Native particle MCMC (SMC/PG), revised: dropping AdvancedPS - #2859
Conversation
`SMC` and `PG`/`CSMC` are implemented directly on Libtask and DynamicPPL, and the AdvancedPS dependency is gone. Each `observe` is one filtering step: under `SMCContext` every likelihood term becomes a `Libtask.produce`, so a particle is a suspended model execution the sweep advances, weights and resamples. The conditional-SMC reference reproduces the retained trajectory by reusing its values through `InitFromParams`, which is what keeps it on that trajectory when another Gibbs component re-conditions the latents; a fork forgets them and samples afresh. Particle state lives in the task's taped globals rather than `task_local_storage`, and `@addlogprob!` reweights through the likelihood accumulator (#1996). Resampling schemes become types -- `StratifiedResampler`, `SystematicResampler`, `MultinomialResampler`, optionally wrapped in `ESSThresholdResampler` -- with stratified as the default, since systematic is order-dependent and can fail to be consistent as the particle count grows. Conditional sweeps draw their ancestors multinomially, because a correct conditional version of stratified or systematic resampling is scheme-specific rather than "pin one draw and keep the rest". Internal seeds derive through a counter-based Philox generator, so a fixed user seed gives the same draws on every Julia version and platform (#2781, AdvancedPS.jl#110). `SMC` resamples once at the end of a sweep so the returned particles are an equal-weight sample, and chains carry `log_normalizing_constant`, with `ess_per_step` for `SMC`. Fixes #2781 Fixes TuringLang/AdvancedPS.jl#110 Fixes TuringLang/AdvancedPS.jl#39 Fixes TuringLang/AdvancedPS.jl#6 Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
…eference The tilde handlers took a varinfo argument and ignored it, reading `particle.varinfo` instead, so anything the model body had accumulated since the previous tilde was discarded. `@addlogprob!` and `:=` reach the varinfo directly rather than through a handler, so a chain's `loglikelihood` omitted every `@addlogprob!` term and `:=` variables never appeared at all, where every other sampler reports both. The handlers now thread the argument through; a trailing update, which has no later handler to carry it, is mirrored where it happens. That mirroring keys on the produce-aware accumulator, so `gibbs_update_state!!` swaps in the plain one for its out-of-task re-evaluation -- outside a task, `get_taped_globals` returns whichever particle was last consumed rather than failing. A reference now carries only the retained trajectory's values. The set of addresses each particle had assumed went with it, and so did the deep copy it incurred on every fork. Its stated justification -- that a slice assume `x[1:2]` cannot be checked against the stored `x[1]`, `x[2]` -- was wrong: `haskey` resolves it. `check_MoGtest_default` compares against the exact posterior means rather than the idealised cluster labels, which sit far enough from the truth to spend most of the tolerance before any Monte Carlo error. Draws are unchanged throughout. On a T=200 state-space model, 100 PG(16) iterations run in 9.1 s against 10.9 s and allocate 9.4 GiB against 11.6 GiB. Co-Authored-By: Claude Code <noreply@anthropic.com>
|
Turing.jl documentation for PR #2859 is available at: |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2859 +/- ##
==========================================
+ Coverage 85.10% 85.90% +0.79%
==========================================
Files 23 23
Lines 1517 1603 +86
==========================================
+ Hits 1291 1377 +86
Misses 226 226 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Drop a work-stage reference that means nothing to a reader with only the repository, and a stray line about nondeterministic evaluation order that sits above `centred_normal` instead of the model it describes. Assisted-by: AI coding agent
Libtask only instruments calls marked with `@might_produce`, and the `produce` that reweights an `@addlogprob!` term sits below `acclogp_and_mirror!!`, which was unmarked. Reweighting worked only because the compiler chose to inline that helper: force `@noinline` and the term stops reweighting with no error of any kind, leaving PG at the prior (mean 0.10 against a target of 0.79). Assisted-by: AI coding agent
`resample_propagate!` ran at the top of the sweep loop, so every sweep resampled once on the equal weights the particles start with. For a multinomial draw, which every conditional sweep uses, that duplicates particles before any data has been seen and costs diversity for nothing. One resample is still spent after the last observation: a particle reveals that it has finished only by producing nothing, one pass later, so the sweep cannot know which observation was the last. Assisted-by: AI coding agent
PG's state was the retained `Particle`, which owns a live `Libtask.TapedTask`. Serialising a chain sampled with `save_state=true` therefore died inside the taped code, in `write(::IOBuffer, ::Core.IntrinsicFunction)`. `PGState` now holds only the retained trajectory's raw values, which is all the next sweep's reference needs, and `Particle` takes those values directly. `gibbs_update_state!!` re-derives the trajectory on a `trajectory_varinfo`, so the produce-aware accumulator is never installed outside a sweep and the swap that used to guard against its inert `produce` goes away. Nothing carries a generator over: the reference consumes no randomness of its own, and every other particle is seeded from the sampler's `rng`. Assisted-by: AI coding agent
PG draws its first sweep from the prior and dropped `initial_params` without a word, where SMC warns. Gibbs passes the keyword to every component's first step, so the same silence applied there. Assisted-by: AI coding agent
Every particle scoring `-Inf` used to surface as `DomainError with [NaN, NaN, ...]: Categorical: vector p is not a probability vector`, naming neither the model nor the cause. Report it from the sweep instead, at the observation no trajectory could account for. Assisted-by: AI coding agent
A sweep whose last act was a resample leaves every weight zero, so the closing resample drew a second time over equal weights and duplicated particles for nothing. Return the sweep's own population instead. Assisted-by: AI coding agent
There was a problem hiding this comment.
The core logic looks good to me.
Agent found two minor issues I think worth a second look. They are unlikely to impact user too much though.
Here are the MWEs.
using Turing, Distributions, StableRNGs, Test
@model function constrained_posterior()
x ~ Normal()
@addlogprob! (; logprior = x >= 0 ? 0.0 : -Inf)
0.0 ~ Normal(x, 1)
end
smc = sample(
StableRNG(1), constrained_posterior(), SMC(), 1_000; progress=false
)
pg = sample(
StableRNG(1), constrained_posterior(), PG(10), 200; progress=false
)
for chain in (smc, pg)
xs = vec(collect(chain[:x]))
logjoints = vec(collect(chain[:logjoint]))
@show count(x -> x < 0, xs)
@test all(x -> x >= 0, xs)
@test all(i -> xs[i] >= 0 || logjoints[i] != -Inf, eachindex(xs))
endSMC: 507 / 1000 negative draws
PG: 95 / 200 negative drawslogprior didn't affect particle weight (loglikelihood is fine otherwise)
using Turing, Distributions, MCMCChains, StableRNGs, Test
@model function one_observation()
x ~ Normal()
0.0 ~ Normal(x, 1)
end
chain = sample(
StableRNG(1),
one_observation(),
SMC(),
32;
chain_type=MCMCChains.Chains,
progress=false,
)
ess = vec(collect(chain[:ess_per_step]))
@show names(chain, :internals)
@show ess
@test any(x -> !ismissing(x), ess)names(chain, :internals) =
[:log_normalizing_constant, :ess_per_step,
:logprior, :loglikelihood, :logjoint]
ess =
Union{Missing, Float64}[missing, missing, ..., missing]agent explained that this is because ess_per_step being vector will be replaced by missing
An explicit `@addlogprob!` prior factor is not part of the prior proposal. Omitting it from particle weights silently changes the target distribution. Assisted-by: Codex <codex@openai.com>
MCMCChains drops vector-valued statistics. Store each filtering step as a scalar internal so the complete diagnostic survives multi-chain concatenation. Assisted-by: Codex <codex@openai.com>
A particle sweep requires at least one particle. Reject non-positive counts at the public boundary instead of failing later during weight reduction. Assisted-by: Codex <codex@openai.com>
`@addlogprob! (; logprior=...)` reaches DynamicPPL's `acclogp!!`, which already routes the field on to `acclogprior!!` -- overridden here anyway. Producing there drops a 27-line override that had to restate DynamicPPL's field-name validation to reach the one line that mattered. Assisted-by: Claude Code <noreply@anthropic.com>
Every other Turing sampler's `sample` defaults to `verbose=true`, so the same call reported differently through the ensemble wrapper than when made directly. Inert today: no `post_sample_hook` is defined for `SMC`. Assisted-by: Claude Code <noreply@anthropic.com>
`particle_rng`'s default argument has no caller, and its `Random.Sampler` wrapper draws the same stream as `rand(rng, UInt64)`. The `basic model` testset asserted nothing on a model four other testsets already sample. `bundle_smc_samples` gains the rationale a reader needs: specialising `AbstractMCMC.bundle_samples` on `SMC` instead would recurse when the extension hands the flattened transitions back. Assisted-by: Claude Code <noreply@anthropic.com>
`chain[:logevidence]` became `chain[:log_normalizing_constant]`, and SMC's per-particle `weight` is gone. Both break user code silently and neither was listed. Assisted-by: Claude Code <noreply@anthropic.com>
DynamicPPL's `acclogp!!` splits a two-field `@addlogprob!` across `acclogprior!!` and `accloglikelihood!!`, so one statement produces twice; the tests covered the two fields only in isolation. One entry of `ess_per_step` is therefore one likelihood term, not one observation as the changelog claimed. Assisted-by: Claude Code <noreply@anthropic.com>
Every particle owns its own `Philox2x`, seeded from the sampler's generator at creation and again at every fork, so a step that skips resampling has nothing to decorrelate: the streams simply continue. The refresh was a vestige of AdvancedPS's `TracedRNG`, which replayed per-step seeds to rebuild the reference trajectory; this implementation replays values instead. Draws therefore change. Seeding `Philox2x(UInt64, seed)` directly, rather than reseeding a generator the bare constructor had already seeded from the OS, leaves the state and the stream identical. Assisted-by: Claude Code <noreply@anthropic.com>
`mirror_onto_particle` had one call site. The comment claiming the RNG
section must come first because `Particle` names the generator type in
its signature was wrong: `Particle{RT<:AbstractRNG,WT<:Real}` does not.
The rest restated its own code, and the changelog carried mechanism that
belongs in the commit log.
Assisted-by: Claude Code <noreply@anthropic.com>
The effective sample size lies in `[1, nparticles]`, so a threshold outside `[0, 1]` silently meant "always resample" or "never". Worse, `Bool <: Real`, so `SMC(SystematicResampler(), true)` read `true` as a threshold. `PG` now takes any `Integer` particle count; `PG(Int32(4))` was a `MethodError`. Assisted-by: Claude Code <noreply@anthropic.com>
`acclogprior!!` produced before checking `ignore_missing_accumulator`, so a varinfo without a `LogPrior` accumulator would have contributed a weight it never accumulated. Unreachable today, since `trajectory_varinfo` always installs one, but it inverts the invariant the produce rests on. The `hasacc` check is repeated rather than pushed into `acclogp_and_mirror!!` behind a keyword: a keyword call lowers through `Core.kwcall`, which `Libtask.@might_produce` does not name, so the produce could go uninstrumented. Assisted-by: Claude Code <noreply@anthropic.com>
`save_state=true` stored `nothing` as the chain's sampler state and `initial_state` was dropped in `kwargs...` without a word, while every other keyword SMC cannot honour already warns. Both presuppose an MCMC loop with state to carry, which one sweep does not have. `SMC` had an `SMCState` before this rewrite, so this is a user-visible change and the changelog now says so. Making room for the new cases also removed three duplicate `sample` calls: `@test_logs` returns the chain. Assisted-by: Claude Code <noreply@anthropic.com>
`E[Oᵏ] = n·Wᵏ` is what makes a resampling scheme unbiased and nothing checked it, only that the indices stayed in bounds. For systematic resampling the count is within one offspring of `n·Wᵏ` on every draw, which is sharper than the mean; that bound does not hold for stratified, where a particle spanning strata can fall outside it. The exact-posterior grids were not exact to the tolerance they are compared against: `range(0.05, 4.0; length=400)` left weight 2e-5 on its last point, shifting the reference mean for `q` by 0.0075 and understating its standard deviation by 6%. Widened until the moments agree to seven digits with a far wider grid, with the coverage asserted rather than assumed. Assisted-by: Claude Code <noreply@anthropic.com>
"both fields of a named addlogprob are weighted" pins the produced value of each field, so the sampling-level testset covering the same ground goes, as does a `@test_throws ErrorException` standing beside an assertion on the same message. The exact-SSM testsets were 4m12s of the file's 6m41s. Their tolerances are batch-means standard errors, so 16 particles rather than 32 widens them to match the worse mixing and the assertions stay as sharp. Assisted-by: Claude Code <noreply@anthropic.com>
|
Thanks, @sunxd3 -- both fixed. |
SMCandPG/CSMCnow run directly on Libtask and DynamicPPL. Resamplers aretypes, stratified by default; conditional sweeps draw ancestors multinomially.
Philox seeding makes draws reproducible across Julia versions. Chains carry
log_normalizing_constant. Roughly twice as quick asmain, allocating half.Builds on #2848. Supersedes #2853.
Fix #2781
Fix TuringLang/AdvancedPS.jl#110
Fix TuringLang/AdvancedPS.jl#39
Fix TuringLang/AdvancedPS.jl#6
Close #2848. Close #2853.