PaperlayerView exact revision on arXiv

Source-linked AI summary

High-Dimensional Simulation-Based Inference in Latent Spaces

Lars Kühmichel, Stefan T. Radev, Bhanu Prasanna Koppolu, Masoumeh Davoudi, Jerry M. Huang, Paul-Christian Bürkner

arXiv:2609.37381v1cs.LGstat.ML

TL;DR

SBI has usually compressed observations rather than increasingly high-dimensional inference targets. This paper learns a low-dimensional parameter representation for posterior inference and decodes samples back to the original space. Across its case studies and generative families, latent inference achieves similar or better accuracy and marginal calibration with faster sampling than standard inference.

  • Problem

    SBI has mainly compressed observations, while increasingly high-dimensional inference targets raise the question of whether targets can also be compressed.

    Source: §1 · p. 2, §2.1 · p. 2

  • Method

    A two-stage approach learns and freezes a low-dimensional parameter autoencoder, fits an amortized posterior over its codes, and decodes posterior samples to parameter space.

    Source: §1 · p. 2, §2.2 · p. 3

  • Results

    Across four case studies and three generative families, latent SBI achieved similar or better accuracy and marginal calibration than standard inference with considerably faster sampling.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The results suggest latent SBI is a competitive choice for very high-dimensional amortized inference, improving sampling speed while retaining or increasing posterior quality.

    Source: §5 · p. 9

    When code capacity is insufficient, the θ-only encoder may preserve high-prior-variance directions whether or not x informs them.

    Source: §2.3 · p. 4

Abstract

from arXiv · show

Neural simulation-based inference (SBI) has been widely successful in inferring a relatively small number of interpretable parameters from potentially high-dimensional observations, such as images or time series. Accordingly, representation learning in SBI has focused almost exclusively on compressing the observations used to condition the posterior. More recently, however, SBI has begun to target increasingly high-dimensional parameter spaces, raising the complementary question of whether the inference target itself should be compressed. Our answer is a practical merger of SBI and latent generative modeling, which learns a low-dimensional representation of the simulator parameters, performs posterior inference directly in this latent space, and maps posterior samples back to the original parameter space. We characterize the conditions under which latent-space inference recovers the desired target posterior and systematically study its empirical trade-offs. Across four case studies and three generative families, we compare latent and standard estimators while controlling for network capacity, regularization, optimization, and training compute. At matched training compute, latent-space inference achieves accuracy and marginal calibration comparable to direct target-space inference while sampling up to more than an order of magnitude faster.

1 INTRODUCTION

The introduction asks whether inference targets, not only observations, can be compressed for high-dimensional SBI. It proposes latent target compression to retain accurate, calibrated inference while speeding sampling.

1 INTRODUCTION

  • Recent SBI work is moving toward very high-dimensional parameter spaces, complementing its traditional focus on low-dimensional parameters inferred from complex or large-scale data.Latent generative models also learn in compressed target spaces, but their representation learning is designed primarily for synthesis.Source: §1 · p. 1
  • SBI has typically compressed high-dimensional observations while leaving the inference target intact, whereas target-reduction methods for inverse problems are often problem- or algorithm-specific.The paper brings target compression into amortized SBI, where earlier inverse-problem reductions are not directly portable.Source: §1 · pp. 1–2
  • At matched training compute, latent SBI maintains accuracy and marginal calibration while sampling faster than otherwise matched target-space estimators.The comparison controls for network capacity, regularization, optimization, and training compute.Source: §1 · p. 2
  • The approach learns an information-maximizing parameter autoencoder, freezes its representation, fits a conditional posterior over latent codes, and decodes samples to parameter space.The figure depicts the encoder mapping θ to ϑ, inference conditioned on s(x), and the decoder mapping ϑ back to θ.Source: §1 · p. 2, Figure 1 · p. 2

2 METHOD

The method compresses simulator parameters into a latent representation, performs amortized posterior inference there, and decodes samples back to parameter space. Its error analysis separates posterior approximation, information lost by the code, and decoder error, while identifying limits of a θ-only encoder.

