"Fixes" for PG-in-Gibbs

[penelopeysm]

This PR:

(1) Adds ProduceLogLikelihoodAccumulator on steps n ≥ 2 of PG, otherwise Libtask.produce is never called.

(2) Removes all uses of set_retained_vns_del!, to avoid errors with PG-within-Gibbs (issue is that it would get a VariableOrderAccumulator from a different component sampler that didn't have the variables that PG cared about).

The role of set_retained_vns_del! was to mark specific variables for resampling by comparing order against num_produce. The idea was that we would only mark variables for resampling if they were 'in the future' i.e. we hadn't reached their tilde-statements in the model yet.

Instead of doing that, we now simply mark everything as to be resampled. This is fine because models can only be propagated in the forward direction i.e. there is really no harm in marking previously seen variables for resampling because we won't see them again! There are two cases where this could lead to different behaviour:

• if the same variable appears on the lhs of tilde twice, which is generally invalid anyway.
• if the same varinfo (with the del flags set) is later re-used. To circumvent this problem, I introduce an inverse function which marks all variables as not being for resampling. It turns out that there is only one place where this inverse function has to be called: specifically, it has to be called before creating the reference particle on the next pMCMC iteration, to ensure that the trajectory of the reference particle is indeed fixed.

Collectively, this completely removes the need to track which variables have already been seen. With these proposed changes there is simply no need to track the order of a variable or the num_produce of the varinfo, meaning that we can remove VariableOrderAccumulator entirely (with concomitant positive implications for performance).

I have not exhaustively tested these changes on breaking, but I have tried making exactly the same changes to the del flag on main and they lead to no changes in the sampled results if the PRNG seed is set.

---

Note

While we aren't quite there yet, this change also paves the way for eventually getting rid of the del flag entirely (TuringLang/DynamicPPL.jl#982).

Notice that the changes in this PR already remove the need for a per-variable del flag; instead we can have a single boolean for the entire VarInfo which marks whether something should be resampled or not. This would already represent a substantial simplification.

However, ultimately we should like to shift this information outside of the VarInfo because then it means that VarInfo doesn't need to track information that only PG needs. The way I would propose to do this is with contexts: one which never resamples variables (which would correspond to all del flags unset), and one which always resamples variables (corresponding to all del flags set). Notice further that these roles are very similar to DefaultContext and the upcoming InitContext. There are probably a ton of subtleties (e.g., in that we have to use the traced rng rather than just naively resampling), but I think there is light at the end of the tunnel.

[github-actions]

Turing.jl documentation for PR #2629 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2629/

[codecov]

Codecov Report

❌ Patch coverage is 0% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 24.43%. Comparing base (c062867) to head (1c424ab).
⚠️ Report is 1 commits behind head on mhauru/dppl-0.37.

Files with missing lines	Patch %	Lines
src/mcmc/particle_mcmc.jl	0.00%	19 Missing ⚠️

Additional details and impacted files

@@                  Coverage Diff                  @@
##           mhauru/dppl-0.37    #2629       +/-   ##
=====================================================
+ Coverage              0.00%   24.43%   +24.43%     
=====================================================
  Files                    22       22               
  Lines                  1485     1502       +17     
=====================================================
+ Hits                      0      367      +367     
+ Misses                 1485     1135      -350     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[penelopeysm]

Since we don't ever use num_produce, we don't need to modify it either.

[penelopeysm]

Changes in this file are all due to the issue with Gibbs-conditioned-variables now triggering resampling. Thus now it's forbidden to have a dynamic number of Gibbs-conditioned-variables.

[penelopeysm]

In general, we now can't handle the different number of Gibbs-conditioned-variables. However, in the specific case of these tests, we can change the model to something equivalent which has the same number of observations:

@model function dynamic_model_for_ess()
    b ~ Bernoulli()
    if b
        x ~ MvNormal([1.0], I)
    else
        x ~ MvNormal([1.0, 2.0], I)
    end
end

The one below this can similarly use filldist.

Do you think it's worth keeping these tests? i.e. was the point to make sure that the sampling results are numerically correct, or was the point of the test to make sure that you could have different numbers of Gibbs-conditioned variables (in which case rewriting it would defeat the purpose)?

[penelopeysm]

(If it's the former, I think we should include a likelihood term, I think being able to sample from the prior correctly is a bit of a low bar to clear for an inference algorithm implementation.)

[mhauru]

It's a regression test, and I don't know what the original issue was that caused it to error. We don't have too many dynamic model tests, so I would keep it with the MvNormal because it doesn't cost us much and might still catch something.

An alternative, though: Any interest in turning this into a test to check whether you can still do dynamic models with loops if you balance out the branches with @addlogprob!!(0.0) calls? Hacky, but I kinda want to know that it works as a workaround if need be. So that would be something like

@model function dynamic_model_for_ess()
    b ~ Bernoulli()
    x_length = b ? 1 : 2
    max_x = 2
    x = Vector{Float64}(undef, x_length)
    for i in 1:x_length
        x[i] ~ Normal(i, 1.0)
    end
    for _ in (x_length+1):max_x:
        @addlogprob!!(0.0)
    end
end

[penelopeysm]

Okay, happy to do that, will add a comment with the intent too

[penelopeysm]

Currently unsure what to do with the tests.

@model function dynamic_bernoulli()
    b ~ Bernoulli()
    if b
        x ~ MvNormal([1.0], I)
    else
        x ~ MvNormal([1.0, 2.0], I)
    end
end

chain = sample(dynamic_bernoulli(), Gibbs(:b => PG(10), :x => ESS()), 2000; discard_initial=100)

This errors because if x in the global varinfo is a length-2 vector and PG selects b = 1, then logpdf(MvNormal([1.0], I), x) will fail.

@model function dynamic_bernoulli_2()
    b ~ Bernoulli()
    x_length = b ? 1 : 2
    max_x = 2
    x = Vector{Float64}(undef, x_length)
    for i in 1:x_length
        x[i] ~ Normal(i, 1.0)
    end
    for _ in (x_length+1):max_x
        @addlogprob!(0.0)
    end
end

chain = sample(dynamic_bernoulli_2(), Gibbs(:b => PG(10), :x => ESS()), 2000; discard_initial=100)

This runs but gives a rather incorrect value of b = 0.75 (ish, depends on number of iterations and stuff).

[penelopeysm]

I think I actually understand why the mean is wrong now (at least when it runs; I don't know why it crashes). Suppose that x[1] = 1.0 and x[2] = 2.0 (on average this is true).

For particles with b = true (i.e. only observe x[1]), Libtask.produce gets called twice: once with logpdf(Normal(1.0), 1.0) = -0.919 and once with a manual logpdf of 0.

For particles with b = false (i.e. observe x[1] and x[2]), Libtask.produce gets called twice, once with
logpdf(Normal(1.0), 1.0) = -0.919 and once with logpdf(Normal(2.0), 2.0) = -0.919.

When PG starts its thing, it just samples b from the prior, so there should be 50% true and 50% false.

On the first reweighting, the logpdf's (i.e., weights) are equal and thus one would expect that the resampling step generates an equal amounts of true and false.

On the second reweighting however the weights are different. The unnormalised weights are exp(0) = 1 and exp(-0.919) = 0.3989, which is a ratio of about 2.5, so when resampling you're approximately 2.5x more likely to get a particle with b = true rather than b = false.

So mean(b) would be expected to be 2.5 / (2.5 + 1) = 0.715, which is sort of in the right region.

Furthermore, in general x[2] will not be equal to 2.0, meaning that logpdf(Normal(2.0), x[2]) will be smaller and the ratio of weights larger than 2.5, and mean(b) will be even more biased towards 1. This is exactly what we observe with mean(b) coming out at 0.772.

[penelopeysm]

Indeed if you make this change:

  for _ in (x_length+1):max_x
-     @addlogprob!(0.0)
+     # same as @addlogprob!(logpdf(Normal(2.0), 2.0))
+     @addlogprob!(logpdf(Normal(), 0.0))
  end

you get back a value of mean(b) = 0.592 which is closer to 0.5 (albeit still larger than 0.5 because of the last paragraph in the previous comment).

And to completely remove the bias, I think you'd have to add a term that is the expectation value of logpdf(Normal(), x) where x ~ Normal(). That's a slightly nasty integral, but we can approximate it with Monte Carlo (or alternatively use WolframAlpha):

julia> mean(logpdf.(Normal(), randn(10_000_000)))
-1.418849291536427

Calling @addlogprob! with this gives... mean(b) = 0.5050!

[penelopeysm]

Finally, none of this actually explains why it used to "work" with the previous non-tilde-observe Gibbs sampler. The current behaviour makes complete sense to me given our tilde-observe-Gibbs implementation, but it still seems to me that the exact same thing should have happened before this PR. To understand that I might have to take another long walk.

[mhauru]

Thanks. That makes sense.

I think I know why it used to work: The problem was masked by another bug. Namely, we would accumulate the logp contributions from the Gibbs-conditioned tilde_assume statements into the VarInfo, and then add it to the score/log likelihood at the next produce/resampling, with the line return score + DynamicPPL.getlogp(varinfo) in advance!. But in the above model there is no next produce/resampling, because there is no likelihood term. Observe instead this, running on latest release:

julia> module MWE

       using Turing

       @model function dynamic_bernoulli()
           b ~ Bernoulli()
           x_length = b ? 1 : 2
           x = Vector{Float64}(undef, x_length)
           for i in 1:x_length
               x[i] ~ Normal(i, 1.0)
           end
           0.0 ~ Normal()
       end

       chain = sample(dynamic_bernoulli(), Gibbs(:b => PG(10), :x => ESS()), 2000; discard_initial=100)
       describe(chain)
       end
WARNING: replacing module MWE.
Sampling 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| Time: 0:00:04
Chains MCMC chain (2000×4×1 reshape(::Matrix{Union{Missing, Float64}}, 2000, 4, 1) with eltype Union{Missing, Float64}):

Iterations        = 101:1:2100
Number of chains  = 1
Samples per chain = 2000
Wall duration     = 4.23 seconds
Compute duration  = 4.23 seconds
parameters        = b, x[1], x[2]
internals         = lp

Summary Statistics
  parameters      mean       std       mcse    ess_bulk    ess_tail       rhat   ess_per_sec
      Symbol   Float64   Float64   Float64?    Float64?    Float64?   Float64?      Float64?

           b    0.7975    0.4020     0.0098   1692.0019         NaN     0.9995      399.7170
        x[1]    0.9588    1.0110     0.0232   1888.4870   1014.7637     0.9999      446.1344
        x[2]    2.0290    0.9868    missing     missing     missing    missing       missing

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

           b    0.0000    1.0000    1.0000    1.0000    1.0000
        x[1]   -1.0495    0.2851    0.9589    1.6325    2.9982
        x[2]    0.1122    1.2627    2.0212    2.6988    3.9777

[mhauru]

On your earlier explanation of how PG works, that's helpful, and matches what I imagined it would be. In that light, I just didn't understand how this could come to pass:

I actually wonder if this is a Libtask thing. Let's say that we have a particle with b == false, so x_length = 2. After it sees x[1], it has to resample. Suppose it then resamples (transforms itself) into a particle with b == true and x_length = 1. What happens then, does it go on to sample x[2] or does it exit the loop because the termination condition x_length == 1 is now hit?

[mhauru]

One comment/proposal, otherwise I think all good to go.

[mhauru]

It's a regression test, and I don't know what the original issue was that caused it to error. We don't have too many dynamic model tests, so I would keep it with the MvNormal because it doesn't cost us much and might still catch something.

An alternative, though: Any interest in turning this into a test to check whether you can still do dynamic models with loops if you balance out the branches with @addlogprob!!(0.0) calls? Hacky, but I kinda want to know that it works as a workaround if need be. So that would be something like

@model function dynamic_model_for_ess()
    b ~ Bernoulli()
    x_length = b ? 1 : 2
    max_x = 2
    x = Vector{Float64}(undef, x_length)
    for i in 1:x_length
        x[i] ~ Normal(i, 1.0)
    end
    for _ in (x_length+1):max_x:
        @addlogprob!!(0.0)
    end
end

[penelopeysm]

Pretty sure that the diff of this PR matches when you reviewed so will merge. Dynamic models to be discussed at next Monday's team meeting...