NamedTuple-valued random variables #289
Red-Portal
started this conversation in
Ideas
Replies: 1 comment
-
So we already have |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Hi, I think it might be interesting to consider whether we ever want to support NamedTuple-valued random variables.
Currently, most of the inference algorithms expect vector-valued random variables.
This is probably fine for MCMC algorithms running on the CPU.
But for more advanced inference algorithms that require some introspection into the model and GPUs,
I think this interface is becoming a liability.
Variational Inference
For example, to support subsampling in variational inference, we need to deal with models that have local variables. In this case, we need to operate on a subset of random variables at each iteration. That is, for a model like this
we wish to subsample over$i = 1, \ldots, n$ . Then, for each batch $j$ , we need to select the axis such that $\left[x,\, y_j\right]$ . This is not very nice to do if we flatten everything. For instance, we receive a sample $\theta$ one way to do this would be to unflatten and obtain the structured representation $\left[\, x, \, y_1, \, \ldots, \, y_n \right]$ , select the variables $\left[\, x, \, y_i \right]$ , and then flatten back.
If the random variables$\theta$ were already structured not flattened, we can entirely skip the redundant flattening-unflattening process.
Subsampling MCMC
The story is similar for implementing subsampling-based MCMC algorithms.
RJMCMC
This also concerns certain MCMC algorithms that need to inspect the structure of the parameters, such as RJMCMC. For example, if we know the name of the model indicator random variables, the algorithm can simply deduce the model order from the length of these variables and do business as usual. However, if the random variables are vector-valued, we need additional machinery to do this. There is also the problem of having flatten-unflatten parameters of varying dimensionality.
GPU Support
There is also a concern for supporting GPUs. As we all know, Bayesian models often involve lots of scalar-valued hyperparameters. Since evaluating the likelihood involves flattening-unflattening, we have to face lots of scalar-indexing, which don't play nice with Julia's GPU ecosystem. With
NamedTuple
-valued RVs, however, we can probably mix the types of the different random variables, and leave some things on the CPU when appropriate.What Needs to Change?
LogDensityProblems
: Probably doesn't need to change except forwith_gradient
andwith_hessian
. We could also makeLogDensityProblems.dimension
to return aNamedTuple
of dimensions.Bijectors
: New interfaces forStacked
andTransformedDistribution
will have to be added.DynamicPPL
: Probably an additional specialization forLogDensityProblems
?StatsBase
: Mayberand
andlogpdf
?I would love to hear yous thoughts on this!
Beta Was this translation helpful? Give feedback.
All reactions