Table of Contents
Fetching ...

Dimension-Free Multimodal Sampling via Preconditioned Annealed Langevin Dynamics

Lorenzo Baldassari, Josselin Garnier, Knut Solna, Maarten V. de Hoop

TL;DR

The paper analyzes continuous-time annealed Langevin dynamics for multimodal targets in a dimension-varying setting, showing that uniform-in-dimension sampling guarantees are achievable by carefully designing the annealing smoothing and preconditioning spectra. It provides explicit dimension-dependent bounds on annealing bias and robustness to score and initialization errors, with clear spectral conditions on the smoothing and diffusion operators that ensure dimension-free performance. The results are complemented by numerical experiments demonstrating that a decaying preconditioner spectrum preserves stability across high dimensions, while improper spectra lead to deterioration. Collectively, the work offers principled guidelines for constructing ALD schemes that remain accurate and stable as problem dimensionality grows, including infinite-dimensional to finite-dimensional transitions and robustness to model misspecification. The findings bridge diffusion-based sampling theory with practical design strategies for high-dimensional multimodal targets.

Abstract

Designing algorithms that can explore multimodal target distributions accurately across successive refinements of an underlying high-dimensional problem is a central challenge in sampling. Annealed Langevin dynamics (ALD) is a widely used alternative to classical Langevin since it often yields much faster mixing on multimodal targets, but there is still a gap between this empirical success and existing theory: when, and under which design choices, can ALD be guaranteed to remain stable as dimension increases? In this paper, we help bridge this gap by providing a uniform-in-dimension analysis of continuous-time ALD for multimodal targets that can be well-approximated by Gaussian mixture models. Along an explicit annealing path obtained by progressively removing Gaussian smoothing of the target, we identify sufficient spectral conditions - linking smoothing covariance and the covariances of the Gaussian components of the mixture - under which ALD achieves a prescribed accuracy within a single, dimension-uniform time horizon. We then establish dimension-robustness to imperfect initialization and score approximation: under a misspecified-mixture score model, we derive explicit conditions showing that preconditioning the ALD algorithm with a sufficiently decaying spectrum is necessary to prevent error terms from accumulating across coordinates and destroying dimension-uniform control. Finally, numerical experiments illustrate and validate the theory.

Dimension-Free Multimodal Sampling via Preconditioned Annealed Langevin Dynamics

TL;DR

The paper analyzes continuous-time annealed Langevin dynamics for multimodal targets in a dimension-varying setting, showing that uniform-in-dimension sampling guarantees are achievable by carefully designing the annealing smoothing and preconditioning spectra. It provides explicit dimension-dependent bounds on annealing bias and robustness to score and initialization errors, with clear spectral conditions on the smoothing and diffusion operators that ensure dimension-free performance. The results are complemented by numerical experiments demonstrating that a decaying preconditioner spectrum preserves stability across high dimensions, while improper spectra lead to deterioration. Collectively, the work offers principled guidelines for constructing ALD schemes that remain accurate and stable as problem dimensionality grows, including infinite-dimensional to finite-dimensional transitions and robustness to model misspecification. The findings bridge diffusion-based sampling theory with practical design strategies for high-dimensional multimodal targets.

Abstract

Designing algorithms that can explore multimodal target distributions accurately across successive refinements of an underlying high-dimensional problem is a central challenge in sampling. Annealed Langevin dynamics (ALD) is a widely used alternative to classical Langevin since it often yields much faster mixing on multimodal targets, but there is still a gap between this empirical success and existing theory: when, and under which design choices, can ALD be guaranteed to remain stable as dimension increases? In this paper, we help bridge this gap by providing a uniform-in-dimension analysis of continuous-time ALD for multimodal targets that can be well-approximated by Gaussian mixture models. Along an explicit annealing path obtained by progressively removing Gaussian smoothing of the target, we identify sufficient spectral conditions - linking smoothing covariance and the covariances of the Gaussian components of the mixture - under which ALD achieves a prescribed accuracy within a single, dimension-uniform time horizon. We then establish dimension-robustness to imperfect initialization and score approximation: under a misspecified-mixture score model, we derive explicit conditions showing that preconditioning the ALD algorithm with a sufficiently decaying spectrum is necessary to prevent error terms from accumulating across coordinates and destroying dimension-uniform control. Finally, numerical experiments illustrate and validate the theory.
Paper Structure (44 sections, 6 theorems, 259 equations, 6 figures)

This paper contains 44 sections, 6 theorems, 259 equations, 6 figures.

Key Result

Theorem 3.1

Fix $d\geq 1$ and $\epsilon >0$. Define Consider the ALD dynamics up to time Then $\mathcal{B}_{\mathrm{ann}}^d(T^d) \le \epsilon.$

Figures (6)

  • Figure 1: ALD step count vs. dimension. Number of ALD time steps required for the empirical $\mathrm{KL}(\rho_\star^d\|\rho^{\mathrm{ALD},d})$ to fall below the prescribed accuracy $\epsilon=0.3$, plotted against the truncation dimension $d$. The red curve corresponds to the default flat-spectrum choice, with preconditioner $\Gamma=I$ and smoothing $C=40I$, while the green curve corresponds to a tailored spectral design: $\Gamma=\mathrm{Diag}(j^{-1.5})_{j\ge1}$ and $C=\mathrm{Diag}(40\cdot j^{-2.7})_{j\ge1}$.
  • Figure 2: Illustration of annealed Langevin dynamics. The ALD scheme considered here starts from a heavily smoothed, and hence tractable, version of the target and simulates a Langevin diffusion whose drift is adapted to the current smoothing level. As time progresses, the smoothing is gradually removed, so the dynamics moves from an almost unimodal distribution back to the original multimodal target through a controlled "complexification" of the landscape.
  • Figure 3: Annealed-induced bias across dimensions. Empirical $\mathrm{KL}(\rho_\star^d\|\rho^{\mathrm{ALD},d})$ as a function of the truncation dimension $d$, shown on a logarithmic $y$-axis. The green curve corresponds to the prescribed spectra $\Gamma=\mathrm{Diag}(j^{-1.5})_{j\ge1}$ and $C=\mathrm{Diag}(40 \cdot j^{-2.7})_{j\ge1}$, while the red curve corresponds to the flat-spectrum choice $\Gamma=I$ and $C=40I$. The $\mathrm{KL}$ is estimated via $k$NN with $k=20$; robustness to $k$ is reported in Appendix \ref{['app:numerics-details']}.
  • Figure 4: Preconditioning and dimension dependence of the KL. Empirical $\mathrm{KL}(\rho_\star^d\|\rho^{\mathrm{ALD},d})$ as a function of the truncation dimension $d$, shown on a logarithmic $y$-axis. The green curve corresponds to runs of the ALD sampler with preconditioner $\Gamma=\mathrm{Diag}(j^{-3.5})_{j\ge1}$, while the red curve corresponds to runs with $\Gamma=\mathrm{Diag}(j^{-1})_{j\ge1}$. The $\mathrm{KL}$ is estimated via $k$NN with $k=20$; robustness to $k$ is reported in Appendix \ref{['app:numerics-details']}.
  • Figure 5: Robustness to $k$. Supplementary plots for Figure \ref{['fig:bias-error-vs-dimension']}: we report the $k$NN estimate for $k\in\{20,50,80\}$ (the main text reports $k=20$), illustrating that the observed trend is stable across these choices of $k$.
  • ...and 1 more figures

Theorems & Definitions (8)

  • Theorem 3.1
  • Proposition 4.1
  • Proposition 4.3
  • Proposition 4.4
  • Lemma 1.1
  • proof
  • Lemma 2.1
  • proof