2.1 SIMULATION-BASED INFERENCE WITH HIGH-DIMENSIONAL TARGETS

  • Ordinary SBI often compresses complex observations while retaining interpretable θ, but applications such as denoising and image recovery can make θ high-dimensional.In ultrasound imaging, θ as a whole may be interpretable only through domain expertise.Source: §2.1 · p. 2

2.2 TWO-STAGE LATENT INFERENCE

  • Effective dimension deff(ε) counts the θ-directions needed for posterior inference, so ambient parameter dimension may overstate inference complexity even when every coordinate must be recovered.It is defined using a map T_k for which the conditional mutual information I(θ; x | T_k(θ)) is at most ε.Source: §2.2 · p. 3
  • Because the encoder sees θ but not x, its code dimension can at best match the intrinsic dimension of p(θ), an upper bound on deff(ε).The posterior network therefore solves inference in the encoded subspace, and the decoder restores samples in the original parameter space.Source: §2.2 · p. 3
  • The frozen InfoVAE encoder and decoder are trained on prior parameter samples; a second stage fits the amortized posterior on encoded training pairs.The encoder is not updated during posterior fitting, and the paper reports that joint fine-tuning was less accurate than a frozen control in every ablation variant.Source: §2.2 · p. 3

2.3 ERROR DECOMPOSITION OF LATENT COMPRESSION

  • Latent inference adds code information loss and decoder inversion error to code-space posterior approximation; Proposition 1 bounds joint posterior error by a posterior, code, and decoder term.The posterior term is the amortization error in code space; the code term measures information about θ carried by x but discarded by the code, and the decoder term measures inversion error.Source: §2.3 · pp. 3–4
  • With sufficiently flexible encoder and decoder families, the first-stage maximizer at fixed I(θ; ϑ) makes the decoder term vanish.The first stage also targets the code term through reconstruction; if code dimension is at least deff(ε), some code keeps that term below ε.Source: §2.3 · p. 4
  • When the code cannot retain all of θ, the θ-only encoder preserves directions that reduce reconstruction error, which for a linear code are the largest-prior-variance directions regardless of x.The paper therefore ablates the compression ratio.Source: §2.3 · p. 4

3 RELATED WORK

SBI has mainly reduced the observation side, while target-side reductions in Bayesian inverse problems are often tailored to a particular forward model, likelihood, or sampling algorithm. Latent generative modeling offers a two-stage compression recipe, but its use in SBI must preserve statistical accuracy rather than only sample quality.

3 RELATED WORK

  • SBI has primarily reduced observations, whereas Bayesian inverse-problem reductions generally do not provide a general parameter-compression mechanism for amortized SBI.Observation representations are justified by preserving posterior information; inverse-problem reductions are typically tailored to a particular forward model, likelihood, or sampling algorithm.Source: §3 · p. 4
  • High-dimensional generation commonly encodes data, learns a generative model in latent space, and decodes samples, but this recipe had not yet entered SBI.In SBI, statistical accuracy—not sample quality—is the ultimate success metric.Source: §3 · p. 4
  • Bayesian inverse-problem methods also compress targets, but often require an explicit likelihood, gradients, or forward operator and observation-specific guided sampling.The paper asks whether learned target compression can systematically improve amortized SBI without compromising statistical accuracy.Source: §3 · pp. 4–5

4 EXPERIMENTS

The experiments compare latent and target-space SBI across four high-dimensional case studies at matched training compute and model size. Latent inference improves several accuracy and calibration outcomes, but results vary by task, compression level, and the autoencoder’s representation.

4.1 CASE STUDIES

  • Four case studies span a 2,080-coordinate Gaussian target with eight effective block-scale degrees of freedom, 32 × 32 random fields, Fashion-MNIST deblurring, and 256 × 256 satellite images.The Gaussian posterior and random-field posterior have analytic references; the image tasks use natural-image redundancy, with the satellite task leaving attributes such as roof material and season unconstrained.Source: §4.1 · p. 5

4.2 GENERAL SETUP

  • Latent and target-space variants are matched on training wall-clock compute and model size, including autoencoder training and the frozen autoencoder in the latent model.Other training details are held fixed; the map-to-satellite figure also reports a matched training budget.Source: §4.2 · pp. 5–6, Figure 4 · p. 7

