Efficient Stochastic Optimisation via Sequential Monte Carlo

AI summary

Overview

Research area: Stochastic optimisation, sequential Monte Carlo (SMC) methods, and gradient estimation for loss functions whose gradients are expectations over intractable, parameter-dependent distributions (stat.ML).

Technical level: Advanced. The paper relies on Feynman–Kac flows, SMC sampler weight recursions, χ²-divergence and Polyak–Łojasiewicz (PŁ) convergence analysis.

Scope: The paper proposes a general SMC-based framework ("SOSMC") for first-order stochastic optimisation with intractable gradients, provides convergence and effective sample size (ESS) analysis for the recursions approximated by the method, and evaluates it on reward-tuning of energy-based models (EBMs).

What This Paper Is About

Many machine learning and statistics problems require minimising a loss whose gradient is an expectation taken under a distribution that depends on the parameters being optimised, such as ∇θℓ(θ) = E_{X∼πθ}[Hθ(X)], where πθ(x) = e^{−Uθ(x)}/Zθ and Zθ is intractable. Existing stochastic approximation methods attack this by running inner Markov chain Monte Carlo (MCMC) sampling loops to estimate the gradient, which becomes computationally expensive. The paper develops sequential Monte Carlo (SMC) samplers that replace those inner sampling loops with efficient SMC approximations, reusing particles across optimiser iterations to obtain gradient estimates.

Key Contributions

1. (C1) A general SMC-based optimisation framework. The authors develop a flexible SMC framework for optimising functions with intractable gradients, and show that existing algorithms (for example the JALA-EM algorithm of Cuin et al. (2025) for maximum marginal likelihood estimation, and the SMC sampler of Carbone et al. (2023) for EBM training) arise as special cases of it.

2. (C2) High-level theory for the approximated recursions. They analyse idealised versions of their recursions, proving exact gradient recovery and a convergence rate under Polyak–Łojasiewicz and smoothness assumptions, and they provide a theoretical discussion of how the construction of the marginal distributions affects the effective sample size (ESS).

3. (C3) Empirical demonstration on reward tuning of energy-based models. They run experiments on reward-tuning of energy-based models across several settings, reporting improvements in computational efficiency and optimisation performance relative to baselines.

4. A practical algorithm (Algorithm 1). Stochastic Optimisation via SMC (SOSMC) interleaves one SMC sampler step per outer optimiser iteration, forms a weighted-particle gradient estimate, updates parameters with any first-order optimiser, and resamples when ESS falls below τN.

Main Findings

• Exact gradient recovery in the idealised scheme. Lemma 1 shows that if the Feynman–Kac expectations are computed exactly, the ratio E[H_{θ_{k−1}}(X_{k−1})W_{k−1}] / E[W_{k−1}] equals π_{θ_{k−1}}(H_{θ_{k−1}}) = ∇ℓ(θ_{k−1}), i.e. the weighted SMC recursion recovers the true gradient.

• Convergence rate for the idealised recursion. Under Assumption A1 (ℓ is μ-PŁ and L-smooth), Proposition 1 gives ℓ(θ_k) − inf_θ ℓ(θ) ≤ (1 − γμ)^k (ℓ(θ_0) − inf_θ ℓ(θ)) for any step-size γ ≤ 1/L, which follows from the fact that the idealised update reduces to exact gradient descent.

• Finite-particle behaviour is left unresolved. Remark 3 states that with finitely many particles the update uses a stochastic, biased, self-normalised SMC estimate; for a fixed target sequence standard SMC theory gives bias and mean-squared error of order O(1/N) for bounded test functions, but because the target sequence here depends on previous particle approximations, proving such bounds requires mean-field/interacting particle tools, which the authors leave to future work.

• ESS degrades with large step-sizes and large gradients. The quantity ρ_k(γ) := N / (1 + χ²(π_{θ_k} ‖ π_{θ_{k−1}})), which relates to ESS when the forward kernel is π_{θ_k}-invariant and the backward kernel is its time reversal (Dai et al., 2022), takes the form ρ_k(γ) = N exp(−γ²‖∇ℓ(θ_{k−1})‖²_{Σ^{−1}}) in the Gaussian case π_θ = N(θ, Σ).

• A general χ² expansion. Proposition 2 shows χ²(π_{θ_k}, π_{θ_{k−1}}) = γ² ∇ℓ(θ_{k−1})ᵀ I_{θ_{k−1}} ∇ℓ(θ_{k−1}) + o(γ²), where I_{θ_{k−1}} is the expected outer product of ∇θ U_{θ_{k−1}}(X); the local geometry of the loss landscape therefore informs step-size choice.

• Adaptive step-size rule. Motivated by the exponential sensitivity of ESS, the authors adapt γ_k multiplicatively: γ_k = c γ_{k−1} when ESS > τ_adapt N and γ_k = γ_{k−1}/c when ESS < τ_adapt N, with c > 1.

• Reward tuning of Langevin processes: SOSMC converges faster than ImpDiff, SOUL is more variable. Against implicit diffusion (ImpDiff) and stochastic optimisation via unadjusted Langevin algorithm (SOUL) of Marion et al. (2025), SOSMC variants and SOUL show swift convergence relative to ImpDiff, but SOUL exhibits greater run-to-run variability and, in one setting, a failure mode where a single chain fails to transition between modes; Metropolis-corrected kernel variants were robust.

