Reinforced sequential Monte Carlo for amortised sampling

AI summary

Overview

Research area: Machine learning — amortised (neural) sampling, sequential Monte Carlo (SMC), and maximum-entropy reinforcement learning (MaxEnt RL).

Technical level: Advanced. The paper assumes familiarity with variational inference, Markov chain Monte Carlo, importance sampling, diffusion models, and RL objectives such as trajectory balance and soft Bellman equations.

Scope: The paper builds a unified mathematical framework linking hierarchical variational inference, entropy-regularised RL, and sequential Monte Carlo, then uses that framework to train neural samplers with off-policy samples produced by SMC, improving both the sampler and the Monte Carlo estimates.

What This Paper Is About

Sampling from a complex, multi-modal probability distribution known only up to a normalising constant is a central problem in statistics, physics, and chemistry. Two families of methods attack it: classical Monte Carlo methods (SMC, MCMC), which are asymptotically correct but slow, and amortised neural samplers (diffusion samplers, autoregressive models), which are fast after training but can collapse onto a subset of modes or fail to match the target. This paper shows that the two families are two views of the same object, then trains the neural sampler using samples produced by SMC, and feeds the trained sampler back into SMC as its proposal distribution.

Key Contributions

1. A unified account of three frameworks. The authors give what they describe as the first single treatment connecting hierarchical variational inference, MaxEnt RL, and SMC with annealed importance sampling. In this view, a learned sampling policy plays the role of an SMC proposal kernel, and a learned flow (value) function plays the role of an intermediate target density or twist function.

2. SMC as a behaviour policy for off-policy sampler training. Because SMC produces weighted samples that better approximate the target, the authors use trajectories derived from SMC as off-policy training data for the amortised sampler — circumventing the importance-weighting variance that plagues naive off-policy training of reverse-KL objectives.

3. Training stabilisation techniques. They propose adaptive importance-weight tempering (raising weights to a power chosen by binary search so that effective sample size stays above a threshold), adaptive resampling, a specific split of objectives (trajectory balance for proposals, subtrajectory balance for flows), and diffusion-specific parameterisations.

4. Importance-weighted experience replay. They derive a principled replay buffer in which each historical batch is weighted by its particle estimate of the normalising constant, multiplied by the within-batch self-normalised weights — a variant of prioritised experience replay grounded in SMC rather than temporal-difference error.

Main Findings

• Large gains in the gradient-free setting. On multi-modal targets, the full method (TB/SubTB + SMC + IW-Buf) reduces the Sinkhorn distance on GMM40 (d=2) from roughly 595–617 (DDS, TB) to 6.46, and the EUBO on GMM40 (d=5) from roughly 3,000 to 2.3. Gains persist at d=32 on ManyWell (Sinkhorn 22.97 versus 29.58 for TB and DDS).

• Gradient-based benchmarks improve most on hard targets. On Funnel (d=10) MMD drops to 0.050 from 0.109 (TB, LV) and 0.116 (DDS). On Robot4 (d=10) MMD falls to 0.103 from about 0.42. On ManyWell (d=64) MMD falls to 0.043 from 0.243 (TB) and 0.260 (LV).

• The replay buffer is what stabilises training. The SMC-augmented method without the importance-weighted buffer shows occasional divergence (an entry marked with × on GMM40 d=50, and high variance such as Sinkhorn 64.48 ± 103.61 on Robot4). Adding the buffer removes the instability and delivers the best numbers on those same targets.

• Separating the two objectives matters. Using trajectory balance only for the policy/proposal parameters and subtrajectory balance only for the flow/target parameters works best; other combinations are reported as suboptimal or unstable.

• Amortisation and particle methods are complementary rather than competing. The policy provides low-variance proposals for SMC, while SMC's weighted population explores regions that the current policy underestimates, closing the loop between the two.

• Improvements extend beyond synthetic targets. The paper also reports results on the Boltzmann distribution of alanine dipeptide conformations and on discrete sequence spaces, where the method improves both target approximation and mode coverage relative to purely amortised or purely Monte Carlo baselines.

Methodology in Plain English

The starting observation is that the logarithm of an annealed importance sampling weight, evaluated on a trajectory, is the same expression that appears inside the square of the trajectory balance loss used to train neural samplers. That algebraic coincidence means an SMC sampler is, in effect, a trained sampler plus a correction — and it means the two can share parameters.