4.4 RESULTS

  • For map-to-satellite, latent draws follow the map’s street layout while standard draws mostly show generic vegetation; latent sampling is 19× faster.Latent inference scores better on every reported accuracy and marginal-calibration metric and has higher SSIM, although C2ST separates both variants from true images.Source: Figure 4 · p. 7, §4.4 · p. 9

4.3 METRICS

  • Metrics assess target-space accuracy, marginal and joint calibration, and within-family sampling speed, using sample-based measures that need no tractable density.The reported measures include NRMSE, CRPS, coverage, calibration error, C2ST, and speedup relative to standard sampling.Source: §4.3 · pp. 7–8

4.4 RESULTS

  • On the correlated Gaussian, latent inference lowers NRMSE, CRPS, and calibration error, brings coverage closer to nominal, and moves C2ST closer to chance.Both variants recover block variances, while latent inference is closer to the exact posterior; latent sampling is 1.8× to 3.6× faster, including 2.3× for normalizing flow.Source: Figure 2 · p. 6, §4.4 · p. 8
  • Across three inference networks, compression ratios at or below the d = 8 optimum meet or beat standard-SBI accuracy and error thresholds, while greater compression degrades quality.Figure 5 compares CRPS, covariance error, and calibration error against standard-SBI baselines.Source: §4.4 · p. 8, Figure 5 · p. 9
  • Fashion-MNIST draws look convincing and have similar NRMSE and CRPS, but standard inference has better C2ST; random-field draws match expected amplitude and smoothness despite pixel-based metrics being far from optimal.The random-field posterior is not conditioned on simulator noise, and latent draws lack the highest frequencies reconstructed by the autoencoder.Source: §4.4 · p. 9
  • Latent inference improves accuracy and marginal calibration on map-to-satellite, while Fashion-MNIST results are mixed and random-field draws lack the highest frequencies.On the correlated Gaussian, latent inference is better calibrated and closer to the exact posterior; beyond d = 8, information loss quickly degrades inferential quality.Source: §4.4 · pp. 8–9

5 DISCUSSION AND CONCLUSION

The authors characterize conditions for latent-space SBI to recover the target posterior and report comparable or better accuracy and marginal calibration with faster sampling. They also identify limits in the autoencoder and in the scope of the empirical comparisons.

5 DISCUSSION AND CONCLUSION

  • Across four case studies and three generative families, latent SBI achieved similar or better accuracy and marginal calibration than standard inference with considerably faster sampling.The framework’s error is bounded by posterior, code, and decoder terms.Source: §5 · p. 9
  • The autoencoder is trained on parameters alone, so its code can discard directions of θ that the observation constrains; decoded draws also remain on its learned manifold.The pixel-wise decoder loses high-frequency detail detectable by C2ST; perceptual or adversarial losses could reduce this gap.Source: §5 · p. 9
  • Comparisons use a single training budget, leaving changes in latent-versus-standard performance with budget unstudied; sequences and graphs are also outside the evaluated target types.The paper notes that these structures call for different autoencoder inductive biases.Source: §5 · p. 10

AI USE STATEMENT

  • The authors report using generative AI across theoretical, experimental, implementation, literature, and writing work, while stating that authors reviewed the assisted work and take responsibility for the final content.Source: AI USE STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

A CASE STUDY DETAILS

  • Case-study details cover correlated-Gaussian covariance targets, Fashion-MNIST deblurring, Gaussian random fields, and map-to-satellite translation.Source: A CASE STUDY DETAILS · pp. 16–17

B MODEL CONFIGURATIONS AND HYPERPARAMETERS

C THE EQUAL-COMPUTE PROTOCOL IN DETAIL

The protocol defines training budget as accumulated training time, with autoencoder training included and validation time excluded. It also distinguishes evaluation FLOPs from wall-clock speed on the deblurring study.

C THE EQUAL-COMPUTE PROTOCOL IN DETAIL

D CALIBRATION DIAGNOSTICS

