Technical Summary: Efficient Stochastic Optimisation via Sequential Monte Carlo

Problem Statement

The paper addresses the challenge of optimizing loss functions $\ell(\theta)$ where the gradient $\nabla_\theta \ell(\theta)$ is intractable. Specifically, the gradient is expressed as an expectation over a parameter-dependent distribution $\pi_\theta$:
$$\nabla_\theta \ell(\theta) = \mathbb{E}_{X \sim \pi_\theta} [H_\theta(X)]$$
where $\pi_\theta(x) = e^{-U_\theta(x)} / Z_\theta$ and the normalizing constant $Z_\theta$ is intractable. This setting is prevalent in machine learning and statistics, including:

• Maximum Marginal Likelihood Estimation (MMLE): In latent variable models where the marginal likelihood gradient requires integrating out latent variables.
• Training Energy-Based Models (EBMs): Where the negative log-likelihood gradient involves an expectation over the model distribution.
• Reward Tuning of Generative Models: Aligning pre-trained models with downstream objectives via KL-regularized optimization.

The primary difficulty lies in the dependence of $\pi_\theta$ on $\theta$. Standard stochastic approximation methods typically rely on inner sampling loops using Markov Chain Monte Carlo (MCMC) to estimate these gradients. These inner loops are computationally expensive, slow to converge, and often require fresh sampling at every optimization step, leading to significant computational overhead.

Methodology

The authors propose Stochastic Optimisation via Sequential Monte Carlo (SOSMC), a framework that replaces expensive inner MCMC loops with efficient Sequential Monte Carlo (SMC) approximations.

Core Framework

Instead of sampling independently from each $\pi_{\theta_k}$, SOSMC treats the sequence of distributions $(\pi_{\theta_k})_{k \geq 0}$ as a Feynman-Kac flow. The method maintains a collection of weighted particles $\{(X^{(i)}_k, w^{(i)}_k)\}_{i=1}^N$ that approximate the current target distribution $\pi_{\theta_k}$.

1. Sequential Sampling: At each optimization iteration $k$, the algorithm propagates particles from the previous distribution $\pi_{\theta_{k-1}}$ to the current $\pi_{\theta_k}$ using a forward Markov kernel $K_k$.
2. Weight Updates: Particles are assigned importance weights based on the ratio of the target densities and the transition kernels. The unnormalized weight update follows:
   $$W^{(i)}_k = W^{(i)}_{k-1} \frac{\Pi_{\theta_k}(\bar{X}^{(i)}_k) L_{k-1}(\bar{X}^{(i)}_k, X^{(i)}_{k-1})}{\Pi_{\theta_{k-1}}(X^{(i)}_{k-1}) K_k(X^{(i)}_{k-1}, \bar{X}^{(i)}_k)}$$
   where $L_{k-1}$ is a backward kernel and $\Pi_\theta$ is the unnormalized density.
3. Gradient Estimation: The intractable gradient is estimated using the weighted particle approximation:
   $$g_k = \sum_{i=1}^N w^{(i)}_k H_{\theta_k}(X^{(i)}_k)$$
4. Resampling: To mitigate weight degeneracy, the algorithm monitors the Effective Sample Size (ESS) and resamples particles when ESS falls below a threshold.
5. Parameter Update: The estimated gradient $g_k$ is fed into a standard first-order optimizer (e.g., SGD, Adam) to update $\theta_{k+1}$.

Kernel Flexibility

The framework is general regarding the choice of kernels. The authors demonstrate that specific choices of forward and backward kernels recover existing algorithms as special cases:

• Using Unadjusted Langevin Algorithm (ULA) kernels recovers the Jarzynski-adjusted Langevin algorithm for EM (JALA-EM) and the SMC sampler proposed by Carbone et al. (2023).
• The framework also supports Metropolis-adjusted kernels (MALA, RWM) and Hamiltonian Monte Carlo (HMC), allowing for more robust exploration of the state space.

Key Contributions

The paper outlines three main contributions:

1. General SMC Framework: Development of a flexible SMC-based framework for optimizing functions with intractable gradients. The authors show that several existing algorithms are special cases of this framework, unifying previous approaches under a single theoretical structure.
2. Theoretical Analysis:
   • Convergence: Convergence rates are established for an idealized version of the algorithm (assuming exact expectations) under standard assumptions (Polyak-Łojasiewicz condition and smoothness).
   • ESS Behavior: A theoretical discussion is provided on the behavior of the Effective Sample Size (ESS) with respect to the step size and gradient magnitude, highlighting the sensitivity of particle degeneracy to the optimization trajectory.
3. Empirical Validation: Extensive experiments demonstrate the method's effectiveness in fine-tuning energy-based models across various settings, showing significant improvements in computational efficiency and optimization performance compared to existing baselines.

Experimental Results

The authors evaluate SOSMC on three primary tasks:

1. Reward Tuning of Langevin Processes:

   • Compared against Implicit Diffusion (IMPDIFF) and Stochastic Optimisation via Unadjusted Langevin (SOUL).
   • Findings: SOSMC variants (particularly those with Metropolis-corrected kernels like MALA and HMC) achieved faster convergence to high-reward regions and more stable Negative Log-Likelihood (NLL) trajectories compared to IMPDIFF. SOUL showed higher run-to-run variability and failure modes in transitioning between modes.

2. Reward Tuning of EBMs on 2D Datasets:

   • Tested on synthetic datasets (Two Moons, Circles, Blobs) with half-plane indicator rewards.
   • Findings: SOSMC-ULA achieved better objective contours (reward vs. KL divergence) than IMPDIFF, especially for small regularization strengths ($\beta_{KL}$). Crucially, the weighted particle rewards in SOSMC closely tracked the true expected reward, whereas IMPDIFF's unweighted particle rewards were poor proxies when parameters changed rapidly, requiring expensive fresh sampling for evaluation.

3. Reward Tuning of EBMs on MNIST:

   • Applied to a high-dimensional convolutional EBM pre-trained on MNIST.
   • Findings: SOSMC successfully tuned the model to satisfy various reward functions (brightness, darkness, half-plane) without degrading optimization behavior or exhibiting "reward hacking," demonstrating applicability to realistic, high-dimensional image settings.

4. Bayesian Image Deblurring (MMLE):

   • Applied to a maximum marginal likelihood estimation problem for image deblurring with a total variation prior.
   • Findings: The SOSMC-MYULA variant outperformed the Moreau-Yosida Particle Gradient Descent (MYPGD) baseline, achieving lower Mean Squared Error (MSE) and higher Structural Similarity Index Measure (SSIM), as well as sharper image reconstructions.

Significance and Claims

The paper claims that SOSMC provides a computationally efficient and theoretically sound alternative to standard MCMC-based stochastic optimization for intractable gradient problems. By leveraging the sequential nature of the optimization process, SOSMC reuses samples from previous iterations, avoiding the cost of "fresh" sampling required by methods like IMPDIFF or standard MCMC-based EM.

The authors emphasize that their framework is general and adaptable, capable of incorporating various proposal kernels (ULA, MALA, HMC) to suit different problem geometries. They note that while the method introduces a bias of order $O(1/N)$ and mean-squared error of order $O(1/N)$ (standard for SMC), the practical gains in convergence speed and stability are significant.

The work is presented as a methodological advancement. The authors state that future work will focus on deriving rigorous convergence rates for the particle approximations in the interacting setting and exploring accelerated schemes. They do not claim to introduce new application-specific objectives but rather a more efficient engine for solving existing classes of optimization problems in probabilistic inference and generative modeling.