Reimplement particle MCMC (SMC and PG) natively, removing AdvancedPS#2853
Reimplement particle MCMC (SMC and PG) natively, removing AdvancedPS#2853yebai wants to merge 34 commits into
Conversation
SMC/PG/CSMC are reimplemented directly on Libtask + DynamicPPL, dropping the AdvancedPS dependency. A particle is a suspended model execution advanced one observation at a time; the conditional-SMC reference trajectory is regenerated by replaying a per-particle TracedRNG (Random123 Philox), so no reference values need to be stored, and a child forked from the reference is simply reseeded to branch off. All particle state lives explicitly on the particle via Libtask's taped globals -- no task_local_storage -- and @addlogprob! reweights by producing inside the likelihood accumulator, so no type piracy is required. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
- Inline traced_rng.jl into particle_mcmc.jl as a section (it must precede Particle/PGState, which name TracedRNG in their types). - Merge the container unit tests into test/mcmc/particle_mcmc.jl and remove the now-empty essential test group. - Drop the set_taped_globals! call in fork: deepcopy already preserves the task-particle back-reference (verified against the full particle suite). - Add WHY notes for the might_produce hints and the reference RNG rewind. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Stratified resampling satisfies the unbiased-offspring condition needed for the particle Gibbs invariance proof (Andrieu, Doucet & Holenstein, 2010) and is consistent as N grows. Systematic resampling shares the expected offspring counts but is order-dependent and not consistent in general (Gerber, Chopin & Whiteley, 2019), so it sits outside the invariance proof. Make stratified the default scheme for ESSResampler (hence for SMC and PG); systematic remains available explicitly. A note in the resampling section records the trade-off. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Chain metadata recomputes the log-prior and Jacobian terms from a particle's raw values, so accumulating them during the sweep was wasted work; particles now carry only the produce-aware likelihood accumulator and the raw values. Metadata stays correct (test_chain_logp_metadata passes for SMC and PG). Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
resample_propagate! deepcopied every resampled slot; deepcopying a Libtask TapedTask is the dominant allocation in a sweep, so this made PG allocate ~40% more than the AdvancedPS implementation (59.9M vs 41.9M allocations; ~22% slower). Reuse each surviving parent's object for its first offspring, deepcopying only additional offspring and any descending from the retained reference (as AdvancedPS does). Both kinds of child are reseeded via the new `reseed!` helper. Allocations now match the old implementation and PG is slightly faster; CSMC reference reproduction is unchanged (0 mismatches over the reference-consistency check, full suite green). Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Open the file with a concise "Key design" overview (the observe-as-filtering-step idea, the RNG-replay reference mechanism, and the fresh-seed-per-step invariant it rests on) plus a section map. Move the threadsafe-eval guard into the model-evaluation section with a note on why particle samplers require it, and note the log-evidence telescoping in the sweep. Comments only; no behaviour change. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Add a HISTORY.md breaking-change entry for the native SMC/PG reimplementation (AdvancedPS removed, resamplers are now types, stratified is the new default), and add test_rng_respected coverage for SMC and PG so RNG determinism is locked in. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Bring across the valuable tests the ck/smc branch added, adapted to the native internals: - resampling schemes: stratified and multinomial both hit the analytic coinflip posterior, and the two schemes produce genuinely different draws; - CSMC reference consistency: over 30 conditional sweeps the reference particle regenerates the retained trajectory exactly and its traced-RNG keys stay aligned with the trajectory length (previously only checked in a scratch script); - @addlogprob! is also respected under MH, not just PG; - a particle advanced without resampling matches a direct init!! evaluation (values and log-likelihood) seeded identically. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
SMC(Systematic()) and friends require Turing.Inference qualification: the scheme types are not exported, and a clean export is blocked by the Multinomial resampler colliding with the re-exported Distributions.Multinomial. Document this in the SMC docstring (whether to export a renamed subset is a separate API decision). This matches the previous AdvancedPS-qualified API, so it is not a regression. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
inc_step! only ever advances by one and set_step! was only ever called as set_step!(_, 1). Remove the unused `n` parameters: inc_step! takes no count, and set_step! becomes rewind! (reset to the first step), which also reads clearer at the call site where the reference RNG is rewound for replay. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
SMC's `sample` override already calls error_if_threadsafe_eval before delegating to mcmcsample, so the same check at the top of SMC's first `step` never fires on a distinct path. Remove it. PG keeps its check, being the only guard there (PG has no `sample` override). Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Use DynamicPPL.LogProbType for particle log-weights and the log-evidence rather than a hardcoded Float64, and make ESSResampler parametric on its threshold type so it keeps whatever Real the user passes. Also bound the previously-open type parameters (TracedRNG's key type as <:Unsigned, SMCState's fields as <:AbstractVector) so every parametric struct carries meaningful subtype bounds. Integer types were already platform-native `Int`, not `Int64`. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
|
Turing.jl documentation for PR #2853 is available at: |
Random123 is only used by the TracedRNG in particle_mcmc.jl, so scope the import there rather than in Inference.jl. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2853 +/- ##
==========================================
+ Coverage 85.09% 86.25% +1.16%
==========================================
Files 23 23
Lines 1516 1586 +70
==========================================
+ Hits 1290 1368 +78
+ Misses 226 218 -8 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
The native SMC/PG reimplementation changes the resampler API and default, so give it its own breaking release: move the changelog entry under a new 0.47.0 heading and bump the version. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
`Random.seed!(trng::TracedRNG, key)` took an untyped `key`, making it ambiguous with `Random.seed!(::AbstractRNG, ::Nothing)` (Random stdlib) for the call `seed!(::TracedRNG, ::Nothing)`. The Aqua ambiguity check flags this on CI (the conflicting Random method isn't present in every local Random, so it didn't reproduce locally). The seed is always an integer, so constrain `key::Integer`, which removes the overlap with `::Nothing`. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
`split_key` derived a fresh seed by re-seeding a stdlib `MersenneTwister`. That is fragile two ways: re-seeding to split a generator can yield correlated streams (Steele et al., OOPSLA 2014), and MersenneTwister streams are not reproducible across Julia versions, so SMC/PG results drifted between versions even under a StableRNG. Both affected the previous AdvancedPS implementation (#2781, AdvancedPS.jl#110). Derive the seed through Philox instead -- a counter-based generator with a fixed, portable algorithm and strong avalanche -- which is both well-decorrelated and version-stable. The full particle suite, including exact CSMC reference reproduction, stays green. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Add the foundational particle MCMC reference to the module's design note, where the short "(ADH 2010)" citations in the resampling notes point. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
The native particle-MCMC RNG (Philox seed derivation) changes the exact draws, and `StableRNG(23)` now lands a ~2.7σ tail draw for E[s] that just exceeds atol=0.1 at 3_000 iterations. The estimator is unbiased -- E[s] scatters tightly around the true 2.042 across seeds -- so the fix is simply more headroom: CSMC mixes the variance slowly, and 10_000 draws bring the error comfortably inside atol on every Julia version and platform. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
`SMC(; threaded=true)` and `PG(n; threaded=true)` now evaluate the particles across threads within each sweep. Only the per-particle model advances (the expensive part) run in parallel; resampling stays serial, and because every particle's RNG is seeded serially in `resample_propagate!` before the parallel region, the threaded run reproduces the serial draws bit for bit. The default (serial) path keeps its scalar tally and allocates nothing extra. `threaded` is a keyword-only field: an inner constructor suppresses the positional default constructor, so it cannot collide with the existing `(scheme, threshold::Real)` forms (`Bool <: Real`). HISTORY.md gains notes on this, on the cross-version RNG reproducibility, and on the weight-diagnostic columns, plus a mention that multiple chains run under `MCMCThreads()` / `MCMCDistributed()`. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
There was a problem hiding this comment.
Overall, I think this is a very nice revision of #2848. Although, there are a few minor choices I intentionally made when drafting that should be reconsidered here.
- Passing
AbstractMCMCEnsemblethrough toreweight!(only for SMC) for consistency - Producing the log score upon
tilde_observe(oraccloglikelihood!!for@addlogprob!) to ensure log score is up to date with log likelihood accumulation.
Addresses Charles's review comment on internal consistency. Emit the per-step `produce` from `tilde_observe!!` and an `accloglikelihood!!` overload -- after the log-likelihood accumulator has been updated -- instead of from inside `acclogp` before it. This removes the one-step lag between a particle's produced weight and its accumulated log-likelihood, and makes `@addlogprob!` terms reach the accumulator (hence the reported log-likelihood), not only the weight. The produced score is the accumulator's increment, so the accumulator stays the single source of truth: `acclogp` now inherits the generic `LogProbAccumulator` method, and `ProduceLogLikelihoodAccumulator` is just a marker type flagging a particle's varinfo so the produce sites know to emit. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
Per Charles's review: SMC is not an MCMC (nor PMCMC) algorithm, so `SMCContext` names the leaf context more accurately. Pure rename, internal to particle_mcmc.jl. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
Per Charles's review: `Particle{RT<:TracedRNG,WT<:Real}`, so `logweight` tracks
`DynamicPPL.LogProbType` and follows it if that is ever changed. Behaviour is
unchanged -- taped-globals access stays type-unstable as before (the `::Particle`
typeassert still holds for the parametric type), so this only concretises the
stored particle's field types.
Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Co-Authored-By: Claude Code <noreply@anthropic.com>
Answers Charles's review question ("is it always InitFromPrior?"). Verified that a
variable is never already present at `tilde_assume!!` across SMC, PG, CSMC and
PG-in-Gibbs: particle varinfos start empty, each variable is assumed exactly once,
and the CSMC reference reproduces its trajectory by replaying RNG seeds rather than
reusing stored values. So `InitFromPrior` is always correct, and a conditional
`InitFromParams` branch would be dead code unless the varinfo were pre-populated
(e.g. initial_params for particle samplers, which is not currently supported).
Documented rather than adding the inert branch.
Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Co-Authored-By: Claude Code <noreply@anthropic.com>
Per Charles's review: the retained particle already carries the reference trajectory's varinfo and rng needed to resume the next conditional sweep, so particle Gibbs needs no separate state struct. `const PGState = Particle` keeps the semantic name at the PG call sites while adding no new type. `gibbs_update_state!!` now updates the reference varinfo in place, which is safe -- Gibbs replaces the state with the returned value and never reads the pre-update one again. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
…ate the two levels Renames the within-sweep parallelism keyword `SMC`/`PG(; threaded=...)` to `multithreaded`, and documents (in HISTORY and a note at `reweight!`) that it is a distinct axis from AbstractMCMC's chain-level ensemble. Within-sweep threads a single sweep's particle evaluations; the ensemble (`MCMCThreads`/`MCMCDistributed`) runs whole chains independently; the two compose. Only threading is offered within a sweep -- particles resample every step (all-to-all) and are live Libtask tasks, so distributing one sweep across processes would be communication-bound rather than a speed-up. Addresses Charles's review that the parallelism description conflated the two levels, and follows the discussion that these are genuinely different axes. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
AGENTS.md is the tool-agnostic convention; CLAUDE.md now just includes it via @AGENTS.md so both entry points resolve to one source. Co-Authored-By: Claude Code <noreply@anthropic.com> Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
… weighting Per Charles's review: SMC is a single weighted sweep, not a Markov chain, so overload `AbstractMCMC.sample` to run the sweep and bundle the whole population in one shot instead of faking iteration through the step loop with an `SMCState` cursor. Removes `SMCState` and both SMC `step` methods; `discard_initial`/`thinning` then have no loop to (not) apply to. A final resampling step makes the returned particles an equal-weight sample, so downstream chain summaries (`mean`, etc.) are correct without weighting; SMC chains no longer carry a per-particle `weight` column. Also drops the `const PGState = Particle` alias -- with SMC carrying no state struct, PG's state is plainly a `Particle`. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
…rmalizing_constant SMC chains now carry `ess_per_step`: the effective sample size (1/Σŵ²) after each observation's reweight, a sweep-level degeneracy diagnostic reported alongside the normalizing-constant estimate. `sweep!` records the trajectory and returns it (PG discards it), and the ESS formula is factored out of `should_resample` into a shared `ess` helper. Renames the marginal-likelihood / normalizing-constant statistic `logevidence` to `log_normalizing_constant` -- the standard SMC/PMCMC term for the estimated p(y) -- including the internal `log_normalizing_constant(particles)` helper. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
|
Thanks, this was a valuable review. What changed, and where I took a different route: Score ↔ accumulator consistency. The SMC output. Parallelism. I diverged here. PG state Done — the retained Naming / parametrization "Is it always |
The `logevidence` -> `log_normalizing_constant` stat rename missed the SMC log-evidence assertion in test/mcmc/Inference.jl (CI caught it on the min-Julia mcmc/Inference shard). Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
`Multinomial`/`Stratified`/`Systematic` -> `MultinomialResampler`/ `StratifiedResampler`/`SystematicResampler`, so the resampling-scheme types no longer shadow common names (notably `Distributions.Multinomial`) inside `Turing.Inference`. Algorithm-description docstrings are unchanged. (Review finding 4.) Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
The particle-Gibbs reference reproduced its trajectory by replaying RNG seeds, which only regenerates the retained values when the model is unchanged. Inside Gibbs the model is re-conditioned between sweeps -- one component's variables become observations parametrised by the others -- so seed-replay re-drew the reference from a *different* prior rather than reproducing the retained trajectory. CSMC then lost its defining invariant and behaved like independent SMC: exact at large N but biased at small N. On a linear-Gaussian state-space model, Gibbs(:x => PG(2), :ρ => MH()) put E[x1] at 0.64 against NUTS's 0.96; the fix closes the gap. Reproduce by value instead: the sampler state carries the retained values, and the reference re-runs the model with InitFromParams, so it stays the retained trajectory regardless of re-conditioning. A fork clears the carried values (reseed!) and thus samples fresh past the fork point. The existing consistency test only covers same-model reproduction; the new test pins the reference under changed conditioning, the Gibbs scenario. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
softmax can return weights summing to a hair under one; the cumulative walk then falls short of the final uniform and the unguarded loop indexes one past the last particle -- a BoundsError. `s < length(weights)` clamps the walk to the last particle. Bounding by the weight count (not the resample count) keeps the last particle reachable in the conditional sweep, where fewer indices are drawn than there are particles. The realistic trigger is astronomically rare -- softmax undersum is a few ULP, so the uniform must land within that of the stratum edge -- but the overrun is a latent array-index bug; the test exaggerates the undersum to hit it deterministically. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
SMC draws its whole population from the prior, so a user's initial_params cannot apply -- yet it was silently swallowed by kwargs. Warn instead, as SMC already does for discard_initial and thinning. (PG ignores it too, but has no custom sample overload to host the warning without risking spurious ones inside Gibbs, so it is left unchanged.) Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
The reference_values comment stated the conclusion -- value-replay, not seed-replay, keeps the reference on the retained trajectory across Gibbs re-conditioning -- but not the mechanism. A draw is x = g(u; θ) of the RNG output and the distribution parameters; seed-replay fixes u and recomputes x' = g(u; θ') while value-replay reuses x' = x, so the two agree only when θ' = θ. Re-conditioning changes a θ that another block owns, so θ' ≠ θ and a seed-replayed reference would move with the shifted prior instead of holding its retained values. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
Opened for discussion.
Reimplements
SMCandPG/CSMCdirectly on Libtask + DynamicPPL and drops the AdvancedPSdependency. A rewrite based on #2848 (thanks @charlesknipp; commits co-authored).
Each
observeis one filtering step: underParticleMCMCContext, every likelihood termbecomes a
Libtask.produce, so a particle is a suspended model execution we advance, weight,and resample. The conditional-SMC reference is reproduced by replaying a per-particle
TracedRNG(no stored trajectory); a fork is reseeded to branch off. Particle state lives inthe task's taped globals (no
task_local_storage), and@addlogprob!reweights via thelikelihood accumulator.
Posteriors, log-evidence, and exact reference reproduction match the old implementation;
gibbs.jl, Aqua, and log-density-metadata tests pass; performance is on par.Fix #2781
Fix TuringLang/AdvancedPS.jl#110
Fix TuringLang/AdvancedPS.jl#39
Fix TuringLang/AdvancedPS.jl#6