Scalable Diffusion SBI for Compositional Inference under Simulator Misspecification

Vincent D. Zaballa; Elliot E. Hui

cs.LGSep 29, 2026

Scalable Diffusion SBI for Compositional Inference under Simulator Misspecification

Authors: Vincent D. Zaballa, Elliot E. Hui

Organizations: University of California, Irvine

Abstract

Simulation-based inference is challenging when many heterogeneous observations must be composed, hierarchical latent structure must be preserved, and the simulator is misspecified relative to observed data. We develop sampling and fine-tuning methods for diffusion-based inference in design-conditional settings, where the same simulator is queried across different experimental conditions $ξ$ . We extend compositional score-based inference with a continuous-time diffusion coefficient that accounts for the number of observations, avoiding Jacobian and auxiliary-covariance corrections. We introduce Hierarchical Blockwise Diffusion Sampling (HBDS), which infers shared parameters and group-specific latent states using a single pretrained model, with the hierarchy specified only at sampling time. Together, these methods support variable observation sets and groupings without retraining. To address misspecification, we introduce path-regularized fine-tuning that adapts the learned likelihood to observations and transfers corrections to posterior inference. Using Girsanov's theorem, we quantify path divergence between pretrained and fine-tuned models across experimental designs and interpret it alongside predictive errors to distinguish candidate misspecification correction from unnecessary adaptation. We evaluate compositional sampling on exact-score Gaussian and Simple Likelihood, Complex Posterior benchmarks, HBDS with analytic and learned scores on a controlled hierarchical model, and fine-tuning and localization on a separate analytic model with known design-dependent discrepancy. Finally, we apply the framework to 940 measurements across four cell lines in a mechanistic Bone Morphogenetic Protein signaling model, where fine-tuning improves posterior-predictive accuracy relative to the pretrained model and shifts posterior marginals toward the least-squares reference while retaining spread.

Figures & tables

Figure 1: Schematic overview of the framework. ( Left ) Compositional sampling combines observations using the derived diffusion coefficient

g_{n}(t)

and optional Langevin correctors, without retraining. ( Middle ) HBDS alternates updates of shared parameters

\theta

and group-specific latent states

R_{k}

conditioned on each group’s data

\mathcal{D}_{k}

. ( Right ) Path divergence

\mathcal{A}(\theta^{\mathrm{ref}},\xi)

measures the magnitude of likelihood adaptation across experimental designs. Comparing this field with predictive improvement helps distinguish candidate misspecification correction from over-adaptation.

Figure 2: Likelihood fine-tuning of the joint diffusion surrogate. At fixed reference parameters and design, generated responses are evaluated against observations with a path-space KL penalty relative to the PT model. Gradients through sampling update the token embeddings, FiLM layers, and score network, which also support posterior queries.

Figure 3: Conditional path-law geometry under fine-tuning.

	Gaussian ( $T=400$ , $n=100$ )	SLCP ( $n=30$ )
Method	max-sW $\downarrow$	sW $\downarrow$	MMD $\downarrow$	C2ST $\rightarrow 0.5$
GAUSS	$0.24\pm 0.23$	$0.45\pm 0.18$	$0.05\pm 0.16$	$0.97\pm 0.04$
JAC	$0.26\pm 0.22$	$0.53\pm 0.19$	$0.05\pm 0.14$	$0.95\pm 0.05$
Langevin	$0.52\pm 0.50$	$0.81\pm 0.28$	$0.13\pm 0.22$	$0.97\pm 0.03$
F-NPSE	$0.29\pm 0.24$	$0.83\pm 0.29$	$0.11\pm 0.13$	$0.98\pm 0.03$
F-NPSE-SDE ( $L=1$ )	$0.29\pm 0.25$	$0.74\pm 0.31$	$0.10\pm 0.15$	$0.98\pm 0.03$

Table 1: Comparison of compositional samplers on the Gaussian benchmark (

n=100

T=400

) and SLCP (

n=30

). Entries are mean

\pm

standard deviation over five seeds for Gaussian and 25 test instances using one network seed for SLCP.

Figure 4: ( Left ) Pooled sampling uses one shared latent. ( Right ) HBDS retains separate group latents with a shared $\theta$ . Colors identify the generating groups; plotted latent states are continuous sampler outputs. Dotted lines mark the generating $\theta=1$ .

