Hierarchical Neural Simulation-Based Inference Over Event Ensembles
Authors: Lukas Heinrich
Siddharth Mishra-Sharma
Chris Pollard
Philipp Windischhofer
Venue: Accepted by TMLR ()

Abstract: When analyzing real-world data it is common to work with event ensembles, which comprise sets of observations that collectively constrain the parameters of an underlying model of interest. Such models often have a hierarchical structure, where ``local'' parameters impact individual events and ``global'' parameters influence the entire dataset. We introduce practical approaches for frequentist and Bayesian dataset-wide probabilistic inference in cases where the likelihood is intractable, but simulations can be realized via a hierarchical forward model. We construct neural estimators for the likelihood(-ratio) or posterior and show that explicitly accounting for the model's hierarchical structure can lead to significantly tighter parameter constraints. We ground our discussion using case studies from the physical sciences, focusing on examples from particle physics and cosmology.

Published in Transactions on Machine Learning Research (02/2024)
Hierarchical Neural Iipullatioo-Basele Inference Over Event Ensembles
Lukas Heinrich Technical University of Munich
Lheinrich@tum.de
Siddharth Mishra-Sharma
smsharma&@mit.c edu
MIT, Harvard University, IAIFI
Chris Pollard
University of Warwick
Philipp Windischhofer
windbischhofrruuchioooo..k
University of Chicago
Reviewed on OpenReview: https: // openreview. net/ forum? id- Jy2IgzjoFH
Abstract
When analyzing real-world data it is common to work with event ensembles, which comprise sets of observations that collectively constrain the parameters of an underlying model of interest. Such models often have a hierarchical structure, where "local" parameters impact individual events and "global" parameters influence the entire dataset. We introduce practical approaches for frequentist. and Baycsian datasct-wide probabílistic inference cases where the likelihood is intractable, but simulations can be realized via hierarchical forward model. or posterior and show that explicitly
We construct neural estimators for the likelihood accounting for the model's hierarchical structure can lead to significantly tighter parameter constraints. We ground our discussion using case studies from the physical scienees, focusing on examples from particle physics and cosmology.
1 Introduction
Datasets composed of multiple samples are ubiquitous in scientific and more broadly real-world data analysis tasks. These datasets are typically governed by underlying models that exhibit rich hierarchical structure, with local parameters shaping individual events while global parameters exert influence across the entire dataset. This layered structure, if appropriately utilized, can greatly augment the efficiency and effectiveness of the inference process.
The complexity of scientific models and digl-diimesionaality of datasets has led to a recent surge in interest in implicitly-speccfeedd models, where the likelihood function intractable but simulations can be realized via mechanistic forward modeling. The paradigm of simulation-basell inference (SBT), augmented using tools from machine lcarning and differentiable optimization more broadly, has emergod a powerful approach for performing inference in such scenarios (Cranmer al., 2020). However, the majority of existing methods for simulation-basee based inference are designed for learning from individual data points and do not fully capitalize on the hierarchical structure of the data-generating process.
To address these gaps, we propose a set of novel approaches that augment existing simulation-based inference techniques, both frequentist and Bayesian paradigms, with the goal of exploiting the hierarchical structure of the governing models. We contextualize our discussion using case studics from the physical seiences, with particular focus on particle physics (particle collider data) and astrophysics (strong gravitational lensing images).

---

