Compositional amortized inference for large-scale hierarchical Bayesian models
Authors: Jonas Arruda
Vikas Pandey
Catherine Sherry
Margarida Barroso
Xavier Intes
Jan Hasenauer
Stefan T. Radev
Venue: ICLR 2026 Poster (2026)
Keywords: Amortized Bayesian Inference
Hierarchical Models
Compositional Modeling
Score Matching
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We extend amortized Bayesian inference to hierarchical models using compositional score matching with adaptive solvers and a novel error-damping estimator.

Abstract: Amortized Bayesian inference (ABI) with neural networks has emerged as a powerful simulation-based approach for estimating complex mechanistic models. 
However, extending ABI to hierarchical models, a cornerstone of modern Bayesian analysis, has been a major hurdle due to the need to simulate and process massive datasets. 
Our study tackles these challenges by extending compositional score matching (CSM), a divide-and-conquer strategy for Bayesian updating using diffusion models. 
We develop a new error-damping estimator to address previous stability issues of CSM when aggregating large numbers of data points. 
We first verified the numerical stability with up to 100,000 data points on a controlled benchmark. 
We then evaluated our method on a hierarchical AR model, achieving competitive performance to direct ABI baselines on smaller problem sizes while using less than one full model simulation for larger problem sizes. 
Finally, we address a large-scale inverse problem in advanced microscopy with over 750,000 parameters, demonstrating its relevance to real scientific applications.

Published as a conference paper at ICLR 2026
COMPOSITIONAL AMORTIZED INFERENCE FOR LARGE- SCALE HIERARCHICAL BAYESIAN MODELS
Jonas Arruda
Vikas Pandey Center for Modeling, Simulation, & Imaging in Medicine (CeMSIM) Rensselaer Polytechnic Institute, NY, USA
Bonn Center for Mathematical Life Sciences, & Life and Medical Sciences Institute (LIMES) University of Bonn, Germany
Catherine Sherry Molecular and Cellular Physiology Albany Medical College, NY, USA
Margarida Barroso
Molecular and Cellular Physiology Albany Medical College, NY, USA
Xavier Intes
Jan Hasenauer
Center for Modeling, Simulation, & Imaging in Medicine (CeMSIM) Rensselaer Polytechnic Institute, NY, USA
Bonn Center for Mathematical Life Sciences, & Life and Medical Sciences Institute (LIMES) University of Bonn, Germany
Stefan T. Radev
Center for Modeling, Simulation, & Imaging in Medicine (CeMSIM) Rensselaer Polytechnic Institute, NY, USA
ABSTRACT
Amortized Bayesian inference (ABI) with neural networks has emerged a pow- erful simulation-basee approach for estimating complex mechanistic models. How- ever, extending ABI to hierarchical models, a cornerstone of modern Bayesian analysis, has been major hurdle due to the need to simulate and process massive datasets. Our study tackles these challenges by extending compositional score matching (CSM), fusion models. We develop a new damping estimator to address previous stability issues of CSM when aggregating large numbers of data points. We first verified the numerical stability with up to 100,000 data points on a controlled benchmark. We then evaluated method hierarchical AR model, achieving competitive performance to direct ABI baselines on smaller problem sizes while using less than one full model simulation for larger problem sizes. Finally, we address large-scale inverse problem in advanced microscopy with over 750,000 parameters, demonstrating its relevance to real scientific applications.
divide-and-conquee strategy for Bayesian updating using dif-
INTRODUCTION
Simulation based inference (SBI; Cranmer et al. 2020) has entered a new era, leveraging deep learning advances deliver markedly more efficient computational statistics. Within this frame- work, amortized Bayesian inference (ABI: Bürkner et al. 2023) now scales Bayesian analysis to high-dimensional, mechanistic models, driving ofae-oofthh--ae art discoveries in fields as diverse as astrophysics (Dax al. 2025) and neuroscience (Tolley et al. 2024).
The core idea of ABI straightforward: train conditional generative model on simulations from parametric Bayesian model p(0, Y) over parameters 0 and (potentially high-dimensional) observables Y. The network can then obtain independent samples from the posterior p(0 Y) in a fraction of the time required by gold-stan standad Markov chain Monte Carlo (MCMC) methods. And simple benchmarking suites have already received extensive attention (Lueckmann et al., 202T), recent research increasingly turns to pressing challenge in Bayesian inference: affording amortized

---

Published as a conference paper at ICLR 2026
inference for hierarchical, mixed-effects, multilevel models (Rodrigues al. 2021 Heinrich et al. 2024 AAruda et al. 2024; Habermann et al. 2024).
Hicrarchical models (HMs) are the default textbook choice in Bayesian modeling (Gelman et al. 2013 McElreath, 2018). They also drive advanced analyses at the frontier of life sciences, such pharmacokineties with large patient cohorts (Arruda et al. 2024) biomedical imaging (Wang) et al. (019). However, their nested structure strains inference algorithms: standard MCMC rarely scales large datasets (Blei et al.] 2017; Margossian and Saul] 2023), and ABI faces both network design and simulation efficiency hurdles. Crucially, direct ABI approaches for estimating HMs require exhaustive simulations for each training sample (cf. Figure I| left). This renders existing ABI
approaches impractical for many real-world hierarchical models, particularly those involving large datasets or expensive simulators.
To overcome this bottleneck, we extend compositional score matching (CSM), a divide-anddconquee strategy originally introduced for Bayesian updating exchangeable models (Geffner al. (023) and recently adapted to complete pooling (Linhart et al.) (026) and time-series models (Gloeckler et al., 2025). Our method enables amortized inference HMs without ever simulating multiple data groups keeping the entire dalaset in memory (cf. Fifure
right). Moreover, it affords modern score-based diffusion models (Song and Ermon 2019: Song et al., 2021b) that have already shown
considerable potential in SBI (Sharrock et al. 2024 Gloeckler et al. 2024)
Data Generation
Compositional Bayeslan Inference for Hierarchical Models
for Simulation-
Setiiion-asse Based Training
Inversr factorization
Example
Naive training
Flat (ours)
Full data (X Yp.... Xy}
Hyperpcior
e.g., 3D stack of 512 512 imagees, where each pixel forms time series
pip)
plaok-wise
Time
Posterior medlaan
local paramotons
Condisional prioe
Globel eeliustor trained on Dla daia
Vlog &o 1
j= I.....J
Uncertaiaty
Data model p(Y| 0..))
f ~ qporl I Y...... Xy)
Sompling with error damping
Local estimator trained on flat data.
An
Training sample
19.0, "......... Yy. Egl
Vlog ay 10 y
8
BU 1) simmlations sarred per training batch with batch size B and groupe
Figure Compositional inference for hierarchical Bayesian models. Overview of our training procedure (left) and inference stages (right) for amortized hierarchical Bayesian modeling. Amortized posterior sampling uses our error-damping compositional score estimator to achieve rapid inference very high-dimensional hierarchical problems.
Despite the conceptual appeal of CSM, observe that current aggregation methods fail, even for simple, non-hierarchical models, as the number of observations grows. Here, we show that these instabilities are due to compounding approximation errors and introduce new compositional estimator that remains stable even in hierarchical models with more than 250,000 data groups and 750,000 parameters. Concretely, we develop and showcase
1. A new simulation-efficient method for estimating large hierarchical Bayesian models with diffusion models;
2. A stable reformulation of compositional score matching with stochastic differential equations (SDEs);
An error-damping mini-batel estimator that enables efficient scaling as the number of groups J becomes very large (e.g., hundreds of thousands).
2

---