Figure 5: Likelihood adaptation and prediction error in the analytic toy at a fixed LSR reference. ( Left to right ) Known discrepancy $h(\xi)$ , PT absolute error $e_{\mathrm{PT}}(\xi)$ , FT absolute error $e_{\mathrm{FT}}(\xi)$ , and path divergence $\mathcal{A}(\theta^{\text{ref}},\xi)$ .

Figure 6: BMP predictions and design-localized adaptation. ( Left ) BMP4/BMP10 competition in NMuMG. ( Middle ) The same competition in BMPR2 KD. ( Right ) BMP4 titration in NMuMG. From top to bottom: LSR and FT predictions against observations; low (L), middle (M), and high (H) LSR cost groups; experiment-normalized path divergence $\mathcal{A}(\theta,\xi)$ ; and reduction in absolute prediction error from LSR to the FT bootstrap median.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Compositional diffusion improves finite-budget Gaussian sampling. ( a–d ) Diffusion schedules, covariance tracking, and matched particle paths. At

n=100

, ordinary-VP covariance errors require a logarithmic scale; path gaps and braces mark off-scale excursions. ( e ) Endpoint sW reduction at matched cost, with positive values favoring

g_{n}^{2}(t)

; bars are pointwise 95% intervals over five paired seeds. ( f ) Step size relative to the local contraction limit. Ordinary VP crosses

1

, indicating numerical amplification, while the derived coefficient stays below it. Accuracy improves even below this threshold.

Figure 8: Joint-state diagnostic for the hierarchical toy with 20 observations across two groups and no ULA correctors. ( Left ) Pooled sampling with one shared latent. ( Right ) HBDS with separate group latents. Rows use analytic (top) and learned (bottom) scores; dotted lines mark the generating $\theta=1$ . This oracle diagnostic conditions HBDS latent updates on the true $\theta$ and parameter updates on the true group signs; pooled sampling receives no such information.

SBI symbol	BMP meaning	Dimension
$\xi$	Experimental design: ligand doses $L_{1:J}$	$J=5$
$\theta$	Shared affinities $K_{ijk}$ and efficiencies $\varepsilon_{ijk}$	$2IJK=60$
$R^{k}$	Cell-line receptor state $A_{1:I},B_{1:K}$	$I+K=5$
$y$	Observation: steady-state SMAD response $S$	$1$

Appendix

Table 2: BMP-to-SBI mapping and dimensions.

Figure 9: Comparison between BMP receptor posterior samples (columns) across LSR best-fit $\theta^{\text{LSR}}$ ( top ), applying a truncated lognormal prior to the BMP receptors’ distributions ( middle ), and applying the marginal hierarchical receptor priors Appendix D ( bottom ). Model identifiability in the $\theta^{\text{LSR}}$ parameter set can be seen in the multi-modal distributions of the $\theta^{\text{LSR}}$ fits as well as the case of BMPR2, where $\theta^{\text{LSR}}$ only works if the ACVR1 KD and BMPR1A KD conditions are 1.5 $\times$ the amount of the parent cell line. Imposing a diffusion prior on the receptors in the BMP simformer avoids this clash of parameters.

Table 12

	$L=0$		$L=1$		$L=3$		$L=5$
$\boldsymbol{\lambda}$	RMSE $\downarrow$	Median $\downarrow$	RMSE $\downarrow$	Median $\downarrow$	RMSE $\downarrow$	Median $\downarrow$	RMSE $\downarrow$	Median $\downarrow$
$0$	0.29	12.16	0.36	12.96	0.36	12.97	0.36	12.97
$10^{-4}$	0.30	28.15	0.32	89.55	0.32	90.63	0.32	90.87
$5\times 10^{-4}$	0.28	14.30	0.34	14.89	0.33	14.62	0.34	14.94
$10^{-3}$	0.34	29.97	0.36	41.20	0.36	41.36	0.36	41.86
$10^{-2}$	0.33	45.26	0.40	93.51	0.42	94.78	0.41	99.30

Appendix

Table 5: BMP regularization and Langevin-correction sweep. Each setting used 250 HBDS draws and one sampling seed.

L

