Deprecate AdvancedPS - #2848
Conversation
I fixed this locally, but forgot to commit to the PR Co-authored-by: Penelope Yong <penelopeysm@gmail.com>
Co-authored-by: Penelope Yong <penelopeysm@gmail.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2848 +/- ##
==========================================
- Coverage 85.09% 84.05% -1.04%
==========================================
Files 23 24 +1
Lines 1516 1643 +127
==========================================
+ Hits 1290 1381 +91
- Misses 226 262 +36 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
the unit tests currently failing are only off by a slightly larger absolute tolerance. Not totally convinced its the fault of the sampler, but I'd appreciate some feedback/sanity checks |
Co-authored-by: Shravan Goswami <123811742+shravanngoswamii@users.noreply.github.com>
yebai
left a comment
There was a problem hiding this comment.
I did a partial review today. Here are a few comments:
It would be helpful to include a high-level overview of the design of the new algorithm. I realize this is an unusual request for a code review, but particle Gibbs is a subtle algorithm, and it's difficult to assess correctness without understanding the overall design and implementation.
I think it's worth making a strong effort to avoid task-local storage and keep the implementation as functional as possible (i.e., by tracking all state explicitly in the TapedTask or particle).
I understand this is a substantial piece of work, so I'm happy to iterate on this over a few review cycles.
The latest commit updates the interface to use taped globals.
Under the hood,
There are a handful of similarities between my implementation and AdvancedPS. I suggest comparing specifically this AdvancedPS code and the highlighted section of my rewrite; which should provide a bit more clarity in terms of correctness of implementation. To ensure this implementation is indeed valid, I wrote a few unit tests I feel are worth highlighting:
|
|
Many thanks, @charlesknipp -- can you help take a look at #2853? #2853 passes all existing and extra tests and also fixes a few other legacy issues that date back to the early days of Turing.jl. We could try to find a way to consolidate this PR and #2853 in the end. |
…Lang#2855) ### The bug `PG` / `CSMC` posteriors have been biased since v0.41.0. `AdvancedPS.fork` marks a fork of the reference particle as no longer retained by calling `delete_retained!` for its side effect and [discarding the return value](https://github.com/TuringLang/AdvancedPS.jl/blob/main/ext/AdvancedPSLibtaskExt.jl#L128). Turing's AdvancedPS integration returned a fresh `TracedModel` rather than mutating the one it was handed, so `resample` never became `true`. That flag is what `tilde_assume!!` consults. With it still `false`, a descendant of the reference found every address present in its copy of the retained varinfo and took the `InitFromParams` branch, replaying the retained values instead of drawing new ones. Every offspring of the reference was therefore a duplicate of it rather than a branch off it, the population lost the diversity conditional SMC depends on, and the retained trajectory was over-selected. The error is minor but statistically detectable. On a two-state HMM with ten observations, `PG(16)` state marginals over 20k draws sat up to 7.0 batch-means standard errors from the exact forward–backward values, and 1.3 after the fix. Under `Gibbs(:p => MH(), :z => CSMC(16))`, with `p` the transition probability, 6.4 and 7.4 standard errors on two seeds, and 1.7 after. ### The regression test This PR adds a targeted regression test. The test reaches the sampler only through `sample`, so it applies to any PG/CSMC implementation, not to this one's internals. The model is a two-state hidden Markov chain of length $T = 8$ whose transition parameter is a second unknown: $$ \begin{aligned} i &\sim \mathrm{Categorical}(1/2, 1/2) \\ z_1 &\sim \mathrm{Categorical}(1/2, 1/2) \\ z_t &\sim \mathrm{Categorical}(P_{z_{t-1}, 1}, P_{z_{t-1}, 2}) \\ y_t &\sim \mathcal{N}(\mu_{z_t}, \sigma^2) \end{aligned} $$ for $t = 2, \dots, T$, where $P_{k,k} = p$ and $P_{k,3-k} = 1 - p$, so the chain stays put with probability $p = (0.35, 0.65)_i$. The emission means are $\mu = (-1, +1)$ with $\sigma = 0.8$. Sampled with `Gibbs(@varname(i) => MH(), @varname(z) => CSMC(8))`. ### Why this model Conditional SMC leaves $p(z \mid i, y)$ invariant only if the reference particle is *exactly* the retained path $z^{(n)}$. There are two ways to lose that, and each ingredient of the model is there to expose one of them. **Descendants that copy the reference.** After resampling, a child assigned the reference as its ancestor must continue with fresh randomness; if it keeps replaying the retained values it *is* the reference again, so the population over-counts one path and the weights stop representing the filtering distribution. Detecting this needs uneven weights, hence means separated by $2$ at $\sigma = 0.8$: putting a state in the wrong place costs $\approx (2/\sigma)^2/2 \approx 3$ nats, so the ESS gate triggers resampling at almost every step. **A reference rebuilt by replaying uniforms.** Replay regenerates the path by pushing the stored uniforms back through the sampling map $z_t = g(u_t; \theta, z_{t-1})$, where $g$ is the inverse CDF of the latent's conditional — here $z_t = z_{t-1}$ iff $u_t < p$. That is faithful only while $p$ is fixed: once the MH step moves $i$, every $u_t$ between $p$ and $p'$ flips its step, and because $z_{t-1}$ is itself an argument of $g$, one flip changes what "stay" means for the rest of the path. Two design consequences follow, in the notes below. ### Exact target The configuration space is finite, $2 \cdot 2^{8} = 512$ points, so the posterior is available in closed form — weighted by the model's own log density $\ell$, not by a reimplementation of it: $$ \pi(i, z \mid y) = \frac{e^{\ell(i,z)}}{\sum_{i', z'} e^{\ell(i', z')}} $$ and the state marginal it implies is $$ \pi_t = \Pr(z_t = 2 \mid y) = \sum_{(i,z) : z_t = 2} \pi(i, z \mid y) $$ The statistic over 6000 draws is the mean absolute error of the state marginals: $$ D = \frac{1}{T}\sum_{t=1}^{T} \left| \hat{\pi}_t - \pi_t \right| < 0.01 $$ Averaging over $t$ rather than taking a maximum is deliberate: the Monte Carlo error at each $t$ has random sign and partly cancels, while both defects shift many marginals in the same direction and add up. ### Result The same test body, three seeds, four implementations: | Implementation | 468 | 469 | 470 | Verdict | |-----------------------------------|--------|--------|--------|----------| | main + this fix | 0.0053 | 0.0045 | 0.0024 | passes | | main as-is (cloned reference) | 0.0354 | 0.0264 | 0.0303 | bias detected | | TuringLang#2848 (replayed reference) | 0.0219 | 0.0244 | 0.0227 | bias detected | | TuringLang#2853 (the rewrite) | 0.0032 | 0.0022 | 0.0049 | passes | Twelve runs, no misclassification. The worst passing value leaves 1.9x headroom under the threshold and the weakest bias signal sits 2.2x above it, with a factor of four between the two groups and nothing in between. The testset costs about 50 s. That TuringLang#2853 passes on the same tolerance is a check on the test rather than on that branch: it shares no code with the AdvancedPS-based implementation, so anything the test measures is a property of the sampler's output distribution. ### Design notes on the regression test *The parameter assigned to the `MH` Gibbs step has to be the transition.* The map $g$ involves only the parameters of the conditional being sampled. An emission parameter such as $\sigma$ appears in $p(y_t \mid z_t)$, a density that is evaluated and never sampled, so $g$ never sees it, and replay reproduces the retained path exactly — the reference comes out right by accident. Measured: with the parameter moved to the emission, the error on TuringLang#2848 falls from 0.017–0.022 to 0.0028, which is noise. *The latents have to be dependent.* With independent latents, the target factorises, $p(z \mid \theta, y) = \prod_t p(z_t \mid \theta, y_t)$, and the reference has no lineage to corrupt: its value at step $t$ enters only its own weight at step $t$. A reference regenerated under $p'$ is then one more prior-like draw among the $N - 1$ fresh ones, exchangeable with them, and the population is a plain importance sampler for the right target. In the Markov chain, the reference *is* a lineage — its descendants inherit the corrupted prefix at every resampling step — so the error compounds along $t$ and again across outer Gibbs iterations, which shows up as a shift in path functionals: the expected number of switches moves by $+0.06$ to $+0.14$. --------- Co-authored-by: Claude Code <noreply@anthropic.com>
To facilitate a more uniform ecosystem, this PR is a means of deprecating AdvancedPS in favor of a Turing specific sequential Monte Carlo interface.
For context, this is a little more than a reorganization.
Design Choices
Libtask.produceno longer occurs in the accumulator to ensure that a varinfo reflects the produces log scoreresampleandvarinfoin the task global storageParticleContainerno longer requires tracked RNG, which should be handled in the samplerAbstractMCMC.sampleinstead of stepping through and bundling with the APIMinor User Facing Changes
Multinomial,Systematic, andESSResampler)Criticisms
DynamicPPLLibtaskExt.jlParticleContaineror keep it as an alias forVector{<:Particle}Feel free to make changes. I am open to suggestions. Otherwise, unit tests pass locally so it should be ready to merge.