Source-linked AI summary
Auto-Encoding Sequential Monte Carlo
Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, Frank Wood
TL;DR
The paper addresses how to learn models and proposals effectively in structured probabilistic models using AESMC. It analyzes ELBO tightness and proposal learning, introduces alternating ELBO training, and reports improvements in model and proposal learning while noting a high-variance gradient-estimation limitation.
Problem
Proposal learning with AESMC is affected by finite optimization, limited proposal expressiveness, and the choice and particle dependence of ELBO objectives.
Method
AESMC uses an SMC-based ELBO with reparameterized gradient estimation and introduces alternating ELBOs to address proposal-learning issues.
Results
Alternating ELBOs improve both model and proposal learning in some cases, while AESMC provides better structured generative-model training than IWAE.
Takeaways & Limitations
SMC-based AESMC offers a scalable approach to simultaneous model learning and proposal adaptation in structured generative models.
Takeaways & Limitations
The REINFORCE term for discrete ancestor choices is omitted because its additional variance is problematically high, introducing a small bias into gradient estimates.
Abstract
from arXiv · showhide
We build on auto-encoding sequential Monte Carlo (AESMC): a method for model and proposal learning based on maximizing the lower bound to the log marginal likelihood in a broad family of structured probabilistic models. Our approach relies on the efficiency of sequential Monte Carlo (SMC) for performing inference in structured probabilistic models and the flexibility of deep neural networks to model complex conditional probability distributions. We develop additional theoretical insights and introduce a new training procedure which improves both model and proposal learning. We demonstrate that our approach provides a fast, easy-to-implement and scalable means for simultaneous model learning and proposal adaptation in deep generative models.
1 INTRODUCTION
The paper builds on AESMC, which combines lower-bound model learning with SMC-based marginal-likelihood estimation, especially for structured sequential models. It adds theoretical analysis of proposal learning and a new training procedure.
- AESMC maximizes a lower bound to the log marginal likelihood while using SMC rather than importance sampling for estimation.
- SMC is particularly powerful for models with sequential structure, motivating its use in auto-encoding model learning.
- The paper expresses the ELBO gap as a KL divergence on an extended sampling space, clarifying when objectives can be optimized perfectly.
- It extends analysis of proposal-learning efficiency and develops an improved training procedure from these theoretical results.
2 BACKGROUND
The background introduces state-space models and approximate inference, then describes SMC and importance-weighted auto-encoders as foundations for AESMC. SMC uses sequential targets, weighted particles, and resampling to estimate posteriors and marginal likelihoods.
- STATE-SPACE MODELS: State-space models define latent and observed sequences through initial, transition, and emission densities.
- STATE-SPACE MODELS: Inference approximates posterior distributions or posterior expectations, which is difficult for nonlinear, nondiscrete, and non-Gaussian models.
- SEQUENTIAL MONTE CARLO: SMC sequentially targets distributions over increasing latent prefixes and uses weighted particles to approximate the final posterior.
- SEQUENTIAL MONTE CARLO: SMC also provides an unbiased marginal-likelihood estimator from intermediate particle weights.
- SEQUENTIAL MONTE CARLO: Resampling focuses computation on particles favored by new observations and helps avoid exponential weight divergence as sequence length increases.
- IMPORTANCE WEIGHTED AUTO-ENCODERS: IWAEs maximize an importance-sampling ELBO, which reduces to the VAE objective when K = 1.
3 AUTO-ENCODING SEQUENTIAL MONTE CARLO
AESMC uses an SMC-based ELBO for amortized model learning and proposal adaptation, with gradients estimated through reparameterized SMC samples. Its KL analysis shows that tightness requires both suitable proposals and intermediate targets, while gradient estimation omits a high-variance REINFORCE term.
- OBJECTIVE FUNCTION: AESMC applies stochastic gradient ascent to a dataset-averaged SMC ELBO for joint model learning, proposal adaptation, and inference amortization.
- OBJECTIVE FUNCTION: The objective uses a sequential proposal family conditioned on latent histories and observations, together with the SMC marginal-likelihood estimator.
- GRADIENT ESTIMATION: Reparameterization expresses latent sampling through auxiliary random variables, allowing gradients of the SMC estimator to be evaluated with stochastic gradient ascent.
- GRADIENT ESTIMATION: The additional REINFORCE term for discrete ancestor choices is omitted because experiments found it had problematically high variance, introducing a small gradient bias.
- BIAS & IMPLICATIONS ON THE PROPOSALS: The ELBOSMC gap from log pθ(y1:T) is represented as KL(QSMC||PSMC), enabling analysis of proposal and target-distribution mismatch.
- BIAS & IMPLICATIONS ON THE PROPOSALS: For AESMC, a tight bound requires learning intermediate target distributions corresponding to marginals of the final target, not only correct individual proposals.
- BIAS & IMPLICATIONS ON THE PROPOSALS: For most convenient model parameterizations, a perfect proposal and tight bound may be impossible, although SMC can still support more ambitious structured inference problems than importance sampling.
4 IMPROVING PROPOSAL LEARNING
Larger particle counts tighten the ELBO but can harm proposal learning because gradient signal shrinks faster than its variability. The proposed ALT procedure uses different objectives and particle counts for model and proposal updates to improve their joint learning.
- Larger K often harms proposal learning for both IWAE and AESMC despite producing tighter bounds.
- As K increases, the expected proposal gradient and its standard deviation both decrease, but the expected gradient decreases faster, raising relative noise.For K > 10, the estimate is roughly equally likely to be positive or negative, degrading stochastic updates.
- Increasing K reduces the proposal-gradient signal-to-noise ratio, causing updates to degrade toward pure noise when K is too high.
- 4.1 ALTERNATING ELBOS: ALT alternates coordinate-descent updates, using separate inference types and particle counts for θ and φ.The algorithm estimates gθ and gφ with their respective ELBO choices before updating the model and proposal parameters.
- 4.1 ALTERNATING ELBOS: The suggested configuration uses SMC with large K for θ and IS with small K for φ, matching their differing learning requirements.Small-K IS supports proposal learning through better SNR, while large-K SMC supports model learning.
5 EXPERIMENTS
Experiments evaluate AESMC for generative-model learning and proposal adaptation, from analytically tractable state-space models to high-dimensional stochastic videos. They find that more particles improve model learning but can worsen proposal learning, while ALT mitigates this trade-off.
- The experiments test whether tighter bounds and improved inference affect proposal learning, whether AESMC outperforms IWAE, and whether ALT improves both.
- 5.1 LINEAR GAUSSIAN STATE SPACE MODEL: In an LGSSM, optimizing ELBOSMC with respect to θ produces better generative models, and using more particles accelerates convergence toward the optimum.The bootstrap proposal prevents gradients with respect to θ from being backpropagated through q.
- 5.2 PROPOSAL LEARNING: For proposal learning, increasing particles worsens φ quality under ELBOIS and similarly worsens proposal learning under ELBOSMC.The Gaussian unknown-mean experiment compares learned φ with the analytic posterior.
- 5.2 PROPOSAL LEARNING: ALT alleviates the particle-induced proposal-learning problem while making model learning faster through a more accurate proposal distribution.
- 5.3 MOVING AGENTS: On Moving Agents, AESMC models outperform IWAE and more particles improve the ELBO for both methods, although ALT is less stable than either.The task uses partially occluded stochastic video sequences and a VRNN generative model and proposal.
20 IWAE -356.6 AESMC -356.1
At 40 particles, AESMC reaches -356.1 versus -356.2 for IWAE. Figure 4 tracks max(ELBOSMC, ELBOIS) over training and reports final values across algorithms and particle counts.
- -356.1: AESMC reaches a final max(ELBOSMC, ELBOIS) of -356.1 at 40 particles, compared with -356.2 for IWAE.
- Figure 4 plots the rolling mean over 5 epochs of max(ELBOSMC, ELBOIS) on the test set.Lines average three random seeds, with shaded areas showing standard deviation.
- Figure 4 encodes particle count by color and algorithm by line style, with a table reporting each learned model's final value.
6 CONCLUSIONS
The paper presents AESMC as an SMC-based ELBO method for model learning, adds theoretical KL-divergence analysis, and introduces ALT to address proposal-learning issues. Experiments report better structured-model training than IWAE and, in some cases, improved joint model and proposal learning.
- AESMC performs model learning with an ELBO based on the SMC marginal-likelihood estimator, optimized using SGA and the reparameterization trick.
- For structured problems, the SMC-based objective yields better generative-model training than IWAE, attributed to superior inference and a tighter bound.
- The paper expresses ELBO bias relative to log marginal likelihood as a KL divergence on an extended space and analyzes its implications for IWAE and AESMC.
- AESMC can achieve a zero KL only when SMC intermediate target distributions corresponding to target marginals are learned.
- Alternating ELBOs uses different objectives for θ and φ and, in some cases, improves both model and proposal learning.
A.2 REINFORCE & REPARAMETERIZATION
The appendix studies a reparameterization-based gradient estimator for ELBOSMC and compares it with an estimator that includes a REINFORCE term. Including that term substantially increases variance.
- The estimator is formed by sampling auxiliary variables, reparameterizing them into particles, and evaluating the resulting ELBOSMC gradient.
- Including the first term in the estimator produces much higher variance in the T = 200 model with K = 16 and 100 samples.
B PROOFS FOR BIAS & IMPLICATIONS ON THE PROPOSALS
The proofs characterize when SMC estimators and proposal distributions can achieve exactness, and connect these conditions to practical neural proposal architectures and predictive evaluation.
- Proof setup: The analysis defines sequential target distributions, unnormalized targets, and normalizing constants used to establish SMC estimator properties.
- Weight conditions: Under the stated estimator identities, the incremental weights must be constant with respect to the particle trajectory.
- Proposal implications: A perfect proposal requires each intermediate proposal marginal to match the corresponding marginal of the final target distribution.
- Proposal implications: The proofs show that exact estimator equality implies q(x1|y) = p(x1|y) for all x1.
- Neural implementation: The VRNN implementation uses shared components between the proposal and generative model, with shared parameters trained using gradients for both θ and φ.
- Predictive evaluation: Using ten training particles and five prediction particles, AESMC models produced more consistent future trajectories than IWAE models in the moving-agents visualization.
C.3 OPTIMIZING ONLY PROPOSAL PARAMETERS
Proposal-only experiments compare importance-sampling and SMC objectives across training and testing particle counts, revealing systematic effects of particle number and algorithm choice on inference quality.
- Experimental design: The experiments optimize ELBOIS and ELBOSMC with K ∈ {10, 100, 1000}, then evaluate proposals using IS and SMC with Ktest ∈ {10, 100, 1000}.
- Evaluation metric: Inference quality is measured using a proxy comparing proposal-based approximate marginal means with true marginal means from Kalman smoothing.
- Results: Training with ELBOSMC and Ktrain = 1000 produced worse SMC inference than ELBOIS training at any tested particle count.
- Particle-count effects: Increasing Ktest improves inference quality, whereas increasing Ktrain worsens inference quality.
- Algorithm combinations: The tested training-testing combinations rank as (IS, SMC) ≻ (SMC, SMC) ≻ (IS, IS) ≻ (SMC, IS).