• Reward tuning of EBMs on 2D datasets. Baseline EBMs π_0(x) ∝ exp(−E_{θ_0}(x)) were trained with persistent contrastive divergence (PCD), with E_θ parameterised by an MLP with smooth activations so that ∇_x E_θ is well-defined for Langevin sampling; the standard contrastive loss was augmented with mild energy and gradient regularisation. Terminal density snapshots are shown across increasing β_KL for the reward R_lower.

• No quantitative benchmark tables are reported in the provided content. Numerical values beyond the theoretical expressions (e.g. exact rewards, NLL values, runtimes, particle counts, dataset sizes) are not reported in the content available; the paper states that complete details are in Appendix E.

Methodology in Plain English

The core idea is to stop throwing away samples between optimiser iterations. Standard approaches treat each gradient step as a fresh problem: to estimate ∇θℓ(θ_k) = E_{X∼π_{θ_k}}[H_{θ_k}(X)], they run an inner MCMC loop that approximately samples from the current π_{θ_k}, which is slow and expensive.

Instead, the authors treat the sequence of distributions π_{θ_0}, π_{θ_1}, π_{θ_2}, … generated as the parameters move as a sequence of targets for an SMC sampler. Particles are propagated forward through a Markov kernel K_k, reweighted using the ratio of unnormalised targets and the backward kernel L_{k−1}, and resampled when the effective sample size drops below a threshold τN. The weighted particle average of H_{θ_k} then serves as the gradient estimate fed to any first-order optimiser.

Crucially, because the particles from the previous iteration are reused rather than discarded, the sampler tracks the slowly drifting target distribution efficiently. The paper shows this framework is flexible: choosing the forward kernel to be the unadjusted Langevin algorithm (ULA) with a particular backward kernel recovers known algorithms as special cases, and other choices (such as π_{θ_k}-invariant Metropolis-adjusted Langevin algorithm (MALA) kernels with their time reversals) fit into the same template.

The theoretical part studies what happens when the expectations in the Feynman–Kac identity are computed exactly — in that idealised limit the scheme is exactly gradient descent, so standard convergence theory applies. The authors then analyse how the χ²-divergence between consecutive targets controls ESS, giving guidance on step-size tuning and motivating their adaptive step-size rule.

Why This Matters

Impact on research. The paper unifies previously separate SMC-based optimisation methods (JALA-EM for maximum marginal likelihood estimation, the Carbone et al. (2023) SMC sampler for EBM training) under one framework, and offers a systematic route to designing new samplers by mixing different forward and backward kernels. It also connects ESS behaviour explicitly to step-size and gradient magnitude, giving a principled basis for tuning.

Real-world applications mentioned in or implied by the paper:

• Maximum marginal likelihood estimation in latent variable models, including stochastic EM-type procedures.
• Training energy-based models and latent energy-based models, where sampling from the model is required at every gradient step.
• Fine-tuning and reward-tuning of generative models (alignment with a potentially non-differentiable reward under a KL-regularised objective).
• State-space model estimation via SMC-EM style algorithms for marginal log-likelihood gradients.

Industry relevance. Any pipeline that fine-tunes generative models or trains EBMs pays for inner sampling loops; replacing them with a recycled particle population is a direct route to lower compute cost per optimisation step. The paper's experiments were run on a personal computer and a Google Colab T4 GPU, and the code is released at https://github.com/akyildiz-group/SOSMC, which lowers the barrier to adoption.

Future Directions

• Convergence theory for the finite-particle, interacting case. The authors explicitly leave to future work the combination of mean-field/interacting particle bounds (Del Moral and Rio, 2011) with biased stochastic gradient descent analyses (e.g. Demidovich et al., 2023; Karimi et al., 2019) to quantify error when the target sequence depends on previous particle approximations.
• Full SMC error analysis of the particle approximations. Standard SMC analysis is said to be applicable but is deferred to future work.
• Refined sufficient conditions for the χ² expansion. The paper points to Appendix B for analogous results with more explicit sufficient conditions.
• Broader kernel and dynamics design. The framework allows alternative choices such as underdamped Langevin dynamics or deterministic dynamics with coherent importance-weight estimation (Schönle et al., 2025a), and different forward/backward kernel parametrisations that coincide in the continuous-time limit (Schönle et al., 2025b).

Target Audience

Researchers and graduate students in machine learning and statistics working on stochastic optimisation with intractable gradients, SMC methods, energy-based models, generative model fine-tuning, or maximum marginal likelihood estimation. It is also relevant to practitioners who implement training pipelines for EBMs or reward-tuned generative models and want to reduce the cost of inner sampling loops. Readers need familiarity with MCMC, importance sampling, and stochastic approximation to follow the theory sections.

Authors’ abstract

The problem of optimising functions with intractable gradients frequently arises in machine learning and statistics, ranging from maximum marginal likelihood estimation procedures to fine-tuning of generative models. Stochastic approximation methods for this class of problems typically require inner sampling loops to obtain (biased) stochastic gradient estimates, which rapidly becomes computationally expensive. In this work, we develop sequential Monte Carlo (SMC) samplers for optimisation of functions with intractable gradients. Our approach replaces expensive inner sampling methods with efficient SMC approximations, which can result in significant computational gains. We establish convergence results for the basic recursions defined by our methodology which SMC samplers approximate. We demonstrate the effectiveness of our approach on the reward-tuning of energy-based models within various settings.