Diffusion posterior sampling for simulation-based inference in tall data settings

1. Problem Statement

The paper addresses the challenge of Simulation-Based Inference (SBI) in "tall data" settings.

• Context: In many scientific fields, models are defined by complex simulators where the likelihood function $p(x|\theta)$ is intractable. Traditional methods like MCMC are often inapplicable.
• Tall Data: The specific setting involves multiple independent and identically distributed (i.i.d.) observations $x_{1:n} = (x_1, \dots, x_n)$ generated from the same underlying parameter $\theta$.
• Goal: To infer the posterior distribution $p(\theta | x_{1:n})$. As $n$ increases, the posterior becomes sharper, improving parameter identifiability.
• Existing Limitations:
  • NPE (Neural Posterior Estimation): Requires retraining for every new number of observations $n$ or fixed input dimensions, leading to high simulation costs and lack of flexibility.
  • F-NPSE (Factorized Neural Posterior Score Estimation): A recent method that composes individual scores from a single-observation network to handle variable $n$. However, because the composed score does not correspond to a known forward diffusion process, it relies on Langevin dynamics (MCMC) for sampling. This is computationally expensive, sensitive to hyperparameters, and unstable, especially as $n$ grows.

2. Methodology

The authors propose a new algorithm that retains the amortized and compositional benefits of F-NPSE but eliminates the need for Langevin dynamics by explicitly approximating the diffusion process of the tall data posterior.

Core Mathematical Insight

The paper derives an exact expression for the score of the diffused tall data posterior $\nabla_{\theta_t} \log p_t(\theta_t | x_{1:n})$.
Using the factorization of the posterior (Equation 1) and Bayes' rule, the score is shown to be:
$$\nabla_{\theta_t} \log p_t(\theta_t | x_{1:n}) = (1-n)\nabla_{\theta_t} \log p^\lambda_t(\theta_t) + \sum_{j=1}^n \nabla_{\theta_t} \log p_t(\theta_t | x_j) + \nabla_{\theta_t} \log \mathcal{L}_\lambda(\theta_t, x_{1:n})$$

• The first two terms are tractable (analytical prior score + learned single-observation scores).
• The third term, $\nabla_{\theta_t} \log \mathcal{L}_\lambda$, is a correction term involving backward kernels. F-NPSE ignores this term, necessitating Langevin corrections.

The Approximation Strategy

To make the correction term tractable and enable deterministic sampling (e.g., DDIM), the authors approximate the backward kernels using Gaussian distributions (following the Tweedie framework).

• They approximate the backward transition kernels $q_{0|t}$ as Gaussians with means $\mu$ and covariances $\Sigma$.
• This allows the correction term to be expressed as a linear combination of Gaussian log-factors.
• The final approximate score formula (Equation 16) becomes a precision-weighted sum of individual scores:
  $$\nabla_{\theta_t} \log p_t(\theta_t | x_{1:n}) \approx \Lambda(\theta_t)^{-1} \left( \sum_{j=1}^n \Sigma_{t,j}^{-1} \nabla_{\theta_t} \log p_t(\theta_t | x_j) + (1-n)\Sigma_{t,\lambda}^{-1} \nabla_{\theta_t} \log p^\lambda_t(\theta_t) \right) + F$$
  Where $F$ vanishes if the covariance matrices are constant.

Proposed Algorithms

Two strategies are introduced to compute the covariance matrices $\Sigma_{t,j}$ to ensure they remain constant (making $F=0$):

1. JAC (Jacobian Approximation): Estimates $\Sigma_{t,j}$ by taking the Jacobian of the learned score network.
   • Pros: Highly accurate in noise-free settings.
   • Cons: Computationally expensive ($O(m^2)$ matrix operations) and unstable with noisy scores.
2. GAUSS (Gaussian Approximation): Assumes a constant Gaussian approximation for the covariance. It estimates the initial covariance $\Sigma_0$ empirically by running a short DDIM chain for each observation and computing the sample covariance.
   • Pros: Significantly faster, highly stable, and robust to noise.
   • Cons: Slightly less precise than JAC in perfect noise-free conditions, but superior in realistic SBI settings.

3. Key Contributions

1. Theoretical Derivation: Derivation of the exact score for the diffused tall data posterior and the identification of the missing correction term in F-NPSE.
2. Novel Sampling Algorithm: Introduction of a method to approximate the tall posterior diffusion process using Gaussian backward kernels, enabling the use of deterministic samplers (DDIM) instead of Langevin dynamics.
3. Two Practical Implementations: The JAC and GAUSS algorithms, with GAUSS emerging as the robust, scalable choice for real-world applications.
4. Elimination of Langevin Bottleneck: The method removes the computational bottleneck and instability associated with Langevin steps in compositional SBI.

4. Experimental Results

The authors evaluated their methods against F-NPSE (using Langevin) on three levels of difficulty:

• Toy Models (Gaussian & GMM):

  • Speed: GAUSS is 2.5x faster and JAC is 1.8x faster than Langevin.
  • Robustness: In noisy settings (simulating imperfect score networks), Langevin and JAC become unstable or diverge. GAUSS maintains high accuracy and stability.
  • Scalability: Performance remains consistent as parameter dimension $m$ and observation count $n$ increase.

• SBI Benchmarks (SLCP, SIR, Lotka-Volterra):

  • Using learned scores (NPSE), GAUSS consistently outperformed Langevin in terms of Sliced Wasserstein (sW), Maximum Mean Discrepancy (MMD), and Classifier Two-Sample Test (C2ST) metrics.
  • Langevin and JAC often diverged as $n$ increased due to error accumulation, whereas GAUSS remained stable.

• Real-World Application (Jansen and Rit Neural Mass Model):

  • Applied to a complex computational neuroscience model with 4 parameters.
  • The method successfully inverted the model, revealing parameter dependencies (indeterminacy) that were hidden in single-observation settings.
  • GAUSS was the only method that passed the local Classifier-Two-Sample Test ($\ell$-C2ST) for both 3D and 4D configurations, confirming statistical consistency with the true posterior.

5. Significance

• Efficiency & Stability: The work provides a significantly faster and more stable alternative to existing tall-data SBI methods, making high-precision inference feasible for complex simulators with large datasets.
• Scalability: By enabling deterministic sampling, the method scales effectively to high-dimensional problems and large numbers of observations ($n$), which is critical for modern scientific applications (e.g., neuroscience, physics).
• Identifiability: The approach demonstrates that aggregating observations via tall data posteriors can resolve parameter indeterminacies (coupling effects) that single-observation inference cannot, offering deeper insights into complex system dynamics.
• Framework: It establishes a new paradigm for compositional inference in score-based models, moving beyond simple score addition to modeling the full diffusion dynamics of the composed posterior.

In conclusion, the paper successfully bridges the gap between the flexibility of compositional score-based inference and the efficiency of deterministic diffusion sampling, offering a state-of-the-art solution for simulation-based inference in tall data regimes.