Technical Summary: Theoretical Guidelines for Annealed Langevin Dynamics in Compositional Simulation-Based Inference

1. Problem Statement

In Simulation-Based Inference (SBI), the objective is to approximate the posterior distribution $p(\theta \mid x)$ of a stochastic model when the likelihood is intractable but the model can be simulated. A specific challenge arises when $n$ independent observations $x_1, \dots, x_n$ share a common unknown parameter $\theta^\star$. As $n$ increases, the multi-observation posterior concentrates, requiring inference methods that sharpen accordingly.

While pooling strategies (training a single network on all $n$ observations) exist, they scale poorly with $n$ and require retraining for new sample sizes. Compositional strategies offer an alternative by combining individual posterior estimators trained once on single-observation data, amortizing cost over $n$.

Current compositional approaches (e.g., Geffner et al. [13], Linhart et al. [16]) aggregate individually learned posterior scores into a composite score. A critical theoretical flaw exists: this composite score does not correspond to the score of any distribution along the true forward diffusion path of the multi-observation posterior. Consequently, sampling from it via a standard reverse Stochastic Differential Equation (SDE) introduces an irreducible bias, as the reverse SDE requires the exact score of the joint forward diffusion process to be valid.

To address this, Annealed Langevin Dynamics has been proposed as a principled alternative. Instead of targeting the true posterior via a reverse SDE, it treats the composite score as the genuine score of a sequence of tractable "bridging densities" ($\pi_t$) and samples from them successively using the Unadjusted Langevin Algorithm (ULA). However, the hyperparameters of this process—step sizes ($h_t$), steps per level ($k_t$), and the number of annealing levels ($T$)—have historically been chosen empirically, lacking theoretical guarantees on sampling accuracy.

2. Methodology

Theoretical Framework

The authors derive Wasserstein error bounds for annealed Langevin dynamics when using approximate scores. Building on Dalalyan and Karagulyan [7], they establish:

1. Intermediate Bounds (Proposition 1): A per-level bound for a single ULA chain that accounts for both the composite score error (approximation bias) and the Wasserstein distance between consecutive bridging densities.
2. Global Bounds (Proposition 2): A global sampling guarantee obtained by composing $T$ intermediate bounds.

The analysis relies on assumptions of log-concavity and smoothness of the target and bridging densities. The authors demonstrate (Lemma 1) that these properties are preserved or strengthened along the Variance-Preserving (VP) diffusion path, even if individual posteriors are only smooth (not log-concave), particularly as $n$ grows (consistent with the Bernstein–von Mises theorem).

Hyperparameter Decision Rule

The paper proposes an explicit decision rule to select hyperparameters ($h_t, k_t$) to guarantee a prescribed sampling accuracy $\gamma$ (Wasserstein distance).

• The rule splits the error budget $\gamma$ into a bias component (controlled by step size $h_t$) and a variance/mixing component (controlled by steps $k_t$) using a weight parameter $\omega \in (0,1)$.
• Step Size ($h_t$): Bounded by the smoothness ($M_t$) and log-concavity ($m_t$) constants of the bridging density and the score matching error $\epsilon_{DSM}$.
• Number of Steps ($k_t$): Determined by the Wasserstein distance between consecutive bridging densities $W_2(\pi_{t_{p+1}}, \pi_{t_p})$ and the log-concavity constant.

Comparative Analysis: Geffner vs. Linhart

The authors apply these guidelines to compare two specific composite score formulations:

1. Geffner et al. [13]: Uses a linear aggregation of scores.
2. Linhart et al. [16]: Incorporates a second-order Gaussian approximation of backward diffusion kernels, utilizing covariance matrices ($\Sigma_{t,i}, \Sigma_{t,\lambda}$).

In a Gaussian setting, the authors derive closed-form expressions for the smoothness and log-concavity constants and the Wasserstein distances between bridging densities for both methods.

3. Key Contributions

1. First Theoretical Guidelines: The paper provides the first theoretically grounded guidelines for selecting hyperparameters in compositional SBI using annealed Langevin dynamics, moving beyond empirical tuning.
2. Wasserstein Error Bounds: It derives explicit bounds linking the sampling error to the properties of the bridging densities and the score approximation error.
3. Decision Rule: An explicit formula is provided to choose step sizes and iteration counts to ensure the Wasserstein error remains below a threshold $\gamma$.
4. Geffner vs. Linhart Comparison:
   • Step Sizes: In the Gaussian setting, Linhart's bridging densities admit larger step sizes than Geffner's (Corollary 1), as the ratio of log-concavity to smoothness constants is more favorable.
   • Efficiency: Near the target distribution ($t \to 0$), Geffner's method may require fewer steps, but near the prior ($t \to 1$), Linhart's method requires significantly fewer steps and allows larger step sizes. Empirically, the "switch point" where Linhart becomes more efficient occurs very early (often near $t=0$), making Linhart's approach more efficient overall in terms of total score evaluations.
5. Generalization: The guidelines derived in the Gaussian setting are shown to generalize effectively to non-Gaussian settings (e.g., GMM priors, SIR, Lotka-Volterra models) via second-order approximations.

4. Results

• Gaussian Setting: Theoretical analysis confirms that Linhart's composite score yields a more efficient sampler, requiring fewer total Langevin steps to achieve the same Wasserstein accuracy compared to Geffner's score.
• Non-Gaussian Setting: Empirical experiments on four inference tasks (GMM prior, GMM likelihood, SIR, Lotka-Volterra) demonstrate that the hyperparameters derived from the Gaussian theory successfully control the final Wasserstein error below the prescribed threshold $\gamma$ across varying dimensions and observation counts.
• Efficiency: Across all tested tasks, Linhart's method consistently required fewer total steps (score function evaluations) than Geffner's method while maintaining comparable error levels.
• Robustness: The method remains robust even when individual posteriors are multimodal, as the multi-observation posterior tends to become unimodal and Gaussian-like with increasing $n$.

5. Significance and Claims

The paper claims to close a theoretical gap in compositional score-based SBI by providing a principled method for tuning annealed Langevin dynamics. Its significance lies in:

• Theoretical Grounding: Offering guarantees on sampling quality (Wasserstein distance) rather than relying on heuristics.
• Practical Utility: Providing a "well-understood and theoretically grounded starting point" for practitioners, eliminating the need for manual, empirical hyperparameter tuning.
• Methodological Preference: Demonstrating that Linhart's composite score formulation is theoretically and empirically superior to Geffner's in terms of sampling efficiency (fewer steps, larger step sizes) in both Gaussian and complex non-Gaussian settings.

The authors remain modest, noting that their theoretical guarantees rely on smoothness and approximate log-concavity assumptions. They acknowledge that extending these guarantees to highly multimodal distributions or those satisfying only log-Sobolev inequalities remains a direction for future work. However, they assert that the derived guidelines are effective for a broad range of practical SBI problems.