Do we need to store linked variables in VarInfo?

[mhauru]

As far as I can see, the only reason to do linking is because samplers need it. Samplers moreover interface with a vector of floats, not a VarInfo. So really the only place when we need linked variables is when we call getindex_internal or something like vi[:]. If I'm mistaken about this, and there's a different need, all of the below may be moot, so do say so.

This lead me to wonder why we bother storing linked variables at all, rather than converting at getindex/setindex time. Currently our philosophy seems to be that samplers give and take vectors of linked floats, and when interfacing with a sampler we typically call vi[:] and unflatten(vi, vec_of_floats) to read/write those vectors of floats from/to the internal storage of the VarInfo, where the variables are linked if link has been called on vi at some point in its history. Then whenever we need the unlinked versions of variables, such as when calling logpdf, we convert the internal, linked storage to the unlinked one.

An alternative approach would be to say that VarInfo only ever stores unlinked variables, and rather than calling link on vi once at the beginning of sampling and using getindex/setindex, if you would want linked variables you would need to call getindex_linked and setindex_linked or unflatten_linked, and the linking would happen on the fly in that call.

Needs linked                                Needs unlinked 
                                                           
┌──────────┐         ┌───────────┐          ┌─────────────┐
│          │         │           │          │             │
│ Sampler  │ ◄─────► │  VarInfo  │  ◄─────► │ logpdf call │
│          │         │           │          │             │
└──────────┘   ▲     └───────────┘     ▲    └─────────────┘
               │                       │                   
               │                       │                   
                                                           
          Alternative:             Current:                
          (Un)link here            (Un)link here           

I'm not sure if the alternative is better, but I do see some benefits to it:

• VarInfos would have less state that they carry. Currently when you call getindex_internal or unflatten you need to know/remember whether this VarInfo has been linked at some point in its history (or check for it) to know what you'll get.
• Sampler interface code would be more explicit, at the call-site, as to whether it wants linked variables or not. This would make the interface code a bit easier to understand.
• VarInfo would be simpler. We could drop the istrans flags, and never worry about things like having some variables linked and others not. Note that we would still need to carry the linking transformations in the VarInfo, to know how to transform when someone calls e.g. getindex_linked.

This could also help us move towards something like SimpleVarInfo or the @parameters struct of TuringLang/Turing.jl#2492, where the storage in VarInfo would more closely match how the model sees the variables. I haven't thought about this carefully, but I could imagine this leading to some performance gains and maybe code simplification.

This is currently more an idea to be discussed than a proposal to implement. Any thoughts? Downsides to linking at getindex/setindex time?

[yebai]

I don't recall a good reason for the current design. It could be something inherited from the good old days. One vague motivation was that bijectors used to be constructed on the fly at the logpdf call site, so that is the only place we can perform linking.

@torfjelde @sunxd3, are you aware of any specific reasons for preferring the current approach?

[sunxd3]

I don't know of a reason for it either. (JuliaBUGS always stores the untransformed values, but JuliaBUGS is simpler than DynamicPPL too.)

[torfjelde]

There are ofc efficiency and numerical accuracy concerns with this (see below), but the main issue, which, I don't quite see a way around atm, is variable constraints where the transformations require knowledge of the realizations of the other variables in the model.

Example is the demo_dynamic_support model in the TestUtils. Here you have the support of one variable depending on the value the other variable takes on, leading to the transformation used in (inv)linking potentially changing between evaluations.

How would we handle this if we don't save the transformed variables in the varinfo?

And regarding efficiency and numerical accuracy.

We have to do two things with transformed variables:

1. Compute constrained -> unconstrained and vice versa.
2. Compute logabsdetjac correction to the logpdf.

What you're suggesting would lead to at least one more call of (1), since we would now have to do unconstrained (in the sampler) -> constrained (in the logpdf) -> unconstrained (in the sampler).

Moreover, with_logabsdet_jacobian can be used to compute both (1) and (2) in one go, which is more efficient than doing them separately.

For numerical accuracy, I believe we ran into some issues with LKJBijector or similar due to unnecessary passes through constrained -> unconstrained. I don't remember the details and so it might be unrelated, but issues on this should be floating around somewhere.

Both of the above is ofc worsened when we start involving AD too.

Whether it's worth it or not, is a different question.

[yebai]

Related: TuringLang/Turing.jl#1583

[mhauru]

