Scalable Diffusion SBI for Compositional Inference under Simulator Misspecification

We start from a Simformer (Gloeckler et al., 2024), a masked diffusion transformer trained with denoising score matching on simulator draws of parameters \(\theta\), observations \(y\), and experimental designs \(\xi\). Because any subset of variables can be conditioned on, the same network answers posterior queries \(p(\theta\mid y,\xi)\) and likelihood queries \(p(y\mid\theta,\xi)\). Everything below happens after pretraining.

The application that motivated this combination is a mechanistic model of Bone Morphogenetic Protein (BMP) signaling, where ligands bind pairs of receptors and the resulting complexes drive a transcriptional response.

Under conditional independence, the posterior given \(n\) observations factorizes into single-observation posteriors, so one score network evaluated \(n\) times can stand in for a model trained on the whole set. F-NPSE (Geffner et al., 2023) uses this in discrete time, and its reverse transitions contract as \(n\) grows. We carry that contraction into continuous time by keeping the variance-preserving drift and prescribing a smaller accumulated noise variance, which fixes the diffusion coefficient.

\[ g_n(t)^2=\beta(t)\,\frac{1+(n-1)\bigl(1-\alpha(t)\bigr)^2}{\bigl[1+(n-1)\bigl(1-\alpha(t)\bigr)\bigr]^2} \]

Here \(\alpha(t)\) is the cumulative squared signal retention and \(\beta(t)\) the usual noise rate. One observation recovers the ordinary VP-SDE. With more, the coefficient falls toward \(\beta(t)/n\) at high noise. The resulting sampler, F-NPSE-SDE, needs only score evaluations. It skips the score Jacobians of JAC and the auxiliary covariance estimates of GAUSS (Linhart et al., 2026).

Pooling every observation into one compositional posterior assumes they all share the same latent variables, and grouped data, such as the receptor states \(R^k\) of the BMP model, break that assumption. Each of the four cell lines has its own \(R^k\), while the biophysical parameters \(\theta\) are common to all of them.

Hierarchical Blockwise Diffusion Sampling (HBDS) keeps one shared parameter state and a separate latent state per group. At every reverse-diffusion time each group state is updated from its own data and the current shared parameters, then the shared parameters are updated from all groups. The pretrained network is reused unchanged, with no hierarchical estimator to train and no need to tokenize the whole dataset.

The toy model below isolates this structure. Each group \(g\) has a sign-valued latent \(z_g\), all groups share one parameter \(\theta\), and \(x_{gi}\) is observation \(i\) in group \(g\).

\[ z_g\sim\mathrm{Rad}\!\left(\tfrac12\right),\qquad \theta\sim\mathcal{N}(0,1),\qquad x_{gi}\mid\theta,z_g\sim\mathcal{N}(5z_g+\theta,\,1). \]

\(\mathrm{Rad}(\tfrac12)\) gives the two signs equal probability, so the two groups sit ten units apart when their signs differ.

A surrogate trained on simulations inherits the simulator's errors. We adapt it to observed data by sampling responses from the learned likelihood at a fixed reference parameter \(\theta^{\mathrm{ref}}\), scoring them against observations with a negative mean-squared-error reward, and differentiating through the sampler. A path-space KL penalty holds the fine-tuned (FT) model near the pretrained (PT) one.

Fine-tuning changes the law of the whole sampling trajectory, and Girsanov's theorem gives that change in closed form.

\[ \mathcal{A}(\theta,\xi):=\mathrm{KL}\bigl(\mathbb{P}_{\phi}^{\theta,\xi}\,\big\|\,\mathbb{P}_{\phi_0}^{\theta,\xi}\bigr)=\tfrac12\,\mathbb{E}_{\mathbb{P}_{\phi}^{\theta,\xi}}\!\left[\int_0^T g(t)^2\,\bigl\|s_{\phi}^{y}-s_{\phi_0}^{y}\bigr\|_2^2\,dt\right] \]

Evaluated across designs at the reference parameter, \(\mathcal{A}\) is a map of how far the model was moved at each experimental condition. It measures the size of the adaptation and says nothing on its own about whether predictions improved, so we always read it next to predictive error.

A linear-Gaussian simulator over a two-dimensional design space omits a localized bump \(h(\xi)\) that is present in the observations. We fine-tune on ten noisy observations and compare predictions over the full design grid.

