PaperlayerView exact revision on arXiv

Source-linked AI summary

Source Distribution Estimation by Posterior Averaging

Trung-Dung Hoang, Lisa M. Koch

arXiv:2609.02622v1cs.LG

TL;DR

Source Distribution Estimation requires finding parameter distributions whose simulated observations match real data, despite fixed-surrogate methods optimizing an approximation that may be inaccurate in explored regions. The paper instead uses posterior-averaging expectation maximization with fresh simulations, and reports consistent gains over fixed and iterated surrogate baselines, especially on Lotka–Volterra.

  • Problem

    Existing SDE methods optimize a likelihood surrogate trained once under a fixed proposal prior, so surrogate agreement need not imply agreement under the true simulator.

  • Method

    Posterior Averaging alternates an E-step that trains a posterior on fresh simulations from the current source with an M-step that fits the source to posterior averages over observed data.

  • Results

    The method consistently outperforms fixed-surrogate baselines and their iterated variants, most clearly on Lotka–Volterra.

  • Takeaways & Limitations

    Posterior averaging targets self-consistency with the true simulator rather than relying on a frozen surrogate likelihood.

  • Takeaways & Limitations

    The evaluation covers three synthetic benchmark simulators, one simulation budget, n = 10,000 observations, and three seeds; behavior on real data, smaller n, and higher dimensions remains open.

Abstract

from arXiv · show

Simulation-based science often requires a distribution over simulator parameters whose push-forward reproduces a set of real observations: this is the source distribution estimation (SDE) problem. Existing methods fit the source against a likelihood surrogate trained once from a fixed proposal prior. Their objective is therefore stated only in terms of the surrogate instead of the true simulator, which may fail for inaccurate areas in parameter space where the surrogate was never trained. We instead solve SDE by expectation maximization: an E-step trains an amortized posterior on fresh simulations from the current source estimate, and an M-step refits the source to the average of that posterior over the observed data. We give two parameterizations, (1) separate source and posterior flows and (2) a single shared conditional flow. We evaluate our method on three benchmark tasks under both broad and misspecified initial priors. Both improve on existing fixed surrogate approaches and on iterated variants of each, most clearly on Lotka--Volterra, where no baseline falls below 0.96 data-space C2ST while our methods reach 0.64-0.68 in three of four initial-prior settings.

1 Introduction

Source Distribution Estimation infers a parameter distribution whose simulator-generated observations match real data, even when the simulator is a black box accessible only through forward sampling.

  • Source Distribution Estimation seeks q(θ) whose simulator-induced observation distribution matches the empirical distribution of real observations.The simulator provides samples x ∼ p(x | θ), while the likelihood density may be intractable.
  • The motivating application is to infer parameter distributions compatible with real observations so simulations remain consistent with reality.Simulators encode system knowledge across domains including astrophysics, epidemiology, and neuroscience.
  • The simulator is treated as a black box: researchers can draw forward samples for any θ but cannot access its implementation or gradients.This setting includes licensed or proprietary simulators and systems with discrete or combinatorial internal stochasticity.

2 Related Work

In the non-differentiable black-box setting, existing SDE methods rely on fixed learned simulator surrogates, with NEB and Sourcerer using surrogate-likelihood approaches.

  • Existing SDE implementations use a fixed learned simulator surrogate rather than direct simulator differentiation in the non-differentiable black-box setting.
  • NEB trains a normalizing-flow likelihood surrogate from simulations generated under a fixed broad prior, then optimizes the source distribution against it.
  • Sourcerer selects the maximum-entropy feasible source distribution, but only its surrogate-likelihood variant applies under the paper’s non-differentiable black-box assumption.
  • Under this assumption, Sourcerer’s surrogate likelihood is trained identically to NEB’s surrogate.

3 Method