Published as a conference paper at ICLR 2026
2 BACKGROUND AND RELATED WORK
Hierarchical Bayesian models Hierarchical Bayesian models are the default choice to model dependencies in nested data, where observations organized into clusters, levels, or groups (Gelman et al. 2013). From a Bayesian perspective, any parametric data model p(Y () can incorporate a multilevel structure via hierarchical prior. For instance, a two-level model defines two stages
Y, PY, 10, 0g p(O 7)), ~ p(n)
via a hyperprior p(n) encoding global variation between groups and conditional prior p(0 (1) encoding local variation within groups. The task of Bayesian estimation is to estimate the full posterior over local and global parameters:
p(n, On. Yı.) x p(n)
IY, 0z)p(0, n),
In),
j=1
where J denotes the number of groups and the data model factorizes over N₃ observations within group j as p 0,) = II p(y,,n f fy.mm-))..
The gold-standard approaches for estimating hierarchical models are Markov chain Monte Carlo (MCMC) methods (Gelman et al. 2020). While MCMC methods offer strong theoretical guarantees, they are typically too slow for real- or big data applications. Moreover, MCMC cannot be trivially
applied to simulation- based models (Sisson and Fan 2011); hence, the appeal of amortized inference.
Amortized Bayesian inference (ABI) In ABI, a generative network learns a global posterior functional Y + q(- Y). Typically, the network minimizes strictly proper scoring rule S (Gneiting) et al. 2007 over a training budget of B Monte Carlo samples from the joint model p(0, Y):
B
Ex(e.) S(a( Y),
~
Σ B
S(a(- Y(
(3)
Using a universal density estimator for q. such as coupling flows (Draxler et al., 2024), ensures that Eq. can, in principle, converge to the correct target for a large simulation budget B - co. Since Y is typically high-dimensional, a summary network h(Y) can be jointly trained to learn data embeddings
the fly (Radev et al.) (2020) or implicitly incorporated into the architecture of (Gloeckler et al. 2024). Crucially, q repays users with zero-shot sampling for any new observation (lnww) compatible with p(0, Y), making ABI an attractive avenue for efficient estimation of complex hierarchical models.
ABI for hierarchical models Previous work has already ported the basic idea of ABI to hicrarchical settings (Habermann al. 2024 Arruda al. 2024; Heinrich al. 2024 Rodrigues et al. (02TT) These works leverage the inverse factorization of Eq. in different ways to design hierarchical neural networks with inductive biases that capture the probabilistic symmetries (e.g.. permutation invariance for exchangeable groups) of HMs. However, these approaches only approximate parts of the joint posterior (Eq. 2 or scale poorly even when the number of groups becomes moderately large.
Scalability issues arise since the expectation in Eq. now runs over p(n,01.... 01....,..),, Yi,... necessitating the simulation of a dataset of datasets {Y...... Y for each batch instance (cf. Fig ure 1] left). Even for J x 1000, single training batch requires tens of thousands of simulations, quickly exceeding typical simulation budgets for non-trivial models. Similar-sized problems can also become practically infeasible for established MCMC samplers (e.g., NUTS, Hottman et al. 2014). even for models with closed-form -fom likelihoods (see Experiment 2).
Building prior work by Geffner et al. (202)) Linhart al. (2026); Gloeckler et al. (2055)
address these efficiency issues in a "divide-and www.idd-anddcooquer" manner via compositional score matching (CSM; cf. Figure right). In addition, stability and scalability. To the best of our knowledge, we provide the first amortized method capable of handling large-scale hierarchical Bayesian models with or without explicit likelihoods.
introduce several key improvements to CSM in terms of
Score matching Score-based modeling (Song and Ermon 2019) and diffusion models (Ho et al. 2020) provide a powerful framework for generative modeling by learning to reverse a noise-adding process. Diffusion models build forward process that gradually corrupts sample 0 into pure
3

---

Published as a conference paper at ICLR 2026
Table 1:
Convergence of sampling methods for Gaussian toy example Experiment 1 with maximum budget of 10,000 steps per sample (v- converges, X- fails).
Method
N=10 N=100
=
==100k
Annealed Langevin sampler (Geffner et al. 2023
x
x
x
Euler-Maruyama sampler
X
EXXi
Probability ODE sampler
V
Adaptive second-orde -oder sampler
V
GAUSS (Linhart et al. 2026)
Any sampler with damping (ours) Any sampler with schedule adjustment (ours)
V
noise at each time step E [0. I]. taking the form 0, a.0 + 016; with E, ~ (0, I). The factors
Or and d are time-dependent functions that satisfy ot of = 1 for variance-preserving processes. These functions can be parameterized by the log signal-to-noise ratio X, = ( o known as the noise schedule (Kingma and Gao 2023).
The conditional denoising score matching loss can be expressed terms of an unconditional score (Vincent 2011 Liet al. 2024):
min Pp (00
[we]l 81 (0,, A, Va, pog 00)
Intuitively, during training, the score time is determined by the initial 0, and endpoint Oo alone and not the condition Y. However, during inference, where the endpoint is not known, the condition Y allows us steer towards an endpoint belonging the posterior distribution. Given a noise schedule, the unconditional score is analytically known as Ve, p(O, 0o) = /o,/o, The weighting function wt 0, often chosen to match the noise schedule X (see (see) Kingma and Gao (2023) for a detailed review of different weighting functions and noise schedules), is instantiated here as the likelihood weighting proposed by [Sng et al. 2021a). Given a trained score model s(O,, Y, At). can sample from the posterior distribution pl YJ by integrating either the reverse stochastic differential equation (Song et al. 2021b, SDE) or the reverse probability flow, which enables posterior sampling using of-te-of-the-ane art SDE and ODE solvers (more details in Appendix A.I).. For an introduction to diffusion models in SBI, see Simons et al. (2023). However, this formulation has neither been used for hierarchical modeling nor explored for compositional score matching in recent prior work (Geffner et al. 2023 Linhart et al. 2026), as discussed next.
3 METHOD
3.1
COMPOSITIONAL. SCORE MATCHING
major challenge in Bayesian inference arises when dealing with varying and potentially large numbers of observations, especially in hicrarchical models. To address this for non-hicra non-hierarchical models, Geffner et al. (023) introduced compositional score matching (CSM), which enables the aggregation of multiple conditionally independent score estimates into a global posterior estimate. In the following, we first introduce a naive extension of CSM for estimating the global parameters r) of hierarchical models. We then propose solution to the stability problems of the naive approach that allows us estimate the full joint posterior (Eq, 2 of large hierarchical models.
Compositional score and bridging density Suppose we have J exchangeable groups of data points, Y1... Then, the compositional global posterior can be written as
p(n YI.) x p(n) 1- II p(n Y,)
((n)
(see Appendix Proposition A.Ip. Let (r( Y,)) be the time- varying density the noise-corrupted parameter 7, for & E [0, 1I. as defined by the forward diffusion process. Then, we can define the bridging densities (r( Yı.)) p(n,) (1-)(-0))) IT Pr(n. Y;). This results in linear
4

---

Published as a conference paper at ICLR 2026
composition of the prior and individual posterior scores:
log YI:)) (1-
+
V n. log Po pe (7: Y2).
(6)
After training, we can sample from the base distribution Pt=1 P = -V (o. 7I) and use the composi- tional score sample from the posterior distribution of 17. The score model can also be conditioned on m groups jointly, rather than on single group. This results in compositional update that involves only k [J/m] scores per posterior evaluation, which can improve the robustness of the score estimation. However, this comes at an increased computational cost because each training iteration requires batch simulation of m full groups of data points rather than just one group.
Sampling with compositional scores Geffner et al. (2023) employ annealed Langevin sampling to invert the diffusion process for posterior inference, which needs many score evaluations for accurate inference dolicoeur-Martinent
al., 2021) and sensitive to the step-size at each sampling iteration. In contrast, Linhart et al. (2026) used a second-order Gaussian approximation of the backward diffusion kernels to bypass Langevin sampling, introducing the need to approximate potentially large covariance matrices and limiting their experiments to only 100 observations.
In the remainder, we demonstrate that is possible instead leverage the SDE formulation by aggregating the compositional scores in the reverse SDE (see Appendix AI.] to sample from the posterior. The rationale that posterior samples can still follow the correct distribution, even if the stochastic process does not replicate the true time reversal of diffusion (Vuong et al. 2025) Crucially, this allows the use of more efficient numerical solvers.
However, regardless of the number of conditioning groups k. increasing the number of score terms leads to error compounding due to potential mismatch of marginal densities pt and the corresponding forward diffusion process, resulting in unstable dynamics and divergent samples (see /egure 5 in the Appendix). Even higher-order solvers require extremely small step sizes, rendering them impractical even for moderately large problems (cf. Table T). The next section introduces three new modifications to the naive CSM approach that stabilize the bridging density (Eq. ( and unlock unprecedented scalability to large datasets.
3.2 STABILIZING AND SCALING COMPOSITIONAL SCORE MATCHING
In contrast to most previous work, adopt the SDE formulation to perform compositional inference with adaptive solvers that automatically adjust the step size during integration. This modification is essential for larger numbers of groups J, where the need for finer granularity (i.e., smaller step sizes) increases and manual tuning becomes infeasible it avoids the need for annealed Langevin sampling, which requires many steps per noise level and becomes prohibitively expensive when error correction is needed.
Marttienar-aaanneuu et al.) (022). Moreover,
However, simply using adaptive solver does not address the two major challenges for scaling the compositional approach to very large datasets: 1) the bridging densities (Eq, (6) become unstable as J increases (see Table Ip. and 2) memory requirements grow substantially when accumulating scores over the full dataset.
Flexible error-damping bridging densities We propose to stabilize the bridging density by intro- ducing a damping factor of the accumulated score. Yet, naively applying damping factor to the compositional score to prevent from diverging would bias the posterior samples. Instead, to mitigate the instability at large J, we introduce a more flexible class of error-damping bridging densities:
J
(r YI.)) x p(n)) (1)))
II Pr, dtt),
1 I). where d(0) = do and d(1) d, < 1, and the latent diffusion prior is ( =NO,
N' Jd.
The key idea is to define monotonic function d(t) that modulates the accumulation of score contributions throughout the diffusion trajectory during inference. In high-noise regimes, we reduce the influence of the individual terms to prevent the score from diverging, whereas for + 0, we allow their contributions to accumulate, recovering the true posterior. This construction motivated by
5

---

Published as a conference paper at ICLR 2026
the observation that adaptive solvers require smaller steps in high-noise regimes avoid numerical instability (see Appendix Figuure 5). As a damping schedule, we propose an exponential decay
d(t) do exp(- (n(o-/n(///// di) with dg and hyperparameter di that can be tuned during inference.
Mini-batch estimation for memory efficiency To address memory constraints in large-data scenar- ios, we propose a mini-batch estimator for the compositional score:
J
Sup (M: YI,1, Ae) (1 J) 1,
Σ
Sp Y Xx),
(8)
M
where j; U(I,..... J and M is the mini-batch size. Proposition 3.1. The mini-batch estimator (Eq. g is an unbiased estimator of the compositional score.
For a short proof, see Appendix A.3 Combining this estimator with the damping function yields our final compositional form:
4 71 YIJ, Ao) = d(t) (1--)(1-0)) - x V,, log
se (m Yis, Aei).
This error- drmping mini-batch estimator scales well with increasing numbers of groups and maintains stability across the reverse diffusion process, as shown in our experiments.
Noise schedule adjustment for sampling Finally, we propose to use different noise schedules for training and inference. As discussed by Karras et al. (2022) and Kingma and Gao (2023), the noise schedule plays a role akin to importance sampling during training. During inference, spending less time in the high-noise regime of the reverse process improves stability and allows for larger step sizes, which is particularly important in the large-J setting. In the case of the cosine schedule X(t) -2 log( (tan(tt/2) + 2s proposed by Nichol and Dhariwal (2021). this can be easily achieved by increasing the shift parameter which effectively compresses the high-noise portion of the schedule.
3.3 COMPOSITIONAL SCORE MATCHING FOR HIERARCHICAL MODELING
To employ our stable compositional formulation for simulation-efficient hierarchical Bayesian model- ing, we represent the posterior at each hierarchical level with its own score estimator. The outputs of these score estimators are then connected via the inverse factorization (see Figure 1p. This design is similar to the frameworks introduced by Habermaan et al. (2024) and Heinrich et al. (2024). ) but avoids the need for hierarchical embeddings and exhaustive simulation. At higher levels of the hierarchy, we use our stable compositional formulation (Eq. 19, enabling the training of the global score model single groups. For example, in two-level model, have
global local (0.,, 17, Y, Ac) = Ve, lo n, YJ), + x
(7. Y,, Ae) = In,logp Y,). (10)
For each group, we can learn a shared summary representation h, h(Y,) ;) via a summary network h. Either the raw data Y, or its summary h, is then used as input to both the global and local score-based models. The design of the summary h should be adapted to the specific data modality (e.g., recurrent networks or transformers for time series). When conditioning on multiple groups, we encode exchangeability via second summary network (e.g., a DeepSet, Zaheer et al. 2017). which aggregates individual summaries into a permutation-invarian global summary.
The global and local score networks can be trained jointly via denoising score matching objectives, (11)
min Ep(O.., Y) 20 |||e + local (O., 17, Y, Ae 1 + || global (n: Y, )
with 7; = am7 + OrEr and 0, a + orE,, where E; ~ NN 1). Having trained the score models, we can sample from the joint posterior via ancestral sampling:
global
( Oy
Jocal (0) n. Y,),
(12)
where we used the compositional score (Eq. 9 to sample the global parameters and then the local parameters conditioned on the global sample using standard score-based diffusion.

---

Published as a conference paper at ICLR 2026
Number of Stops
KL Divergence
100
1.00
0.75
Postertor
Contraction
Callbr Callbrationn elton Error
04
10
RMSS
0.50
0..5
10"
0.25
10a
10
Number of Stops
Divergence
100
1.00
0.75
Postertor
Contraction
Callbration Error
04
10a
RMSSE
0.50
101
02-
10a
0.25
107
10
Data Size
Data Size
Data Size
Data Size
Data Size
d1
0.003
0.1
Figure Evaluation of the error-damping estimator for the Gaussian toy example. Different evaluation metrics are shown for different dataset sizes and damping factors dj or cosine shifts s. The mini-batch size was set to 10% of the dataset size, and for each step, 10 runs were performed. The median and median absolute deviation are reported, besides for those runs in which none converged.
EXPERIMENTS
To systematically evaluate the proposed methods, we consider three case studies.
Gaussian toy example: An analytically tractable Gaussian model with up to 100,000 synthetic data points, used to assess the accuracy and stability of compositional score estimation.
Hierarchical time series model: A grid of AR(I) processes with global and local parameters, used to evaluate hierarchical estimation against gold-standard MCMC.
Real-world application: Time-resolved Bayesian decay analysis in Fluorescence Lifetime Imaging (FLI), used to demonstrate scalability to nigh-dimensional hicrarchical real data.
For the two synthetic examples, assessed convergence across varying data sizes by recording the number of sampling iterations of the adaptive sampler. In the Gaussian toy example, we can calculate the KL divergence between the compositional and the posteriors, relative mean squared error (RMSE) normalized by the known variance, posterior contraction, and calibration error (Ap- pendix A.4). For the hierarchical models, computed these metries separately at both the global and local levels. Appendix A.5 provides further details about the architecture.
4.1 EXPERIMENT 1: SCALING AND STABILIZING CSM WITH ERROR-DAMPING
Setup and baseline This first experiment serves both as a sanity check and as demonstration of the stabilizing effects of our error-damping estimator, highlighting the accuracy and scalability of compositional score matching in controlled setting. We consider a Gaussian model of dimension D=10 with conditionally independent groups and global latent variable (see Appendix[A.d Since the posterior is analytically tractable, it enables exact measurement of accuracy and convergence. We scale the number of observations up to 100,000 to test the effect of dataset size on the error accumulation of the individual scores. Below, we summarize our results and provide practical recommendations.
Damping factor We find that the optimal damping factor d, depends on the number of composed groups: larger datasets require smaller damping factors for convergence (Figure 2). However, overly small factors can prevent posterior contraction, worsen calibration, and even hinder convergence. With initial factor of 0.1, we can successfully compose 100,000 scores. At this scale, the analytical posterior becomes nearly a point estimate, so even slight deviations in our estimate can significantly increase the Kl--divergence, but the RMSE remains negligible. The damping factor is a tunable hyperparameter, and a value on the order of 1 //n often serves as good starting point.
Mini-batching Our mini-batch estimator alone reduces computational cost per sampling step but does not resolve instability due to егтоr accumulation, which prevents convergence beyond 1000 data points (see Appendix Figure 6). Using smaller batches instead the full dataset lowers both the KL-ddivereence and posterior calibration emror, albeit with slight increase in RMSE. We
7

---

Published as a conference paper at ICLR 2026
A
a
B 0.921
log a
B
ECDF Difference
Stacked ECDFs
1.000
0.913
Estimate
0.1
0.0
-01
2.5 Ground truth
0.0
0.2
0.4
0.6
0.8
1.0
Ground truth
Ground truth
Fractional statistic
Figure 3: Assessing inference for high-resolution grids (128x 128). A Global parameter recovery across 100 datasets, showing the posterior median and median absolute deviation. B Posterior calibration plot for the global parameters using SBC (Säilynoja et al. 2022).
attribute this to a smoothing effect on the accumulated score errors. In practice, we recommend using mini-batches of about 10% of the data to balance accuracy and computational demands.
Noise schedule shifting Adjusting the noise schedule improves stability and mitigates error accumu- lation (Appendix| Figure 6). large shift of s=10 enables scaling to 100,000 groups and improves the KI--divergence, RMSE, and calibration error. As expected, both Kl.-divergence and calibration degrade with larger datasets due to increased error accumulation, but the shifted schedule helps to mitigate this effect. Moreover, a linear schedule appears suboptimal for compositional score matching, failing to converge even on smaller datasets (see Appendix Figure 7
Number of conditions Scaling 100,000 groups also becomes feasible by conditioning the diffusion model on subsets of 100 groups (Appendix Figure 6]. However, increasing the number of conditioning groups does not necessarily lead to better posterior contraction or lower RMSE. Notably, the number of conditions has to be chosen before training, and conditioning simulations, since each training sample incorporates multiple groups. The choice of the number of groups per subset introduces a trade-off between scalability and accuracy. While larger subsets can reduce the variance in the compositional score estimation, they require more expressive networks to compose group-leve information into accurate score estimates. In practice, using few conditions can yield performance gains without incurring major training costs.
Appendix[Fiurre
more groups requires additional
In summary, experiments with the analytically tractable Gaussian toy example demonstrate that the arror-damping mini-batch estimator affords scalable compositional inference for up to 100,000 units of information. While mini-batching alone insufficient to ensure convergence, combining with damping and noise schedule shifting reduces score accumulation errors and computational cost.
4.2 EXPERIMENT 2: SCALING HIERARCHICAL INFERENCE
Setup and baseline Our second experiment evaluates whether our approach can accurately infer the joint posterior for non-trivial hierarchical Bayesian model. We simulate a grid of local AR(1) processes with shared global drift and local variation parameters (see Appendix A.7). We increased the grid size up 128x 128 test the scalability of the method, resulting in up 16,384 local parameter vectors. For this grid of AR(1) processes, a direct comparison to NUTS (as implemented in Stan, Carpenter al. 2017) possible, which is widely regarded as the gold standard for Bayesian inference and provides the most reliable benchmark for evaluating how well our method captures the correct shrinkage in the local parameters. Additionally, we benchmark against direct hierarchical ABI approach (Habermann et al. 22024 Heinnich et al. (002)) using a normalizing flow (ABI-NF) or a modernized version with a diffusion model (ABI-DM), sharing the same settings as ours and the same training budget.
Results Our results support the carlier findings regarding the role of the damping factor: tuning the damping function essential to balance posterior contraction and estimation error (Appendix Figure 8): however, a too large cosine shift might hinder calibration. Moreover, we find that neither the damping factor nor the cosine shift alone is sufficient to ensure convergence on high- grids (e.g., 128x x 128), but their combination stabilizes inference (Figure 3). However, for these large-scale settings, achieving well-calibrated posteriors often comes at the cost of reduced accuracy in parameter recovery. This difficulty arises due to the strong contraction of the global posterior and compounding errors while solving the reverse SDE. As a result, calibration becomes challenging in the high-resolution regime.
8

---

Published as a conference paper at ICLR 2026
Table 2: Benchmarking against NUTS (gold-standard MCMC) and direct ABI methods for the hierarchical AR(I) model. The simulation budget is reported as the equivalent number of full grids. We show the median and median absolute deviation over 100 datasets for global parameters.
Grid size Metric
NUTS
ABI-NE
ABI-DM
ABI-NF-10 ABI-DM-I0
Ours
4x4
Simulation budget
625
625
625
RMSE
0.07 (0.01) 0.01 (0.01)
0.08 (0.01)
0.08 (0.01)
Contraction
0.96 (0.03) 0.87 (0.07)
0.95 (0.03)
0.96 (0.03)
16x16
Simulation budget RMSE
40 0.03 0.19 (0.01) 0.12 (0.01) 0.09 (0.01)
40
400
400
40
0.04 (0.01)
0.08 (0.01)
Contraction
1.00 1.00 (0.00) 0.51 (0.13) 0.71 (0.05) 0.88 (0.04)
0.99 (0.01)
1.00 (0.00)
128
Simulation budget RMSE
<1 ()
0.01 (0.01)
0.09 (0.05)
Contraction
1.0 (0.00)
1.0 (0.01)
In terms of precision, we observe that our method yields results comparable to those of NUTS at both the global and local parameter levels (Table 2). with slightly higher local RMSE (Appendix Table 3). Crucially, our method scales effortlessly to significantly larger grid sizes, such as 128 128. In contrast, NUTS requires approximately hours on a perffrrnreee cluster with 64 CPU cores, whereas our likelihood-fe free approach completes inference within few minutes
single GPU. Moreover, already at a resolution of 32 x2, posterior sampling with NUTS for 100 datasets takes a similar amount of time as training
score-based model and performing amortized inference.
Furthermore, direct hierarchical approaches do not scale to larger grid sizes and already show worse precision 16x 16 grids due to the need for larger training budgets Table 2). Only for a 10-fold increase in the training budget do they achieve similar accuracy. Importantly, our total simulation budget amounts to less than single grid of size 128 128, rendering amortized methods that train on full grids completely infeasible with this low number of training samples.
Inference-time hyperparameter optimization Importantly, because inference with our method is amortized, we can perform grid-based or even Bayesian hyperparameter optimization. We tune the damping factor and noise shift by selecting the best configuration based on the sum of the RMSE and calibration error. To generalize beyond the proposed decay damping, we introduce the following flexible decay function: d(tt) = do + (dl do) - (1 - )²)))),) which adds two hyperparameters (a and (3) that enable smooth interpolation between linear, exponential, and cosine-like behaviors. We perform Bayesian optimization over 3 E [0.3, 2]. d1 E [10-5, 10-1], and dg [10-3, 1]. This increases the runtime on a 32 x 32 grid from 3 to 7 minutes, primarily due to a reduction in early failures during sampling. The best configuration yields di = 0.005, do 0.94, 3.53, a 0.39, and B 1.97, suggesting that the learned schedule strongly favors a sharp, exponential decay (Figure 3).
In summary, our experiment with the hicrarchical AR(I) model revealed that compositional score matching, when combined with damping and noise schedule shifting, enables accurate and scalable inference in hierarchical models with thousands of groups. Even though NUTS is competitive on small grids, its cost and requirement for a tractable likelihood can make impractical for estimating complex models from large datasets, whereas our compositional approach remains viable.
4.3 EXPERIMENT 3: APPLICATION TO FLUORESCENCE LIFETIME IMAGING (FLI)
Practical relevance Our final experiment demonstrates the practical utility of our approach for real-world data, enabling scalable posterior estimation fluorescence lifetime imaging (FL.I, where existing methods struggle with noise and high dimensionality. FLI is an important tool in pre-clinical cancer imaging, particularly for in vivo drug-target analysis (Verma et al., 2025). However, FLI remains challenging as it requires sub-nanosecond acquisition, computationally heavy pixel-wise curve fitting, and must deal with noisy decay profiles from wow-quantum-yinn -ield dyes, leading to high uncertainty (Yuan et al.] [2024] Trinh and Esposito| (021). Bayesian approaches have been explored in prior work (Wang et al.) 2019: Rowley et al. 2016), but, to the best of our knowledge, we present the first application of fully Bayesian hierarchical model to FLI data.

---

Published as a conference paper at ICLR 2026
A
of
[ne]
Af [X]
Servin Freies Cure
Pocvier sintion Pusteior vits
Monn Wiin
wwwvioe Pretee Cutt
D
3372
Figure 4: Inference for fluorescence lifetime imaging. A Mean intensity across time for each pixel, representing the fluorescence data. B Time series data and fitted posterior median for representative pixels. C Spatial map of the fitted local posteriors (medians) per pixel. D Spatial map of R for each pixel, comparing our results with flat Bayesian model and popular baseline (MLE).
Setup and baseline We analyze time-resolved fluorescence decay data (Figure 4A-B), where each pixel in measured series of 512x512 images is modeled using bi-exponential decay with local decay rates o and Th and mixture weights A Each local parameter has a global mean and standard deviation, resulting in hierarchical inference problem with over 250,000 groups (see Appendix A. Unlike amortized methods, which train on full-image simulations to generalize across spatial structures, our approach trains single pixels, requiring only the equivalent 350 full images for training. We compare our approach with the field's gold standard method based on maximum likelihood estimation (MLE; assuming Gaussian noise make MLE feasible) and non-hicrane non-bierarchical diffusion model trained on single-pixel time series.
Results To assess the performance of the baseline non-hierarchical approach and our proposed method, first consider 100 held-out synthetic images. We find that per-pixel MLE fails to recover the ground truth due to photon-limited noise. In contrast, our hierarchical approach accurately captures both global and local structures (Appendix Figure 9 [Op. Nevertheless, estimating global variances remains challenging under very high noise conditions. Finally, we apply our method to real FLI data (Appendix A. Using the trained score-based hierarchical model, we fit over 750, 000 local parameters efficiently (Figure 4C).
Qualitatively, the inferred mean lifetime closely matches the standard MLE fit (Appendix IT). Our approach achieves excellent image-wide fits, with a mean R -0.96 (s.d., 0.017) for posterior predictive medians, versus 0.871 (s.d., 0. 10) for MLEE and 0.921 (s.d., 0.12) for the pixel-wise approach (Figure 4)), illustrated in Figure 4B. Across pixels, the mean posterior predictive p-value 0.20 (s.d., 0.337), indicating slight underdispersion; masking the final third of the decay tail increases the mean valuue 0.40 (s.d., 0.38), confirming that our model captures the core dynamics.
CONCLUSION
Hierarchical Bayesian models (HBMs) of utmost importance in statistics, but their estimation remains challenging. Here, we demonstrated that compositional score matching (CSM) provides scalable and flexible framework for estimating large HBMs. Moreover, we introduced an error- damping mini-batch estimator that resolves the inherent instability of CSM and scales up to hundreds of thousands of data points. As a notable limitation, we observed that posterior calibration becomes difficult under extreme contraction, leaving room for further improvement. Future work could further explore temporal aggregation (Gloeckler al. 2025), systematically test the trade-off introduced by different damping schedules, refine mini-batch selection using informativeness criteria (Peng et al. 2019) and generalize our experiments to more than two levels.
ACKNOWLEDGMENNS
This work was supported by the German Federal Ministry of Education and Research (BMBF) RMMUNE/031L0293C the European Union via the ERC grant INTEGRATE, grant agreement number 101 126146, and under Germany's 's Excellence Strategy by the Deutsche Forschungsgemein- schaft (DFG, German Research Foundation) (EXC -447-39065813 EXC 2151- 155--39073004)
10

---

Published as a conference paper at ICLR 2026
the University of Bonn via the Schlegel Professorship of J.H, and the National Science Foundation under Grant No. 2448380. J.A. thanks the Global Math Exchange Program of the Hausdorff Center for Mathematics for their financial support. We also thank Niels Bracher for his helpful insights diffusion models. Views and opinions expressed are however those the authors only and do not necessarily reflect those of the funding agencies.
AUTHOR CREDIT
J.A.: Conceptualization, Methodology, Software, Formal analysis, Validation, Visualization, Funding acquisition, Writing original draft, Writing review & editing. V.P.: Data curation, Software, Writing - review & editing. C.S.: Investigation, Data curation, Writing review & editing. M.B.: Supervision, Investigation, Data curation, Writing review & editing. X.I. Supervision, Project administration, Resources, Writing review & editing. J.H.: Supervision, Project administration, Funding acquisition, Writing review & editing. S.T.R.: Conceptualization, Methodology, Software, Supervision, Resources, Funding acquisition, Writing - original draft, Writing review & editing.
REFERENCES
T. Akiba, S. Sano, T. Yanase, T. Ohta, and M. Koyama. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2019.
Arruda, Schälte, C. Peiter, O. Teplytska, U. Jaehde, and Hasenauer. An amortized approach to non-linear mixed-effects modeling based on neural posterior estimation. In International Conference on Machine Learning, pages 1865-1901. PMLR, 2024.
M. Betancourt and M. Girolami. Hamiltonian monte carlo for hierarchical models. Current trends in Bayesian methodology with applications, 79(30):2-4, 2015.
D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112((18):859-877) 877, Apr. 2017. ISSN 0162-1459, doi: 10. 00.10000016214592211...1 1287777
Boelts, M. Deistler, M. Gloeckler, Á. Tejero-Cantero, J.-M. Lueckmann, G. Moss, P. Steinbach, T. Moreau, F. Muratore, J. Linhart, et al. shi reloaded: toolkit for simulation-bast inference workflows. Journal of Open Source Software, 10(108):7754, 2025.
P.-C. Bürkner, M. Scholz, and S. T. Radev. Some models are useful, but how do know which ones? towards a unified Bayesian model taxonomy. Statstic Surveys, 17:216-310, 2023.
B. Carpenter, A. Gelman, M. D. Hoffman, D. Lee, B. Goodrich, M. Betancourt, M. Brubaker, J. Guo, P. Li, and A. Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76:1-32, 2017.
K. Cranmer, J. Brehmer, and G. Louppe. The frontier of simulation saaseattoo--bseddd inference. Proceedings of the National Academy Sciences, 117(48)330555-3066 2020.
M. Dax, S. R. Green, J. Gair, N. Gupte, M. Pürrer, V. Raymond, J. Wildberger, J. H. Macke, A. Buonanno, and B. Schölkopf. Real-time inference for binary neutron star mergers using machine learning. Nature, 639(8053):49 2025.
F. Draxler, S. Wahl, C. Schnörr, and U. Köthe. On the universality of volume-preservee and coupling- based normalizing flows. In Proceedings of the 41st International Conference on Machine Learning, pages 1 11613-11641, 2024.
C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios. Neural spline flows. Advances in neural information processing systems, 32, 2019.
T. Geffner, G. Papamakarios, and A. Mnih. Compositional score modeling for simulation based inference. In International Conference on Machine Learning, pages 1 11098-1111 1116. PMLR, 2023.

---

Published as a conference paper at ICLR 2026
A. Gelman, J. B. Carlin, H. S. Stern, D. B. Dunson, A. Vehtari, and D. B. Rubin. Bayesian Data Analysis (3rd Edition). Chapman and Hall/CRC, 2013.
A. Gelman, A. Vehtari, D. Simpson, C. C. Margossian, B. Carpenter, Y. Yao, L. Kennedy, J. Gabry, P.-C. Bürkner, and M. Modrák. Bayesian workflow. arXiv preprint urXiv:2011.0133 2020.
M. Gloeckler, M. Deistler, C. Weilbach, F. Wood, and J. H. Macke. All-in-one simulation-bail based inference. In Proceedings of the 41st International Conference on Machine Learning, pages 15735-15766, 2024.
M. Gloeckler, S. Toyota, K. Fukumizu, and H. Macke. Compositional simulation-based inference for time series. In The Thirleenth International Conference on Learning Representations, 2025.
T. Gneiting, F. Balabdaoui, and A. E. Raftery. Probabilistic Forecasts, Calibration and Sharpness. Journal of the Royal Statistical Society Series B: Statistical Methodology, 69(2):243 2007. doi: 10.1 1.1111/j... 1467-9868.788
D. Habermann, M. Schmitt, L. Kühmichel, A. Bulling, S. Radev, and P.-C. Bürkner. Amortized bayesian multilevel models. CoRR, abs/2408. 2024.
L. Heinrich, Aishra-Sharma, C. Pollard, and Windischhofer. Hierarchical neural simulation- based inference over event ensembles. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 68500,6881,111 2020.
M. D. Hoffman, A. Gelman, et al. The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1593-162,, 2014.
A. DolicocurMaartnnaal K. Li, R. Piché- Taillefer, T. Kachman, and I. Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint a0S.v:2 14080, 2021.
Karras, M. Aittala, T. Aila, and Laine. Elucidating the design space of diffusion-baase generative models. Advances in neural information processing systems, 35:26565-26577. 2022.
D. Kingma and R. Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36:65484-666 65516, 2023.
Z. Li, H. Yuan, K. Huang, C. Ni, Y. Ye, M. Chen, and M. Wang. Diffusion model for data-driven black-box optmizzttion. arXiv preprint arXiv:2403.13211 2024.
J. Linhart, G. Cardoso, A. Gramfort, S. L. Corff. and P. L. C. Rodrigues. Diffusion posterior sampling for simulation based inference in tall data settings. Transactions on Machine Learning Research, 2026. ISSN 2835-8856.
J.-M. Lueckmann, J. Boelts, Greenberg, P. Goncalves, and Macke. Benchmarking simulation- based inference. In International conference artificial intelligence and statistics, pages 343-351. PMLR, 2021.
C. C. Margossian and L. K. Saul. The shrinkage-delinkage trade-off: An analysis of factorized gaussian approximations for variational inference. Uncertainty in Artificial Intelligence, pages 1358-1367. PMLR, 2023.
R. McElreath. Statistical rethinking: A Bayesian course with examples in R and Stan. Chapman and Hall/CRC, 2018.
A. Q Nichol and Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pages 8162-8171. PMLR, 2021.
V. Pandey, I. Erbas, X. Michalet, A. Ulku, C. Bruschini, E. Charbon, M. Barroso, and X. Intes. Deep learning- based basedd temporal deconvolution for photon üime-of-fligh distribution retrieval. Optics letters, 49422).4577-4466 2024.
12

---

Published as a conference paper at ICLR 2026
V. Pandey, E. Millar, I. Erbas, L. Chavez, J. Radford, Crosbourne, M. Madhusudan, G. G. Taylor, N. Yuan, C. Bruschini, et al. Real-time wide-field fluorescence lifetime imaging via acquisition for biomedical upplications. PhotoniX, 6(1):58, 2025.
X. Peng, L. Li, and F.-Y. Wang. Accelerating minibatch stochastic gradient descent using typicality sampling. IEEE transactions neural networks and learning systems, 31(11):4649 4659, 2019.
S. Т. Radev, U. K. Mertens, A. Voss, L. Ardizzone, and U. Kothe. BayesFlow: Learning complex stochastic models with invertible neural networks. IEEE Transactions on Neural Networks and Learning Systems, 33(4): 4522-1466, Apr. 2020.
S. T. Radev, M. Schmitt, L. Schumacher, L.. Elsemüller, V. Pratz, Y. Schälte, U. Köthe, and P.-C. Bürkner. BayesFlow: Amortized Bayesian workflows with neural networks. Journal of Open Source Software, 8(89):5702, 2023.
Rodrigues, T. Moreau, G. Louppe, and A. Gramfort. Hnpe: Leveraging global parameters for neural posterior estimation. In M. Ranzato, A. BBeyglzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 13432-13443. Curran Associates, Inc., 2021.
M. I. Rowley, A. C. Coolen, B. Vojnovic, and P. R. Barber. Robust Bayesian fluorescence lifetime estimation, decay model selection and instrument response determination for low-intensity flim imaging. PLoS one, (1(6))0158404 2016.
Säilynoja, P.-C. Bürkner, and A. Vehtari. Graphical test for discrete uniformity and its applications in goodness-of-fi fit evaluation and multiple sample comparison. Statistics and Computing. 32(2):32, 2022.
T. Salimans and J. Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022.
L. Sharrock, J. Simons, S. Liu, and M. Beaumont. Sequential neural score estimation: Likelihood free inference with conditional score based diffusion models. In International Conference on Machine Learning, pages 44565-44603 PMLR, 2024.
Simons, Sharrock, S. Liu, and M. Beaumont. Neural score estimation: Likelihood-free inference with conditional score based diffusion models. In Fifth Symposium on Advances in Approximate Bayesian Inference, 2023.
S. A. Sisson and Y. Fan. Likelihood-free meme. Handbook of Markov Chain Monte Carlo, pages 313-335, 2011.
J. T. Smith, R. Yao, N. Sinsuebphon, A. Rudkouskaya, N. Un, J. Mazurkiewicz, M. Barroso, P. Yan, and Intes. Fast fit-free analysis of fluorescence lifetime imaging via deep learning. Proceedings of the national academy of sciences, 116(48):2409-2444
2019.
Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Y. Song, C. Durkan, Murray, and S. Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34: 1111441114288 2021a.
Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021b.
N. Tolley, P. L. Rodrigues, A. Gramfort, and S. R. Jones. Methods and considerations for estimating parameters biophysically detailed neural models with simulation based inference. PLOS Computational Biology, 20(2);e101 108, 2024.
B. Torrado, B. Pannunzio, L. Malacrida, and M. A. Digman. Fluorescence lifetime imaging mi- croscopy. Nature Reviews Methods Primers, 40 4(1):80, 2024.
13

---

Published as a conference paper at ICLR 2026
A. L. Trinh and A. Esposito. Biochemical resolving power of fluorescence lifetime imaging: tangling the roles of the instrument response function and photon-statistics. Biomedical optics express, 12(7):3775-377 377888 2021.
A. Verma, V. Pandey, C. Sherry, T. Humphrey, C. James, K. Matteson, J. T. Smith, A. Rudkouskaya, X. Intes, and M. Barroso. Fluorescence lifetime imaging for quantification of targeted drug delivery in varying tumor microenvironments. Advanced Science, 12(3):2403533 2025.
Vincent. connection between score matching and denoising autoencoders. Neural computation, 23(7):1611-177, 2011.
A. Vuong, M. T. McCann, J. E. Santos, and Y. T. Lin. Are really learning the score function? reinterpreting diffusion models through wasserstein gradient flow matching. Transactions on Machine Learning Research, 2025. ISSN 2835-8856.
Wang, V. Chacko, A. K. Sagar, K. W. Eliceiri, and M. Yuan. Nonparametric empirical Bayesian framework for snorescence-lifeting imaging microscopy. Biomedical optics express, 10(11): 5497-5517, 2019.
N. Yuan, Pandey, A. Verma, T. Humphrey, M. Barroso, X. Intes, and Michalet. Quantifying lifetime uncertainty in gated-iced fluorescence lifetime imaging. Optica Open. Preprint, 2024. doi: 10. 11..36//oppiccope2228894944434
M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Y. Zhang and Mikelsons. Solving stochastic inverse problems with stochastic BayesFlow. In 2023 IEEE/ASME International Conference on Advanced Intelligent Mechatronics (AlM), pages 966-972 2023.
14

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1
STOCHASTIC DIFFERENTIAL EQUATION FORMULATION OF THE DIFFUSION PROCESS
The forward diffuusion process for t [0, 1] can be specified as stochastic differential equation Song et al. (
d0.- f(0,, + g(t) dW.-
For known variance- process, the drift and diffusion coefficients are given by
((0,t) - हे log(1 )) 0, g(1)²
4 log(1 + +
dt
with of = sigmoid(A:) and of sigmoid - X₂) as discussed Kingma and Gao 2023). Time can be reversed via the reverse-time SDE
d0: = [ t) - g(t) Vo, p: Y) dt + g() dW
which enables posterior sampling using state-of-the-art SDE solvers. The corresponding probability ODE is
d0. = f(Or.)) - 1
20()) x) dt.
A.2 FACTORIZATION OF THE GLOBAL POSTERIOR IN HIERARCHICAL MODELS
Proposition A.1. Consider two-level hierarchical model with global parameter n, local parameters
for = 1,....JJ J, and ii.d. observations Y, for each group. The marginal posterior of the global parameter satisfies
p(n Yı:) p(n)
II p(n Y,).
Proof. Using Bayes' rule and the conditional independence we get:
p(n Yı:)) x p(n))(((ee) p(Y |n) pp(n) IY, (n).
The posterior based on group alone is p(n Y,) x ppp()))) In)
using again Bayes' rule. Therefore,
J
II 1=1
r(n (Y) x p(n)" II p(Y, m).
5=1
Combining results gives
p(n Yı:)) x p(n) II M(Y, p()) p(n)' I n) p(n)--J II p(n Y,).
A.3 MINI-BATCH ESTIMATOR IS UNBIASED Proposition A.2. The mini-batch estimator
J M
(nh Yo) = =
+
Σ sp Y, Ac)
M j=1
with M samples, where each sample Y; sampled uniformly from the set (Y,,... }) is an unbiased estimator of the full compositional score.
15

