Technical Summary: Reinforced Sequential Monte Carlo for Amortised Sampling

Problem Statement

The paper addresses the challenge of sampling from high-dimensional, unnormalised probability distributions $\pi(x) = R(x)/Z$, where the normalising constant $Z$ is unknown. This problem is central to Bayesian inference and computational chemistry (e.g., sampling molecular conformations). Classical approaches like Markov Chain Monte Carlo (MCMC) and Sequential Monte Carlo (SMC) are robust but often require many steps to converge. Conversely, amortised methods (e.g., diffusion models, GFlowNets) learn a parametric sampler to generate samples in finite time but often suffer from mode collapse, insufficient exploration of multi-modal landscapes, or instability when trained with on-policy objectives. The authors seek a synergy that combines the exploration capabilities of particle-based methods with the efficiency of amortised sampling.

Methodology

The proposed framework, Reinforced Sequential Monte Carlo (Reinforced SMC), unifies Hierarchical Variational Inference (HVI), Maximum-Entropy Reinforcement Learning (MaxEnt RL), and SMC/Annealed Importance Sampling (AIS).

1. Theoretical Unification

The authors establish a mathematical connection between these three fields:

• HVI & MaxEnt RL: Sampling from a target distribution is framed as a deterministic Markov Decision Process (MDP) where states are $(n, x_n)$ and actions are transitions $x_{n+1}$. The objective is to learn a policy $\vec{p}_\theta$ that maximises expected rewards (log-probabilities) and entropy. This is equivalent to minimising the reverse Kullback-Leibler (KL) divergence between the joint trajectory distribution and the target.
• SMC Connection: The importance weights in SMC correspond to the terms in the Trajectory Balance (TB) and Subtrajectory Balance (SubTB) objectives used in GFlowNets. Specifically, if the SubTB loss is minimised to zero, the proposal kernels satisfy detailed balance with intermediate targets, rendering importance weights uniform.

2. Core Algorithm: Off-Policy Training with SMC

The central innovation is using SMC as a behaviour policy to train the amortised sampler off-policy.

• Mutual Benefit: The amortised sampler (policy $\vec{p}_\theta$ and flow functions $F^\phi_n$) provides proposal kernels and intermediate targets for SMC. In return, SMC generates samples that better explore the target distribution (including high-probability regions the current sampler might miss). These SMC samples are used to update the amortised sampler.
• Training Objectives:
  • The policy parameters $\theta$ are updated using the Trajectory Balance (TB) loss.
  • The flow (intermediate target) parameters $\phi$ are updated using the Subtrajectory Balance (SubTB) loss.
  • The authors empirically find that using TB for $\theta$ and SubTB for $\phi$ yields the most stable training.

3. Key Technical Components

• Importance-Weighted Experience Replay (IW-Buf): To leverage historical data, the authors introduce a replay buffer where samples are prioritised not by temporal difference error (as in standard RL) but by their importance weights. This allows the buffer to maintain a particle approximation of the target distribution that converges weakly to $\pi(x)$.
• Adaptive Weight Tempering: To mitigate the high variance of importance weights (which causes training instability), an adaptive inverse temperature parameter $\lambda$ is introduced. $\lambda$ is dynamically adjusted to ensure the Effective Sample Size (ESS) remains above a user-defined threshold, smoothing the weights without introducing excessive bias.
• Adaptive Resampling: Resampling in the SMC loop is performed only when the ESS drops below a threshold, balancing computational cost with weight degeneracy.
• Flow Learning: Instead of fixed geometric annealing schedules, the intermediate targets $F^\phi_n$ are learned. The authors propose a parameterisation combining a learnable annealing schedule $\beta^\phi_n$ with a learnable correction term $\tilde{F}^\phi_n$, allowing the model to approximate the ground-truth intermediate targets more closely.

Key Contributions

1. Unified Framework: The paper provides the first unified treatment of HVI, MaxEnt RL, and SMC/AIS, demonstrating that optimal policies in MaxEnt RL correspond to optimal proposals in SMC.
2. Off-Policy Training Scheme: It introduces a method to train amortised samplers using off-policy samples generated by SMC, creating a feedback loop where the sampler improves the SMC proposals, and SMC improves the sampler's exploration.
3. Stabilisation Techniques: The derivation of an importance-weighted experience replay buffer and an adaptive weight tempering scheme to reduce gradient variance and improve training stability.
4. Learned Intermediate Targets: A principled approach to learning the intermediate distributions (flows) in SMC, moving beyond fixed annealing schedules.

Experimental Results

The method was evaluated on synthetic multi-modal targets in both continuous and discrete spaces, as well as on the Boltzmann distribution of alanine dipeptide conformations.

• Continuous Space (Diffusion Samplers):
  • Gradient-Free Setting: On benchmarks like GMM40, Funnel, and ManyWell, the proposed method (TB/SubTB + SMC + IW-Buf) significantly outperformed on-policy baselines (DDS, LV, TB) in terms of EUBO and Sinkhorn distance, indicating superior mode coverage and reduced mode collapse.
  • Gradient-Based Setting: The method improved performance on complex targets like Robot4 and GMM40 compared to advanced baselines like SCLD and CMCD, particularly in scenarios with sharp, separated modes.
• Discrete Space (Sequence Generation):
  • Applied to biochemical sequence design (QM9, sEH, TFbind8, L14-RNA1), the IW-Buf approach consistently improved mode coverage (lower EUBO) and Pearson correlation with the target reward compared to standard on-policy training and other buffer prioritisation schemes.
• Molecular Conformation (Alanine Dipeptide):
  • The method successfully approximated the Boltzmann distribution, outperforming baselines that failed to converge. The inclusion of SMC and experience replay was critical for navigating the complex energy landscape.

Significance and Claims

The paper claims to advance the field of machine learning by bridging the gap between amortised sampling and Monte Carlo methods.

• Synergy: It demonstrates that the strengths of MC methods (exploration, robustness) and amortised methods (speed, generalisation) can be combined. The MC method benefits from the learned kernels of the amortised sampler, while the amortised sampler benefits from the off-policy exploration of MC.
• Generality: The framework does not assume a specific form for the training loss (e.g., it works with general off-policy objectives) and is applicable to both continuous and discrete spaces.
• Stability: The proposed techniques (adaptive tempering, learned flows) address common instabilities in training neural samplers, particularly in high-dimensional, multi-modal settings.

The authors note that while their approach shows substantial improvements over on-policy training, future work could explore better target distributions for SMC, alternative objectives for learning intermediate targets, and applications to high-dimensional Bayesian posteriors and discrete biological discovery problems.