Two controls trade correction against preservation. The penalty strength \(\lambda\) acts during training. Autoguidance acts at sampling time and needs no retraining.

Fine-tuning corrects the model near the discrepancy and adds error elsewhere, while the pretrained model does the opposite. Autoguidance (Karras et al., 2024) combines the two during sampling with the score

\[ s_w^{y}=s_{\mathrm{PT}}^{y}+w\,\bigl(s_{\mathrm{FT}}^{y}-s_{\mathrm{PT}}^{y}\bigr), \]

where \(w=0\) recovers the PT model and \(w=1\) the FT model. An intermediate weight keeps the correction near the bump and reduces the error that fine-tuning introduced at designs far from the observations, which lowers error over the whole design grid. Neither model is retrained, and the only added cost is one more score evaluation per sampling step.

We combine fine-tuning with HBDS on 940 measurements from a parent cell line (NMuMG) and three receptor knockdowns. The model has 60 shared biophysical parameters and five receptor quantities per cell line. The fine-tuning reference is the least-squares fit (LSR) that is standard for this model, used as an anchor and not as ground truth.

With parameters held at the reference, FT improves on the simulator throughout the BMPR2-knockdown competition series and at lower BMP4 doses. In the titration those gains generally coincide with larger path divergence. In the NMuMG competition series adaptation often increases error, which is the case the table above labels over-correction.

We benchmark fine-tuning against the most relevant existing method, Flow Matching Corrected Posterior Estimation (FMCPE). FMCPE trains a neural posterior estimator on simulations and then corrects it with a smaller calibration set, which we build here from LSR fits. It needs a dedicated posterior model and correction flows for each observation set, so every row with a different number of observations below was retrained. The Simformer reuses one set of weights across observation sets through masking and compositional sampling. FMCPE still gives a direct comparison, because both methods can be scored by the same predictive errors through the BMP simulator.

Likelihood fine-tuning is anchored at a single reference \(\theta^{\mathrm{ref}}\), the LSR fit, so one might expect posterior inference to reduce to that point estimate. It does not. The FT posterior marginals generally shift toward the reference relative to pretraining while retaining substantial spread, which indicates that the likelihood-side correction transfers to posterior queries without collapsing inference onto the fine-tuning anchor.

Path divergence compares two models at the same \((\theta,\xi)\). The same construction also compares one model with itself at nearby conditioning values. Writing \(z=(\theta,\xi)\), the KL divergence between path laws at \(z\) and \(z+\delta z\) is, to second order, \(\tfrac12\,\delta z^{\top}\mathcal{I}_{\phi}^{\mathrm{path}}(z)\,\delta z\) with

\[ \mathcal{I}_{\phi}^{\mathrm{path}}(z)=\mathbb{E}_{Y\sim\mathbb{P}_{\phi}^{z}}\!\left[\int_0^T g(t)^2\,J_z s_{\phi}^{y}(Y_t,t;z)^{\top}J_z s_{\phi}^{y}(Y_t,t;z)\,dt\right]=\begin{bmatrix} I_{\theta\theta} & I_{\theta\xi}\\ I_{\xi\theta} & I_{\xi\xi}\end{bmatrix}. \]

This is a Fisher metric on the family of path laws \(\mathcal{M}_{\phi}\). It can be singular where the conditioning variables are not locally identifiable, for example when different parameter combinations give indistinguishable responses.

Each block answers a different question.

• \(I_{\theta\theta}\) measures how much the path law changes with the mechanistic parameters at a fixed design. Designs where it is large are locally informative about \(\theta\).
• \(I_{\xi\xi}\) measures how quickly the path law changes across designs. For the pretrained model, large values mark regions where denser simulator coverage would help most.
• \(I_{\theta\xi}\) measures whether a parameter change and a design change have the same first-order effect.

The paper proposes using these blocks to move a population of reference parameters during fine-tuning, to choose designs, and to weight a fine-tuning curriculum. These are proposed extensions. None of them is implemented or tested in the experiments above.

What are the limitations?

Compositional sampling and HBDS both rely on approximations to intermediate-time posteriors. Fine-tuning at a fixed reference can improve misspecified regions while degrading regions that were already well modeled, and on the BMP data the comparison with the least-squares fit remains mixed. These results point toward adaptation that varies across experimental designs instead of one global regularization strength.