---

Published as a conference paper at ICLR 2026
Proof. By linearity of expectation, we have
E
Spp Te Xe)
Σ Ey, (1 Yx, [e)]
M
M
==1
Since each y is sampled uniformly from {YI..... YJ).
1
Ey, 1a (N.. Y, Xc)
Σ
Su (Te Ae)
J
j=1
J
M s
E
M Σ
S (7) Yj Ax) = > S y, Ae).
M
M
j=1
j=1
Adding the constant prior term (1 J) V) llg ((n:) yields the full compositional score. Hence, the estimator is unbiased.
A.4 EVALUATION METRICS
All experiments were repeated times, and the median and median absolute deviation from the following standard metrics reported:
Root mean squared error (RMSE) RMSE measures the deviation between posterior samples and the ground-truth parameters. Given posterior samples ( (local or global) for parameters, in the dataset i, and true parameters Ozj. the RMSE is defined as:
S Σ S
RMSE,
01g
aggregated over datasets via median and over the parameters j via the mean. We normalize RMSE by dividing the empirical range of the ground-truth parameters.
Calibration error Calibration Error measures how well the empirical coverage of posterior credible intervals matches their nominal level. For level E [0.005, 0.995), we compute the a-credible interval for each parameter and check whether the ground-truth value falls within it. Let I denote the indicator that the true value lies within the interval:
N
CalibrationError, = median
a
2=1
where aggregation is across grid of a values. We calculated the mean calibration enror over the parameters j. This metric is sensitive to both over- and under-confidence in the posteriors.
Posterior contraction We define posterior contraction the relative reduction in variance from prior to posterior:
Var posterior (,))
Contraction,
Var (0,)
where values are clipped to [0, 1). This reflects how much uncertainty has been reduced due to conditioning on the data, with values near I indicating strong learning.
KL Divergence (Gaussian Case) In the Gaussian toy example, where the true posterior is analyti- cally tractable and Gaussian, we compute the KL divergence between the empirical posterior (()) (estimated from samples) and the true Gaussian posterior p()):
Σ Σg
TH (H, m.p) p,p) "e," 2 11))
(l) 2 log
where Mg E, are the empirical mean and covariance of posterior samples, and pep Ep are the parameters of the analytical posterior.
16