is the number of corrector steps per diffusion step.

Method	Feedback	Update mechanism	Sampler gradients	Learned auxiliaries	Reference regularization
DPOK ( Fan et al., 2023 )	$R$	Policy gradient	–	Optional value baseline	Discrete path KL
DDPO ( Black et al., 2024 )	$R$	Policy gradient	–	–	No explicit PT KL
ReFL ( Xu et al., 2023 )	$\nabla_{y}R$	Intermediate clean-prediction gradient	1 selected step	–	Pretraining denoising loss
AlignProp ( Prabhudesai et al., 2023 )	$\nabla_{y}R$	Randomly truncated rollout backpropagation	Random trailing steps	–	No explicit PT KL
DRaFT- $K$ ( Clark et al., 2024 )	$\nabla_{y}R$	Truncated rollout backpropagation	Last $K$ steps	–	LoRA weight decay (implementation)
ELEGANT ( Uehara et al., 2024b )	$\nabla_{y}R$	SOC objective via rollout gradients	Full	Initial-time value function and initial-law sampler	Drift-control cost and initial-law KL

Appendix

Table 6: Representative approaches to reward- and constraint-based generative-model adaptation, compared by the feedback they require, how model updates are obtained, whether gradients propagate through sample generation, any additional learned objects, and how deviation from a reference model is controlled. In the feedback column,

R

denotes reward values without sample derivatives,

\nabla_{y}R

denotes differentiable reward or loss feedback,

(h,h^{\star})

denotes constraint statistics and target moments, and

\nabla_{y}E

denotes gradients of a differentiable unnormalized target energy

E(y)

, with

p^{\star}(y)\propto\exp[-E(y)]

; when an energy is instead defined relative to a reference distribution

p_{0}

, the equivalent tilted target is

p^{\star}(y)\propto p_{0}(y)\exp[-E(y)]

. In the sampler-gradient column, “Full” denotes differentiation through the complete sampled trajectory and explicit step labels indicate partial differentiation. In the learned-auxiliary column, “–” means that no additional function or parameter is fitted beyond the adapted generator or control and any supplied reward model; computed adjoint states, trajectory queues, and replay buffers are not counted. PT denotes a pretrained model. For Implicit Diffusion, Full ∗ denotes full finite-horizon sampler differentiation. Our implementation uses checkpointed backpropagation through a complete discretized rollout rather than the original queued joint sampling–optimization scheme, so that each gradient update corresponds to a complete rollout under a single selected design.

Method	Observations	LSR fits	RMSE $\downarrow$	Median distance $\downarrow$	$\ell$ -C2ST $\downarrow$
FMCPE-LSR	235	1,204	$0.242\pm 0.004$	$203.661\pm 38.307$	$0.198\pm 0.023$
FMCPE-LSR	235	2,408	$0.231\pm 0.026$	$203.451\pm 70.019$	$0.150\pm 0.039$
FMCPE-LSR	235	3,612	$0.184\pm 0.010$	$77.887\pm 19.521$	$0.194\pm 0.050$
FMCPE-LSR	235	4,816	$0.168\pm 0.011$	$136.460\pm 77.855$	$0.157\pm 0.049$
FMCPE-LSR	470	1,204	$0.256\pm 0.008$	$283.935\pm 26.484$	$0.137\pm 0.020$
FMCPE-LSR	470	2,408	$0.202\pm 0.006$	$108.803\pm 19.193$	$0.081\pm 0.013$

Appendix

Table 7: Predictive errors through the original BMP simulator on all 940 measurements, using 250 parameter draws per run. Values are mean

\pm

SE over three FMCPE training seeds or three sampling seeds of each fixed PT or FT simformer. FMCPE

\ell

-C2ST uses separate shared-normalization correction refits and a joint 80-dimensional parameter target.

Metric	PT	FT	Improvement (PT/FT)
Posterior $W_{2}$ $\downarrow$	$0.09$	$0.07$	$1.35\times$
Posterior mean error $\downarrow$	$0.08$	$0.03$	$2.88\times$
Correct-region likelihood RMSE $\downarrow$	$0.05$	$0.26$	$0.18\times$
Misspecified-region likelihood RMSE $\downarrow$	$0.45$	$0.30$	$1.49\times$

Appendix