Posterior Averaging alternates posterior estimation under the current source with source refitting to posterior averages over real data, directly targeting SDE self-consistency. Under exact steps, this update is monotone in the induced data-space KL, while practical implementations use neural approximations and require finite initial KL for informative guarantees.

  • General algorithm: The E-step fits an amortized posterior from simulations generated under the initial or current source, and the M-step refits the source to posterior samples for each real observation.The general procedure repeats these steps until a fixed number of rounds or convergence.
  • Variants: PA uses separate unconditional source and conditional posterior flows, while PA-Shared uses one conditional flow with null and data-dependent contexts for both modes.Context dropout enables one objective to train both modes, but shared parameters can couple posterior fitting to source updates.
  • Convergence: Under full support and exact E- and M-steps, each update cannot increase the induced marginal KL to the observed-data distribution.Equality holds only when consecutive sources agree almost everywhere, assuming the displayed KL quantities are finite.
  • Convergence: The monotonicity guarantee is informative only when the initial induced marginal has finite KL relative to the real-data distribution.If the initial marginal assigns zero density to a region with positive real-data mass, the initial KL is infinite and the bound is vacuous.
  • Self-consistency: The M-step targets the average posterior over real data, implementing the self-consistency operator that fixed-surrogate baselines target only indirectly.The practical update is an approximation because the posterior is represented by an NPE model and the average is estimated from observed samples.

4 Experiments

The experiments compare PA and PA-Shared with fixed-surrogate and iterated baselines across three tasks, broad-support priors, and shifted-support priors. PA methods generally perform best, especially on Lotka–Volterra, while baseline iteration is unreliable.

  • Tasks and methods: The evaluation uses Two Moons, SLCP, and Lotka–Volterra, with true task priors generating 10,000 observations per task.Methods are compared using data-space C2ST, where 0.5 indicates indistinguishable distributions.
  • Broad initial prior: Under broad-support priors, PA methods ranked best across nearly all task-prior combinations, except Two Moons, where NEB-it marginally outperformed PA.The Two Moons comparisons were 0.513 vs. 0.515 under Gaussian and 0.498 vs. 0.524 under wide-box.
  • Broad initial prior: 0.655 and 0.675 were PA-Shared’s Lotka–Volterra C2ST values under Gaussian and widebox, while all baselines exceeded 0.964.PA reached 0.914 and 0.892 under those priors, still outperforming every baseline.
  • Misspecified initial prior: Under shifted-support priors, PA and PA-Shared both outperformed baselines on Lotka–Volterra under shift_half, reaching 0.683 and 0.643.Under shift_full, PA reached 0.903 while PA-Shared reached 0.984, similarly to the baselines.
  • Convergence across rounds: Across rounds, NEB-it and both PA methods generally descended, whereas Sourcerer-it improved on Two Moons, stayed constant on Lotka–Volterra, and deteriorated on SLCP.Figure 1 reports per-round C2ST for the four iterative methods and includes one-shot baselines for reference.

5 Discussion and Conclusion

The paper replaces fixed surrogate-likelihood optimization with posterior averaging, using fresh simulations and alternating E- and M-steps. It reports monotonic KL reduction under idealized conditions and stronger empirical performance, while identifying important recovery and evaluation limitations.

  • Method: Posterior averaging trains a posterior on fresh simulations from the current source and refits the source to its average over observed data.The method has separate-flow and shared-flow parameterizations.
  • Theory: Under idealized conditions, posterior averaging monotonically decreases the induced-data KL.The guarantee concerns exact E- and M-step behavior.
  • Empirical findings: Empirically, posterior averaging consistently outperforms fixed-surrogate baselines and their iterated variants, most clearly on Lotka–Volterra.The discussion attributes this distinction to iteration alone not closing the surrogate gap.
  • Limitations: The evaluation measures reproducing po rather than recovering a unique π⋆, because multiple source distributions can induce the same marginal.A source that performs well in data space may still place mass where a domain scientist would not.
  • Limitations: The study is limited to three synthetic benchmark simulators, one simulation budget, n = 10,000 observations, and three seeds.Behavior on real observations, smaller n, and higher-dimensional simulators remains open.

A.1 Full KL Derivation