The diagnostics assess conditional posterior fit using a classifier two-sample test and compare sampled covariance against an analytic reference. They also show that marginal calibration can miss decoder-range and high-frequency deficits.

D CALIBRATION DIAGNOSTICS

  • Marginal coverage cannot detect that latent draws lie in the decoder’s lower-dimensional range, which has zero volume under the true posterior.Source: D CALIBRATION DIAGNOSTICS · p. 19
  • The classifier two-sample test separates true target-observation pairs from posterior-draw pairs through their conditional distribution; exact posterior draws score 0.50.Controls test decoder reconstructions and posterior draws in latent space separately.Source: D CALIBRATION DIAGNOSTICS · p. 19
  • For Gaussian random fields, covariance relative Frobenius error is 0.071 (0.082) for standard and 0.076 (0.097) for latent inference with flow matching (diffusion), against a 0.064 exact-draw floor.The error is dominated by low frequencies and hides the latent draws’ high-frequency deficit.Source: D CALIBRATION DIAGNOSTICS · p. 19
  • The latent-space two-sample test is at its null on correlated Gaussian and random-field studies, and just above it on deblurring, indicating that posterior networks reproduce the encoded conditional.The information-term estimates have caveats: some are upper bounds, and on correlated Gaussian the entropy terms are dominated by encoder variance collapse.Source: D CALIBRATION DIAGNOSTICS · p. 20

D.1 CAVEATS OF THE AUTOENCODER

  • The VAE’s pixel-wise reconstruction loss and small KL penalty favor directions of largest prior variance, and latent draws retain the autoencoder’s reconstruction deficits.For random fields, latent draws retain 9% of analytic power above wavenumber 13, a band holding 0.2% of variance.Source: §D.1 · p. 20
  • For map-to-satellite, C2ST separates latent draws at 0.98, about as well as autoencoder reconstructions at 0.99.For random fields, the radially averaged spectrum relative error is 0.46 for latent versus 0.012 for standard inference, with a 0.004 floor.Source: §D.1 · p. 21

E COMPUTATIONAL COMPLEXITY AND HARDWARE

Latent sampling reduces per-evaluation cost through fewer convolutional grid positions, while its decoder adds some cost per draw. Reported speedups depend on the timing protocol and account for autoencoder training within the matched budget.

Where the speedup comes from

  • The 4 × 4 × 32 deblurring code has 64 times fewer grid positions than the 32 × 32 target, and the narrower latent network uses over two orders of magnitude fewer FLOPs per evaluation.The decoder adds 5–14% to latent-variant FLOPs per draw; an 8 × 8 × 8 code with the same number of entries would save more than an order of magnitude in FLOPs at the same widths.Source: E COMPUTATIONAL COMPLEXITY AND HARDWARE · p. 21

Amortizing the first stage

  • About 2.7 × 10^5 draws repay the first-stage cost when the autoencoder is treated as additional on the deblurring study with flow matching.The estimate divides 25% of the 6-hour budget by the per-1,000-draw timing difference of 20,099 and 322 ms.Source: E COMPUTATIONAL COMPLEXITY AND HARDWARE · p. 21

Solver steps

  • Cost timing fixes the solver and step count at 30 in both spaces, yielding 180 network evaluations per flow-matching draw and 60 for the two-step diffusion sampler; correlated Gaussian is timed on CPU with adaptive solvers.All reported posterior draws use adaptive solvers with data-dependent steps, which are not part of the reported speedup. Speedups are comparable within a family, but absolute costs are not comparable across families.Source: E COMPUTATIONAL COMPLEXITY AND HARDWARE · p. 21

Sampling batch size

  • The measured speedup is smaller than the FLOP ratio, and matching sampling batch sizes makes the reported throughput ratio a lower bound on attainable throughput.At the selected batch sizes, target-space networks are compute-bound while latent networks are not; larger latent batches are not counted.Source: E COMPUTATIONAL COMPLEXITY AND HARDWARE · p. 21

Hardware and comparison

  • For k = 8 block variances, latent flow matching scatters around exact posterior widths, while standard widths are too narrow for most conditions and too wide for some.Figure 6 compares posterior standard deviation and W1 distance to exact draws across 1,000 test conditions with 100 draws each.Source: Figure 6 · p. 22

