cs.LGSep 29, 2026
SaveScalable Diffusion SBI for Compositional Inference under Simulator Misspecification
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 and optional Langevin correctors, without retraining. ( Middle ) HBDS alternates updates of shared parameters and group-specific latent states conditioned on each group’s data . ( Right ) Path divergence 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 ( , ) | SLCP ( ) | |||
| Method | max-sW | sW | MMD | C2ST |
| GAUSS | ||||
| JAC | ||||
| Langevin | ||||
| F-NPSE | ||||
| F-NPSE-SDE ( ) | ||||
Table 1: Comparison of compositional samplers on the Gaussian benchmark ( , ) and SLCP ( ). Entries are mean 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 . Colors identify the generating groups; plotted latent states are continuous sampler outputs. Dotted lines mark the generating .
Figure 5: Likelihood adaptation and prediction error in the analytic toy at a fixed LSR reference. ( Left to right ) Known discrepancy , PT absolute error , FT absolute error , and path divergence .
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 ; 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 , 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 ; bars are pointwise 95% intervals over five paired seeds. ( f ) Step size relative to the local contraction limit. Ordinary VP crosses , 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 . This oracle diagnostic conditions HBDS latent updates on the true and parameter updates on the true group signs; pooled sampling receives no such information.
| SBI symbol | BMP meaning | Dimension |
|---|---|---|
| Experimental design: ligand doses | ||
| Shared affinities and efficiencies | ||
| Cell-line receptor state | ||
| Observation: steady-state SMAD response |
Table 2: BMP-to-SBI mapping and dimensions.
Figure 9: Comparison between BMP receptor posterior samples (columns) across LSR best-fit ( 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 parameter set can be seen in the multi-modal distributions of the fits as well as the case of BMPR2, where only works if the ACVR1 KD and BMPR1A KD conditions are 1.5 the amount of the parent cell line. Imposing a diffusion prior on the receptors in the BMP simformer avoids this clash of parameters.
| RMSE | Median | RMSE | Median | RMSE | Median | RMSE | Median | |
|---|---|---|---|---|---|---|---|---|
| 0.29 | 12.16 | 0.36 | 12.96 | 0.36 | 12.97 | 0.36 | 12.97 | |
| 0.30 | 28.15 | 0.32 | 89.55 | 0.32 | 90.63 | 0.32 | 90.87 | |
| 0.28 | 14.30 | 0.34 | 14.89 | 0.33 | 14.62 | 0.34 | 14.94 | |
| 0.34 | 29.97 | 0.36 | 41.20 | 0.36 | 41.36 | 0.36 | 41.86 | |
| 0.33 | 45.26 | 0.40 | 93.51 | 0.42 | 94.78 | 0.41 | 99.30 | |
Table 5: BMP regularization and Langevin-correction sweep. Each setting used 250 HBDS draws and one sampling seed. is the number of corrector steps per diffusion step.
| Method | Feedback | Update mechanism | Sampler gradients | Learned auxiliaries | Reference regularization |
|---|---|---|---|---|---|
| DPOK ( Fan et al., 2023 ) | Policy gradient | – | Optional value baseline | Discrete path KL | |
| DDPO ( Black et al., 2024 ) | Policy gradient | – | – | No explicit PT KL | |
| ReFL ( Xu et al., 2023 ) | Intermediate clean-prediction gradient | 1 selected step | – | Pretraining denoising loss | |
| AlignProp ( Prabhudesai et al., 2023 ) | Randomly truncated rollout backpropagation | Random trailing steps | – | No explicit PT KL | |
| DRaFT- ( Clark et al., 2024 ) | Truncated rollout backpropagation | Last steps | – | LoRA weight decay (implementation) | |
| ELEGANT ( Uehara et al., 2024b ) | SOC objective via rollout gradients | Full | Initial-time value function and initial-law sampler | Drift-control cost and initial-law KL |
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, denotes reward values without sample derivatives, denotes differentiable reward or loss feedback, denotes constraint statistics and target moments, and denotes gradients of a differentiable unnormalized target energy , with ; when an energy is instead defined relative to a reference distribution , the equivalent tilted target is . 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 | Median distance | -C2ST |
|---|---|---|---|---|---|
| FMCPE-LSR | 235 | 1,204 | |||
| FMCPE-LSR | 235 | 2,408 | |||
| FMCPE-LSR | 235 | 3,612 | |||
| FMCPE-LSR | 235 | 4,816 | |||
| FMCPE-LSR | 470 | 1,204 | |||
| FMCPE-LSR | 470 | 2,408 |
Table 7: Predictive errors through the original BMP simulator on all 940 measurements, using 250 parameter draws per run. Values are mean SE over three FMCPE training seeds or three sampling seeds of each fixed PT or FT simformer. FMCPE -C2ST uses separate shared-normalization correction refits and a joint 80-dimensional parameter target.
| Metric | PT | FT | Improvement (PT/FT) |
|---|---|---|---|
| Posterior | |||
| Posterior mean error | |||
| Correct-region likelihood RMSE | |||
| Misspecified-region likelihood RMSE |
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 favor fine-tuning and values below indicate degradation.
Figure 11: Regularization sweep on the joint-RBF misspecification toy. Columns increase from to . ( Top ) FT absolute prediction error , with full-grid MAE annotated. ( Bottom ) Path divergence . Stars mark observed designs; dashed contours enclose . 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 from PT ( ) to FT ( ), 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.
| Error improves | Error does not improve | |
|---|---|---|
| Small | Adequate simulator or mild correction | Under-correction |
| Large | Corrected misspecification | Over-correction or forgetting |
Table 9: Interpretation of design-localized fine-tuning. Path divergence measures departure from the PT model at fixed , while predictive error change is evaluated against observations at the same design .
Figure 13: LSR solution structure and cost groups across retained fits. ( Left ) PCA coordinates colored by total LSR cost. ( Right ) The same projection colored by low, middle, and high cost groups, containing , , and fits. PC1 and PC2 explain and 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 , 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 and 30 phosphorylation efficiencies , 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.
More in this topic
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