Thanks @torfjelde, that's really helpful. I'll need to give both dynamic models and with_logabsdet_jacobian a think. For dynamic models I wonder if this could be handled by storing in the VarInfo not a bijector, but a function that returns a bijector given the values of the other variables. This might require also storing the order in which variables appear in the model (related to #833) so that the aforementioned function would only depend variables that came before. This might get clunky.

[mhauru]

This has also made me realise that what we do now with dynamic models is a unintuitive as well. Namely,

@model function demo_dynamic_constraint()
    m ~ Normal()
    x ~ truncated(Normal(), m, Inf)
end

when sampled in unconstrained space is such that the value of x may change even if the sampler does not change it, if m is changed. I've probably had this same realisation around 3 times now since I started on the project, but I seem to not internalise it, and I think it's because I don't fundamentally understand what a VarInfo really is.

[torfjelde]

For dynamic models I wonder if this could be handled by storing in the VarInfo not a bijector, but a function that returns a bijector given the values of the other variables

Hmm, not sure I'm fully on board with this being simpler than allowing the internal representation of VarInfo to be different from that of the model 😕

when sampled in unconstrained space is such that the value of x may change even if the sampler does not change it, if m is changed

This what you want to happen though, right?

[yebai]

when sampled in unconstrained space is such that the value of x may change even if the sampler does not change it, if m is changed

This what you want to happen though, right?

Unfortunately, no. This is desirable from a user perspective, but it creates difficulties for HMC: warm-up in HMC generally doesn't work very well when the support of model parameters changes during inference. That said, we can keep this feature, but better to provide an informational message to warn the modeller to modify their model to avoid such uses.

[torfjelde]

Unfortunately, no.

I meant from the perspective of functionality; this will ofc make things different for samplers, but one could definitively implement samplers which took these things into account, and so it seems a bit drastic to not allow such a model to be used at all.

But agree with informing the user 👍 Think this is the sort of case where people should only be doing it if they really know what they're doing

[mhauru]

This what you want to happen though, right?

I wouldn't say that it's what I want, but I can't think of a better alternative. I find that this sort of stuff makes it hard to reason about how things work, but maybe it's irreducible complexity (rather than reducible, i.e. fixable with a better design).

[penelopeysm]

I feel like I have come back to this question about 10 times now, and if it keeps bringing me back, I feel like there must be at least a pretty strong argument for at least trying to get rid of linked values in VarInfo. I don't immediately see the issue with the dynamic-support distribution, since the transform is performed in tilde_assume!! (which sees dist, so is always correct) and thus the correct unlinked value can be stored in the VarInfo.

The main suboptimal situation I can see happening, which is similar to Tor's comment above, is if the output VarInfo of evaluate!! needs to be reused to extract the linked parameters again. That would incur a linked->unlinked->linked cost. But:

1. Anything that uses LogDensityFunction doesn't even see a VarInfo (technically it sees an OnlyAccsVarInfo), so by definition the VarInfo doesn't need to contain linked values.

2. Surely, in some cases the sampler already has that information? Consider MH, which might generate a new proposal in linked space, and then has to call evaluate!! with a set of linked params to find out whether it's a good step or not. But in that case, no new params need to be extracted from the VarInfo, since the sampler already has access to them. In fact you should just use OnlyAccsVarInfo for this.

3. The remaining possibility is sometimes you need to sample new values in linked space, i.e., essentially what InitContext with a linked VarInfo. In this case, you need to get back values that are linked. I legitimately don't know when this happens -- probably somewhere in Gibbs.

   • But you could do that with a (hypothetical) LinkedValueAccumulator, which would only be used when needed, and there's no reason why VarInfo / VNT itself needs to store the linked values?
   • Alternatively, why can't you just store the unlinked values, and then re-link at the end using one extra model evaluation? We already have an extra model evaluation in Gibbs so if we could find a way to lump it in, then we don't even need the accumulator.

Note that in tilde_assume!! we already have access to both the linked and unlinked values. To avoid the link-unlink-link cost, we just need to change the signature of accumulate_assume!! to pass in both the linked and unlinked values. The individual accumulator can decide what to do with them. ValuesAsInModelAccumulator, and the log-prob calculators, can use the unlinked values. LinkedValueAccumulator can use the linked values.

Of course when I say 'linked values' what I really mean is a tuple of (potentially linked value, invlink transform), where the latter might be identity.

With Metadata-backed VarInfo, there was less motivation to get rid of linking, because the data structure with and without linked variables is pretty much the same: the only thing that changes is the length of the internally stored vectors.

With VNT-backed VarInfo, getting rid of linking makes things MUCH, MUCH simpler, since we don't have to muck around with this TransformedValue stuff. We would still need the general ArrayLikeBlock mechanism for other purposes, such as extract_priors and also the LinkedValueAccumulator, but we wouldn't need to use it inside VarInfo itself. In particular, this would make VarInfo's metadata entirely equivalent to ValuesAsInModelAccumulator, which presents an obvious path to getting rid of VarInfo's metadata by delegating it to VAIMAcc.

[penelopeysm]

I am pretty sure I can implement this shortly after we tidy up VNT's data structure.