Concretely, the sampler is a hierarchical latent variable model: a chain of conditional distributions that turns a simple initial distribution into a sample. Each conditional is a policy; each intermediate unnormalised density is a flow. The authors train the policies with trajectory balance and the flows with subtrajectory balance, using gradients that do not require importance sampling corrections — a key property of these objectives that makes off-policy learning possible.

The bidirectional loop works as follows. Forward, the trained policies serve as SMC proposal kernels and the trained flows serve as intermediate target densities, giving weighted particle approximations to the target. Backward, those weighted particles are used as training data: a trajectory is drawn from the replay buffer with probability proportional to its importance weight, and a backward trajectory is generated using the reverse kernel. The sampler is then updated on this mixture of on-policy and off-policy data.

Two practical safeguards keep training from collapsing. First, importance weights are raised to a power λ, and λ is chosen as the largest value keeping effective sample size above a fraction γ of the particle count, found by binary search. Second, resampling is performed only when effective sample size falls below a threshold. For diffusion samplers specifically, the flows are parameterised as corrections to a temperature annealing schedule and the proposals are written in terms of the gradient of the target energy.

Why This Matters

Research impact. The paper supplies a conceptual bridge between three literatures that have largely developed in parallel: variational inference with hierarchical latent variables, entropy-regularised RL (and its GFlowNet cousins), and particle filtering. It also reframes experience replay for neural samplers, replacing temporal-difference prioritisation with importance weights derived from SMC, which is likely to influence how off-policy training data is weighted in future sampler work.

Real-world applications.

• Bayesian inference of parameters in scientific models, where the posterior is known only up to a constant.
• Molecular conformation sampling and free-energy estimation in computational chemistry and drug discovery.
• Statistical physics, where Boltzmann distributions over many-body systems are intractable to normalise.
• Generative modelling and reinforcement learning for structured discrete outputs such as sequences, molecules, and programs.

Industry relevance. Amortised samplers promise fast inference at deployment time, but their failure modes — mode collapse, poor coverage — are exactly the failures that matter in fields like drug design and risk modelling. A training procedure that uses particle methods to police the sampler, while keeping the sampler's fixed-cost rollout, is attractive to any team that needs both speed and distributional fidelity.

Future Directions

• Scaling to higher dimensions and longer trajectories. The experiments stop at d=64 and a small peptide; whether the loop remains stable for protein-scale or image-scale targets is open.
• Better objectives for flows. Subtrajectory balance is known to underperform trajectory balance in some diffusion settings due to instability. The authors suggest the improved behaviour policies they introduce may unlock its credit-assignment advantages; testing that claim rigorously is a natural next step.
• Learned reverse kernels. The paper fixes the backward kernel, but recent work makes it trainable. Jointly learning both directions could tighten the SMC proposals further.
• Theory for the replay buffer. The buffer is shown to converge weakly to the target as the number of batches times particles grows, but finite-buffer bias and the interaction between stale weights and current parameters are not fully characterised.

Target Audience

Researchers and graduate students working on Monte Carlo methods, variational inference, diffusion models, or reinforcement-learning-based generative modelling. Practitioners who already use SMC or diffusion samplers and want a principled way to combine them, or who need a stabilised off-policy training recipe for neural samplers, will get the most direct value. The unified-framework section is also useful to readers from any one of the three constituent fields who want a compact statement of how the others relate.

Authors’ abstract

This paper proposes a synergy of amortised and particle-based methods for sampling from distributions defined by unnormalised density functions. We state a connection between sequential Monte Carlo (SMC) and neural sequential samplers trained by maximum-entropy reinforcement learning (MaxEnt RL), wherein learnt sampling policies and value functions define proposal kernels and twist functions. Exploiting this connection, we introduce an off-policy RL training procedure for the sampler that uses samples from SMC -- using the learnt sampler as a proposal -- as a behaviour policy that better explores the target distribution. We describe techniques for stable joint training of proposals and twist functions and an adaptive weight tempering scheme to reduce training signal variance. Furthermore, building upon past attempts to use experience replay to guide the training of neural samplers, we derive a way to combine historical samples with annealed importance sampling weights within a replay buffer. On synthetic multi-modal targets (in both continuous and discrete spaces) and the Boltzmann distribution of alanine dipeptide conformations, we demonstrate improvements in approximating the true distribution as well as training stability compared to both amortised and Monte Carlo methods.