Table 8: Fine-tuning on the analytic design-misspecification toy. Posterior errors are measured against the exact true posterior; likelihood RMSE is evaluated in correctly specified and misspecified regions of the design space. The improvement factor is the PT error divided by the FT error, so values above

1\times

favor fine-tuning and values below

1\times

indicate degradation.

Figure 11: Regularization sweep on the joint-RBF misspecification toy. Columns increase $\lambda$ from $10^{-4}$ to $10$ . ( Top ) FT absolute prediction error $e_{\mathrm{FT}}(\xi)$ , with full-grid MAE annotated. ( Bottom ) Path divergence $\mathcal{A}(\xi)$ . Stars mark observed designs; dashed contours enclose $h(\xi)\geq 0.1$ . Color scales are shared within rows. The sweep used one training seed and 16 Monte Carlo draws per design to estimate response means.

Figure 12: PT–FT score blending on the misspecification toy. Columns vary the sampling weight $w$ from PT ( $0$ ) to FT ( $1$ ), without retraining. ( Top ) Absolute error of the predictive mean against the noise-free response, with full-grid MAE annotated. ( Bottom ) Estimated path divergence from PT. Stars mark observed designs; color scales are shared within rows. Intermediate weights reduce average error while limiting departure from PT.

$\mathcal{A}(\theta,\xi)$	Error improves	Error does not improve
Small	Adequate simulator or mild correction	Under-correction
Large	Corrected misspecification	Over-correction or forgetting

Appendix

Table 9: Interpretation of design-localized fine-tuning. Path divergence

\mathcal{A}(\theta,\xi)

measures departure from the PT model at fixed

\theta

, while predictive error change is evaluated against observations at the same design

\xi

Panel of Figure 13 — Figure 13: LSR solution structure and cost groups across $4{,}816$ retained fits. ( Left ) PCA coordinates colored by $\log_{10}$ total LSR cost. ( Right ) The same projection colored by low, middle, and high cost groups, containing $1{,}296$ , $2{,}245$ , and $1{,}275$ fits. PC1 and PC2 explain $8.9\%$ and $5.3\%$ of the standardized feature variance. Grouping also uses fitting cost, so overlap in this two-dimensional projection is expected.

Figure 14: Design-localized adaptation in the BMP model across four cell lines (rows) and experimental conditions (columns). Thin black lines separate titration series; the thick black line separates ligand-competition (“Rim”) from gradient series (“Gradient”). Green boxes mark the conditions illustrated in Figure 6 . ( A ) Best-fitting LSR parameter group; colors identify groups, not error magnitudes. ( B ) Path divergence $\mathcal{A}(\theta,\xi)$ , with yellow indicating greater adaptation. ( C ) Fine-tuning distance improvement relative to the selected LSR anchor. ( D ) Fine-tuning distance improvement relative to the PT bootstrap median. In C and D, blue indicates reduced distance to observations and red indicates increased distance relative to the comparator.

Figure 15: BMP posterior alignment with the LSR reference. Purple and green show PT and FT marginals; red lines mark the fitted LSR reference. The upper and lower blocks show 30 binding affinities

K

and 30 phosphorylation efficiencies

\varepsilon

, with ligand rows and receptor-complex columns. Type-I receptors label the column groups above; type-II receptors label individual columns beneath them, identifying each receptor pair. FT marginals generally move closer to the reference while retaining spread rather than collapsing to it.

Simulation-Based Inference

Simulation-based inference estimates parameters and uncertainty from simulators when evaluating a tractable likelihood is difficult or impossible.

See all 87 papers

Dec 22, 2025

Diffusion Models in Simulation-Based Inference: A Tutorial Review

Diffusion models have recently emerged as powerful learners for simulation-based inference (SBI), enabling fast and accurate estimation of latent parameters from simulated and real data. Their score-based formulation offers a flexible way to learn conditional or joint distributions over parameters and observations, thereby providing a versatile solution to various modeling problems. In this tutorial review, we synthesize recent developments on diffusion models for SBI, covering design choices for training, inference, and evaluation. We highlight opportunities created by various concepts such as guidance, score composition, flow matching, consistency models, and joint modeling. Furthermore, we discuss how efficiency and statistical accuracy are affected by noise schedules, parameterizations, and samplers. Finally, we illustrate these concepts with case studies across parameter dimensionalities, simulation budgets, and model types, and outline open questions for future research.