Published in Transactions on Machine Learning Research (02/2024)
In the present work we make several methodological contributions that necessary for optimal deployment of simulation-based inference methods on large-scale scientific datasets. Our primary contributions are the following:
We substantiate theoretically well empirically the fact that hierarchy-aware inference in many implicit models with hierarchical structure requires dataset-wide approach, contrasted with the more common paradigm of combining implicit likelihood or posterior estimators associated with individual observations. By "hierarchy-aware inference", we refer to building a still approzimate postcrior or likelihood-ratio estimator that nevertheless takes into account the full (in our case) hierarchical structure of the assumed forward model. We systematically derive conditions under which such hierarchical approach is necessary for correct inference, depending on how the model parameters are partitioned into local and global parameters of interest and nuisance parameters;
We introduce frequentist as well as Bayesian methods for dataset-wide learning that can be used lo perform simulation-based inference over event ensembles and can deal with datasels of varying cardinality. We connect our insights to several common use cases in the physical sciences and show how popular simulation-based inference paradigms can be adapted for optimal datasct-wide lcarning. In particular, we introduce the first approach for end-to-end frequentist, simulation--aes inference targeting hierarchical set-values data;
We show that our machine learning-based inference methods are generically substantially faster than traditional approaches, such as Markov Chain Monte Carlo (MCMC) methods, even when the likelihood tractable, while giving consistent results. They also allow performing inference in "streaming" mode where, e.g., posterior estimates are efficiently updated in real-time as new observations are made without having to perform a re-analysis of the entire updated dataset.
The paper is organized as follows. In Sec. we provide a background discussion of hierarchical frequentist as well as Bayesian inference, and describe the class of forward models we target. We discuss related work in context of this background Sec. 3. In Sec. we describe the methods and architectures used in this paper. Section 5 presents several case studies including simple toy example, partiele physics example, and example from astrophysics, in increasing order of complexity. We conclude in Sec. 6.
2 Theory and background
We use the following notation throughout: an refers to an observation associated with an individual event, 0 are the global (dataset-wide) parameters, and 2, are the local (per-event) parameters. Ensembles of either obscrvations or variables are denoted with curly braces, {x} = {x) We will also differentiate muisance parameters, which are parameters associated with the data-generating process that we are generally not interested in inferring; these will be treated through either profiling (defined later) or marginalization. Nuisance parameters will be denoted by v, or Og and 2p when we wish to distinguish between global and local muisance paramcters.
2.1 Hierarchical models
The joint probability distribution of a set of events with cardinality N, with global parameters of interest 0, global muisance parameters 0, as well as local parameters 2; can be written as:
N p(x) (2).0000) = Ip(a, 0,,,,,,
(1)
In many scientific applications, the number events in a dataset may not be known a priori. This could occur events are obscrved sequentially (such as gravitational wave events observed over time), we aim apply the same model to datasets with different event counts (c.g., astronomical observations of sky patches containing different. numbers of stars/gald sassgaaaxies,, if the rate of events is itself model parameter that may
2

---

Published in Transactions on Machine Learning Research (02/2024)
be interesting to measure (e.g., interaction cross section particle collider experiment). In these cases, we can generalize the joint probability to an extended model:
80
(=3,000) Σ P(
I z,, 0,,,,...
(2)
N=0
Here, p(N f) represents the probability of observing N events given the global parameters A. Often, p(N () is Poisson distribution, and the full model is known marked Poisson process. Our proposed techniques apply to both situations, and we present example applications described by equation as well as equation 2 below.
2.2 Bayesian inference
0, 02) and wish to target the posterior distribution over the global as well as local parameters of interest (and marginalizations thereof) given a dataset.
In the Bayesian paradigm, we introduce a prior p.,
In special cases, it possible combine event-level inferences construct the dataset-wide inference. For example, in the non-exlended case one can construct the dataset-wide full posterior from the event-level posteriors p(0, 0u, 24),
N II p(0, 04 x.), 1=1
p(0, 00, {2) {x}) (0,02)N-1
where the prior factor in the denominator ensures that the prior density is not counted multiple times. Similarly, if one has access to the per-event posteriors marginalized over local parameters 2,, and the target quantity is the wide posterior marginalized over local muisance parameters, a combination is possible due to the independence of the events and the assumption that the various priors also factorize:
p(0 (x}) = / d(zu)p(0,
p(0, (a0) [x}} = f d( p([x} 0, (z(zz) p(O, {2n})
(4)
p((x))
II
/ džpz p(x;;)
p(æ; 0, 0,,2))
pp(2, 1 10) p(0) p(O)--
p(0 xx4)
However in the general case such combinations do not hold; e.g., it is not possible to combine marginal per- event postcriors over global muisance paramcters p(0 z,) = d', p(0, 0, xx) to construct the dataset-wide posterior p(0 {a+),, marginalized over global and local nuisance parameters. Therefore, a general solution for marginalized dataset-wide posteriors necessitates a dataset-wide approach to inference. This be done either by inferring global parameters at a dataset-wide level (in particular if the dimensionality of the resulting posterior is sufficiently wieldy), or by performing letaset-level marginalization.
2.3 Frequentist inference
In the frequentist setting, inference parameters is lypically split into procedures for point estimation and interval estimation. A popular choice for point estimator is the maximum maximum-likelihood estimator (MLE), due to its favorable asymptotic property of being the minimum-variance unbiased estimator. Interval estimation and hypothesis testing are based on a subjective choice of a test statistic ta,e, (x) and a desired confidence level:
(5) where L 0,) ) is a function proportional to the likelihood p({x} 0, 0,) and 18 is the value of the quantile function at - a cumulative probability and serves as the threshold value for the interval. Here, a indicates the size hypothesis test, with a typical value being a = 0.05.
0,0, argming,o. log L(x) (0,0)) (0,0v)l;
18.0. (0 t0.0, ({x)) < t8.0..
As the Bayesian case, one often interested in an inference that only considers the parameters of interest 0 and removes any dependence nuisance parameters 0,- The frequentist. analogue to (partial) marginalivation
3

---

Published in Transactions on Machine Learning Research (02/2024)
is (partial) optimization (also referred to as "profiling"). For interval estimation and hypothesis testing the test statistic is thus modified to read
to((=)) -log -lo (((}) 10,6,) p({x} 0,00)
(6)
where 02 = ((x)),)) is the optimal value for 02 when 0 is fixed to particular value. Similarly to the Bayesian case, it is in general not possible to reconstruct the data-sct wide nuisance-free inferences from per-event nuisance free inferences: In general the dataset-wide profile likelihood ratios cannot be constructed from per-event profile likelihood ratios.
2.4 Simulation-base inference
When the likelihood is not tractable but simulations can be realized via forward modeling, neural simulation- based inference methods can be leveraged to perform inference. In cases where event-level inference quantities can be combined to yield correct dataset-wide quantity, neural posterior or likelihood-ratio estimators trained on individual events can be combined. For example, marginal posteriors p(0, 02 a;;) can be trained and combined using equation 4.
Alternatively, a per-event likelihood ratio estimator parameterized by all parameters, including muisance parameters, can be learned (Cranmer et al., 2015; Nachman & Thaler, 2021) to train neural estimate for p(2,10,0..)
the per-event likelihood ratio sa(x,, 0 02) = p(:: 0o.0w,0]: These can be used to build a dataset-wide likelihood
ratio II p(2.)0.0..) that can be an input for dataset-wide statistical analysis to yield posteriors or likelihood p(
ratios which can then be marginalized or profiled over to yield dataset-level quantities. While this simplifies the training procedure as the networks are only performing inference for single events, disadvantage of this approach is that the downstream statistical treatment can be very costly if the number of nuisance parameters 0, is large, so that the amortization gains from the neural estimates are not fully exploited.
Given these nuances, methods for set-wide simulation- based inference necessitate architectures which can efficiently model joint posteriors or likelihood ratio over global and local parameters of interest, for sets of varying cardinality, without requiring explicit parametcrization over a potentially large number of nuisance parameters (which would otherwise be necessary for hierarchy-aware inference with event-level estimators). In the following, we describe two such architectures for amortized dataset-wide inferences.
3 Related work
The methods presented in this paper build upon previous studies across several disciplines. A likelihood-ffre algorithm for obtaining frequentist, profile likelihood ratios corresponding to individual observations was previously presented in Heinrich (2022). Also in the context of frequentist inference, Nachman & Thaler (2021) explored dataset-wide inference in the context of particle collider studies, without distinguishing nuisance parameters and parameters of interest.
Hierarchical inference for implicit models has been studied using Gaussian posterior density estimators in Wagner-Care Carena et al. (2021; 2022), where again individual (per-event) posteriors are learned and combined to obtain estimates for the global parameters. Agrawal & Domke (2021) considered amortized variational inference for a simple class of hicrarchical models. While similar spirit, we consider hicrarchical inference in the language of modern scientific simulation-bassd inference methods, and discuss subtleties associated with treating classes of variables in different hierarchies (local and global) as muisance parameters.
Hicrarchical Neural Postcrior Estimation (Rodrigues et al., 2021) is a recently-introdneed simulation-based inference method closely related to the deep set-based Bayesian inference model used here. We complement this approach by introducing frequentist analogues for dataset- wide learning; by including amortized estimators that can then target datasets of varying sizes; and by comparing the proposed approach to more traditional methods in terms of inference computational time. Finally, Geffner et al. (2022) introduced a method for efficient sed-wide learning by combining score estimators targeting a sub-set of events. Our method is
4

---

Published in Transactions on Machine Learning Research (02/2024)
complementary, focusing the unexplored setting of hicrarchical models and bridging applications to domain sciences.
4 Methods
Our goal is to target the graphical models described by equation and equation 2 in a simulation-based inference setting. We wish to construct amortized estimators for the likelihood ratio or posterior corresponding to these models such that they (1) sensitive to the full hierarchical structure of the forward model i.e. take dataset-wide view, treating local and global parameters separately, (2) distinguish between nuisance parameters and parameters of interest at various hierarchy levels, and (3) can process datasels of varying cardinality. While previous related works, discussed in Sec. 3, have considered subsets of these requirements, our goal is to target all three with a particular eye towards scientific applications where event ensembles are common. We consider the following architectures, which satisfy these desiderata.
Deep set-based model Here we consider a simple deep scts-inspired (Zaheer et al., 2017) architocture, gglob (2;) are aggregated via permutation-invarian pooling and passed where per-event encoded features &e
through a decoder network gy that outputs the parameters of a variational posterior distribution for the global parameters (8 [ {z}) (Papamakarios & Murray, 2016). At the same time, local per-event embeddings (x;) are used to condition density estimator for the local parameters, do (= {x), 0), including a conditional dependence on the global parameter 0. In practice, the local and global embeddings are obtained by chunking a feature vector obtained with sg. For the variational ansatz we consider either a multivariate normal distribution or a normalizing flow (Rezende Mohamed, 2015), with the decoder outputting either the mcan-covariance of the normal distribution or the conditioning context for normalizing flow. The (negative) sum of the local and global posterior log-densities is used as the minimization objective,
-f = = log glob (o (Bl08 g
N
glob (ar) ))
N
Σ
log Ohp" (zs
8 (I.). 8
(7)
We derive and further motivate this loss function in App. B. For models of the form of equation 2, the cardinality of the input set is drawn as N ~ p(N | 0) in the training. The density estimator parameters (Pp } as well as the embedding network and decoder parameters @ and & are optimized simultancously. We note that this is similar to the set-up used Rodrigues et al. (2021). Fig. shows a schematic illustration of this deep set-based architecture.
-ransformer-based model As an alternative to deep set- bet-based aggregation with varying cardinality at training time, we consider a decoder-only transformer (Vaswani et al., 2017; Phuong & Hutter, 2022) which takes feature embeddings from individual events and produces cumulative posterior embeddings using blocks of self-attention and dense layers, applying causal mask to ensure that output embeddings are only sensitive to preceding events in the sequence. The embeddings are then used to predict the posterior mean and covariance of the variational Gaussian ansatz. A potential advantage of the transformer approach is that, assuming a uniform distribution on the cardinality N, we do not need to vary the cardinality of the input set during training. Since the loss consists of a sum of log-likelihoods corresponding to cumulative posteriors, one for each cardinality N up to some specified maximun value, all cardinalities are considered together; gloob( {m] E, log log 2,,0), leaving the data-point encodings implicit.
--=EN log o
Networks that used LSTM cells to iteratively update posterior embedding based on incoming data embeddings were also considered, but their convergence properties were noticeably poorer than those of the deep set- and transformer-inspitn architectures described above.
5 Experiments
We describe several case studies of dadeset-wide learning, ranging from illustrative "toy" experiments to more prototypical examples representing problems in particle physics and astrophysics. We refer to App. A for
5

---

Published in Transactions on Machine Learning Research (02/2024)
Event-level (local) posteriors
ap (2g Teg O)
Frequentist sctting
Globel parameters
Likelihod--ratin test statistic Eup (f
Joc (ar²) &
Observations
Simulator
Vj
SØ Pro-poent embeedding network
glob (24)
4u (0 {x}}
Zi
Agg
Set-level
Per-event (local) parameters
network
Dataset-level (glohal)
posterior
Figure 1: Schematic illustration of the deep set-based architecture used this work. The red lines tiees/arrwws show the path used only in the frequentist sctting for training global test-statistic estimator while profiling over global nuisance parameters.
additional details on the experiments, including details on training. All experiments are implemented using PyTorch (Paszke et al., 2019).
5.1 Simple vurti-variate normal likelihood
The forward model To verify the ability of our method to recover the true posterior distribution for sets of varying cardinality, we consider simple multivariale normal likelihood with known covariance matrix E; p([x) () Io No(x, (,
The model parameter is the mean vector 0, and Po, Eo are the hyperparameters of the prior multivariate normal dist p(0) = No(uo, En). In this case, the posterior distribution p(0 {x)) is also multivariats normal - the prior and posterior are conjugate with mean Ppost and covariance Mpost of an updated posterior given by Ppoa = (En + -E-') (Eo'l No - -F) and Epost = (Ep +NN- where is the sample mean of z; and N the total cardinality of the dataset. We choose diag(2, 4, 6), with each sample consisting of 5 draws from the multi-variate normal distribution; each individual data point then consists of 15 features.
Inference We train the deep set and transformer architectures described in Sec. on 50,000 samples drawn from this likelihood with prior p(a) = No(0, 3) and a maximum sequence length of Nmax = 200. The cardinality of the training set is randomly varying as N ~ (1, Further details on the training procedure are provided in App. A.
The model is then tested on 500 new sequences, and the distribution (median and middle-68% containment) over inferred standard deviation o for each of the parameters is shown Fig. (left: deep set, middle: transformer) compared to the true expected scaling of the parameters (dashed lines). We show 01, 02, and d3, which are the diagonal entries of Epsst the covariance of the posterior the target mean parameters 8. The deep set model typically gives more faithful posterior widths compared to the transformer, potentially due to the relatively simple nature of the problem combined with the more data-hungry nature of the transformer architecture. Given this, we restrict our experiments in subsequent sections to use the deep set-based architecture, noting that the transformer approach may nevertheless yield advantages depending on the specific structure of the forward model. The right plot shows the evolution posterior for a specific sequence using the deep set model, illustrating convergence of the posterior mass around the true point more data points included. We note that although this is simple, analytically tractable example, the final posterior is obtained by analyzing dalaset of dimensionality 3000 far from trivial task.
Results

---

Published in Transactions on Machine Learning Research (02/2024)
Predicted posterior Deep Set
Predicted posterior width: Transformer
Predicted posterior evoluttoo: Deep Set
cardinding
Ground truth
10a
10
8
150
pte (x))
100
d1
d1
62
10
50
Ground truth
Ground truth
10"
10'
10
10a
10
10
0.5
1.0
1.5
2.0
2.5
Sequence length
Sequence length
Figure 2: The median and middle-68% containment of the inferred posterior width over 500 Lest samples as a function the size of the dataset. Results for the deep sets (left) and transformer (middle) architectures are shown. The expected scaling (dashed lines) is observed in cases, but the deep set architecture shows a narrower spread in outcomes. The right plot shows the evolution of posterior for specific sequence using the deep set model, illustrating convergence of the posterior mass around the true point.
5.2 Mixture models in particle physics: frequentist treatment
We demonstrate the dataset level profiling by adapting the approach described in Heinrich (2022) and compare the trained test statistic to the true profile likelihood ratio (not to be confused by the saelihood-to-cviddne ratio (Miller et al., 2022)).
The forward model In this example we consider a situation that is common in both particle and astrophysics. Here, the observed data can often originates from multiple disparate physics processes of varying intensity (for example a signal process and background process), such that the overall model is a mixture of the per-process models:
N p((x) 0.0,) = [N X(0, 02)) II (c.pa 0,0,) + c. Pata 10,0,); with,
(8)
with c, co, 02) being the relative signal strength and pa and ps the per-event densities for signal and background, respectively. Fast inference for such models is especially important in settings that perform real-time monitoring of the quality of the data stream delivered by the experiment. The observed per- event densities are typically not tractable, but high-fidelity simulators exist. In this work we focus on Gaussian mixture components, since the empirical distributions typically considered in the domain use-cases fcature multi-modal densitics (c.g., in scarches for "bump" fcatures). To demonstrate the simulation-basel Irequentist inference, we consider the special case where ps and Ps are Normal distributions No(μ, or) with fixed hyperparameters Ps/b. S//bs but the component coefficients are a function of both the parameter of interest (the "signal strength") as well as the nuisance parameter (i.e. "the background strength"); C, 0/(0 0,). The overall expected rate is then X(e, 0,) = 0 ty.
The model is again based on the deep set-inspired architecture described in Sec. 4, with a small variation after aggregating the per-event embeddings, these are concatenated with the parameters of interest 0 order to obtain a parameterized neural network classifier one-to-one with the desired test statistic (Cranmer et al., 2015), so([z), F (o(()).. Following the procedure in Heinrich (2022) we aim to train test statistic with best average power across alternatives irrespective of the muisance parameters and show that the result. good approximation of the profile likelihood ratio, i.e., the dataset-wide test. statistic typically used for interval estimation and hypothesis testing in frequentist inference:
Inference
so((x),0) F to([=))=+)))) = -22lo p((x) 0 10,
(9)
p((x} 0 0,))
Details the training procedure well hyperparameter choices provided in App. A.
7

---

Published in Transactions on Machine Learning Research (02/2024)
1.0
0.5 -
0.5
0.0
2.5 5.0 7.5
10
10
10
7.5
7.5
5
ARRL
5
5.0
2.5
2.5
0
10
2.5
5.0
7.5
0.0
0.5
0ML
Figure 3: (Left) True profile likelihood ratio for the mixture model scenario in Sec. 5.2 (bottom) compared to the neural test statistic (top). (Center) Neural test statistic and profile likelihood ratio scatter plot, showing bijective relationship between the two. (Right) Maximum-likelihonn est estimate obtained from the neural Lest. statistic bML compared with the true value pLLR-
As the ground truth is available in this model, can compare the performance of the learned test. statistic directly. In Fig. 3 (left) we show the ground truth profile likelihood ratio fpLR as function of the parameter of interest 0 for 1000 datasets {z} ~ p({x} 02) with average cardinality of 110 events. The profile likelihoods have the typical parabolic structure with its minimum at the maximum-likelihood estimate PpLn = argming {(o) (0). The neural test statistic MLL exhibits similar features, however as it an output of classification lask is constrained to the interval [0,1]. The relationship between neural and profile likelihood ratio test statistic is shown in Fig. 3 (center), where a clear bijective mapping is recognizable. The network therefore learns an equivalent test statistic in terms of inferential power. Similar to the ground truth case, the best parameter can found by minimizing the test statistic 0ML = argming tMI, likelihood-free point estimate. In Fig. 3 (right) the two point estimates can be scen to correspond well to each other.
Results
[x] (0) to produce a fully
One of the major advantages of the neural test statistic is its inference time: While explicit profiling based on either a known or Icarned test statistic (0, 0,) requires multiple optimizations for cach data instance, the neural network statistic allows for fast vectorized computation in single forward pass. We observe the neural network based inference for 1000 datasets to be over two orders of magnitude faster than the standard approach, i.e. explicit construction of the profile likelihood ratio computation of 0, Ō., O via optimization.
5.3 Mixture models in particle physics: Bayesian inference for a narrow resonance
The forward model We also demonstrate a Baycsian datasct-wide inference for the mixture model case. Here the parameters of the model affect the mean of the signal component and the mixture coefficient, whereas the signal width and all background parameters are fixed hyperparameters. Using the notation in equation 8, we have
(x 0.) = No(ar Ou, ox); pb(2.)) = No(a; | (b, co( = 0,
(10)
i.c. the parameter of interest is the relative signal strength and the location of the signal in the observed feature space is a nuisance parameter. In typical applications in particle physics, the signal often very narrow, whereas the background distribution is more diffuse. All parameter choices and the priors p(0) and p(0,) are specified in App. A. Instead of parameterizing the overall Poisson rate X, here we used a fixed batch of events of size No, corresponding to the situation one might encounter in real-time inference task. In summary, the wide joint probability is:
No
p([x} | 0, (,), 0,) = p(0)p(0,) II (0) ps (8,) + (1 cc) pb(x;))
(11)
8

---

Published in Transactions on Machine Learning Research (02/2024)
Background
No - 100
Posterior mean
Posterior standard deviation
Signal + background
Oru = 0.2
0.40
Truth
No 100
0.08
8, B nom
Density
1.0
Signal
2.0
0.35
Deep set
MCMC on (#1og)
Gcu = 0.2
0.06
0.5
med
0.30
MCMC on Soom
med(aal
MCMC Smarg
0.0
0.04
0.25
Deep
Density
1.0
Signal
MCMC on x
0.5
0, true 0.0
0.20
0.02
MCMC on Syom
N₂ = 100 Bcue 0.2
MCMC on Smarg
0.15
0.0
0.00
-2
2
-1
2 True signal position 6n true
0
X
True signal position tutr
Figure 4: Signal- and background densities ps and ps in a simplified particle physies analysis for two different true values of the muisance parameter 8, (left), additionally showing sample datasets {z) of size No- 100. The median mean no (center) and median standard deviation do (right) of the posterior p(0 {a:)), obtained with different inference methods for different true nuisance parameter values 0 Oytrrue The median is computed on ensembles of 400 datasets for each parameter point.
Inference We wish extract the posterior p(0) {x}} characterizing the prevalence of the signal process given a set of No events, marginalizing over 0y. The mixture fraction 0 is fundamental property of nature; tracing its evolution over time is thus excellent indicator of the health of the experimental apparatus. We use deep set work to perform fast, amortized inference in this setting, trained as described in Sec. 4 directly on the set of observations {x}. Using the likelihood of equation 11, the posterior is also accessible through MCMC inference (as described in App. A), which can serve as a point of comparison. In many cases, the individual events x are too high-dimensional for direct inference to be possible. A fundamental problem in the analysis of collider data thus consists the construction of a low-dimensional summary statistic x + s(x), designed to retain the relevant information contained in the data, subsequently used in the inference step. It is common practice in contemporary work to use the signal-to-backgronne density (x) evaluated at particular "nominal" nuisance parameter value fomnom. This observable is known to be a sufficient statistic for 0 under the hypothesis 6,- 0v.nom (Neyman & ES, 1933). Another popular choice is to neglect the hierarchical structure of the data-generating process; in the Bayesian paradigm this involves marginalizing over the nuisance parameters at a per-event level and using (m) = I d0, (=(x | 6z) p(8,)/p6(a)) as summary statistic. Both quantities are known to not be informatioon-presrviing in general. Our example admits analytic likelihoods also for 8nom and Smarg: allowing their performance to be evaluated through the MCMC procedure described above.
ratio &pom (1; (m.mom) ) = (ar 0r,nomm
Results
The center and right panes in Fig. 4 compare the neural posterior estimate (black lines) to MCMC posterior estimates based on æ (blue markers). The width of the neural posterior closely matches the MCMC results, demonstrating the ability of the deep set to extract the entire information about the parameter O contained in the data. Also shown are MCMC posterior estimates for 8nom (red markers), and 8marg (green markers). As expected, these generate wider posteriors! and, hence, weaker parameter constraints, illustrating the importance of a proper treatment of global nuisance parameters.
Using the fully parameterized density ratio p=(x 0p)/Po(z) for hierarchy-aware inference is, as mentioned in Sec. 2.3, complicated by the fact that often many hundred muisance parameters are present contemporary particle physics problems (The ATLAS Collaboration, 2022; The CMS Collaboration, 2022). However, muisance parameters inform the the deep set through their presence in the training dataset and do not require any explicit parameterization of the network, suggesting a much better scaling behavior.
While we do not explore this direction further here, we note that the event-embedding
s sglob (x) learned by the deep set also plays the role of an informatioon-reserving summary statistic that can serve as input to traditional (non-amortized) inference pipelines.
In our simple example, Smarg (2) related Suom By mm for mm R -0.3, leading to a sharp improvement. sensitivity for this model point. This is an artifact that is not expected occur in more complicated environments.

---

Published in Transactions on Machine Learning Research (02/2024)
5.4 Astrophysics example: Strong gravitational lensing
We test. our procedure for hierarchical sinulatioon-assd inference on more complicated problem from the domain of astrophysics, where the likelihood associated with the forward model is intractable the characterization of the distribution of dark matter clumps (subhalos) in galaxics using images of galaxies gravitationally lensed by the chumps as well as a larger smooth mass distribution. The model contains both per-event (i.e., local) as well as dataset-wide (i.e.; global) parameters. Estimating global lensing parameters with more sophisticated forward models was studied in, e-g., Anau Montel et al. (2023); Montel & Weniger (2022); Coogan et al. (2022); Wagner-Carena Carena al. (2022); Brehmer et al. (2019); note that our aim here is instead to showcase the ability to do hierarchical simulation-based inference over local and global parameters with an amortized set-wide estimator.
Different latent realizations Zsub ~ p(Zuub @glob, 0loc)
Vary glab and Iloc
Figure 5: Illustrative samples from the lensing model. The rows show two different choices of local (per-event) parameters, while the different columns show variations on the global (sct-wide) parameters for the fixed choice of local parameters. the abundance of subhalo population the lens, can be seen. The scatter points shows the location of individual subhalos in each image.
variation induced by the global parameters, which control
The forward model The substructure mass function, which describes the mass distribution of clumps, is dn = a-m -8 where the normalization parameterizes the overall abundance of subhalos, and / is the spectral index. Oglab {a, B} are the global parameters in the hierarchical model. The number of subhalos follows a Poisson rate, Naub ~ Pois()subb, where we have the expected number of du In addition we have per-image (local) parameters fjoe {0 Ax. Ay, q), where 0g is the Einstein radius characterizing the overall scale of the lensing ring, (Ar, Ay} the offset vector between the centers of the background source and the lens galaxy, and g E [0, 1] is the axis ratio with corresponding to spherical lens. The subhalo radial positions are drawn from a uniform distribution close to the Einstein radius, ) Nauh (n 0.2", 0g + 0.2") and the azimuthal angle is drawn uniformly from the interval [0, 2x]. Given realization of the masses and positions of the clumps { Yaub. Maub Neah (which otherwise act latent variables, collectively denoted {zsub)), the expected image Pimage can be obtained using a gravitational lensing simulator (Brehmer et al., 2019). The actual image is then a Poisson realization of this expectation, xr Pois(/amagee). We choose our simulated lensed images be 64 64 pixcls in size; see App. A for more details on the lensing forward model. The per-image likclihood is intractable bocause of the high dimensionality of the latent space.
parameterized as power law Amaat
subhalos Hsub J dm dmaat"
p(x 0)
Σ
(MMun (Zeub) Pois(æ [imagee p(Mimage (Zeub ,(oo)) Pois(Neub (sub).
(12)
N1ab=1
10

---

Published in Transactions on Machine Learning Research (02/2024)
Some illustrative samples from the lensing forward model are shown in Fig. 5, showing variation with local (per-event) and global (set-wide) parameters across rows, and different realizations of the subhalo latent parameters (mass and location) along columns. The scatter points show the location of individual sub-clumps in each image.
Inference We would like to infer the posterior over the local as well as global parameters p(0, 2 {a}) given an ensemble of images. As discussed in Sec. 2, building per-image posterior estimators ignores the joint influence of the global parameters over the entire dataset, potentially reducing the identifiability of the local parameters. If a combined posterior for the global parameters is desired, combining per-image estimators into global posterior can also be error-prone as per-observation errors accumulate, especially if flexible density estimators (e.g., normalizing flows) are used, as we do here. We train a hierarchical neural posterior estimator with datasets containing to 25 lensed images, using a ResNet-18 (He et al., 2016) CNN as the per-event feature extractor. App. A describes further details related to model training.
Results Figure 6 shows exemplary posterior inference results on one of the local parameters, the Einstein radius 0g (left) and one of the global paramcters, the amplitude of the clump mass function a (right), on test datasets. The red lines show posteriors using the hierarchical estimator, while the blue dotted lines show results of using per-image estimator. We can see that the posterior mass concentrates around the true value (vertical dashed) as more images are analyzed. On the other hand, the local parameter values are very similar to those obtained without simultaneously constraining dataset-wide global parameters. This is because, in this case, sub-clumps have a minimal (sub-percent level) effect on the lensing image, and the degeneracy between the local and global parameters is weak.
Local parameter inference
Global parameter Inference
25
25
20
Local only retwork
mieearchicl network
0.5
21 0.5
21
immgge
15
Thuth
20
pla (e))
0.4
17 0.4
17
p0h,
0.3
13 0.3
13
30
0.2
0.1
1.50
1.75
2.00
2.25
2.50
1.6
1.8
2.0
2.2
0.0
0
10
0.0
10
20
0.5
0.4
images
piby [x]
20
15
ple (29)
0.4
1.5
0.3
13
10
0.2
0.2
0.1
6.1
0.0
0.0
Og. Olocal param.)
8: (local param.)
(9lobbal param.)
(global param.)
Figure 6: Example results of inference on local (left) as well as global (right) gravitational lensing parameters using the hicrarchical simulation-basee inference model, with the posterior on global parameter shown for various cardinalities as up to 25 images are analyzed. Concentration of the posterior mass around the true global parameter value (vertical dashed lines) is seen as more lensed images are analyzed, and the postcrior distributions on the local parameters are seen to be qualitatively consistent with the true parameter values used for simulation.
Statistical coverage We perform a test of posterior statistical coverage (described in detail in Hermans et al. (2021)) for the gravitational lensing example, serving as a validation check of the learned posterior estimator across cardinalities. Using the same notation as in the main text, let 0 be the global parameter of interest and {x) be datasct. Denoting our learned postcrior density estimator as p(0 {x)). For a confidence level 1 - a, the expected coverage probability is
Big, (el [1e (0 (1 -a))]
(13)
11

---

Published in Transactions on Machine Learning Research (02/2024)
where ppC (x)) (1 gives the - highest posterior density interval (HDPI) of the estimator {x]] and le() is an indicator function mapping samples that fall within the HDPI to one. Given N sampled datasets from the joint distribution (0", (xr) ~ p(0,x), x), the empirical expected coverage for the posterior estimator p(0) {x}) is
N 1e (0" E (10)t))(1aa))
Σ 2=1
(14)
N
The nominal expected coverage is the expected coverage when p(01=)) -(0)=)) and is equal to the confidence level 1 a. In Fig. we show the results of a lest of statistical coverage for the gravitational lensing example, showing empirical against nominal coverage probability for the same posterior estimator evaluated over different cardinalities (corresponding to the different lines shown), using 100 test samples. If an estimator produces perfectly calibrated posteriors, its empirical expected coverage probability and nominal expected coverage probability match (dashed diagonal line). The amortized posteriors obtained shown to be well-calibrated across different cardinalities.
Coverage for substructure fraction a
1.0
Empirical coverage
Conservative
0.6
0.4
Overconfident
0.2
Images in set
5
20
10
25
0.0
0.0
0.2
0.4
0.6
0.8
1.0
Nominal coverage
Figure 7: Test of statistical coverage for the gravitational lensing example, showing empirical against nominal coverage probability. If an estimator produces perfoctly calibrated posteriors, its empirical expected coverage probability and nominal expected coverage probability match (dashed diagonal line). The amortized posteriors obtained are shown be well-calibrated across different cardinalities (corresponding to the different lines).
6 Conclusions
We studied the problem simulation- based inference over event ensembles in the presence of hierarchical structures in the underlying forward model. In this sctting, common in the physical scienoes, hicrarchy-awand parameter inference requires dataset-wide approach. For this purpose, introduced neural estimators for likeihhod-rrttio and posteriors, illustrated their use for Bayesian as well frequentist analysis in situations typical of astro- and particle physics, and substantiated correct inference performance in these scenarios. Our methods do not require explicit parameterization in terms of marginalized or profiled muisance parameters, making them complementary to existing techniques for dataset-wide inference. Once trained, they provide fully amortized inference over datasets of varying cardinality and offer significant speedups over traditional frequentist and Bayesian methods, making them suitable also high-throughput applications.
Our methods apply situations where the cardinality N of the observed datasct follows an arbitrary, but fixed, distribution p(N 8) (cf. equation 2). It is currently an open question to what extent estimators trained pı (N(N] 0) achieve statistically sound performance when applied to datasets
Limitations
12

---

Published in Transactions on Machine Learning Research (02/2024)
corresponding to different distribution P2 0), encoding, for example, the accumulation of additional data. We leave this extension for future investigation, but note that insights from recent work Berman et al. (2022); Chen et al. (2019) describing Bayesian updating as a dynamical system may offer a promising way to address this important problem.
Acknowledgments
We made extensive use of the Einops (Rogozhnikov, 2022), Jax (Bradbury et al., 2018), Jupyter (Kluyver et al,, 2016), Matplotlib (Hunter, 2007), nflows (Durkan et al., 2020), Numpy (Harris et al., 2020), PyMC5 (Salvatier et al., 2016), PyTorch (Paszke et al., 2019), PyTorch Lightning (Falcon et al., 2020), and Scipy (Virtanen et al., 2020) packages.
LH is supported by the Excellence Cluster ORIGINS, which is funded by the Deutsche Forschungsgemeennschallll (DFG, German Research Foundation) under Germany's Excellence Strategy EXC-2094-390783311 SM is supported by the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, is based upon work supported by the U.S. Department of Energy, Office of Science, Office of High Energy Physics of U.S. Department of Energy under grant Contract Number DE-SC0012567. CP acknowledges support through STFC consolidated grant ST/W ST/W000571/L. PW is supported by the Grainger Fellowship at the University of Chicago. This work was initiated at the Aspen Center for Physics, which supported by National Science Foundation grant PHY-1607611. We thank the Munich Institute for Astro-, Particle and BioPhysics (MIAPbP) for their hospitality during the completion of this project.
This material
References
Abhinav Agrawal and Justin Domke. Amortizod variational inference for simple hicrarchical models. Advances in Neural Information Processing Systems, 34:21388- 01111888-1199,,,,, 2021.
Noemi Anau Montel, Adam Coogan, Camila Correa, Konstantin Karchev, and Christoph Weniger. Estimating the warm dark matter mass from strong lensing images with truncated marginal neural ratio estimation. Monthly Notices of the Royal Astronomical Society, 518(2))2746 2760, 2023.
David S Berman, Jonathan J Heckman, and Mare Klinger. On the dynamics of inference and learning. arXiv preprint XrXe::2204. 12939, 2022.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnn and Qiao Zhang, JAX: composable transformations of Python+NumPy programs, 2018. URL http://gi com/ mom/google//arr
Johann Brehmer, Siddharth Mishra-Sharma, Joeri Hermans, Gilles Louppe, and Kyle Cranmer. Mining for dark matter substructure: Inferring subhalo population properties from strong lenses with machine learning. The Astrophysical Journal, 886(1):49, 2019.
Xinshi Chen, Hanjun Dai, and Le Song. Particle flow Bayes' rule. In International Conference on Machine Learning, pp. 1022-1031. PMLR, 2019.
Adam Coogan, Noemi Anau Montel, Konstantin Karchev, Meiert W Grootes, Francesco Nattino, and Christoph Weniger. One never walks alone: the effect of the perturber population subhalo measurements in strong gravitational lenses. arXiv preprint arXi::2209.099: 09911. 2022.
Kyle Cranmer, Juan Pavez, and Gilles Louppe. Approximating likelihood ratios with calibrated discriminative classifiers. arXiv preprint arXiv:1506 500.00.., 2015.
Kylc Cranmer, Johann Brchmer, and Gilles Louppe. The frontier of simulation-basee inference. Proceedings of the National Academy of Sciences, 117(44):30055530000 2020.
Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. nflows: normalizing flows in PyTorch, November 2020. URL btttps://doi.rrgeee orw.iiorr...... http:/ooorgo....//eeeeeeeee 4296287.
13

---

Published in Transactions on Machine Learning Research (02/2024)
William Falcon et al. swwordliihhiinng/ptoorollgggggii 0.7.6 release, May 2020. URL https://doi.org 10. .0.528//zeeooo..330 3828935..
Tomas Geffner, George Papamakarios, and Andriy Mnih. Score modeling for simulation-based inference. arXiv preprint arXit:2209.12223 2022.
Charles R. Harris et al. Array programming with NumPy. Nature, (385(785))3377333 September 2020. doi: 10.10888/411566-022-4444444 URL https://doi. ort...rg/oo. 11110//24188-0002264999
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and patlern recognition, pp. 770 778, 2016.
Lukas Hcinrich. Learning optimal test statistics in the presence of nuisance paramcters. arXiv preprint arxiv:2203. 13079, 2022.
Joeri Hermans, Arnaud Delaunoy, François Rozet, Antoine Wehenkel, Volodimir Begy, and Gilles Louppe. A trust crisis in simulation based inference? your posterior approximations can be unfaithful. arXiv preprint ar NKi:211000655) 2021.
J. D. Hunter. Matplotlib: A 2d graphics environment. Computing In Science & Engincering, 9(3):90 -50 2007.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv:14122 11..1112.688888 6980, 2014.
ar preprint
Thomas Kluyver et al. Jupyter notebooks- - publishing format for reproducible computational workflows. In ELPUB, 2016.
Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
Benjamin Miller, Christoph Weniger, and Patrick Forré Contrastive neural ratio estimation. Advances in Neural Information Processing Systems, 35:3262 3278, 2022.
Noemi Anau Montel and Christoph Weniger. Detection is truncation: studying source populations with truncated marginal neural ratio estimation. arXie preprint arXiv:2211.04291 2022.
Benjamin Nachman and Jesse Thaler. Learning from many collider events at once. Physical Review D, 103 (11):116013, 2021.
J Neyman and Pearson ES. On the problem of the most efficient tests of statistical hypotheses. Philosophical Transactions of the Royal Society of London. Series A, Containing Papers of Mathematical or Physical Character231:289 337, 1933.
George Papamakarios and lain Murray, Fast e-free inference of simulation models with bayesian conditional deusity estimation. Advances in neural information processing systems, 29, 2016.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. Advances in neurol information processing systems, 30, 2017.
Adam Paszke et al. Pytorch: H. Larochelle, A. Beygelzimer, F. d' Alché-Buc, F. Fox, and R. Garnett (eds.), Advances in Neural Informa- tion Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019. URL http://papers. neurips. cc/paper/9015 1atpaa//9001-yytcch-haa
An imperative style, high-performance deep learning library. In H. Wallach,
perforranrrrreeeeeeee deep- learning- .pff.
Mary Phuong and Marcus Hutter. Formal algorithms for transformers. arXiv preprint arXiu:2207.09233. 2022.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference machine learning, pp. 1530 1538.. PMIRR, 2015.
14

---

Published in Transactions on Machine Learning Research (02/2024)
Pedro Rodrigues, Thomas Moreau, Gilles Louppe, and Alexandre Gramfort. HNPE: Leveraging global
parameters for neural posterior estimation. Advances in Newral Information. Processing Systems, 34: 13432 13443, 2021.
Alex Rogozhnikov. Einops: Clear and reliable tensor manipulations with einstein-like notation. In International Conference on Learning Representations, 2022. URL https://openrevieee http:////ppaaeeeeeeeeeee//ooraaa
J Salvatier, TV Wiecki, and Fonnesbeck C. Probabilistic programming in Python using PyMC3. PeerJ Computer Science 2:e55, 2016.
The ATLAS Collaboration. detailed map of Higgs boson interactions by the ATLAS experiment ten years after the discovery. Nature 607, 52-58 55, 2022.
The CMS Collaboration. A portrait of the Higgs boson by the CMS experiment ten years after the discovery. Nature 607, 60-68, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in nearal information processing systems, 30, 2017.
Pauli Virtanen et al. SciPy 1.0: Fundamental Algorithus for Scientific Computing in Python. Nature Methods, 2020. doi: https://doi.org/
019-0686-2.
Sebastian Carner- Ji Won Park, Simon Birrer, Philip J Marshall, Aaron Roodman, Risa H Wechsler, LSST Dark Energy Science Collaboration, et al. Hierarchical inference with bayesian neural networks: An application to strong gravitational lensing. The Astrophysical Journal, 909(2):187, 2021.
Sebastian Wagner-Carena Jelle Aalbers, Simon Birrer, Ethan O Nadler, Elise Darragh-Ford, Philip J Marshall, and Risa H Wechsler. From images to dark matter: End-to-end inference of substructure from hundreds of strong gravitational lenses. arXiv preprint vrXiv:2203.00690 2022.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
A Additional details on experiments
A.1 Details of 'Simple varti-v normal likelihood' experiments
The deep set as well as transformer multivariate normal posterior estimators trained on 50,000 sequences {a) with batch size of 128, withholding 10% of the samples for validation. Sequences consist of up to 200 elements, with each element having 15 features (3 multivariate normal dimensions with 5 draws per dimension). The means across the 3 dimensions of the multivariate normal are the parameters of interest 0, and we target a posterior distribution over these.
glob A 3-layer MLP with RELU activations and hidden size 512 is used as the per-event encoder se which generates 256-dimensional per-event embeddings which are either aggregated (in the deep set model) or fed through transformer layers. An analogous decoder or prediction network Is then outputs the means and log-standard deviations, which are used along with the ground truth values to train the Gaussian density estimator (0 {x)).
The estimators are trained using the AdamW (Loshchilov & Hutter, 2019; Kingma & Ba, 2014) optimizer with initial learning rate 10-4 and cosine anncaling over 100 epochs. The checkpoint corresponding to the lowest validation loss used. Evaluation performed on 500 new test samples, evaluating the posterior estimator for sequences with cardinality varied from 1 to 200 for each test sequence.
15

---

Published in Transactions on Machine Learning Research (02/2024)
A.2 Details of 'Mixture models in particle physics: frequentist treatment' experiments
Hyperparameters The hyperparameters chosen such that at the nominal parameter values f 8, the total number of expected background events is n, = 100 events, and the total number of expected signal events is n, = 10. The hyperparameters of the Gaussian mixture components are μ4 = -7,0, = 2 and N4 = 0, 0.
Minibatch sampling The training algorithi for learning the test statistic for frequentist use requires choosing a weighting function in the (0, 8, 0y)-space. For each minibatch the (y = 0)-labeled instances are sampled from the simulator {a} ~ p6x 0.) at (0 = 8g 02 = 0u,0). The parameter of interest. chosen randomly from a uniform distribution, 0 Unif(3, 7), and a random nuisance parameter value is sampled from Bg (ni 2.0). The (y 1)-labeled instances are sampled from (0 80 = A,8, = 0v, where A Unif 2.0) and 0n Cnif 2.0). Overall, each minibatch has N = 100 positive instances and
N = 100 negative instances and the former are split equally between the two 0 = 6u = A cases.
Training For each minibatch sampled for 0 = to. the instances {z); are evaluated with the deep set conditioning parameter set to 0 = 0n and the loss is evaluated as standard binary cross-entropy loss averaged over the minibatch. The network parameters are optimized with the Adam optimizer with a learning rate of dolem = 10-4 4 for the per-element network parameters and Aget- = 10-3 for the set-wide network parameters. Training procoods for 30000 steps.
A.3 Details of 'Mixture models in particle physics: Bayesian inference for a narrow resonance' experiments
Choice of distributions and priors The model hyperparameters are chosen to be a, = 0.1 for the signal and (4 = 0, σb = 1 for the background. The mixture coofficient 0 is subject to a uniform prior, O ~ Unif(O, 1), and the nuisance parameter 0, has a Gaussian prior p(0.) = No(O2/u = 1, = 2).
Training of the deep set The networks implementing the per-event encoder gglob and decoder oglo are comprised of two dense layers with 128 hidden units and GELU activation functions; the event encoding consists of 64 features. 500k example datasets are generated from the prior distributions over 6g and 0; these datasets are used to train for 300 epochs with the Adaml optimizer, a batch-size 256 and an initial learning rate of 10-3 Cosine-annealing is used over 300 epochs.
The evaluation is performed on independent datasets not used during training.
Markov Chain Monte Carlo We use the implementation of the NUTS algorithm available in PyMC5 Salvatier et al. (2016). For each inference run, two Markov Chains of 50k samples each are sampled. The first 12k samples in each chain are discarded and not used for the estimation of the posterior.
A.4 Details of 'Astrophysics example: Strong gravitational lensing' experiments
We generate 5000 ensembles (drawn from the same global parameters ggabb = {a, each containing lensed images with varying local parameters 0soc = (0x, Ar, Ay.g). We use a ResNet-18 He et al. (2016) convolutional neural network as the per-event encoder s to extract a 128-dim embedding vector from each image; half of the fcatures are concatenated with the true global parameters and used to condition masked autoregressive normalizing flows Papamakarios al. (2017); Rezende & Mohamed (2015) gloe characterizing the local parameters for each image. The other half of the features are averaged, concatenated with the (randomly-varying at training-timor cardinality of the image dataset, and passed through a 3-layer ML. decoder, subsequently conditioning a normalizing flow golub modeling the global parameters of interest.
The experiment is trained with a batch size 16 (each batch containing 16 x 25 400 images), using the Adamk optimizer with cosine-annealed learning rate starting at 10-4 for up to 100 epochs with early stopping, using the model with the best validation loss for evaluation. The sum of global- and normaiionn flow log- lieliieiisoods is used as the optimization objective.
16

---

Published in Transactions on Machine Learning Research (02/2024)
B Motivation for loss function for joint posterior inference
We motivate the loss function used in equation 7, which factorizes the joint posterior over the local and global parameters and 0 into global posterior given the set {o} and a product of posterions over local parameters additionally conditioned on the global parameters:
-6 f = log golob (a (0l00 g+
gloh (mr)) ))
log Uip (zs 8
Jnc (x.),
(15)
Given a hierarchical model with global parameters 0 and local parameters 2s and associated observed data {a), we aim to demonstrate the factorization
p(0, {z)H)z) p(H[z)) II 2(2))0,
(16)
We make two key assumptions, given the nature of common hierarchical models. (a) Given the observations (x), the global parameters 0 are conditionally independent of the local parameters {zi), which gives us
p00, {++)))) ((] {x)).
(17)
and (b) Each local parameter is conditionally independent of any other 2j given 0 and its respective observation Zi, which gives
1(( {x}) II p(z(00xx.).
(18)
Substituting the result from (b) into the expression from (a), obtain our targeted factorized form:
{({{))))))) = (x) II p(zz00zz))
(19)
justifying the loss function used in the hierarchical model.
17


Reviews:
[{"review_id": "oBzWaLO9Dx", "reviewer": "Reviewer_Rc6M", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- Well-motivated problems with clear use cases in particle physics and astrophysics.\n- Good empirical results showing the proposed methods can maintain inference quality while being significantly faster via amortized inference.\n\nWeaknesses\n- The paper does not contain a self-sustaining explanation of the proposed methods. This makes it hard to follow for the readers who have not worked on this problem.\n- There is a concern regarding the derivation of (4).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a neural simulation based inference method for hierarchical probabilistic models over event ensembles. The authors demonstrate that probabilistic inference in many implicit models with a hierarchical structure requires a dataset-wide approach. They introduce neural estimators for likelihood-ratios and posteriors for Bayesian as well as frequentist analysis. With experiments on synthetic, astro- and particle physics experiments. inference methods are generically substantially faster than traditional approaches,", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- The paper heavily relies on previous work. It is hard to follow because the relevant details are not included in the paper. Please provide a self-sustaining explanation of the method in the Section 2 and 3.\n- I could not follow the derivation of (4). For the third equality to hold, it requires $p(\\{x\\}) = \\prod_i p(x_i)$. This does not hold due to the existence of a dataset-wide latent variable $\\theta$. Marginalizing out $\\theta$ from $p(\\{x\\}, \\theta)$ will let different $x_i$ correlate with each other. Did I miss anything in the derivation?\n- On Page 3, it says that “it is not possible to combine marginal per-event posteriors…” From the probabilistic formulation, it is clearly possible to marginalize $\\theta_v$ out from the distribution. Could you provide more details about why it could not happen?\n- For the synthetic experiment in Section 3.1, it is no clear from the text what is inferred from data. What are $\\sigma_1$, $\\sigma_2$ and $\\sigma_3$?", "broader_impact_concerns": "This paper proposed neural simulation based inference methods for hierarchical probabilistic models over event ensembles. There are no obvious ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QV8MANhB5k", "reviewer": "Reviewer_sC5G", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The setup and method are quite clear, and appear sound to me. The experiments clearly list both the forward model and the inference procedure.\n\nThe method doesn't appear particularly novel, and I'm not really sure what the claims are. Below I comment on the three bullet points listed at the end of the introduction.\n1. \"We substantiate theoretically as well as empirically the fact that optimal inference in many implicit models with a hierarchical structure requires a dataset-wide approach, contrasted with the more common paradigm of combining implicit likelihood or posterior estimators associated with individual observations. We systematically derive conditions under which such a hierarchical approach is necessary for optimal inference, depending on how the model parameters are partitioned into local and global parameters of interest and nuisance parameters\"\nIf the authors mean that the variational objective for the global parameter does not factorize over the observations in the presence of global nuisance variables, I would consider this well-known. If there's more to it, I can't find this analysis in the paper. I assume that \"optimal inference\" is some frequentist term of art, rather than a claim that this inference is somehow the best.\n2. \"We introduce frequentist as well as Bayesian methods for dataset-wide learning that can be used to perform simulation-based inference over event ensembles and can deal with datasets of varying cardinality. We connect our insights to several common use cases in the physical sciences and show how popular simulation-based inference paradigms can be adapted for optimal dataset-wide learning. In particular, we introduce the first approach for end-to-end frequentist simulation-based inference targeting hierarchical set-valued data\"\nI don't see a falsifiable statement here, that could be supported by experiments. As for novelty, amortized inference for datasets of unbounded cardinality has been studied before, such as by Le at. al. [1]. I can't comment on the novelty of the frequentist approach.\n3. \"We show that our machine learning-based inference methods are generically substantially faster than traditional approaches, such as Markov Chain Monte Carlo (MCMC) methods, even when the likelihood is tractable, while giving consistent results. They also allow performing inference in “streaming” mode where, e.g., posterior estimates are efficiently updated in real-time as new observations are made without having to perform a re-analysis of the entire updated dataset.\"\nThat variational methods are generally faster than MCMC is true, but hardly surprising. If history serves me well, that was the original motivation for studying them. The proposed method can be applied in a streaming setting, but I don't see any experiments showing that it performs well in such circumstances.\n\nAdditionally, the authors use a set encoder architecture as their main method, and a transformer network as a baseline which factorizes the global variational objective over the observations. The writing suggests that somehow this is a result of using a transformer, but I don't see why there couldn't be a transformer inference network using a global objective that does not factorize.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose an amortized inference procedure for hierarchical models, It defines an inference network for individual observations, conditional on the global parameter, and a permutation-invariant global inference network. The posterior is factorized as the marginal posterior of the global parameter given all observations, and individual local latent variable posteriors conditional on the global parameter and the single local observation. The variational objective has a global term and a sum of local terms. The papers includes experiments on toy models, and various real models used in physics.\n\nI'm having a hard time understanding what exactly are the claims of this paper and how the experiments support them.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Could the authors please point out specifically what are the falsifiable claims in the paper and how their experiments support those?", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "aZGn6eru41", "reviewer": "Reviewer_pDeX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\nThe paper proposes a new method that seems well-motivated, and the experiments\nindicate that the proposed method is indeed able to provide meaningful\nposteriors over parameters in a range of hierarchical inference tasks. The\nexperiments also seem relevant and thereby the paper provides a good motivation\nfor the settings in which the method could be applied in.\n\n## Weaknesses\nThe biggest weakness to me is the structure of the paper. I listed more details\nunder \"Requested Changes\" below, but overall the text can in some parts feel\nvery unstrucutred, e.g. due to the many very small paragraphs that could be one,\nthe usage of paragraph headings as section headings, the\nseemingly-clear-yet-unclear structure in the experiments (e.g. where \"inference\"\ncan include \"training\", but might also not), or the very separate section 4.\nAnd more broadly, the presentation made it difficult for me to properly\nunderstand which parts of section 2 are background, and which parts are the\nactual proposed method.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces an SBI-based method for hierarchical inference of both local and global variables. The authors then evaluate the method on a range of experiments and show that the method is able to recover meaningful posteriors in such hierarchical settings.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- What is meant with \"optimal inference\"? This is used in many parts of the\n  paper and since I am not familiar with the specific usage of \"optimal\n  inferece\" I was quite confused by it's meaning here, as SBI in general is an\n  approximate method that does not provide the true posterior. The term is also\n  not introduced in the background section. Please explain what exactly is meant\n  with \"optimal inference\" and define it properly in the paper.\n- The structure of the paper could be more clear and it is difficult to distinguish\n  background from the proposed method. Maybe having a section 2 purely for the\n  background, and a new section 3 for the method, would be an option?\n- Why is section 4 a separate section, not an experiment? It does feel a bit out\n  of place in the current form and I was unsure how to interpret its meaning.\n- The experiments are all somewhat differently structured, even though the\n  paragraph headers suggest similarity. For example, some experiments describe the actual\n  training procedure, i.e. what data the model is trained on, others do not. And while this can be found in the appendix I believe, Section 3.2 does not link to to it, so by the content\n  of the section itself all training information is missing.\n- Visualizing not just one of the proposed methods would be helpful. This could be done either by\n  having a second figure, or maybe a visualization that includes both methods\n  and thereby also highlights their similarities and difference is possible.\n- The usage of paragraph headings: I believe they should only be used if the\n  title actually belongs to the paragraph itself, but in the paper the text\n  following the paragraph title consists of multiple paragraphs. So I strongly\n  suggest either merging the paragraphs into a single paragraph, or actually\n  using section headings. I believe that the former would be great in the\n  experiments section, as often text was split into multiple paragraphs even\n  though it clearly belongs together (e.g. section 3.3 \"Inference\"), or if it\n  does not then the second paragraph should get a separate heading.\n  On the other hand having actual subsections could be helpful e.g. in the\n  background section to make the structure clearer, and since TMLR does\n  not have a strict page requirement I do not really see any downside (also keep\n  in mind that this way the table of contents displayed by many PDF readers\n  becomes more accurate).\n- In many instances the paper writes \"Eq. equation [X].\". Please fix the\n  doubling.\n- Section 2.2, last paragraph: In the current form the comment on LSTMs feels\n  very out of place. I believe one way to fix this would be to move the \"Deep\n  sets in figure 1\" sentence into the first paragraph as then the overall\n  message of this paragraph becomes more clear.\n- On the transformer-based architecture: It is only used in the first experiment\n  in section 3.1, not in any other, and in this first experiment its performance\n  is worse than the deep sets-based mode. So I am wondering if having the\n  transformer in the paper is actually a helpful improvement, or if it is rather\n  an unneccessary part that distracts from the main messages?\n- Section 3.1: $\\mu_0, \\Sigma_0$ should be introduced more formally, e.g. by explicitly\n  writing (I presume) $p(\\theta) = N(\\mu_0, \\Sigma_0)$.\n- Section 3.3: What do you mean with \"we used a fixed batch of events of size\n  $N_0$, corresponding to the situation one might encounter in a real-time\n  inference task\". Why is this a more realistic scenario? How does this relate\n  to \"real-time\"? This was confusing to me.\n- Figure 5: What is meant with \"the local parameters are correctly recovered\"?\n  Looking at the left-most curve in the top left plot, this does not seem like a\n  true statement. Overall, \"correctly recovered\" seems like a difficult\n  statement in the context of Bayesian inference as we obtain a distribution\n  over parameters, so making this more precise would be good.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors have engaged in a thoughtful conversation with the reviewers and have improved the presentation of their manuscript. The reviewers agree that the work is correct and that the improved presentation clarifies the authors' claims. The only remaining point is the choice of the term `optimal` in the abstract and introduction. It also feels to me that the authors recognize that this can be misconstrued, hence their footnote on the first page. Since readers of the abstract will not have access to this footnote (and my general dislike for footnotes to define such critical aspects of an idea), I am recommending that the authors come up with a small modification that replaces the word optimal directly with what they mean — the effect leading to the footnote becoming superfluous. The definition of what the authors mean by optimal should be what is directly written in the text. I would suggest terms along the lines of "full-covariance approximate inference" or "hierarchy-aware posterior structure" but naturally leave this decision with the authors. I look forward to receiving their updated manuscript.
Decision: Accept (poster)