---

Published as a conference paper at ICLR 2026
Posterior predictive p-value The posterior predictive p-value evaluates how well the observed data are covered by the posterior predictive distribution. In a well-specified model, these p-values are approximately uniform on [0, 1]; thus, their expectations should be approximately 0.5. For S posterior samples, let
f+) median [v1 (v) p(Y 10))--1),
(0),
(c))(6) yt
f((1))
S
For each posterior draw 0!*), define the discrepancy
T
(yt f())
(y,
=
Var,(0)
and then posterior predictive p-value is
Prrc 1/2 E S
((yrtep(())
D(yobs, 0))
RMSE, calibration error, posterior contraction, and empirical CDFs plots are computed using the diagnostics provided in the BayesFlow toolbox (Radev et al. 2023
A.5
SCORE MODEL ARCHITECTURES & TRAINING
MLPs:
Fully connected networks with hidden layers and 256 units per layer, using Mish activations.
Residual local conditioning: Local networks receive a projection of the global latent variables and learn a residual update. Otherwise, global and local networks are simple MLPs.
permutation-invariinn aggregation: To handle multiple condition sets or observations per group, shallow permutation-invariint encoder architecture based on the Deep Set framework Zaheer et al. (2017):
An encoder MLP with 4 layers of 128 hidden units and ReLU activations,
Mean pooling
the set dimension ensure permutation invariance,
decoder MLP with hidden layers (each of size 128) and ReLU activations, projecting to the final output dimension.
Time series summary network: For structured input data such time series (as in the FLI application), we use a hybrid convolutional-recurree architecture. The model begins with stack of ID convolutional layers followed by skipping recurrent path, implemented in (Zhang and
Mikelsons 2023):
A standard recurrent path (bidirectional GRU with hidden size 256).
A skip-convolution path, which downsamples the sequence via strided convolution and feeds the result into a parallel recurrent layer,
Final representations from both paths are concatenated to produce summary embedding, which are then projected by linear layer to fixed summary dimension of size 18.
We parameterize score models to predict the more stable velocity v, = a,e o,0,. and then transform the output to noise é, as it has been shown that this parameterization more stable for all whereas noise-prediction becomes harder for l. close to O where the signal increases and noise decreases Salimans and Ho (2022). Furthermore, we condition the score network on the signal-to- noise ratio (SNR), normalized the interval [-1, 1] similar to the preconditioning introduced in (Karras et al. 2022). The data and parameters were always standardized, and the prior scores were adjusted accordingly by multiplying them by the standard deviation of the parameters.
Noise schedules We employed the following schedules:
Cosine schedule by Nichol and Dhariwal (2021) (with s=0 during training)
((t) -2 log(tan(<1122) 2s,
17

---

Published as a conference paper at ICLR 2026
and Linear schedule by Ho et al. (2020)
X(t) = - log(e - 1).
All our noise schedules are truncated such that the log signal-tor senal-to-noise ratio is A E [-11,, 15] to avoid instabilities in sampling, as detailed in (Kingma and Gao 2023).
As the weighting function for the loss, we employed the likelihood weighting wa - g(13/0o2 proposed by [Song et al. (2021a) for the linear and cosine schedules.
Training We trained all models using AdamW with cosine annealing learning rate schedule. The initial learning rate was set to 5x 5x10-3. The models were trained for 1000 epochs. For the Gaussian toy example, in each epoch, we generated 10,000 new training samples on the fly, as simulations are inexpensive. For the AR(1) experiment, we used fixed budget of 10,000 simulations, and for the FLI application, we used 30,000 samples per epoch, as we found that more training data was needed. For reference, training single score estimator for the FLI task completes in 7.6 hours on a single GPU, while for the AR(1) model takes 0.83 hours.
All models were trained on high-performance computing cluster using an AMD EPYC "Milan" CPU (2.00 GHz), 100 GB DDR4 3200 MHz RAM, and an NVIDIA A40 GPU with 48 GB of memory. Each experiment required 1-2 days for all repeated runs high-performance computing infrastructure with up to 50 parallel jobs.
Sampling For our experiments, we used the adaptive second-order sampler with maximal 10,000 network evaluations and the default settings proposed by [olicoeur cally, we set the absolute error tolerance to Cabs 0.002576 and the relative tolerance Gre = 0.1. To solve the probability ODE, we used an Euler scheme. For annealed Langevin dynamics, we followed the setup from Geffner et al. (2023), using Langevin steps per iteration, a maximum of 2000 iterations, and step size factor of 0.1. For GAUSS, we used the implementation provided by the sbi toolbox (Boelts et al. 2025). ) with the same diffusion model and training settings as those in own implementation.
et al. (2021) Specifi-
To determinc the optimal damping factor dj and shift s for a certain task, we ran Bayesian optimization with untuna (Akiba et al., 2019), using the sum of the average RMSE and expected calibration error as an optimization criterion. We used search grids & E [0. 4] and dj E x 10-5, 0..1]. We chose this simple criterion because the hierarchical structure and shrinkage effects in our experiments encourage unimodal behavior by borrowing strength across observations. More expressive criteria can be used in cases where the posteriors exhibit multiple modes. We also considered do < 1 and found that this can sometimes improve RMSE and calibration. If nothing else is stated, we use mini-batch size of 10 in our experiments.
Code The software code and data for the experiments are available in public GitHub repository: ht https: http:::///ttub.... wtt.ua...coomaaaa com/ wwwpsss.wwwwgg//eeeeeeeeee...eee
A.6 EXPERIMENT 1: GAUSSIAN TOY MODEL
The Gaussian toy model defined as follows:
Y, N(n o²I)
with = 0.1 and n E R00 We observe {Y, 1 with varying J and compute the posterior
p(n (Y,)_-)) Given a normal prior for 17. m N/ o²I), the posterior also Gaussian, and we can calculate it analytically:
pnn (Y, (/)) exp (-1 Tn mJ)).
where μ. 3 E/= Y, and Σラ IL Here, did not employ summary network.
18

---

Published as a conference paper at ICLR 2026
10
Vat Shipe
Number of groups
Number of Steps
0.006
10a
100
Stip s1e
0.001
302
201
0.000
10f
10r
102
102
0
3000
2000
3203
9000
Number of subsets groups
500p
(a) Number of sampling steps.
Adaptive step size for 100 groups.
Figure 5: Assessing the adaptive sampling scheme for compositional inference in the toy model. (a) Increasing numbers sampling steps are needed for increasing number of subsets groups. (b) The adaptive step size is adaptively increased towards the end of the sampling (low noise region).
Number of Steps
Kl Divergence
20*
Calibration tion Error
Posterior
Contraction 0.5
0.4
2.0"
RMSE
102
102
101
0.0
a.
0.0
Data Size
Data Size
Data Size
Data Size
Data Size
100
1000
10000
100000
(a) Varying mini-hatch sizes.
Number of Stops
10"
TASSSSEES
Divergence
200
1.0
Posterror
Contraction
Callbr Calibration ebion Error
RUISE
0.5
20a
0.2
10
0.0
10a
10a
10a
Data Size
Data Size
Data
Size
Data Size
Data Size
10
20
50
100
(b) Varying number of subsets of groups during training (score model trained with DeepSet second summary network).
Figure 6: Evaluation of the error-damping estimator for the toy model. Different evaluation metrics are shown for different mini-batch sizes or varying numbers of subsets of groups. For each experiment, 10 runs were performed. The median and median absolute deviation are reported, besides for those runs where none converged.
Number of Steps
200
10a
201
RMSE
1.0
0.5
Posterior
Contraction
0.5
1.0
Callbration Error
0.4
0.2
1.7
109
107
10a
10a
10
10a
10a
107
10
10a
10'
10a
20a
10'
10'
10a
Data
Data 100000
Data Size
100
1000
10000
Figure 7: Evaluation of the linear noise schedules for the toy model. For each experiment, 10 runs were performed. The median and median absolute deviation is reported, besides for those runs where converged.
19

---

Published as a conference paper at ICLR 2026
A.7 EXPERIMENT 2: HIERARCHICAL AR(1) MODEL
Our hyper-priors are defined as follows:
a (0, 1), 8 N~ 1), loga~ NO0 1).
The local parameters are different for each grid point:
0, ~ ~N(,
0, = 2 sigmoid 4 -0,) 1.
In each grid point j, we have time series of T observations,
3.0 0.1.). Y ~ N(a + 02 1.-11...11). ( 1=1...
On the local level, we perform inference on and afterward transform ô to 0 as NUTS (as imple- mented in Stan Carpenter et al. 2017) performs better on non-centered parameterizations (Betancouri and Girolami 2013). We used parallel chains, each generating 1,000 samples with default settings in Stan. Here, we do not employ summary network.
For the direct hierarchical ABI methods (Heinrich et al. 2024 Habermann al. 2024). we employ
ABI-NF: Normalizing flow with 2 coupling layers using spline transformation (Drkkan al. 2019), trained for 100 epochs,
ABI-DM: Diffusion model of the same size as ours with velocity prediction, cosine schedule, trained on 1000 epochs and Maryuama sampling with 300 steps.
We generated 10,000 pairs of global and local parameters train with compositional approach. This means, for 4 x 4 grids, we use the equivalent of 625 full hierarchical datasets, for 16 x 16 grids, we use 40 full hierarchical datasets, and for 128 128 grids, we used not even full hierarchical dataset.
Table 3: Benchmarking against NUTS (gold-standard MCMC) and direct neural methods for the hierarchical AR(I) model. We show median and median absolute deviation over datasets for local parameters.
Grid size
Method
RMSE
Contraction
4x4
NUTS
0.08 (0.01)
0.99 (0.01)
Ours
0.13 (0.01)
0.97 (0.03)
16x16
NUTS
0.08 (0.01)
0.99 (0.01)
Ours
0.10 (0.02)
0.98 (0.02)
128x128
NUTS
0.01 (0.01)
1.0 (0.00)
Ours
0.09 (0.05)
1.0 (0.01)
20

---

Published as a conference paper at ICLR 2026
Number Stops
10
Stapp
AMSE Global
1.0
Callbrrtion Error
1.0
1.0
Glebal
0.4
10?
0.5
0.5
ROSSE
Contractionn
0.5
1.0
10a
10a
10a
10
10
16>
10c
10a
30
16
10:
10
Size
0.003
0.1
d:
Varying damping factors d1.
Number of Stops
1.0
Callbration Erro
Contraction Global
1.0
Contraction Local
(MSRR Seops
RMSE Global
RMSE Local
1.0
Global
0.5
100
101
20
Data Size
Data Size
Data Size
Data
Size
Data Size 16384
Data Size
10
16
180
-
256
4096
(b) Varying mini-batch sizes. Convergence is achieved only for the smallest dataset.
Number of steens
10a
YMSX grip
RMSE Global
1.0
Calibration Emrrr
Contraction Global
1.0
Contraction Local
1.0
20º
0.5
Global
0.2
RMSE
102
13a
10
19a
10
10a
10a
10
10'
101
Data
Data
0
5= 10
(c) Varying cosine shifts s.
Number of Stepe
10f.
AMSE Global
Call bratton ErTor
Contraction Global
Contraction Local
1.0
0.4
1.00
0.5
Globall
0.5
RMSS
e.
0.5
10P
10a
1.a
10
12
12a
10a
10:
10a
10a
10
10
30'
Size
Data
Size
Data Size
Data
(d) Varying number of subsets groups during training (score trained with a Deep Sct),
Figure 8: Evaluation of the error-damping estimator for the hierarchical AR(I) model. For each experiment, 10 runs were performed. The median and median absolute deviation is reported, besides for those runs where models were trained on a single group.
converged. A mini-batch size of 10% of the data was employed, and score
21

---

Published as a conference paper at ICLR 2026
A.8 EXPERIMENT 3: FLUORESCENCE LIFETIME IMAGING (FLI) MODEL
Model The observed time-resolved fluorescence signal at each pixel is modeled using a bi- exponential function, following the work of Pandey al. (2024) and Smith al. (2019). This approach captures the fluorescence decay dynamics of individual fluorophores, accounting for both the fast and slow decay components associated with different molecular states. By fitting decay model, we can extract information about the characteristic lifetimes of the fluorophores, which is essential for studying molecular interactions and dynamics. The time-dependent fluorescence signal is given
s(t) = I. A (-)==I(Aae'll e + (1--4) ( -/-1// * (RF + (t),
(13)
where of and TI are the fluorescence lifetimes and 4 is mixture parameter. Here, / [0, 1024 denotes the pixel intensity for 10-bit images, (R)(() is the instrument response function, and (7()) represents additive noise. The symbol denotes convolution. For each simulation, we independently sampled time series from the recorded IRF and system generated noise (which makes the likelihood intractable). The maximal photon count in each time series was then normalized to The real data also normalized to 1 on pixel-wise level.
Instrument response function (IRF) The emitted signals are recorded using multiple instruments (detectors, electronics, etc.) which have a characteristic response E() to an instantaneous signal 8(1) (e.g., single photon). The recorded signals from the T Tpperiodic emitted signal can be written as convolution of periodic do, and non non- oon-perioiic E(t):
(l)
F(()) * 51,
E(() (FO, Fo,T)
(14)
(E(t) TO.T) Fo.T
-IRF * Fo,T-
Equation 14 introduces the T-periodic instrument response function IRFor-- The IRF can be measured using the excitation signal from diffused white paper. The FLI experimental details in microscopy, mesoscopy, and macroscopy can be found in Pandey et al. (2025).
The traditional methods of fitting these types of models reviewed in Torrado et al. 2024
Priors The prior distributions were designed with domain knowledge:
1, (l.og(0.2), 0.7²),
1 ~ (((-10011),
mean (log(())) 0.5,)
AT G ((
mean NN.44, 12),
G ~ (((1,0.55)
Local parameters are then sampled from the corresponding global means and standard deviations, as follows:
H (T (),,
Arj ~ N'AT mean ()))))) of ~ (a (a ?))
The local parameters can then be converted to linear scale for simulation:
न exp(log T)),
T
exp(log AT), A
I exp(-a)
This ensures that 72 > both the global and local levels, and that the mixture fulfills A [0, 1]. Additionally, we can compute the average lifetime Timam = Ar + (1- (1-A)
Here, employed time-series summary network. For comparison, model the same size as ours on the flat model using the same prior and simulation budget, but only targeting the local per pixel parameters without conditioning on global parameters.
also trained a diffusion
Data AU565 (HER2+ human breast carcinoma) cells, incubated for 24h with 20 pg/mL. TZM-Alexa Fluor 700 (Donor, D) and 40 pg/mL TZM-Alexa Fluor 750 (Acceptor, A), were imaged using Förster resonance energy transfer (FRET) microscopy quantify trastuzumab (TZM) binding. AU565 cells exhibit relative low level of HER2 heterodimerization that correlate with reduced TZM uptake and sensitivity, which is also influenced by culture conditions (2D vs. 3D). FLI-FRET analysis allows for the quantification of these deneriettion-deeenddent variations in live cells by assessing the proximity of the donor and ncceptor-labeled T7M.
22

---

Published as a conference paper at ICLR 2026
Hi 1p
(a) Recovery of global parameters with hierarchical score based approach (medians and median absolute deviation of the posterior samples).
(b) Global posterior calibration assessed with simulation saadio based calibration diagnostics.
Figure 9: Assessing inference of global parameters for the FLI model. Synthetic data 32x32 grid was generated.
4
0,908
(- 0.967
<- 0.973
3.2
Estinate
64-
2.5
05
2.0
84
02
375
2.6
335
230
235
235 adc Groand mith
258
0.4
0.5
0.5
2.2
Ground Irihh
Gronnd Irath
(a) Recovery of transformed local parameters for one 32x32 grid with hierarchical based approach (medians and median absolute deviation the posterior samples). Deviations from ground truth can be due to the expected shrinkage of the local posteriors.
34
Ground truth
Ground bruhh
Ground bruth
(h) Recovery of transformed local parameters with MLLE.
Figure 10: Assessing inference of local parameters for the FLI model. Synthetic data on a x2x grid was generated. We compared our hierarchical approach against the standard non- con-hicrarchical pixel-wise MLE.
23

---

Published as a conference paper at ICLR 2026
Posterior Median pmean
Estimates gmeen
2.00
1.75
1.50
0.901
1.25
1.25
1.00
11.00
0.75
0.75
0.50
0.25
0.00
0.00
MIE Extimates
(a) Hicrarchical Tmean estimates.
(b) MLE Timean estimates.
(c) Correlation of Timen estimates.
1.0
Coverage: 0.97
Coverage e: 0.99
Coverage: 0.95
Coverage: 0.92
Photon Count
0.8-
Normalized
0.4
0.2
1.0
Coverage: 0.95
Coverage: d: 0.98
Coverage: 1.00
Coverage: 0.99
Photon Count
Normalized
0.6-
0.4
0.2
0.0
100
100
2.00
100
200
100
Time [s]
Time
Time
Time [s]
Data
Posterlor Median
Posterior 95% CI
(d) Local simulations from the hierarchical posterior.
Hierarchical Bayeslan
Density
2
Posterior Predictive value
(e) Quality of posterior prediction measured Bayesian p-value.
Figure 11: Assessing inference of local parameters for the FLI model on real data. We compared our hierarchical approach with the standard an-hierarchical pixel-wise MLE. Owing the low photon count, the average lifetime mean is the most reliable quantity for this non-hierarchical method. Furthermore, we show additional random simulations from hierarchical posterior (median and 95% confidence region out of 100 simulations) and quantitative evaluation of the posterior predictive quality.
24

---

Published as a conference paper at ICLR 2026
C
8
Prior Posterlor
1007
201
100/7
tog
logoud
10
Figure 12: Global posteriors for the real FLI data.
25


Reviews:
[{"review_id": "VEDV90l6mJ", "reviewer": "Reviewer_MQk2", "summary": "The paper introduces an approach to hierarchical modeling for ABI that builds on compositional score matching. Existing approaches struggle to scale with the data requirements for complex hierarchical settings, making it difficult to capture the full, amortized posterior without hitting computational bounds. The proposed method incorporates a CSM reformulation that makes smaller step sizes feasible (needed for larger numbers of groups) and leverages inverse factorization to model different hierarchical levels with separate score estimators. The approach is evaluated on three tasks that include complex hierarchical relationships and high data/parameter counts.", "questions": "- Why the lack of baseline comparisons to the many existing SBI/ABI methods? Can it be reasonably assumed that the proposed method performs on par with CSM (or other competitive baselines) before hitting typical scaling boundaries, or does one sacrifice performance on smaller problems in exchange for a more scalable approach?\n- In Figures 6 and 7 provided in the Appendix, what explains the fairly consistent upward trend in KL divergence and calibration error?\n- For the hierarchical AR(1) model (experiment 2), is there deeper justification behind why NUTS consistently outperforms the proposed method, the large runtime disparity between the two methods notwithstanding?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is well-organized, includes a comprehensive literature review, and provides high quality figures/diagrams.\n- The paper's stated contributions are clear and address a difficult problem in the SBI/ABI space. Scalability to large parameter spaces is a critical difficulty even for state of the art methods, and common workarounds often require heavy summarization or tradeoffs in generality (e.g., using non-amortized methods). Formulations that help stabilize score matching or flatten hierarchical modeling to handle or lighten large data burdens seem like an important step in addressing existing drawbacks in the field.\n- I found the real-world experiment (FLI) to be an intriguing evaluation setting, and the ability to scale to 750k parameters to be an encouraging demonstration of the computational feasibility of the approach.", "weaknesses": "- The empirical evaluation of the proposed method is lacking with respect to existing approaches. While the proposed experiments are larger than many benchmarks common in SBI/ABI, there is virtually no mention of competing SBI/ABI methods, of which there are many in the literature (e.g., variants SNPE/SNLE/SNRE and score-based NPSE/FMPE). Recent papers often push the number of simulated samples to the realm of 100k+, certainly for simpler toy models (e.g., experiment 1), so comparison is feasible when it comes to evaluating posterior accuracy. Without these points of reference, it is very difficult to assess the real-world impact of the proposed contributions.\n\n  The only mention of competing methods is that of NUTS in the hierarchical AR(1) setting, and while the computation time is an order of magnitude larger for this MCMC baseline, it attains consistently better RMSE scores in both global and local settings over the proposed method.\n- Similar to the above point, the paper is lacking an analysis of the training-to-performance tradeoffs of the proposed method. The approach flattens the hierarchical modeling setup to reduce the sample size burden, but there is little in the way of characterizing how this might affect performance or posterior accuracy. Figure 1 portrays the naive hierarchical training alternative; it would be nice to see an analysis of performance vs simulation sample size between these two approaches to better highlight the practical benefits of the proposed method.\n- The presentation of Table 1 is somewhat confusing, in that the number of compositional sampling steps and max runtime feel arbitrarily chosen. Plotting or listing the time-to-convergence of each method instead (if they in fact converge) at each of the listed sample sizes would be far more insightful, making it easier to see the complexity of each method at a glance as $N$ grows.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FDxK8NwF6i", "reviewer": "Reviewer_Pwaw", "summary": "The proposed method aims to perform amortized Bayesian inference for hierarchical models. It does this by extending the framework of Geffner et al., Factorized Neural Posterior Score Estimation (F-NPSE) to hierarchical models. The proposed method separately considers local and global variables of the hierarchical model. Global variables can be targeted in the same way as F-NPSE, but conditioning on groups of datapoints rather than individual observations (S3.1-S3.2). The authors additionally propose a dampening scheme to improve the estimation of the intermediate densities. For targeting inference on local variables, the authors use separate F-NPSE score-based models: for these models, the conditioning set also includes global parameters $\\eta$. Both the inference models on the local and global variables can be trained jointly within the same score-matching objective.", "questions": "- A standard two-layer hierarchical model might be something like: draw global mean $\\eta \\sim p_\\eta$, then draw group means $\\theta_1, \\dots, \\theta_J$ in a region nearby the global mean $\\eta$. Thereafter, draw observations $Y$ around each group mean. In a setting like this, learning posteriors on the $\\theta_j$ by conditioning on its group of points $Y_j$ seems reasonable. Question: did the authors consider or experiment with learning a posterior $q_\\psi^{local}(\\theta \\mid Y_j)$ that does not require conditioning on $\\eta$? This is a question about the utility of the inverse factorization.\n- To that end, say we have a hierarchical Gaussian model like the above, but the group memberships are unknown (the EM algorithm is typically illustrated on such problems). Can your approach be applied when the group memberships of the $Y$’s are unknown? In other words, we don’t know precisely which $Y$’s depend on each local $\\theta_j$\n- Can you give examples of problem settings there the number of groups $J$ becomes overly large (motivating lines 220 and 221 more?)", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Many problems can be cast within the hierarchical modeling framework, making the proposed methodology applicable to a wide variety of settings potentially\n- To my knowledge, the adaptive solver contribution is novel (cf. the gains illustrated in Table 1), as is the use of the dampening adjustment (Eq. 7)\n- I find the work fairly straightforward to understand, provided the reader is familiar with score-based generative modeling. The paper is generally well-written overall\n- The use of mini-batching is indeed a substantial gain in memory for large-scale problems if this has not been done before", "weaknesses": "- When analyzing the contribution of the authors, the F-NPSE and score-based generative modeling components, while critical to the method, are both established work. Rather than a “new method”, I consider the work to be more of an application of the F-NPSE to the hierarchical setting, albeit with numerical stability improvements via the damping and minibatching. This is still a contribution, but has implications to overall novelty and the efficacy of the experimental results.\n- Competing methods for simultaneously handling local and global latent variables are not adequately discussed. For models with local and global variables, neural posterior estimation (NPE), a competing approach to conditional generative modeling that is based on e.g. normalizing flows rather than score-based sampling, is known for its marginalization capabilities (see, e.g. Forward Amortized Variational Inference (Ambrogioni et al.), and works that cite this one). Paired with permutation-invariant neural network architectures, this seems like a powerful competing method that could be used to target arbitrary local or global variables.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0GiEKttafk", "reviewer": "Reviewer_8bUg", "summary": "This paper targets large-scale hierarchical Bayesian inference where many conditionally independent groups share global hyperparameters. The authors build on compositional score matching and propose an SDE-based composition of group-wise posteriors that enables amortized inference at scale. To address the numerical instabilities of compositional approaches in high-noise regimes and with many groups J, they introduce a set of practical improvements: (i) an error-damping bridge that down-weights group contributions in the high-noise part of the diffusion path. (ii) An unbiased mini-batch estimator of the compositional score for large J. (iii) An inference-time noise-schedule shift to shorten the unstable region. Experiments on both toy models and a real-world fluorescence lifetime imaging application with over 750,000 parameters demonstrate that the approach is scalable, fast, and produces practically useful posterior summaries.", "questions": "* For the composed posterior, why does the reverse SDE start from N(0,I/J)? I’m not a diffusion expert—an accessible derivation would help readers (this is not critical, just a clarity request).\n* Have you tried to run any hierarchical ABI baselines? What were the blockers (memory, instability, code availability)? Could you include one controlled comparison?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The main strength of the paper is that it extends the amortized SBI methods to large-J hierarchical models, where previous compositional methods often failed numerically. The combination of SDE-based composition + error-damping bridge + unbiased mini-batching + inference-time schedule shift is well-motivated and empirically effective. These design choices are clearly ablated and explained.\n* The experiments are well-designed, toy cases provide insight, while the FLI case shows real-world value and impressive runtime.\n* The paper provides detailed code and settings, aiding verification and reuse.", "weaknesses": "* Writing is generally clear, but for SBI audience, it would be great if the authors could add more background about diffusion/SDE and the introduction about previous score-based SBI methods to make the paper more accessible. \n* Line 51, The paper lists SBI and ABI as parallel notions. Consider rephrasing to avoid implying ABI is separate from SBI, and the cited ABI reference (Glöckler et al., 2024a) is itself an SBI paper. \n* The paper cites multiple amortized hierarchical approaches but does not compare against them. Even if some implementations are hard to scale, I would appreciate it if the authors could (a) explain the obstacles and (b) add at least one baseline in a moderate-scale scenario to anchor performance and calibration.\n* The error-damping bridge alters the diffusion path by down-weighting scores in the high-noise regime; while highly effective, it likely changes the target path measure. Could you add a theoretical discussion (e.g., a bound on bias in expectations, or conditions under which the deformed path still converges to the correct posterior in the small-noise limit). Even a formal statement plus intuition would strengthen the contribution.\n* The paper limits its scope to two levels. It would help to outline how the method might extend to deeper hierarchies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aZGWjXz9bD", "reviewer": "Reviewer_MSGt", "summary": "The paper proposes a compositional diffusion model technique for tackling inference in hierarchical Bayesian models with large amounts of variables, which could be especially useful in the case where generating data from the model is a significant bottleneck. Further, we may not have access to an explicit likelihood function. Given a hyperprior $p(\\eta)$, priors $p(\\theta_j | \\eta)$ and likelihoods $p(y_j | \\theta_j)$, where $j \\in {1,...,J}$, the task is to get the posterior $p(\\eta, \\theta_1,...\\theta_J | y_1,...y_J)$. The paper tackles this by training diffusion models to approximate the conditional distributions $p(\\eta|y_1,...y_J)$  and $p(\\theta_j| \\eta, y_j)$. In a similar situation, previous literature replaces the score function for noise level t with $\\nabla \\log p(\\eta_t|y_1,...y_J) \\rightarrow (1-J)(1-t)\\nabla \\log p(\\eta_t)  + \\sum_j \\nabla \\log p(\\eta_t|y_j)$, and corrects errors with, e.g., Langevin dynamics. In the hierarchical model context, this necessitates only training score functions for $\\nabla \\log p(\\eta_t|y_j)$, which is beneficial because we can get much more data per simulation than the joint score if we use the same neural network to represent the denoiser for all $j$. This is also more computationally scalable. The paper shows that this naive \"bridging score\" is unstable in the case of very large hierarchical models, however, and proposes a stabler alternative that scales better through a damping schedule and minibatching on the compositional score, and a sampling noise schedule adjustment that puts less emphasis on high noise levels.", "questions": "- It seems that interleaving Langevin dynamics with the SDE steps would help with the theoretical guarantees. Do you think that this would be an aspect of the approach that is worth highlighting? Would this require some different design choices w.r.t. the damping factor and noise schedule? If so, it seems like adding analysis on how to apply the idea to the Langevin dynamics case would strengthen the paper.\n- The baseline where we literally train $\\nabla \\log p(\\eta | y_1, ... y_J)$ is also missing (e.g., with some variant of the Deepset architecture). I understand that this will likely be less scalable than the proposed method and the subset examples are, and in practice the GPU memory may easily run out with a naive implementation. But it seems that the denoiser architecture could be optimized to handle large amounts of inputs $y_j$. Do you think this would be possible? \n- Small detail: In the Appendix, it is mentioned that the noise-prediction becomes unstable for t close to 0. Is it not the opposite, that it becomes unstable at t close to 1? The v-parameterization is equivalent to the epsilon-parameterization at low noise levels since alpha_t=1 and sigma_t=0.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The problem itself is well motivated, and seems to address limitations with previous diffusion methods that could be applied in this domain. Scalability of compositional techniques seems to be a clear problem and the paper proposes sensible ideas for addressing it. \n- While the SDE sampler used in the experiments does not seem to have guarantees of converging to the correct posterior in some limit, the ideas presented here seem to be applicable more broadly to methods with guarantees that include, e.g., MCMC correction steps in the generative process. \n- The inference-time hyperparameter optimization is a pragmatic approach that seems to work quite well in practice. \n- The paper is clearly written and is easy to follow. \n- Ablations for different aspects of the method are in place, and increase confidence that they are useful in tackling the issue.", "weaknesses": "While I like the idea and the experiments that are present in the paper, it seems that the paper could be improved significantly:\n\n*Baselines.* The quantitative baselines in the paper are quite scarce, and it seems to be that the few baselines chosen are not particularly strong. I am not an expert on hierarchical Bayesian models, however, and it may be that I missing some wider context. Problems that occur to me: \n- For the Gaussian toy example, comparisons to other compositional diffusion methods are presented, but we do not have quantitative results. If the main contribution of the paper is the scalable compositional diffusion method, it would seem appropriate to have quantitative comparisons on wall-clock time and inference accuracy for all the data sets against these methods.\n- On the other hand, if the focus is on improving hierarchical Bayesian models in general, then it seems that baselines from that literature should be present. It seems that the cited papers of Habermann, Arruda, Heinrich and Rodrigues would be relevant baselines. It is claimed that they are not scalable to large J, but it would be good to show this concretely at least with some of the methods, comparing to the proposed method with a given simulation budget. \n- For the AR(1) model, comparison to an MCMC method (NUTS) as a gold standard is presented. The key improvement compared to MCMC is highlighted to be improvement in wall-clock time, but a key reason for why the wall-clock time is good is the minibatching. It seems to me that this is a technique also applicable in the real of MCMC methods, i.e., SGLD [1]. This is not directly applicable to the likelihood-free setting, but the experiments also do not consider cases where the likelihood is intractable. \n- For the Fluorescence Lifetime Imaging experiment, the baseline is a non-hierarchical maximum likelihood estimation model. Could one not also do, e.g., MAP estimation with the full hierarchical model? It seems that the likelihood can again be minibatched for efficiency in a similar manner as in the proposed method. \n\n*Datasets.* \n- The paper positions itself as a simulation-based inference paper, but it does not seem to have any inference tasks where the likelihood is intractable. I understand that the method may be more scalable than many alternatives even with likelihoods accessible and extending the method to likelihood-free cases it trivial, but it seems that this would solidify the results in the paper. I do not see this as a major concern, however.\n\n*Theoretical guarantees.* The following are not major concerns, as I think they can be alleviated with methods that are mostly orthogonal to this work. See also question regarding this. \n- The proposed method does not seem to have a guarantee of converging to the correct posterior distribution in some limit, as opposed to the standard approach of training the score $\\nabla \\log p(\\eta_t | y_1, ... y_J)$ directly. It seems that this may make it more difficult to work with the model in practice, although the inference-time hyperparameter optimization does help here. \n- Further, the minibatching introduces some noise into the obtained combined score functions, even at low noise levels at which the full compositional score becomes accurate. \n\nReferences:\n\n[1] Welling, Max, and Yee W. Teh. \"Bayesian learning via stochastic gradient Langevin dynamics.\" Proceedings of the 28th international conference on machine learning (ICML-11). 2011.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)