F PROOFS

The proofs decompose posterior error into latent and decoder contributions, then relate the first-stage InfoVAE objective to matching the encoder-induced conditional distribution under stated assumptions.

Proof of Proposition 1

  • Proposition 1 applies the KL chain rule to split joint posterior error into a latent-code term and a conditional decoder term, with marginalization not increasing divergence.The proof identifies the conditional information term as I(θ; x | ϑ) and the decoder mismatch as Eϑ KL(qξ(θ | ϑ) ∥ qψ(θ | ϑ)).Source: F PROOFS · p. 22

Proof of Proposition 2

  • At fixed I(θ; ϑ), Proposition 2's InfoVAE objective is maximized when the decoder matches qξ(θ | ϑ) and the aggregate code distribution matches the prior, under its stated assumptions.The proposition assumes continuous spaces and sufficiently flexible encoder and decoder families; the Gaussian fixed-variance decoder condition is approached rather than attained in practice.Source: F PROOFS · p. 22
  • When α + λ − 1 = 0, the marginal penalty becomes −β KL(qξ(ϑ) ∥ p(ϑ)); β > 0 alone forces the aggregate code distribution to equal the prior.The proof notes that this preserves the proposition's conclusion for the experimental setting.Source: F PROOFS · p. 22

G ADDITIONAL RESULTS

The additional results compare latent and standard flow matching against exact posterior widths on the correlated Gaussian and describe a separate training-set-size ablation.

Additional results

  • Along the block variances of the correlated Gaussian, standard flow matching is overconfident for most conditions, while latent flow matching balances over- and underconfidence (Figure 6).Source: G ADDITIONAL RESULTS · p. 22
  • Table 6 reports a Fashion-MNIST deblurring training-set-size ablation at the main-comparison epoch counts, with means and standard errors over three seeds.The caption identifies the ablation scope and aggregation but does not provide metric values.Source: Table 6 · p. 23

H ABLATIONS

The ablations examine joint autoencoder fine-tuning and reduced training-set sizes, showing a code-scale collapse under joint training and reconstruction errors that vary with dataset size.

Joint fine-tuning of the autoencoder

  • Joint fine-tuning yields NRMSE 0.065 versus 0.033 for the frozen control on 1,000 test conditions with 64 draws each.Both continuations lasted 300 further epochs and neither was budget-matched to the main comparison.Source: H ABLATIONS · p. 23
  • With fixed code standardization, joint training shrinks standardized-code standard deviation to 0.007 times its initial value while reconstruction error grows 1.5-fold.KL weights of 10^-3 and 10^-2 keep scale at 0.32 and 0.87 times its initial value, but raise reconstruction error 7.0- and 13.6-fold and yield NRMSE 0.049 and 0.104.Source: H ABLATIONS · p. 23

Training-set size

  • Held-out reconstruction RMSE is 0.028, 0.011, and 0.008 with 1,000, 10,000, and 60,000 Fashion-MNIST training targets, respectively.Each training-set size uses its own autoencoder, and both variants train for the main-comparison epoch count at each size.Source: H ABLATIONS · p. 23

I ADDITIONAL RANDOM SAMPLES

The figures show posterior samples for randomly drawn test conditions across Fashion-MNIST deblurring, Gaussian random fields, and map-to-satellite translation. A smoothness sweep shows increasing smoothness with α in both variants.

I ADDITIONAL RANDOM SAMPLES

  • Four randomly drawn Fashion-MNIST deblurring test conditions show posterior samples generated with flow matching.Source: Figure 7 · p. 24
  • Four randomly drawn Gaussian random-field test conditions show posterior samples, with colors on one symmetric scale per condition.Source: Figure 8 · p. 25
  • Smoothness increases with α in both Gaussian random-field variants, matching the ground truth across four equally spaced α values at log σ = 0.The figure shows posterior samples across the sweep.Source: Figure 9 · p. 26
  • Six randomly drawn map-to-satellite test conditions show posterior samples, as in Figure 4.Source: Figure 10 · p. 27

Related papers

Loading 2609.37381v1…