Jonas Arruda, Niels Bracher, Ullrich Köthe +2

Bonn Center for Mathematical Life Sciences, Life & Medical Sciences Institute, University of Bonn, Germany · Center for Modeling, Simulation, & Imaging in Medicine (CeMSIM), Rensselaer Polytechnic Institute, NY, USA · Computer Vision and Learning Lab, Heidelberg University, Germany

Simulation-Based Inference Parameter Estimation Score-Based Generative Modeling Diffusion Models

Ask PDF

Oct 6, 2026

Spectra: Exact Component Transport for Test-Time Prior Adaptation in Simulation-Based Inference

Simulation-based inference (SBI) has become a powerful approach to Bayesian inference in complex scientific models whose likelihoods are difficult or impossible to evaluate. Amortized SBI learns reusable inference models from simulated data, enabling rapid posterior inference for new observations, and modern generative models have made these models increasingly expressive. However, this reuse is limited to the prior distribution chosen during training, whereas scientific analyses often need revised priors as knowledge accumulates or alternative assumptions are tested. We introduce Spectra, a test-time adaptation method for diffusion-based SBI. Spectra uses an exact score-transport identity to obtain the adapted score from a frozen diffusion model in closed form for structured prior changes, without additional simulation or training. Across six SBI benchmarks, Spectra achieves accurate adaptation under strong prior shifts at low online sampling cost. This enables pretrained SBI models to incorporate updated prior information at test time.

Xin Zhao, Nico Scherf, Robert Trampel +2

Department of Neurophysics, Max Planck Institute for Human Cognitive and Brain Sciences, Leipzig, Germany · International Max Planck Research School on Cognitive NeuroImaging, Leipzig, Germany · Methods and Development Group Neural Data Science and Statistical Computing, Max Planck Institute for Human Cognitive and Brain Sciences, Leipzig, Germany +3

Simulation-Based Inference Test-Time Adaptation Diffusion Models

Ask PDF

May 20, 2026

Theoretical guidelines for annealed Langevin dynamics in compositional simulation-based inference

Compositional score-based approaches to simulation-based inference (SBI) approximate the posterior over a shared parameter given

n

independent observations by aggregating individually learned posterior scores: currently, there are two main propositions of such methods (Geffner et al. (2023), Linhart et al. (2026)). As the resulting composite score does not correspond to the score of any distribution along the forward diffusion path of the true multi-observation posterior, sampling from it via a reverse SDE leads to an irreducible bias. Annealed Langevin dynamics provides a principled alternative: it treats the composite score as the genuine score of a sequence of tractable bridging densities and samples from them in succession. When properly tuned, it could lead to a controllable bias. However, its hyperparameters, namely step sizes, the number of steps per level, and the number of annealing levels, have so far been chosen empirically. We derive Wasserstein bounds for annealed Langevin with approximate scores and translate them into explicit decision rules for these hyperparameters that guarantee a prescribed sampling accuracy, while highlighting different theoretical aspects of each composite score formulation. In the Gaussian setting, we obtain closed-form expressions for all relevant quantities and prove that the bridging densities of Linhart et al. (2026) consistently admit larger step sizes and require fewer total Langevin steps than those of Geffner et al. (2023). Furthermore, we show empirically that the tuning obtained in the Gaussian setting generalizes to more complex problems, thus providing a well-understood and theoretically grounded starting point for practitioners using compositional score-based approaches.

Camille Touron, Gabriel V. Cardoso, Julyan Arbel +1

Univ. Grenoble Alpes, Inria CNRS, Grenoble INP, LJK, France · Geostatistics team, Centre for geosciences and geoengineering Mines Paris, PSL University, Fontaineableau, France

Simulation-Based Inference Langevin Dynamics Posterior Sampling

Ask PDF

Abstract

Simulation-Based Inference

Explore similar work

Diffusion Models in Simulation-Based Inference: A Tutorial Review

Spectra: Exact Component Transport for Test-Time Prior Adaptation in Simulation-Based Inference

Theoretical guidelines for annealed Langevin dynamics in compositional simulation-based inference