The derivation establishes the posterior-averaging update and its monotonicity under stated assumptions. The appendix also presents separate-flow and shared-flow algorithmic realizations of the E- and M-steps.

  • Well-definedness: The derivation verifies that the posterior expression is well-defined wherever the relevant source and likelihood factors are positive.Assumption (A1) supplies the needed positivity conditions without requiring a support relation between successive sources.
  • KL derivation: Jensen’s inequality, log-product splitting, Fubini, and Gibbs’ inequality produce the successive steps of the KL bound.The exact-posterior identity makes πt(θ)p(x | θ)/pt(θ | x) constant in θ.
  • Equality condition: KL equality holds exactly when πt+1 = πt almost everywhere, provided the displayed KL quantities are finite.This follows from the strictness of Gibbs’ inequality.
  • PA with separate flows: In the separate-flow algorithm, the E-step trains a posterior from simulations under the current source, and the M-step fits the source using posterior samples for observed data.The procedure repeats until a fixed round limit or convergence.
  • PA with shared flow: The shared-flow algorithm trains source and conditional-posterior components jointly in the E-step, then embeds observations and samples from the shared flow in the M-step.Classifier-free-guidance dropout is applied during the E-step.

C Task Definitions

The experiments use three simulation-based inference tasks as source-distribution benchmarks. Their parameter dimensions and observation-generation mechanisms range from low-dimensional geometric data to a noisy predator-prey dynamical system.

  • Benchmark suite: The benchmark suite contains Two Moons, SLCP, and Lotka–Volterra, using each task’s true prior as a source distribution.The simulators come from the sbibm benchmark suite.
  • Two Moons: Two Moons has two parameters with a uniform prior on [−1, 1] per dimension, and observations combine a fixed-radius arc with parameter-dependent translation.
  • SLCP: SLCP has five parameters with a uniform prior on [−3, 3] per dimension, generating four independent draws from a parameter-dependent two-dimensional Gaussian.Concatenating the draws produces an eight-dimensional observation.
  • Lotka–Volterra: Lotka–Volterra is a predator-prey ODE model with four rate parameters, fixed initial condition (X(0), Y(0)) = (30, 1), and integration horizon T = 20.It is the evaluation’s one real, non-toy simulator.
  • Lotka–Volterra: Lotka–Volterra observations are noisy log-population counts for both species at 10 evenly spaced time points, forming a 20-dimensional observation.

D Architecture and Optimization

PA uses separate source and posterior flows, while PA-Shared uses one conditional flow for both distributions. Both optimize the source and posterior through alternating E- and M-steps, with task-independent hyperparameters unless noted.

  • Configuration: All hyperparameters remain fixed across tasks, seeds, and both experiments unless explicitly noted.The optimization configuration is summarized in Table 4 for PA and PA-Shared.
  • Architectures: PA parameterizes separate source and posterior normalizing flows, whereas PA-Shared uses a single shared conditional flow.PA-Shared evaluates the source with a learnable null context and the posterior with an observation-dependent context.
  • M-step: The M-step fits the source by maximum likelihood to posterior samples averaged over the observed data.PA uses 300 posterior samples per observation, while PA-Shared uses 100 samples per observation for its null-context branch.
  • Shared-flow optimization: PA-Shared snapshots the shared flow before each E-step so M-step samples combine the frozen source with the updated embedder before restoring E-step weights.This separates source-sampling behavior from the jointly updated conditional model during the M-step.

D.4 Protocol

The evaluation uses matched simulation budgets across eight rounds, a fixed observed set, three benchmark tasks, and data-space C2ST measured against the real observations.

  • Budget: Eight rounds use 4,000 fresh (θ, x) pairs per E-step, totaling 32,000 simulations without accumulation across rounds.Each E-step draws fresh simulations from the current source estimate.
  • Observed data: The observed set contains N = 10,000 observations drawn once from the true prior.
  • Tasks: The benchmark tasks are two_moons, slcp, and lotka_volterra, with parameter and observation dimensions (2, 2), (5, 8), and (4, 20), respectively.Experiments use seeds 0, 1, and 2.
  • Metric: Data-space C2ST compares simulated observations from each round's learned source with xreal using 1,000 samples per side, with 0.5 indicating indistinguishability.Closer-to-0.5 C2ST values are better.
Loading 2609.02622v1…