Skip to content

Deprecate AdvancedPS - #2848

Open
charlesknipp wants to merge 32 commits into
TuringLang:mainfrom
charlesknipp:ck/smc
Open

Deprecate AdvancedPS#2848
charlesknipp wants to merge 32 commits into
TuringLang:mainfrom
charlesknipp:ck/smc

Conversation

@charlesknipp

@charlesknipp charlesknipp commented Jul 16, 2026

Copy link
Copy Markdown

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.produce no longer occurs in the accumulator to ensure that a varinfo reflects the produces log score
  • it is no longer necessary to include resample and varinfo in the task global storage
  • ParticleContainer no longer requires tracked RNG, which should be handled in the sampler
  • resampling looks a lot different, but behaves as it did from AdvancedPS
  • SMC sampler overloads AbstractMCMC.sample instead of stepping through and bundling with the API

Minor User Facing Changes

  • resamplers are now structs (Multinomial, Systematic, and ESSResampler)
  • sampling with SMC now supports the ensemble key word argument to parallelize the re-weight step

Criticisms

  • I wonder if the model traces and Libtask utilities belong in an extension like DynamicPPLLibtaskExt.jl
  • I am split on whether I should create a new struct for ParticleContainer or keep it as an alias for Vector{<:Particle}

Feel free to make changes. I am open to suggestions. Otherwise, unit tests pass locally so it should be ready to merge.

@codecov

codecov Bot commented Jul 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 94.36620% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 84.05%. Comparing base (98e9773) to head (1c0be67).
⚠️ Report is 2 commits behind head on main.

Files with missing lines Patch % Lines
src/mcmc/particle_mcmc.jl 93.96% 16 Missing ⚠️
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.
📢 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.

@charlesknipp

Copy link
Copy Markdown
Author

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

Comment thread src/mcmc/particle_mcmc.jl Outdated
Comment thread src/mcmc/particle_mcmc.jl
Comment thread src/mcmc/particle_mcmc.jl Outdated
Comment thread src/mcmc/particle_mcmc.jl Outdated
Comment thread src/mcmc/particle_mcmc.jl Outdated
Comment thread src/mcmc/particle_mcmc.jl Outdated

@yebai yebai left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/mcmc/particle_mcmc.jl Outdated
@charlesknipp

charlesknipp commented Jul 21, 2026

Copy link
Copy Markdown
Author

@yebai

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).

The latest commit updates the interface to use taped globals.

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.

Under the hood, smcsample does the exact same thing as AdvancedPS.sweep!, the only differences are (1) when Libtask produces and (2) how traced RNG works.

  1. Producing in the accumulator implies that the tasks varinfo does not yet apply the current log score. As a solution, I produce later in the stack after the likelihood accumulates; using the taped globals to store the log score until I need to produce.
  2. When resampling, split! ensures that each child has a different seed per RNG trace and prunes dead branches when splitting from the reference. Otherwise, we call update_keys! to generate new seeds according to the samplers RNG (not counting the reference), as opposed to storing another trace in the particle container.

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:

@yebai

yebai commented Jul 22, 2026

Copy link
Copy Markdown
Member

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.

avmi pushed a commit to avmi/Turing.jl that referenced this pull request Aug 24, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants