ADAVI: Automatic Dual Amortized Variational Inference Applied To Pyramidal Bayesian Models
Authors: Louis Rouillard
Demian Wassermann
Venue: ICLR 2022 Poster (2022)
Keywords: Bayesian inference
Hierarchical Bayesian Models
structured Variational Inference
Simulation Based Inference
Inference amortization
Neuroimaging

Abstract: Frequently, population studies feature pyramidally-organized data represented using Hierarchical Bayesian Models (HBM) enriched with plates. These models can become prohibitively large in settings such as neuroimaging, where a sample is composed of a functional MRI signal measured on 300 brain locations, across 4 measurement sessions, and 30 subjects, resulting in around 1 million latent parameters.

Such high dimensionality hampers the usage of modern, expressive flow-based techniques.

To infer parameter posterior distributions in this challenging class of problems, we designed a novel methodology that automatically produces a variational family dual to a target HBM. This variational family, represented as a neural network, consists in the combination of an attention-based hierarchical encoder feeding summary statistics to a set of normalizing flows. Our automatically-derived neural network exploits exchangeability in the plate-enriched HBM and factorizes its parameter space. The resulting architecture reduces by orders of magnitude its parameterization with respect to that of a typical flow-based representation, while maintaining expressivity.

Our method performs inference on the specified HBM in an amortized setup: once trained, it can readily be applied to a new data sample to compute the parameters' full posterior.

We demonstrate the capability and scalability of our method on simulated data, as well as a challenging high-dimensional brain parcellation experiment. We also open up several questions that lie at the intersection between normalizing flows, SBI, structured Variational Inference, and inference amortization.

ADAVI: AUTOMATIC DUAL AMORTIZED
VARIATIONAL INFERENCE
APPLIED To PYRAMIDAL BAYESIAN MODELS
Louis Rouillard
Demian Wassermann
Université Paris-Saclay, Inria, CEA Palaiseau, 91120, France louis. Douis.rruullardeeeaoo
Université Paris-Saclay. Inria, CEA Palaiscau, 91120, France
demi an wasserranne inriiiiiriiii nri inria. fr
ABSTRACT
Frequently, population studies feature oyramidlly--ogganized data represented us- ing Hierarchical Bayesian Models (HBM) enriched with plates. These models can become prohibitively large in settings such as neuroimaging, where sample is composed of a functional MRI signal measured on 300 brain locations, across 4 measurement sessions, and 30 subjects, resulting in around million latent param- eters. Such high dimensionality hampers the usage of modern, expressive flow- based techniques. To infer parameter posterior distributions in this challenging class of problems, we designed a novel methodology that automatically produces a variational family dual to target HBM. This variational family, represented as a neural network, consists in the combination of an attention- based hierarchical en- coder feeding summary statisties to set of normalizing flows. Our automatically- derived neural network exploits exchangeability in the plate-enriched HBM and factorizes its parameter space. The resulting architecture reduces by orders magnitude its parameterization with respect to that of typical flow-based rep- resentation, while maintaining expressivity. Our method performs inference on the specified HBM in amortized setup: to a new data sample to compute the parameters' full posterior. We demonstrate the capability and scalability of our method on simulated data, well chal- lenging high-dimensional brain parcellation experiment. We also open several questions that lie the intersection between normalizing flows, SBI, structured Variational Inference, and inference amortization.
trained, it readily be applied
INTRODUCTION
Inference aims obtaining the posterior distribution X()]X) of latent model parameters 0 given the observed data X. In the context of Hicrarchical Bayesian Models (HBM), x(0 usually has no known analytical form, and can be of complex shape diffferent from the prior's (Gelman et al.) 2004). Modern normalizing-flot flows based techniques -universal density estimators- can overcome this difficulty (Papamakarios et al., 2019a| Ambrogioni et al.) 2021b). Yet. in setups such as neu- roimaging, featuring HBMs representing large population studies (Kong et al. 2019 Bonkhoff et al.) 2021) the dimensionality of 0 can go over the million. This high dimensionality hinders the usage of normalizing flows, since their parameterization usually scales quadratically with the size of the parameter space (e.g. Dinh al. 2017 Papamakarios et al., 2018 Grathwohl et 2018). Popu- lation studies with large dimensional features are therefore inaccessible to off-the- flow-based techniques and their superior expressivity. This can in turn lead to complex, problem-specife deriva- tions: for instance Kong et al.) (2019) rely technique. Such an analytical complexity constitutes a strong barrier to entry, and limits the wide and fruitful usage of Bayesian modelling in fields such neuroimaging. Our main aim is to meet that experimental need: how can we derive a technique both automatic and efficient in the context of very large, hierarchically-organii
Expectation Maximization (EM)
data?
Approximate inference features a large corpus of methods including Monte Carlo methods (Koller & Friedman 2009) and Variational Auto Encoders (Zhang al. 2019). We take particular inspiration

---

ADAVI working principle
X
Sirpciee
4. FF(M)
Grouna
81)
X
E
Simulated Data
Dual Amoriized Variational Family
4 *
Generative Pyramitil HBM
Automatic Derivston
Figure I: Automatic Dual Amortized Variational Inference (ADAVI) working principle. On the left is generative HBM, with 2 alternative representations: a graph template featuring 2 plates Po. Pi of cardinality 2, and the equivalent ground graph depicting a typical pyramidal shape. We note B1 = (P) the batch shape due to the cardinality of P1. The model features 3 latent RV X, and
T= >²1, and one observed RV = [ xx..2]. [x.21 , 22 We analyse automatically the structure of the HBM produce its dual amortized variational family (on the right). The hierarchical encoder HE processes the observed data X through 2 successive set transformers ST to produce encodings E aggregating summary statistics at different hierarchies. Those encodings then used to condition density estimators -the combination of normalizing flow F and link function I- producing the variational distributions for each latent RV.
from the field of Variational Inference (VI) (Blei et al. 2017), deemed to be most adapted to large parameter spaces. In VI, the experimenter posits variational family Q so approximate g(0) x X(0)X). In practice, deriving an expressive, yet computationally attractive variational family can be challenging (Blei al. 2017). This triggered trend towards the derivation of automatic VI techniques (Kucukelbir et I 2016;) Ranganath et al. (2013) Ambrogioni et al. (2021b). We follow that logic and present methodology that automatically derives variational family Q. Fig. [I] from the HBM on the left we derive automatically neural network architecture on the right. We aim at deriving our variational family Q the context of amortized inference (Rezende Mohamed, 2016 Cranmer et al.] (020). Amortization is usually obtained at the cost of an amortization gap from the true posterior, that accumulates on top of a approximation gap dependent on the expressivity of the variational family Q (Cremer et al. 2018). However, once an initial training overhead has been "paid for". amortization means that our technique can be applied to any number of data points to perform inference in few seconds.
Due to the very large parameter spaces presented above, our target applications aren't 't amenable to the generic flow-hased techniques described in Cranmer et al. (2020) Ambrogioni al. (20215). We therefore differentiate ourselves in exploiting the invariance of the problem not only through the design of an adapted encoder, but down to the very architecture of our density estimator. Specifi- cally, we focus on the inference problem for Hierarchical Bayesian Models (HBMs) (Gelman et al.) 2004 Rodrigues et al., 202T). The idea to condition the architecture of a density estimator by an analysis of the dependency structure of an HBM has been studied in (Wehenkel & Louppe 2020 Weilbach al., 2020), in the form of the masking of a single normalizing flow. With Ambrogioni et al. ()021b), we instead share the idea to combine multiple separate flows. More generally, our statie analysis of a generative model can associated with structured VI (Hoffman & Blei 2014] Ambrogioni et al., 222 b). Yet our working principles rather orthogonal: structured VI usually aims at exploiting model structure to augment the expressivity of a variational family, whereas we aim at reducing its parameterization.
Our objective is therefore to derive an automatic methodology that takes as input generative HBM and generates a dual variational family able to perform amortized parameter inference. This vari- ational family exploits the exchangeability in the HBM to reduce its parameterization by orders of magnitude compared to generic methods (Papamakarios al.] 2019b; Greenberg et al. 2019
2

---

Consequently, method can be applied in the context of large, Pyramidally-stuctuunt data, a challenging setup inaccessible to existing flow-based methods and their superior expressivity. We apply our method such large pyramidal setup in the context of neuroimaging (section (.5). but demonstrate the benefit of our method beyond that scope. Our general scheme is visible in Fig. [ figure that we will explain throughout the course of the next section.
Ambrogioni et al. 2021b).
2 METHODS
2.1 PYRAMIDAL BAYESIAN MODELS
We are interested in experimental setups modelled using saarihen Hierarchical Bayesian Mod- els (HBMs) (Kong al., 2019: Bonkhoff al. 202T). These models feature independent sampling from common conditional distribution at multiple levels, translating the graphical notion of plates Gilks al. 1994). This nested structure, combined with large measurements -such as the ones in IMRI- can result in massive latent parameter spaces. For instance the population study in Kong et al.) (2019) features multiple subjects, with multiple measures per subject, and multiple brain vertices per measure, for a latent space of around 0.4 million parameters. Our method aims at performing infer- ence in the context of those large plate-enriched HBMs.
Such HBMs can be represented with Directed Acyclic Graphs (DAG) templates (Koller & Fried- man 2009) with vertices -corresponding to RVs- {0.) !to...L and plates (Pplp=o...> We denote Card(P) the -fixed- candinality of the plate P, i.e. the number of independent draws from common conditional distribution corresponds to. In a template DAG, a given RV A can belong to multiple plates PA.... Pp. When grounding the template DAG into ground graph -instantiating the repeated structure symbolized by the plates P- 0 would correspond to multiple RVs of similar parametric form (0'h......p ). with in 1... Card(Ph). ...,Ip I... Card(Pp). This equiva- lence visible on the left on Fig. where the template RV I corresponds to the ground RVs [y), ???] We wish exploit this plate-ine ilatueinduced exchangeability.
We define the sub-class of models we specialize upon as pyramidal models, which are plate-enriched DAG templates with the following differentiating properties. First, we consider a single stack of This means that any RV 0 belonging to plate Pp also belongs to plates (Palosp- We thus don't I treat in this work the case of colliding plates (Koller & Friedman 2009). Second, we consider single observed RV Oo. with observed value X, belonging to the plate Po (with no other -latent- RV belonging to Po). The obtained graph follows a typical pyramidal struc- ture, with the observed RV the basis of the pyramid, as seen in Fig, D This figure features 2 plates Po and P,, the observed RV is X, at the basis of the pyramid, and latent RVs are T, and x at upper levels of the pyramid. Pyramidal HBMs delineate models that typically arise as part of population studies -for instance in neuroimaging- featuring a nested group structure and data observed at the subject level only (Kong et al. 2019: Bonkhoff et al. 2021).
the plates
The fact that we consider a single pyramid of plates allows us to define the hierarchy of an RV 0 denoted (0,r. An RV's hierarchy is the level of the pyramid is placed Due our pyramidal structure, the observed RV will systematically be at hierarchy 0 and latent RVs at hierarchies > O. For instance, in the example in Fig. the observed RV X is hierarchy 0, I is at hierarchy 1 and both X and x are at hierarchy 2.
Our methodology designed to process generative models whose dependency structure follows a pyramidal graph, and to scale favorably when the plate cardinality in such models augments. Given the observed data X, we wish to obtain the posterior density for latent parameters 0,, ... ,0, exploiting the exchangeability induced by the plates Po. ..., Pp.
2.2 AUTOMATIC DERIVATION OF A DUAL AMORTIZED VARIATIONAL FAMILY
In this section, we derive our main methodological contribution. We aim at obtaining posterior distributions for a generative model of pyramidal structure. For this purpose, we construct family of variational distributions Q dual to the model. This architecture consists in the combination of items. First, a Hierarchical Encoder (HE) that aggregates summary statistics from the data. Second, of conditional density estimators.
3

---

Tensor functions We first introduce the notations for tensor functions which we define in the spirit of Magnus & Neudecker 1999). We leverage tensor functions throughout our entire architecture to reduce its parameterization. Consider function f F - G, and a tensor Tr E F of shape B. We denote the tensor To E G resulting from the element-wise application of [over T as To ((15) (Tp) riffernnce to the programming notion of vectorization in Harris et al. (2020)). In Fig. ST(3)) and L, o F. (B.) are examples of tensor functions. At multiple points in our architecture, we will translate the repeated structure in the HBM induced by plates into the repeated usage of functions across plates.
Hierarchical Encoder For our encoder, our goal is to learn function HE that takes input the observed data X and successively exploits the permutation invariance across plates Po....,Pp In doing so, HE produces encodings E different hicrarchy levels. Through those encodings, our goal is to learn summary statistics from the observed data, that will condition our amortized inference. For instance in Fig. [
the application of HE over X produces the encodings E, and Eg-
To build HE, we need at multiple hierarchies to collect summary statistics across i.i.d samples from a common distribution. To this end we leverage SerTransformers (Lee et al.) 2019): an attention- based, ionnutaaionniivariatt architecture. We use SefTransformers to derive encodings across given plate, repeating their usage for all larger-rank plates. We cast the observed data X as the encoding Eg- Then, recursively for every hierarchy h = 1. ... P + 1, we define the encoding E, as the application to the encoding En-1 of the tensor function corresponding to the set transformer STA-1- HE( then corresponds to the set of encodings (EI.... Ep41} } obtained from the successive application of {ST, I.=....... If we denote the batch shape B, = CPrd x x... X Card
E = (Bx) (E..))
HE(X)= = (E,,.. Er+1}
In collecting summary statistics across the i.i.d. samples in plate Pa-1. we decrease the order of the encoding tensor En We repeat this operation in parallel on every plate of larger rank than the rank of the contracted plate. We consequently produce an encoding tensor E, with the batch shape B,, which is the batch shape of every RV of hierarchy h. In that line, successively summarizing plates Po. .....,pp. Pp,, of increasing rank results in encoding tensors EI...... Ep+1 of decreasing order. In
Fig. there are 2 plates Po and Pi. hence 2 encodings E, = ST (X) and E, = ST, (E,). E
is an order 2 tensor: it has batch shape of B, = Card(P)) -similar to T- whereas E2 is an order tensor. We can decompose E, [e] + [STo([X1.,, x12))), ST, ((,1,
Conditional density estimators We now will use the encodings E, gathering hierarchical sum- mary statistics the data X, to condition the inference on the parameters 0. The encodings (EA. will respectively condition the density estimators for the posterior distribution of parameters sharing their hierarchy {{0; Hicr(0.) = h) h
Consider a latent RV 8 of hierarchy h, = Hier (6,). Due to the plate structure of the graph, 6, can be decomposed in a batch of shape B,, Card(Pn.) x x... CPrd of multiple similar, conditionally independent RVs of individual size So.- This decomposition is akin to the grounding of the considered graph template (Koller & Friedman] 2009). A conditional density estimator is a 2-step diffcomorphism from a latent space onto the event space in which the RV 0, lives. We initially parameterize every variational density as a standard normal distribution in the latent space RSS, First, this latent distribution is reparameterized by conditional normalizing flow F; (Rezende) & Mohamed 2016; Papamakarios et al. 2019a) into distribution of more complex density in the space R0; The flow F; is a diffeomorphism in the space Se; conditioned by the encoding E Second, the obtained latent distribution is projected onto the event space in which 0, lives by the application of a link function diffeomorphism 4, For instance, if 0; is a variance parameter, the link function would map R onto R+* (4, = Exp as an example). The usage of F, and the link function 4 is repeated on plates of larger rank than the hicrarchy h, of 0,. The resulting conditional density estimator 4i for the posterior distribution XX) is given by:
Ui N ( Br, So, IB,, IB,,xSo) xSe.) 0, = 4, of (B., (B.) Eu.:En.) q. (0,: En.)
In Fig. r hn ",',"' > is associated to the diffeomorphism L F (1:.) This diffeomorphism is con- ditioned by the encoding E, Both I and E, share the batch shape 13 Card(P.). Decomposing the encoding E, = ef], ef is used to condition the inference on y', and eî for u X is associated to the diffeomorphism/ 2 o Fx, and to 4 Fx. both conditioned by Eg-
4

---

M
Non amortized (MF-/I)
Total wall time (s)
Amortized (ADAVI)
10
P2
M
Mr49
E
10a
Pi
M
ITC
Arho
Σ
10a
Po
Po
Po
II
10
102
10
examples
(a) Cumulative GPU training and inference time for non-amortized (MF-VI) and amortized (ADAVI) method.
(b) Experiment's HBMs from left right: Non-conjugate (NC) (section 3.2). Gaussian random effects (GRE) (3.1 Gaus-
sian mixture (GM) B.4). Multi-scale (MSHBM) (355).
Figure 2: panel (a): inference amortization the Gaussian random effects example defined in eq. (6): as the number of examples rises, the amortized method becomes more attractive; panel (b): graph templates corresponding to the HBMs presented as part of our experiments.
Parsimonious parameterization Our approach produces a parameterization effectively indepen- dent from plate cardinalities. Consider the latent RVs .1...., Oz. estimators have a parameterization quadratic with respect to the size of the space they are applied to (e.g- |aapamakarios et al. 2018). Applying single normalizing flow to the total event space of 01.-.. ......... ................................ Pg would thus result in O Se. Iy Card(Pp)] weights. But since we instead apply multiple flows on the spaces of size Sp, and repeat their usage across all plates Phu ... Pp, we effectively reduce this parameterization to:
Normalizing flow-based density
weights O
Σ S2.
As consequence, our method can be applied to HBMs featuring large plate cardinalities without scaling up its parameterization to impractical ranges, preventing computer memory blow-up.
2.3 VARIATIONAL DISTRIBUTION AND TRAINING
Given the encodings Ep provided by HE, and the conditioned density estimators q. define parametric amortized variational distribution as a mean field approximation (Blei et al. 2017):
(0 = (0: (X)) = II
(0,; Eng, ②)
(4)
i=l...
In Fig. we factorize (T, K, A|X) = 9, (T; x (A: E2) gr (m; E2). Grouping parameters as
V (x Φ), our objective is to have do X) X). Our loss an amortized version of the
classical ELBO expression (Blei et al. 2017 Rezende & Mohamed 2016):
V/"" arg min W
1 M
Σ m==1
log qu (0"| logp(X ('), X p(X), TTh qu (0
(5)
p(2) the sampling of according to the distribution p(z). We jointly train HE and qr, i = 1... to minimize the amortized ELBO. The resulting architecture performs amor- tized inference latent parameters. Furthermore, since parameterization is invariant to plate cardinalities, our architecture is suited for population studies with large-dimensional feature space.
Where denote
EXPERIMENTS
In the following experiments, we consider a variety inference problems on pyramidal HBMs. We first illustrate the notion of amortization (section We then test the expressivity (section 3.2 3.4
scalability (section 3.3) our architecture, as well as its practicality on a challenging neuroimaging experiment (section 3.5
5

---

Baseline choice In our experiments use as baselines: Mean Field VI (MF-VI) (Blei et al. 2017) is a common- practtcceecceeeeeeee method; (Sequential) Neural Posterior Estimation INPE-C. SNPE-C) Green berg et al. 2019) is a carructur-unwwaat likelihood-free method: SNPE-C results from the sequential -and no longer amortized- usage of NPEAC; C Total Latent Space Flow (TLSF) (Rezende & Mohamed (016) is reverse-KL. counterpoint to SNPE-C: both fit single normalizing flow to the entirety of the latent parameter space but SNPE-C uses a forward KL loss while TLSF uses a reverse KL loss; Cascading Flows (CF) (Ambrogioni al. 2021b) is our main point of comparison in this section. For relevant methods, the suffix -(N)A designates the (non) amortized implementation. More details related to the choice and implementation of those baselines can be found in our supplemental material.
structure-aware, prior-aware method: CF-A
3.1 INFERENCE AMORTIZATION
In this experiment we illustrate the trade-off between amortized versus non-amortized techniques (Cranmer et al. 2020 For this, we define the following Gaussian random effects HBM (Gelman
ct al.
2004) (see Fig. 2b GRE):
D, N = 2, 50
p~ (((DD,
G
1e Iu N(u,
M = l1ogo=1....
(6)
g=1...G
0a Gg. Ox 1.0, 0.2, 0.05
1901/p0 Nu oz) x= [2..m]
In Fig. 2a we compare the cumulative time to perform inference upon a batch of examples drawn from this generative HBM. For single example, non-amortized technique
be faster -and deliver a posterior closer to the ground truth- than an amortized technique. This is because the non- amortized technique fits solution for this specific example, and can tune it extensively. In terms of ELBO, on top of an approximation gap an amortized technique will add an amortization gap (Cremer et al., 2018). On the other hand, when presented with a new example, the amortized tech- nique can infer directly whereas the optimization of the non-amortized technique has to be repeated. As the number of examples rises, an amortized technique becomes more and more attractive. This result puts in perspective the quantitative comparison later performed between amortized and non-amortized techniques, that are qualitatively distinct.
3.2 EXPRESSIVITY IN A NON-CONJUGATE CASE
In this experiment, underline the superior expressivity gained from using normalizing flows used by ADAVI or CF- instead of distributions of fixed parametric form -used by MF-VI. For this consider the following HBM (see Fig. 2b
N, D = 10, 2
Ta: Tg = 0.5, 0.3
Gamma (Tp,ra)
8a Laplace ab) B = [b*]a=1...N
This example is voluntarily non-canonical: place ourselves in setup where the posterior distri- bution of a given an observed value from B has no known parametric form, and in particular is not of the same parametric form as the prior. Such example is called non-conjugate in Gelman al. (2004). Results are visible in table [.NC: MF-VI is limited in its ability to approximate the correct distribution it attempts fit the posterior distribution of the same parametric form as the prior. As a consequence, contrary to the experiments in section 3.3 and section 3.4 -where MF-VI stands strong baseline- here both ADAVI and CF-A are able surpass its performance.
Proxy to the ground truth posterior MF-VI plays the role of an ELBO upper bound in our experi- ments GRE (section 3.3). GM (section 3.4) and MS- (section 3.5). We crafted those examples to be conjugate: MF-VI thus doesn't it feature any approximation gap. meaning KL (L(()))lp(()))) ~ (g(( (f) )
()) approximately equal to the evidence of the observed data. As a consequence, any inference method with the same ELBO value ccaccllated over the same examples- as MF-VI would yield an approximate posterior with low KL divergence to the true posterior. Our main focus in this work are flow-based methods, whose performance would be maintained in non-conjugate cases, contrary to MF-VI (Papamakarios et al. (0119)). We further focus amortized methods, providing faster inference for multiplicity of problem instances, e.g. section 3.I VI is therefore not to be taken as part of a benchmark but as a proxy to the unknown ground truth posterior.
As such, its FL)OO = p( X)

---

Table Expressivity comparison on the non-conjugate (NC) and Gaussian mixture (GM) examples. NC: both CF-A and ADAVI show higher ELBO than MF-VI. GM: TLSF-A and ADAVI show higher ELBO than CF-A, but do not reach the ELBO levels of MF-VI, TLSF-NA and CF-NA. Are compared from left to right: ELBO median (larger is better) and standard deviation; for non- amortized techniques: CPU inference time for one example (seconds); for amortized techniques: CPU amortization time (seconds). Methods are ran over 20 random seeds, except for SNPE-C and TLSF-NA who were ran on 5 seeds per sample, for a number of effective runs of 100. For CF, the ELBO designates the numerically comparable augmented ELBO (Ranganath et al. 2016).
HBM
Type
Method
ELBO
Inf. (8)
Amo. (8)
NC
Fixed param. form
MF-VI
-21.0 (+: 0.2)
17
(section 3.2
Flow-based
CF-A
-17.5 (# 0.1)
220
ADAVI
-17.6 (# 0.3)
1,000
GM
Ground truth proxy
MF-VI
171 (# 970)
23
(section 3.4
Non amortized
SNPE-C
-14,800 (# 15,000)
70,000
TLSF-NA
181 (#t 680)
330
CF-NA
191 (+: 390)
240
Amortized
NPE-C
-27,000 (+ 19,000)
1300
TLSF-A
-530 (+ 980)
360,000
CF-A
-7,000 (# 640)
23,000
ADAVI
-494 (+ 430)
150,000
3.3 PERFORMANCE SCALING WITH RESPECT TO PLATE CARDINALITY
In this experiment, we illustrate our plate cardinality independent parameterization defined in sec- tion [.2] We consider 3 instances of the Gaussian random effects model presented in eq. (6). in- creasing the number of groups from G
G 30 and G 300. In doing so, augment the total size of the latent parametric space from to 62 to 602 parameters, and the observed data size from 300 3, 000 30, 000 values. Results for this experiment are visible Fig. (see also table On this example we note that amortized techniques only feature a small amortization gap (Cremer et al., 2018), reaching the performance of non-amortized techniques -as measured by the ELBO, using MF-VI& as an upper bound- at the cost of large amortization times. We note that the performance of (S)NPE-C quickly degrades as the plate dimensionality augments, while TLSF's performance is maintained, hinting towards the advantages of using the likelihood function when available. As the HBM's plate cardinality augments, we match the performance and amortization time of state-of-te thetthe-rrt methods, but do maintaining constant parameterization.
3.4 EXPRESSIVITY IN A CHALLENGING SETUP
In this experiment, we test our architecture challenging setup in inference: mixture model. Mixture models notably suffer from the label switching issue and from a loss landscape with multiple strong local minima (Jasra et al. 2005). We consider the following mixture HBM (see Fig. GM):
K, pr Og. o2 = 1, 1.0, 0.2, 0.05 G,L.D, N 3, 3,3.2,50
~ NN(,,
MD [mn)==1... 1=1...
(8a)
1 1 ~ N(u, (((1,,0)) o²))
M4.G [ng] 3 1
T E [0, 11 ~ Dir([K] x L)
11 1 I...G
(8b)
zan 9 [uq 1 Mix(= (2
X [. [1
9=1..6
Where Mix(TT, [pt...., PNl) denotes the finite mixture of the densities [p1...., .PN] with n the mix- ture weights. The results are visible in table T In this complex example, similar to TLSF-A we obtain significantly higher ELBO than CF-A, but we do feature an amortization gap, not not reaching the ELBOO level of non-amortized techniques. We also note that despite efforts (S)N failed to
7

---

Parameterization floer is butuer)
100
Inference quality (anger is betuw)
Time Newar is butter)
104
90
mehhode:
10a
101
CF-A
204
72 72.mm
300
ADAM
Non amortized
methods:
SXPE-C
10a
TLSF
CF-NA
107
10a
60
pracys
162
50
350
Figure 3: Scaling comparison on the Gaussian random effects example. ADAVI -in red- maintains constant parameterization as the plates cardinality goes up (first panel); it does so while maintaining its inference quality (second panel) and a comparable amortization time (third panel). Are compared from left to right: number of weights in the model; closeness of the approximate posterior to the ground truth via the ELBO median -that allows for comparable numerical range G augments; CPU amortization + inference time (s) for a single example -this metric advantages non- non-amortized methods. Non-amortized techniques using plain lines. MF-VI VI, in dotted lines, plays the role of the upper bound for the ELBO. Results for SNPE-C and NPE-C have to be put perspective, from G 30 and G 300 respectively both methods reach data regimes in which the inference quality is very degraded (see table Implementation details are shared with table
represented using dashed lines, and amortized techniques
reach the ELBO level of other techniques. We interpret this result the consequence of forward- KL-bbased training taking the full blunt of the label switching problem, as seen in appendix D.2 Fig. D.3 shows how our higher ELBO translates into results of greater experimental value.
3.5 NEUROIMAGING: MODELLING MULTI-SCALE VARIABILITY IN BROCA'S AREA FUNCTIONAL PARCELLATION
To show the practicality of our method in sigh-dimensional context, we consider the model pro- posed by Kong et al. (2019). We apply this HBM to parcel the human brain's Inferior Frontal Gyrus in functional MRI (iMRI)-based connectivity networks. Data is extracted from the Human Con- nectome Project dataset (Van Essen et al., 2012). The HBM models population study with 30 subjects and 4 large fMRI measures per subject, as seen in Fig. [6] MSHBM: this nested structure creates large latent space of = 0.4 million parameters and even larger observed data size of = 50 million values. Due to our parsimonious parameterization, described in eq. 44) we can nonetheless tackle this parameter range without memory blow-up, contrary to all other presented flow-based methods -C,, TLSF, NPE-C. Resulting population connectivity profiles can seen in Fig. 4 We in addition interested in the stability of the recovered population connectivity considering sub- sets the population. For this we are to sample without replacement hundreds of pouplopulations of 5 subjects from our population. On GPU, the inference wall time for MF-VI 160 seconds per sub-population, for mean log( ELBB)) of 28.6( 6.2.)) (across 20 examples, seeds per example). MF-VI can again be considered as ELBO upper bound. Indeed the MSHBM can be considered as a 3-level (subject, session, vertex) Gaussian mixture with random effects, and therefore features conjugacy. For multiple sub-populations, the total inference time for MF-V VI reaches several hours. On the contrary, ADAVI is an amortized technique, and as such features amortization time of 550 seconds, after which can infer on any number of sub-populations in a few seconds. The posterior quality similar: mean log(- (LBB)) of 29 29.0(=0.01). As shown in supplemental material -as a more meaningful comparison- the resulting difference in the downstream parcellation task is marginal (Fig. E.7
We therefore bring the expressivity of flow- methods and the speed amortized techniques to parameter ranges previously unreachable. This is due our plate inddnalltt-inddeendeenn pa- rameterization. What's more, our automatic method only necessitates a practitioner to declare the generative HBM, therefore reducing the analytical barrier entry there exists in fields such
8

---

Population Networks μr
Population Parcellation
Network 0
Network
Figure 4: Results for neuroimaging experiment. On the left, networks show the top 1% con- nected components. Network (in blue) agrees with current knowledge in semantic/phonologic pro- cessing while network (in red) agrees with current networks known in language production (Heim) iet al. 2009) Zhang et al. 2020). Our soft parcellation, where coloring lightens as the cortical point less probably associated with one of the networks, also agrees with current knowledge where more posterior parts are involved in language production while more anterior ones in seman- tic/phonological processing (Heim et al. 2009 Zhang et al. 2020).
roimaging for large-scale Bayesian analysis. Details about this experiment, along with subject-level results, can be found in our supplemental material.
DISCUSSION
Exploiting structure in inference In the SBI and VAE setups, data structure can be exploited through learnable data embedders (Zhang et al. 2019 Radev et al. 2020). We go one step be- yond and also use the problem structure to shape our density estimator: we factorize the parameter space of a problem into smaller components, and share network parameterization across tasks we know to be equivalent (see section 2.2 and 3.33) In essence, we construct our architecture not based on ground HBM graph, but onto its template, principle that could be generalized to other types of templates, such as temporal models (Koller & Friedman 2009). Contrary to the notion of black box, we argue that experimenters oftentimes can identify properties such as exchangeability in their experiments (Gelman et al., 2004). As our experiments illustrate (section 3.4] section (33) there much value in exploiting this structure. Beyond the sole notion of plates, static analysis of forward model could automatically identify other desirable properties that could be then leveraged
for efficient inference. This concept points towards fruitful connections be made with the field of lifted inference (Broeck et al. 2021 Chen et al. 2020).
Mean-Fielh approximation A limitation in our work is that our posterior distribution is akin mean field approximation (Blei et al. 2017): with the current design, no statistical dependencies can be modelled between the RV blocks over which we normalizing flows (see section 2.3 Re- grouping RV templates, we could model more dependencies at given hierarchy. On the contrary, our method prevents the direct modelling of dependencies between ground RVs corresponding to repeated instances of the same template. Those dependencies can arise as part of inference (Webb) (ct al., 2018). We made the choice the Mcan Field approximation to streamline our contribution, and allow for clear delineation of the advantages of our methods, not tying them up to method augmenting a variational family with statistical dependencies, an open research subject (Ambrogioni eet al. 2021b; Weilbach et al.) 2020 Though computationally attractive, the mean field approxima- tion nonetheless limits the expressivity of our variational family (Ranganath et al.) 2016 Hoffman & Blei 20144) We ponder the possibility to leverage VI architectures such the one derived by Ranganath et al. (2016); Ambrogioni et al. (2021b) and their augmented variational objectives for structured populations of normalizing flows such ours.
Conclusion For the delineated yet expressive class of pyramidal Bayesian models, we have intro- duced a potent, automatically derived architecture able perform amortized parameter inference. Through a Hierarchical Encoder, our method conditions a network of normalizing flows that stands as variational family dual the forward HBM. To demonstrate the expressivity and scalability of our method, we successfully applied it to a challenging neuroimaging setup. Our work stands as an original attempt to leverage exchangeability in a generative model.

---

ACKNOWLEDGMENTS
This work was supported by the ERC-SIG Nnu ID:757672.
We would like to warmly thank Dr. Thomas Yeo and Dr. Ru Kong (CBIG) who made pre- processed HCP functional connectivity data available us.
We also would like to thank Dr. Majd Abdallah (Inria) for his insights and perspectives regarding our functional connectivity results.
REPRODUCIBILIIY STATEMENT
All experiments were performed on a computational cluster with 16 Intel(R) Xeon(R) CPU E5-2660 v2 @ 2.20GHz (256Mb RAM), 16 AMD EPYC 7742 64-Core Processor (512Mb RAM) CPUs and NVIDIA Quadro RTX 6000 (22Gb). Tesla V100 (32Gb) GPUs.
All methods were implemented in Python. We implemented most methods using Tensorflow Prob- ability (Dillon et al. 2017). and SBI methods using the SBI Python library (Tejero-Canters Cantero et al. 2020
As part of our submission we release the code associated to our experiments. Our supplemental ma- terial furthermore contains entire section dedicated to the implementation details of the baseline methods presented as part of our experiments. For our neuromimaging experiment, we also provide section dedicated to pre-processing and post-processing steps
REFERENCES
Luca Ambrogioni, Kate Lin, Emily Fertig, Sharad Vikram, Max Hinne, Dave Moore, and Marcel van Gerven. Automatic structured variational inference. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 676-684 PMLR, 13-15 Apr htm.hmml
2021a. URL
mwr
Luca Ambrogioni, Gianluigi Silvestri, and Marcel van Gerven. Automatic variational inference with cascading flows. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 254-263. PMLR, 18-24 4 Jul 2021b. URL https: /ttp:e:://proceddinga. www.ediiinni.........mmmmmmmmmmm mlr prw.spesssseo...
ambrogioni21a. html
Eli Bingham, Jonathan P. Chen. Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D. Goodman. Pyro: Deep universal probabilistic programming. J. Mach. Learn. Res., 30:28:1-28:: 2019. URL htth:::
403. html
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational Inference: A Review for Statis- ticians. Journal of the American Statistical Association, 112(518):859 877, April 2017. ISSN 0162-1459, 1537-274X. doi: 10. (088/0062145992011111 abs/ / /1/./161..000 60670
URL
arXiv: 1601.00670.
Anna K Bonkhoff, Jae-Sung Lim, Hee-Joon Bae, Nick A Weaver, Hugo J Kuijf, Matthijs Biesbrock, Natalia Rost, and Danilo Bzdok. Generative lesion pattern decomposition of cognitive impairment after stroke. Brain Communications, 05 2021. ISSN 2632-1297.
doi: 10. 1005/oriiccom///ccal1100 URL Ecabl10 fcab110.
Guy van den Broeck, Kristian Kersting, Sriraam Natarajan, and David Poole (eds.). An introduction to lifted probabilistic inference. Neural information processing series. The MIT Press, Cambridge, Massachusetts, 2021. ISBN 978-0- 2-88--0222-4459-3
Yuqiao Chen, Yibo Yang. Sriraam Natarajan, and Nicholas Ruozzi. Lifted hybrid variational in- ference. In Christian Bessiere (ed.), Proceedings of the Twenty-Ninte International Joint Con- ference Artificial Intelligence, IJCAI-20, pp. 4237 4244. International Joint Conferences
10

---

Artificial Intelligence Organization, 7 2020. doi: 10.24963/ijcai.2220/558
URL
https:
//doi. 2422///0.2966331aaaa.222222
Main track.
Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-baset inference.
Proceedings of the National Academy of Sciencex, pp. 201912789, May 2020. ISSN 0027-8424, 1091-6490. doi:
Utthttp:// www pnas www.ppasss.gggoooooo//o//o/
10.1073/ pnas. 1912789117
Chris Cremer, Xuechen Li, and David Duvenaud. Inference Suboptimality in Variational Autoen-
coders. arXiv: 0.0555558 fes, stat], May 2018. URL http:::/aaxxiiv.rrr/aa////////// 03558
arXiv: 1801 .03558.
Joshua V. Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A. Saurous. TensorFlow Distributions. arXiv:1711.1066 [cs, stat], November 2017. URL http://axxxVvorrg/ab/////7 10604 arXiv: 1711.10604.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP.
1605 1605.000500033 fcs, stat), February 2017. URL. http::/arrxiv..orgbbbb///tteeee
08803
arXiv: 1605.08803.
David L. Donoho. High- nigh-dimmnssioal data analysis: The curses and blessings of dimensionality. In AMS CONFERENCE ON MATH CHALLENGES OF THE 21ST CENTURY, 2000.
Andrew Gelman, John B. Carlin, Hal S. Stern, and Donald B. Rubin. Bayesian Data Analysis. Chapman and Hal//CRC, 2nd ed. edition, 2004.
W. R. Gilks, A. Thomas, and J. Spiegelhalter. A Language and Program for Complex Bayesian Modelling. The Statistician, 43(1):169, 1994. ISSN 00390526. doi: 10,2307/2348941. URL
atttps:/vve www. ott.issortsssooeeeeoooooooeeeees
org/
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duve- naud. FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models. 10X10:1810.01367 [cs, stat], October 2018. htRhhtp::/ /ttp///aavvvv///a///aieeaaeaaaaa wttpg//as/111..001111
arXiv: 1810.01367.
David S. Greenberg, Marcel Nonnenmacher, and Jakob H. Macke. Automatic Posterior Transfor- mation for Likelihood--Free Inference. arXiv: /905. 0004877488 [cs, staf], May 2019. URL http:
// wwwpais/rrgaa////o./.. /tt//999500000
arXiv: 1905.07488.
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kem, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Ar- ray programming with NumPy. Nature, 585(7825):357-36, September 2020. doi: 1038/ ///////15566-2200 141138821111800026--------
Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre
s41586-026 020-2649-2222
URL
Stefan Heim, Simon B. Eickhoff, Anja K. Ischebeck, Angela D. Friederici, Klaas E. Stephan, and Katrin Amunts. Effective connectivity of the left BA 44, BA 45, and inferior temporal gyrus during lexical and phonological decisions identified with DCM. Human Brain Mapping, 30(2): 392-402, February 2009. ISSN 10659471. doi: 10.1002/hbm.20522
Matthew D. Hoffman and David M. Blei. Structured Stochastic Variational Inference. 4114
1rXiv:1404.411 [cs], November 2014. URL arXiv: 1404.4114.
Ekaterina Iakovleva, Jakob Verbeek, and Karteek Alahari. Meta-learning with shared amortized vari- ational inference. Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th International Conference Machine Learning, volume 1 119 of Proceedings of Machine Learning Research, pp. 4572 4582. PMLR, 13-18 Jul 2020. URL https: http::/:///poceeeiggg.... mlr. press/v119/ iakovl eva2 html

---

A. Jasra, C. C. Holmes, and D. A. Stephens. Markov Chain Monte Carlo Methods and the Label Switching Problem in Bayesian Mixture Modeling. Statistical Science, 20(1), February 2005.
ISSN 0883-4237. doi: 10.22440083223050000001
URL https: http:a/://pppoocteuueccccc
Ma oeteaoa.haaaattooooeeeeaaeeeeea-
10. 00.1217/885323300000000
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA. May 7-9, 2015, Conference Track Proceedings. 2015. URL http::
412. 6980
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. arXiv: [cs, arXiv: 1312.6114.
stat], May 2014. URL http:://aaxiiii httpaaawww...rrgrsaaaa/eaaa wwtpor/g///a////// 6114
Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. Adap- tive computation and machine learning. MIT Press, Cambridge, MA, 2009. ISBN 978-0-262- 01319-2.
Ru Kong, Jingwei Li. Csaba Orban, Mert Rory Sabuneu, Hesheng Liu, Alexander Schaefer, Nanbo Sun, Xi-Nian Zuo, Avram Holmes, Simon B. Fickhoff, and B. T. Thomas Yeo. Spatial topogra- phy of -ndivicual-ppecffic cortical networks predicts human cognition, personality, and emotion. Cerebral cortex, 29 6:2533-2551, 2019.
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M. Blei. Automatic Differentiation Variational Inference. arXiv: 103.0008830088888888 [cs, stat], March 2016. URL http: arXiv: 1603.00788.
// //t.xii.cgg/ab/////..aaa.a
708
Olivier Ledoit and Michael Wolf. well-conditioned estimator for large dimensimmniional covariance matrices. Journal of Multivariate Analysis, 88(2):365 411, 2004. ISSN 0047-259X. doi: https: (0bi:o///0..0065500772xX11000000
URL https: /twp///wwww//..// www.wwiieeeeedeeeeeet.t...ceeeee com/
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for iitentiin-bbsed permutation-invaate invariant neural networks. In Kama- lika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Confer- ence on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3744 3753. PMLR, 09-15 Jun 2019, URL httpttt:///ppceeddiiie wwweeediiinsmm......mmmmmmmm ptr.ppress/v9771eee11 html
Jan R. Magnus and Heinz Neudecker. Matrix Differential Calculus with Applications in Statis- tics and Econometrics. John Wiley, second edition, 1999. ISBN 0471986321 9780471986324 047198633X 9780471986331
George Papamakarios and lain Murray. Fast S' S'eesilloSS-free free Inference of Simulation Models with
Bayesian Conditional Density Estimation. arXiv; 1605.06376 [cs, stat], April 2018. htLttt
//arxiv.org/abas 605. 06376 arXiv: 1605.06376.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked Autoregressive Flow for Density Estimation. 0rXiv:/705.07055 [cs, stat], June 2018. UtU.htp://aarxii.oor//b/////
07C57 arXiv: 1705.07057.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. 19/2:099/2.020 [cs, star), December 2019a. URL
02762
arXiv:
1912.02762.
George Papamakarios, David C. Sterratt, and lain Murray. Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows. 1rXiv:180..0720 [cs, stat), January 2019b. URL httt:://aaxxxii wwt.ix/ww///.//........ abs//aas..11.. 07226
arXiv: 1805.07226.
Stefan T. Radev, Ulf K. Mertens, Andreass Voss, Lynton Ardizzone, and Ullrich Köthe. BayesFlow: Learning complex stochastic models with invertible neural networks. arXiv:2003.0 06281 [cs, stat), April 2020. _ pottt://arxivee http://awwwww..ogra/////3303003
arXiv: 2003.06281.
12

---

Rajesh Ranganath, Sean Gerrish, and David M. Blei. Black Box Variational Inference. 1rXiv:1401.00 [cs, stat], December 2013. URL http:///axii...//////// wwtp://gggg......//o. 0411.0118
arXiv: 1401.0118.
Rajesh Ranganath, Dustin Tran, and David M. Blei. Hierarchical Variational Models.
arXiv:1511.02560 [cs, stat], May 2016. URL http:p//aaaav.orrga/b///51.....1
arXiv: 1511.02388
Danilo Jimenez Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows.
wrXiv:1505.05777 [cs, stat], June 2016. URL http://aaxxv..org/aa///5533
arXiv: 1505.05770.
Pedro L. C. Rodrigues, Thomas Moreau, Gilles Louppe, and Alexandre Gramfort. Leveraging Global Parameters for Flow-based Neural Posterior Estimation. arXiv:2102.0644 102.06477 [cs, q-bio, star], April 2021.
06477
arXiv: 2102.06477.
Rui Shu, Hung Bui, Shengjia Zhao, Mykel Kochenderfer, and Stefano Ermon. Amortized infer- ence regularization. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gareett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran As- sociates, Inc., 2018. htthttss:: http:/::///pcceeeeeeeeeee wwwp:e/wwd..................eeee http://uuusssccpppppp///////////
181 (819932ff5cf44 44044444444444477744444440402020 Papee224444002222ppp.p...e...... pdf
Alvaro Tejero-Canters Jan Boelts, Michael Deistler, Jan-Matthis Lueckmann, Conor Durkan, Pe- dro J. Gonçalves, David S. Greenberg, and Jakob H. Macke. SBI toolkit for simulation-based inference. arXiv:2007. 091 [cs, q-bio, stat], July 2020. URL
2007. 09114
arXiv: 2007.09114.
Owen Thomas, Ritabrata Dutta, Jukka Corander, Samuel Kaski, and Michael U. Gutmann. Likelihood- inference by ratio estimation. arXiv:1611.1024 [stat], September 2020. URL http::/
org/ awwaabaslllll 10242 arXiv: 1611.10242.
D. C. Van Essen, K. Ugurbil, E. Auerbach, D. Barch, T. E. Behrens, R. Bucholz, A. Chang, L. Chen, M. Corbetta, S. W. Curtiss, S. Della Penna, D. Feinberg, M. F. Glasser, N. Harel, A. C. Heath, L. Larson-Prior, D. Marcus, G. Michalareas, S. Moeller, R. Oostenveld, S. E. Petersen, F. Prior, B. L. Schlaggar, S. M. Smith, A. Z. Snyder, J. Xu, and E. Yacoub. The Human Connectome Project: data acquisition perspective. Neuroimage, 62(4):2222-22) Oct 2012.
Stefan Webb, Adam Golinski, Robert Zinkov, N. Siddharth, Tom Rainforth, Yee Whye Teh, and Frank Wood. Faithful Inversion of Generative Models for Effective Amortized Infer- ence. 1712:: [cs, stat], November 2018. URL/ittp::/arxiivorr///ass
00287
arXiv: 1712.00287.
Antoine Wehenkel and Gilles Louppe. Graphical Normalizing Flows, arXiv:2006.02548 [cs, stat), October 2020. URL http:://aaxxv.ooo// http///as///00../ 02548
arXiv: 2006.02548.
Christian Weilbach, Boyan Beronov, Frank Wood, and William Harvey. Structured conditional con- tinuous normalizing flows for efficient amortized inference in graphical models. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Arti- ficial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 4441 14511 PMLR, 26-28 Aug 2020. URL httpt: httpss://proceedinggg mlr. http://cwe...ioo.........rrrraaa
wei Reilbach20a.110 html
Mike Wu, Kristy Choi, Noah Goodman, and Stefano Ermon. Meta-Amortized Variational Inference and Learning. arXiv: (ttp::/90..0/// 01950 [cs, stat), September 2019. URL http:://axxiv.rrrr////
abs/1902. 01950
arXiv: 1902.01950.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep Sets. arXiv: / 7703.06 14 [cs, stat], April 2018. URL URL http:://axxv.
Srg/abs/1703 06114 arXiv: 1703.061 14.
Cheng Zhang, Judith Butepage, Hedvig Kjellstrom, and Stephan Mandt. Advances in Variational Inference. IEEE Transactions Pattern Analysis and Machine Intelligence, 21088)200822022 August 2019. ISSN 0162-8828, 2160-9292, 1939-3539. doi: 10.1109/TPAMM22288888887777
URL URI
wttpt:ttp//............/p pttttp::///eeexplorrt wew.eeeeee..... wwtpeeeeer.oocce..oe/////e/eeeee /
13

---

Yizhen Zhang, Kuan Han, Robert Worth, and Zhongming Liu. Connecting concepts in the brain by mapping cortical representations of semantic relations. Nature Communications, 11(1):1877, April 2020. ISSN 2041-1723, doi: 10. 10.1008/:444667022-11
15804...
14

---

SUPPLEMENTAL MATERIAL
This supplemental material complements our main work both with theoretical points and experi- ments:
A complements methods section 2 We present the HBM descriptors needed for the
automatic derivation of our dual architecture;
B complements discussion section
3 We elaborate on various points including amorti-
zation;
complements the Gaussian random effects experiment described in eq. 6 We present
results mostly related to hyperparameter analysis;
complements the Gaussian mixture with random effects experiment (section 3.4) We explore the complexity of the example at hand;
complements to the MS-HBM experiments (section 3.5 We present some context for the experiment, toy dimensions experiment and implementation details.
justification and implementation details for the baseline architectures used in our experi- ments;
A COMPLEMENTS TO THE METHODS: MODEL DESCRIPTORS FOR AUTOMATIC VARIATIONAL FAMILY DERIVATION
This section is complement to section 2 HBM needed for our method to derive its dual architecture. This information is of experimental value, since those descriptors need to be available in any API designed implement method.
We formalize explicitly the descriptors of the generative
we denote pfa the plates the RV / belongs to, then the following HBM descriptors are the needed input to derive automatically our ADAVI dual architecture:
{=....
== (Pp).o....P
Card = {P, + #Pplo=o...
Hier {0, h4 min(p
Pp E plates(0.)))
(A.1)
Shape = {0: + So. Ji=0...
sevent
Link = {0: + (I: So, + Spyen) )i=
Where:
V lists the RVs in the HBM (vertices in the HBM's corresponding graph template);
P lists the plates in the HBM's graph template;
Card maps a plate to its cardinality, that to say the number of independent draws from common conditional density it corresponds to;
Hier maps a RV to its hierarchy, that to say the level of the pyramid placed at, or equivalently the smallest rank for the plates it belongs to;
Shape maps a RV to its event shape Sgyent Consider the plate-enriched graph template representing the HBM. A single graph template RV belonging to plates corresponds to multiple similar RVs when grounding this graph template. gevent is the potentially high- order shape for any of those multiple ground RVs.
Link maps RV 0 to its link function I. The Link function projects the latent space for the RV 0 onto the event space in which 0 lives. For instance, 0 is variance parameter, the link function would map IR onto * (1 = Exp as an example). Note that the latent space of shape Sn is necessary an order unbounded real space. therefore potentially implies reshaping to the high-order shape Sevent
Those descriptors be readily obtained from static analysis of generative model, especially when the latter is expressed in a modern probabilistic programming framework (Dillon et al. 2017 Bingham et al. 2019
15

---

B COMPLEMENTS TO OUR DISCUSSION
B.1 AMORTIZATION
Contrary to traditional VI, we aim at deriving variational family Q in the context of amortized in- ference (Rezende & Mohamed] 2016 Cranmer et al.) 2020). This means that, once an initial training overhead has been "paid for", our technique can readily be applied to a new data point. Amortized inference is an active area of research in the context of Variational Auto Encoders (VAE) (Kingma)
& Welling/ 2014 Wu al. 2019 (Shu et al. 2018 Takovleva al. 2020). It is also the original setup of normalizing flows (NF) (Rezende & Mohamed| [0166] Radev et al.) (2020), our technology of choice. From this amortized starting point, Cranmer et al. (2020); PPapamakarios et al. (2019b); Thomas et al. 2220);; Greenberg et al. (2019) have notably developed sequential techniques, refining a posterior -and losing amortization- across several rounds of simulation. To streamline our contri- bution, we chose not build upon that research, and rather focus on the amortized implementation of normalizing flows. But we argue that our contribution actually rather orthogonal to those: similar to to/Ambrogioni et al. (202bb) we propose a principled and automated way to combine several density estimators in a hierarchical structure. As such, our methods could be applied to a different class of estimators such VAEs (Kingma & Welling 2014]. We could leverage the SBI techniques and ex- tend our work into a sequential version through the reparametcrization of our conditional estimators q (see section 2.2 Ultimately, our method is not meant as an alternative to SBI, but a complement to it for the pyramidal class of problems described in section 2.1
We choose to posit ourselves eet al.] (2019) (see section (.5), the inference is performed on a specific data point. An amortized method could therefore appear more natural option. What's more, it is generally admitted that amortized inference implies an amortization gap from the true posterior, which accumulates on top of the approximation gap that depends on the expressivity of the considered variational family. This amortization gap further reduces the quality of the approximate posterior for a given data point.
amortized technique. Yet, in target experiment from Kong
Our experimental experience on the example in section 3.4 however makes us put forth the value that can be obtained from sharing learning across multiple examples, as amortization entitles (Cranmer al. 2020). Specifically. we encountered less issues related to local minima of the loss, canonical issue for MF-VI (Blei et al.) 2017) that is for instance illustrated in our supplemental material. We would therefore argue against the intuition that a (locally) amortized technique is necessarily waste- ful in the context of single data point. However, the results in table aable[and and underline there is much work to be done for amortized technique to reach the performance consistency and train- ing time of amortized techniques, especially in high dimension, where exponentially more training examples can be necessary to estimate densities properly (Donoho 2000).
Specializing for local parameter regime ass sequential method entitles (Cranmer al. 2020 could therefore make us benefit from amortization without too steep an upfront training cost.
B.2 EXTENSION TO A BROADER CLASS OF SIMULATORS
The presence of exchangeability in a problem's 's data structure is not tied to the explicit modelling of a problem as HBM: ZZheerr et al. 2018) rather describe this property as an permutation invariance present in the studied data. As a consequence, though our derivation is based on HBMs, we believe that the working principle of our method could be applied to a broader class of simulators featuring exchangeability. Our reliance on HBMs is in fact only tied to our usage of the reverse KI. loss (see section 2.3 readily modifiable implementation detail.
In this work, we restrict ourselves to the pyramidal class of Bayesian networks (see section 2. Going further, this class of models could be extended to cover more and more use-cases. This bottom-up approach stands at opposite ends from the generic approach of SBI techniques (Cranmer et al. 2020). But, our target experiment in 3.5 demonstrates, we argue that in the long run this bottom-ur upp approach could result in more scalable and efficient architectures, applicable to challeng- ing setups such as neuroimaging.
16

---

B.3 RELEVANCE OF LIKELIHOOD- FREE METHODS IN THE PRESENCE OF A LIKELIHOOD
As part of our benchmark, we made the choice to include likelihood-free methods (NPE-C and SNPE-C), based on forward KL loss. In our supplemental material (appendix C.4) also study the implementation of our method using a forward KL loss.
There is a general belief in the research community that likelihood-free methods are not intended to be as competitive as likelihood--aase methods in the presence of a likelihood (Cranmer et al. 2020). In this manuscript, we tried provide quantitative results to nourish this debate. We would argue that likelihood-fii methods generally scaled poorly to high dimensions (section (.3). The result of the Gaussian Mixture experiment also shows poorer performance in multi- modltl case, but we would argue that the performance drop of likelihood-free methods is actually largely due to the label switching problem (see appendix D.2). the other hand, likelihood-fret methods are dramatically faster train and can perform par with likelihood-basee methods in examples such as the Gaussian Random Effects for = (sce table 2). Depending on the problem at hand, is therefore not straightforward to systematically disregard likelihood-fri free methods.
As an opening, there maybe is more at the intersection between likelihood- free and likelihood- baseed methods than meets the eye. The symmetric loss introduced by Weilbach al. (2020) stands fruitful example of that connection.
B.4 INFERENCE OVER A SUBSET OF THE LATENT PARAMETERS
Depending the downstream tasks, out of all the parameters 8, experimenter could only be
interested in the inference of a subset ,, Decomposing 0 = O U O2 the goal would be to derive a variational distribution qi (01) instead of the distribution q (O), O2).
Reverse KL setup We first consider the reverse KL setup. The original ELBO maximized as part of inference is equal to:
ELBO(q) = log p(X) - 02)lp(),, 02|X]
(B.2)
loo [o Q θ2) log ((0,,
To keep working with normalized distributions, we get similar expression for the inference of θ only via:
EL)O( = E,g [lg (X,,((X,, θ1) log (01)
(B.3)
In this expression, p(X, θ1) ) is unknown: it results from the marginalization of O₂ in p. which is non-trivial to obtain, even via Monte Carlo scheme. As a consequence, working with the reverse KL does not allow for the inference over subset of latent parameters.
Forward KL setup Contrary reverse KL, in the forward KL setup the evaluation of p is required. Instead, the variational family is trained using samples (0, X) from the joint distribution p. In this setup, inference can be directly restricted over the parameter subset θ₁. Effectively, one wouldn't 't have to construct density estimators for the parameters 22, and the latter would be marginalized in the obtained variational distribution g(O)). However, as our experiments point out (section [3.3, section 3.4), likelihood-free training can be less competitive in large data regimes or complex inference problems. As a consequence, even if this permits inference over only the parameters of interest, switching forward KL loss can be inconvenient.
B.5 EMBEDDING SIZE FOR THE HIERARCHICAL ENCODER
in our architecture is the embedding size for the Set Transformer (ST) architecture (Lee al. 2019). The impact of the embedding size for single ST has already been studied in Lee et al. (2019). 9). as consequence we didn't devote any experiments to the study of the impact of this typee-paarameeer
An important
However, our architecture stacks multiple ST networks, and the evolution of the embedding size with the hierarchy could be an interesting subject:
17

---

it is our understanding that the embedding size for the encoding E,, should be increasing with:
the number of latent RVs 0 whose inference depends on Ex, i.e. the latent RVs of hierarchy
the dimensionality of the latent RVs 0 of hierarchy h the complexity of the inference problem at hand, for instance how many statistical moments need to be computed from i.i.d data points
experimentally, kept the embedding size constant across hierarchies, and fixed this stant value based on the aforementioned criteria (see appendix F2). This approach is prob- ably conservative and drives up the number of weights in HE
higher-hierarchy encodings are constructed from sets of lower-hierarchy encodings. Should the embedding size vary, it would be important not "bottleneck" the information col- lected at low hierarchies, even if the aforementioned criteria would argue for a low embed- ding size.
There would be probably experimental interest in deriving algorithms estimating the optimal em- bedding size at different hierarchies. We leave this to future work.
B.6 BOUNDS FOR ADAVI'S INFERENCE PERFORMANCE
When considering an amortized variational family, the non-amortized family with the same para- metric form can be considered as an upper bound for the inference performance -as measured by the ELBO. Indeed, considering the fixed parametric family g(0; V) for given data point X the best performance can be obtained by freely setting up the v parameters. Instead setting vrr f( -amortizing the inference- can only result worst performance. On the other hand the parameters for another data point X can then readily be obtained via 2 fx²) (Cremer al. 2018
In a similar fashion, it can be useful to look for upper bounds for ADAVT's performance. This is notably useful to compare ADAVI to traditional MF-VI (Blei et al. (017):
1. Base scenario: traditional MF-VI In traditional MF-VI, the variational distribution is MI-VI
Il Ui
Prior's 's parametric forn (0:):
Prior's parametric form can for instance be a Gaussian with parametric mean and variance;
4,
in non-conjugate cases, using the prior's parametric form can result in poor perfor- mance due an approximation gap, seen in section 3.2
due to the difference in expressivity introduced by normalizing flows, except in con- jugate cases, MF-VI is not upper bound for ADAVT's performance.
2. Superior upper limit scenario: normalizing flows using the Mean Field approximation
A family more expressive then MI- can be obtained via a collection of normalizing flows
combined using the mean field approximation: MI-NF = II Za
Normalizing Howw (0:):
every individual
more expressive than the corresponding Prior's parametric form. in non-conjugate case it would provide better performance Paa
4:
Normalising flow
pamakarios al.) 2019a):
since the mean field approximation treats the inference over each 0, as a separate problem, the resulting distribution q MF-NF is more expressive than q MF.
consider a plate-enriched DAG (Koller & Friedman 2009), template RV 0,, and of with j = 1... Car the corresponding ground RVs. In Mi-NE every 0f would be associated to separate normalizing flow;
consequently, the parametcrization of g MF-NF is linear with respect to Card (P). This is less than the quadratic scaling of TLSF or NPE-C -as explained in section 2.2 and ap- pendix F.2 But this scaling still makes q MP-NF not adapted to large plate cardinalities, all the more since the added number of weights -coresspondigg to full normalizing flow- per 0 is high;
18

---

this scaling is similar to the of Cascading Flows (Ambrogioni al. 2021b): CF
can be considered as the improvement of MF-F with statistical dependencies between the q;:
as far as we know the literature doesn't 't feature instances of the q MF-NF architecture. Though straightforward, the introduction of normalizing flows in variational family non-trivial, and for instance marks the main difference between CF and its prede- cessor ASVI (Ambrogioni al. 2021a).
3. Inferior upper limit scenario: non-amortized ADAVI At this point, is useful to con- sider the non-existent architecture A
compared to q ME.NF considering the ground RVs 0/ corresponding to the template RV Og each 0f would longer correspond different normalizing flow, but the same conditional normalizing flow;
each 8 would then be associated to a separate independent encoding vector. There wouldn't be need for our Hierarchical Encoder anymore referenced to in sec-
tion 2.2
for ME-NF the parameterization of ADAVI-NA would scale linearly with Card Each new 0f would only necessitate an additional embedding vector, which would make ADAVINNA more adapted to high plate cardinalities than MF-NF CF;
using separate flows for the of instead of shared conditional flow, ME-NF can be considered upper bound for gADAVI-NA''S performance;
due to the amortization gap, ADAVI-NA can be considered as an upper bound for ADAVT's performance. By transitivity, M is then an even higher bound for ADAVI's performance.
is to be noted that amongst the architectures presented above, ADAVI is the only architecture with parameterization invariant the plate cardinalities. This brings the advantage to theoretically being able to use ADAVI on plates of any cardinality, as seen in eq. (4). In that sense, our main claim is tied to the amortization of our variational family, though the linear scaling of AAAAVI-NA could probably be acceptable for reasonable plate cardinalities.
C COMPLEMENTS TO THE GAUSSIAN RANDOM EFFECTS EXPERIMENT: HYPERPARAMETER ANALYSIS
This section features additional results the experiment described in eq. (6 with G 3 groups. We present results of practical value, mostly related to hyperparameters.
C.1 DESCRIPTORS, INPUTS ADAVI
We can analyse the model described in eq. (6) using the descriptors defined in eq. [A.l Those descriptors constitute the inputs our methodology needs to automatically derive the dual architecture from the generative HBM:
y = [1, MG ,X}
P (Po,
Card (Po + N,Pi G}
(C.4)
Hier {1+ 2, 2, - 11
Shape {u+ (D.), + (D.), + (D.)}
Link = { + Identity, M Identity, X + Identity)
C.2 TABULAR RESULTS FOR THE SCALING EXPERIMENT
A tabular representation of the results presented in Fig. 3 can be seen in table
19

---

Type
Method
ELBO (10%)
weights
Inf. (s)
Amo. (s)
3
Grd truth proxy
MF-VI
2,45 (# 0.15)
10
5
Non amortized
SNPE-C
2.17 (#: 33)
45,000
53,000
TLSF-NA
2.33 (#: 0.20)
18,000
80
CF-NA
2.12 0.15)
15,000
190
Amortized
NPE-C
2.33 (+: 0.15)
12,000
920
TLSF-A
2.37 (+: 0.072)
12,000
9,400
CF-A
2.36 (+: 0.029)
16,000
7,400
ADAVI
2.25 (+ 0.14)
12,000
11,000
30
Grd truth proxy
MF-VI
24.4 0.41)
64
18
Non amortized
SNPE-C
-187 110)
140,000
320,000
TLSF-NA
24.0 0.49)
63,000
400
CF-NA
21,2 (+: 0.40)
150,000
1,800
Amortized
NPE-C
23.6 50)
68,000
6,000
TLSF-A
22.7 (+ 13)
68,000
130,000
CF-A
23.8 (+ 0.06)
490,000
68,000
ADAVI
23.2 (# 0.89)
12,000
140,000
300
Grd truth proxy
MF-VI
244 (# 1.3)
600
240
Non amortized
SNPE-C
-9,630 (# 3,500)
1,100,000
3,100,000
TLSF-NA
243 (# 1.8)
960,000
5,300
CR-NA
212 (+ 1.5)
1,500,000
30,000
Amortized
NPE-C
195 (#3 1
3,200,000
72,000
TLSF-A
202 (+: 120)
3,200,000
2,800,000
CF-A
238 (= 0.1)
4,900,000
580,000
ADAVI
224 (# 9.4)
12,000
1,300,000
Table 2: Scaling comparison the Gaussian random effects example (see Fig. 2b GRE). Methods are ran over 20 random seeds (Except for SNPE-C and TLSF: to limit computational resources us- age, those non-amortized computationally intensive methods were only ran 5 seeds per sample, for a number of effective runs of 100). Are compared: from left to right ELBO median (higher better) and standard deviation (ELBO for all techniques except for Cascading Flows, for which ELBO is the numerically comparable augmented ELBO (Ranganath et al.] (0166); number of train- able parameters (weights) in the model; non-amortized techniques: CPU inference time for one example (seconds); for amortized techniques: CPU amortization time (scconds). Results for NPE-C are extremely unstable, with multiple NaN results: the median value rather random and not necessarily indicative of good performance
C.3 DERIVATION ANALYTIC POSTERIOR
To have ground truth to which compare methods results, derive the following analytic posterior distributions. Assuming we know opr Tg. ozi
N
1
20
(C.5a)
N
120
NV
Tdp
(C.5b)
1
1: = T
Σ
(C.5c)
#=1
G p
p1 ~ N
9
C
Idp
p
(C.dd)
2
02
20

---

Where in equation C.5b we neglect the influence of the prior (against the evidence) on the posterior in light of the large number of points drawn from the distribution. We note that this analytical posterior is conjugate, argued in section 3.2
C.4 TRAINING LOSSES FULL DERIVATION AND COMPARISON
Full formal derivation Following the nomenclature introduced in Papamakarios et al. (2019a) there are 2 different ways in which we could train our variational distribution:
using a forward KL divergence, benefiting from the fact that we can sample from our generative model to produce a dataset (("" )m=L. 0 ~ p(0), X p(X (0). This is the loss used in most of the SBI literature (Cranmer al. (020). those based around the possibility to be likelihood-free, and have a target density p only implicitly defined by simulator:
V = arg min E Ex~ x) (p((
= arg min
Ex ~p 10- (1og log ou [0]
arg min E x()) B ( log v X]]
/ |xX) p(( log go (0) Xd0 dX
arg min
(C.6)
/
arg min
p(X, 8) log qu (d
arg min V
log qu (8mh x")
M
where 4Th p(0) (0), X ((1)))
using a reverse KL divergence, benefiting from the access to a target joint density p(X, 0). The reverse KL loss is amortized version of the classical ELBO expression (Blei al. 2017). For training, one only needs to have access to a dataset {X'm) m=11... X ~ p(X) of points drawn from the generative HBM of interest. Indeed, the 8m points sampled from the variational distribution:
V = arg min Ex~p(x) [KL (g) (1)X)]lp(00
V
= arg min Ex
E (0x) du (8) X) - log pl ((]
W
arg min E xp x) E [xx go
9u (0 - log p(X, 0) + log ((X)]
arg min Ex (x) [E - 10
(0 lox, 0)]]
| /
arg min
ou (0) log 9v (0) log p(X, (x,))]d0ddl dx
(C.7)
/ I
arg min arg min
(0))) qu (0 - p(o (0)xX,0)) 0)
1
x
log qv (0 X - log p(x (mmm 0m
M
m=m
where X p(X), 8m
qu (0
As it more uniquely fits our setup and provided better results experimentally, we chose to focus on the usage of the reverse KL divergence. During our experiments, we also tested the usage of the unregularized ELBO loss:
V arg min
x Σ
1og
M
(C.8)
m=1
where X ~ p(X), nT qo (0)x)
21

---

Mean of analytical KL divergences from the theoretical posterior (low is good)
Early stopping
After convergence
Loss
NaN runs
Mean
Std
Mean
Std
forward KL
0
3847.7
5210.4
2855.3
4248.1
unregularized ELBO
0
6.6
0.7
6.2
0.9
reverse KL
2
12.3
19.8
3.0
4.1
Table 3: Convergence of the variational posterior to the analytical posterior over early stopped
training (200 batches) and after convergence (1000 batches) for the Gaussian random effects exam- ple
This formula differs from the one of the reverse KL loss by the absence of the term qu (0 IXm"), ) and is converse formula to the one of the forward KL (in the sense that it permutes the roles of and p).
Intuitively, it posits our architecture as a pure sampling distribution that aims at producing points 8t in regions of high joint density p. In that sense, it acts as a first moment approximation for the target posterior distribution (akin to MAP parameter regression). Experimentally, the usage of the unregularized ELBO loss provided fast convergence to mode of the posterior distribution, with very low variance for the variational approximation.
We argue the possibility to use the unregularized ELBO loss as a warm-up before switching to the reverse KL loss, with the latter considered here regularization of the former. We introduce this training strategy as an example of the modularity of our approach, where one could transfer the rapid learning from one task (amortized mode finding) to another task (amortized posterior estimation).
Graphical comparison In Figure C.I we analyse the influence these 3 different losses on the training of our posterior distribution, compared to the analytical ground truth. This example is typical the relative behaviors induced on the variational distributions by each loss:
The forward KL provides very erratic training, and results after several dozen epochs (sev- eral minutes) with a careful early stopping in posteriors with too large variance.
The unregularized ELBO loss converges in less then 3 epochs (a couple dozen seconds), and provides posteriors with very low variance, concentrated on the MAP estimates their respective parameters.
The reverse KL converges in less 10 epochs (less than 3 minutes) and provides relevant variance.
Losses convergence speed We analyse the relative convergence speed of our variational posterior to the analytical one (derived in eq. (C.5a) when using the 3 aforementioned losses for training. To measure the convergence, we compute analytically the KL divergence between the variational posterior and the analytical one (every distribution being Gaussian), summed for every distribution, and averaged over a validation dataset size 2000.
We use a training dataset of size 2000, and for each loss repeated the training 20 times (batch size 10, 10 Hm samples per X for 10 epochs, resulting in 200 optimizer calls. This voluntary low number allows us to asses how close is the variational posterior to the analytical posterior after only brief training. Results are visible in appendix C.4] showing a faster convergence for the unregularized ELBO. After 800 more optimizer calls, the tendency gets inverted and the reverse KL loss appears as the superior loss (though we still notice a larger variance).
The large variance in the results may point towards the need for adapted training strategies involving Learning rate decay and/or scheduling (Kucukelbir et al. 2016), an extension that we leave for future work.
22

---

Daba
Ireand
sreesuatzee EURO
0.0
0.0
-0.3
E
-1.0
-1.0
-1.5
-15
-1.
0.0
Examel
33
83
13
00
05
15
03
13
-0.5
Eumpee
-10
-13
05
0.0
03
-20
1.0
05
0.
53
-20
10
0.5
0.0
Figure C.I: Graphical results on the Gaussian random effects example for our architecture trained using 3 different losses. Rows represent different data points. Left column represents data, with colors representing different groups. Other columns represent posterior samples for μ (black) and ut, 12, n². Visually, posterior samples μ² n²² n should be concentrated around the mean of the data points with the same color, and the black points should be repartitioned around the mean of the 3 group means (with a shift towards 0 due the prior). Associated with the posterior samples are analytical solutions (thin black circles), centered on the analytical MAP point, and whose radius correspond to 2 times the standard deviation of the analytical posterior: of the draws from posterior should fall within the corresponding circle.
23

---

C.5 MONTE CARLO APPROXIMATION FOR THE GRADIENTS AND COMPUTATIONAL BUDGET COMPARISON
In section 2.3 for the reverse KL loss, we approximate expectations using Monte Carlo integration. We further train our architecture using minibatch gradient descent, opposed stochastic gradient descent as proposed by Kucukelbir et al. (2016). An interesting syper-paramettizztion of our system resides in the effective batch size of our training, that depends upon:
the size of the mini batches, determining the number of X points considered in parallel
the number of 0m draws per X point, that approximate the gradient in the ELBO
More formally, we define computational budget as the relative allocation of a constant effective batch size batch size 9 draws per X between batch size and 0 draws per X.
To analyse the effect of the computational budget on training, we use dataset of size 1000, and run experiment 20 times over the same number of optimizer calls with the same effective batch size per call. Results can be seen in Fig. C.2 From this experiment we can draw the following conclusions:
we didn't 't witness massive difference in the global convergence speed across computational budgets
the bigger the budget we allocate to the sampling of multiple t point X (effectively going towards stochastic training in terms of the points X"), the more erratic the loss evolution
the bigger the budget we allocate to the X batch size, the more stable the loss evolution, but our interpretation is that the resulting reduced number of 6m draws per X augments the risk of an instability resulting NaN run
Experimentally, obtained the best results by evenly allocating our budget to the X batch size and the number of 1 draws per X point (typically, 32 and 32 respectively for an effective batch size of 1024). Overall, in the amortized setup, experiment stand counterpoint to those of Kucukelbir et al. (2016) who pointed towards the case of a single 4T draw per point X as their preferred papeepparametrization..
D COMPLEMENTS TO THE GAUSSIAN MIXTURE WITH RANDOM EFFECTS EXPERIMENT: FURTHER POSTERIOR ANALYSIS
This section is complement to the experiment described in section 3.4 we thus consider the model described in eq. (8a). We explore the complexity of the theoretical posterior for this experiment.
D.I DESCRIPTORS, INPUTS TO ADAVI
We can analyse the model described eq. Saj using the descriptors defined in eq. AA.I. Those descriptors constitute the inputs methodology needs automatically derive the dual architecture
24

---

butch: theta crows: 1000. Nati runs.
33000
batch: 10, thete crows: 100, Kahd rurse
batcle 100, thete 10, Nati
reversa losa
20000
15000
10000
batch
Training loss
70000
batch:
theta drows: 1000, Nahl runs
batch: 10, thete drawn: 120, Nahi rune
60000
batch: 100, thota draws: 10, NaN runs
2000, theta draws:
50000
Dontrs XL, loss
10000
batch
Figure C.2: Loss evolution across batches for different computational budgets. All experiments designed so that to have the same number of optimizer calls (meaning that batch size x epochs = 1000) and the same effective batch size (meaning that batch size 0 draws per X 1000). Every experiment is run 20 times, error bands showing the standard deviation of the loss at the given time point. Note that the blue line (batch size 1000 draws per X) is more erratic than the other ones (even after a large number of batches). On the other hand, the red line (batch size 1000, 0 draws per X) is more stable, but 19 out of 20 ultimately resulted in instability
25

---

from the generative HBM:
V = 1M,,
ML. ,11 x
PP).P1}
Card = (Po + N, G]
Hier (M + 2, M + 1, 1, +1,X
Shape (M + (T,, D), M (I,, D), IT (T,),X + (D,))
Link =
(D.9)
M > Reshape (I.D,) - (T,, D));
M4.G Reshape((LD)) - (T., D));
11 + SoftmaxCenteredd (L- 1.) + (L.)),
X + Identity
For the definition of the SoftmaxCentered link function, see Dillon et al. 2017
D.2 THEORETICAL POSTERIOR RECOVERY IN THE GAUSSIAN MIXTURE RANDOM EFFECTS MODEL.
We further analyse the complexity of model described in section 3.4
mixture com- Due to the label switching problem (Jasra et al.) 2005), the relative position of the L mixture com- ponents in the D space is arbitrary. Consider a non-degenerate example like the one in Fig. D.3 where the data points are well separated in blobs (likely corresponding the ponents). Since there is no deterministic way to assign component = 1 unequivocally to a blob of points, the marginalized posterior distribution for the position of the component be multi-modal, with -roughly- mode placed at each one of the 3 blobs of points. This posterior would be the same for the components than this 3-mode simplification, especially when the mixture components closer to each other in 2D (and the grouping of points into draws from a common component is less evident).
1 should
and = 3. In truth, the posterior even more complex
In Fig. D.3 note that technique doesn't recover this multi-modality in its posterior, and instead assigns different posterior components to different blobs of points. Indeed, when plotting only the posterior samples for the first recovered component / 1, all points are concentrated around the bottom-most blob, and not around each blob like the theoretical posterior would entail (see Fig. D.3 second row).
This behavior most likely represents a local minimum in the reverse KL loss that common to many inference techniques (for instance consider multiple non-mixing chains for MCMC in Jasra et al. 2005). We note that training in forward KL wouldn't I provide such a flexibility in that setup, as it would enforce the multi-modality of the posterior, even the cost of an overall worst result (as is the case for NPE-C and SNPE-C in table T Indeed, let's imagine that our training dataset features M' draws similar to the one in Fig. D.3 Out of randomness, the labelling of the blobs of points would be permuted across those M' examples. A forward- -L--triid-t density estimator would then most likely attempt to model a multi-modal posterior.
Though it is not similar to the theoretical result, we argue that our result is of experimental value, and close to the intuition one forms of the problem: using our results one can readily estimate the original components for the mixture. Indeed, for argument's sake, say we would recover the theoretical, roughly trimodal posterior. To recover the original mixture components, one would need to split the 3 modes of the posterior and arbitrarily assign a label to each one of the modes. In that sense, our posterior naturally features this splitting, and can be used directly to estimate the I, 3 mixture components.
E COMPLEMENTS TO THE NEUROIMAGING EXPERIMENT
This section is a complement to the experiment described in section 3.3 we thus consider the model described in eq. (E. 10a) and eq. (0. We present toy dimension version of our experiment,
26

---

*/u
Rall
-1
samp
-2
Pasterfor
camp compament
-2.2
15
-2.4
-2.8
samps
0.0
35
(zoom)
componont
-2.0
-0.5
-3.0
0.5
175
1.5
2.0
-35 0.5
100
08
0.8
0.4
Figure D.3: Graphical comparison for various methods on the Gaussian mixture with random effects example. First column represents non-degenerate data point, with colored points corresponding to [x ...., x I.N. [22.1 _,..,, x 22. [13. ...., 23.N1. Note the distribution of the points around multi-colored groups (population components), and 3 colored sub-groups per group (group com- ponents). All other columns represent the posterior samples for population mixture components f1, ... 44 (black) and group mixture components μ1 ... 1433 /45 M 1a ... ...... Sec- ond column represents the results of non-amortized MF-VI (best ELBO score across all random seeds): results are typical local minimum for the loss. Third column represents the result of an amortized CF-A (best amortized EL ELBO). Last column represents our amortized ADAVI technique (best amortized ELBO). First row represents the full posterior samples. Second and third row only represents the first mixture component samples (third row zooms in on the data). Notice how neither technique recovers the actual multi-modality the theoretical posterior. We obtain results of good experimental value, usable to estimate likely population mixture components.
27

---

ful to build an intuition of the problem. We also present implementation details for experiment, and additional neuroimaging results.
E.1 NEUROIMAGING CONTEXT
The main goal of Kong et al. (2019) is to address the classical problem in neuroscience of estimating population commonalities along with individual characteristics. In our experiment, we are interested parcelling the region left inferior frontal gyrus (IFG). Anatomically, the IFG is decomposed in 2 parts: pars opercularis and triangularis. Our aim is to reproduce this binary split from functional connectivity point of view, an open problem in neuroscience (see c.g. Heim et al. 2009).
As Kong et al. (2019), consider a population of S=30 subjects, each with 7 = 4 acquisition ses- sions, from the Human Connectome Project dataset (Van Essen et al. 2012). The fMRI connectivity between a cortical point and the rest of the brain, split in 1) 1, 4S3 regions, is represented vector of length D with each component quantifying the temporal correlation of blood-oxygenation between the point and region. A main hypothesis Kfng et al. (2019), and the fMRI field, is that the fMRI connectivity of points belonging to the same parcel share a similar connectivity pattern correlation vector. Following Kong al. (20T9), we represent D-dimensional correlation vectors as RVs on the positive quadrant of the D-ddimensional unit-sphere. We do this efficiently assum- ing they have a C-normal distribution, Gaussian under the transformation of the link function (xx) SoftmaxCentered(() (Dillon et al. 2017):
3 st = =2 -00, 8
C- (w)) N(DD-1, 2gg) log(e) ~ s*)
M¹ [ L
EL [e] 1
L (mi) 14. ~~ N(C - (mq),
ML. [n4 ==
(E.
log(ai) ~ =(ss ,
Σ = [o1|==1..1
J- (1(t)) pi, ~ N(L -1 (ui),
M. = |n
log(k) ~ U($ , I1 Dir( [T] x L)
L (xxi)))))
1 K. MI ( .N( (1Z'), k2)1
X X
(E. 10b)
Our aim is therefore identify L 2 functional networks that would produce functional par- cellation of the studied IFG section. In this setting, the parameters of interest are the networks μ. Instead of the complex EM computation derived Kong al. (2019), we perform full post-posterior inference for those parameters using our automatically derived architecture
E.2 DESCRIPTORS, INPUTS TO ADAVI
We can analyse the model described in eq. E. 10a and eq. (E. using the descriptors defined in eq. A.T. Those descriptors constitute the inputs methodology needs automatically derive the
28

---

dual architecture from the generative HBM:
V = (M-, EL, S EL, ML.S.T K. II, }
(Po. P2,
Card (Po + N,P1 + T, 5+
Hier (M- 3, 33 M 2, +3,M4.5. -1
Shape = {
M + (L,
E + (L,),
ML. + (),
L + (L.),
LT + (),
k- (1.);
11 (4,),
x (D,)
Link
M.. + L o (Lsh) - D))
E + Exp,
MI.8 + Lo Reshaape (I.D,) - (T. D)).
_ Exp,
ML.ST + L (aahhapo(LL,)). > D)).
+ Exp,
11 + SoftmaxCentered 1,) + (L.)),
X + L
(E.11)
E.3 EXPERIMENT ON HB-H MODEL ON TOY DIMENSIONS
To get an intuition of the behavior of our architecture on the MS-HBM model, we consider following toy dimensions reproduction of the model:
N,T,S,D,L= 2,
9 ,,9 = 444
K kt, Aek o
ct = =4, -4, -3, 3, 2, -2, -1
==
C-1 u1-1 g gt)
log(a) U1( ,c))
(B.12)
-- --1 q ~ N(C ((C))(())))))
log(or)
L- (HIt))) o1 N/L -N(C-((A) (ff),of)
log(R) ~ U(n (+) π Dir([T] x L.)
M st K, 11 Mix(11, (N(C
-(c- ((1)))x)))1)
29

---

Data and posterior sample
1.0
0.8
0.6
Example 0
0.4
0.2
0.0
0.00
0.25
0.50
0.75
1.00
1.0
0.8
0.6
Example 1
0.4
0.2
0.0
0.00
0.25
0.50
0.75
1.00
1.0
0.8
0.6
Example 2
0.4
0.2
0.0
0.00
0.25
0.50
0.75
1.00
Figure E.4: Visual representation of our results synthetic MS-HBM example. Data is rep- resented as colored points on the unit positive quadrant, each color corresponding to a subject x session. Samples from posterior distributions below the data points are 1²* samples. Then samples of μz" Then samples of μ (black lines). No- tice how the μ posteriors distributed around the angle bisector of the are covered by the points the subsequent plate.
represented concentric colored markings. Just
The results can be visualized on Fig. usage of link functions.
E.4 This experiment shows the expressivity we gain from the
30

---

E.4 IMPLEMENTATION DETAILS FOR THE NEUROIMAGING MS-HBBM EXAMPLE
Main implementation differences with the original MS-HBM model Our implementation of the MS-HBB (eq. (E. 10ah and eq. (E.10b) contains several notable differences with the original one from Kong et al. (2019):
we model i distributions as Gaussians linked to the positive quadrant of the unit sphere via the function L. the orignal model, RVs are modelled using Von Mises Fisher distribu- tions. Our choice allows to express the entirety of the connectivity vectors (that only lie on a portion of the unit sphere). However, we also acknowledge that the densities incurred by the distributions on the positive quadrant of the unit sphere are different.
we forgo any spatial regularization, and also the assumption that the parcellation of a given subject s should be constant across sessions This is streamline our implementation. Adding components to the loss optimized at training could inject those constraints back into the model, but this was not the subject of our experiment, so we left those for future work.
Data pre-processing and dimensionality reduction Our model was able to run on the full di- mensionality of the connectivity, D 1483. However, we obtained better results experimentally when further pre-processing the used data down the dimension D 141. The displayed results in Fig. 4 are the ones resulting from this dimensionality reduction:
1. we projected the (S, D X connectome (ying the D unit sphere) to the un- bounded 1- space using the function C
2. in this D space, we performed a Principal Component Analysis (PCA) to bring us down to D
140 dimensions responsible for 80% of the explained data variance
3. in the resulting D space, caleulated the of all the connectivity points, and their standard deviation, and used both to whiten the data
4. from the whitened data, calculated the Ledoit-Wolf regularised covariance (Ledoit & Wolf 2004), that we used to construct the Σg matrix used in eq. (E.
5. finally, we projected the whitened data onto the unit sphere in D 141 dimensions via the function C
To project our results back to the original D space, simply ran back all the aforementioned steps. Our prior for μ has been carefully designed so has to sample connectivity points in the
vicinity of the data point of interest, Our implementation is therefore in spirit close SBI (Cranmer et al., 2020; Paaamakarios et al. 2019b Greenberg et al./ 2019 Thomas et al. (020) that aims at obtaining an amortized posterior only in the relevant data regime.
Mutli-step training strategy In appendix F22 we describe our conditional density estimators as the stacking of a MAF (Papamakarios et al. 2018) on top of a diagonal-scale affine block. To accelerate the training of our architecture and minimize numerical instability (resulting in NaN evaluations of the loss) we used the following 3-step training strategy:
1. we only trained the shift part of our affine block into Maximum A Posteriori regression setup. This can be viewed as the amortized fitting of the first moment of our posterior distribution
2. we trained both the shift and scale of our affine block using unregularized ELBO loss. This to rapidly bring the variance of our posterior to relevant values
3. we then trained our full posterior (shift and scale of our affine block, in addition to our MAF block) using the reverse KL loss.
This training strategy shows the modularity of our approach and the transfer learning capabilities already introduced in appendix C.4 C.4 Loss evolution be seen in Fig. H.5
31

---

Training
train
30
(og
25-
500
1000
1500
2000
2.500
3000
calls
Figure E.5: 3-step loss evolution across epochs for the MS-HBM ADAVI training. Losses switch are visible at epochs 1000 and 2000. Training was run for a longer period after epoch 3000, with no significant results difference.
Soft labelling In eq. 1. and eq.
(0 we define variables as following Gaussian distribu- tions in the latent space TR This means that, considering vertex X, and session network a the squared Euclidean distance between (X and C- (x 5.) in the latent space RD
proportional the log- likelihood of the point (X ) for the mixture component k:
( (xa')
1 log p( (X k) + C(s)
(E.13)
Note that x is the same for both networks. Additionally, considering Bayes theorem:
log p(1 k( log (x x k)) + log p(1 k) log p( ((a))))))
p(1 x logp( X = 0) - log b (X 1) log p(l = 0) p(1 1)
(E.14)
log
p(1 1
F. 13 and log p(1 = k)) via draws from the posterior of (see eq. (E.10a)). To integrate those equations, we used a Monte Carlo procedure.
Where log [x = k)) can be obtained through eq.
E.5 ADDITIONAL NEUROIMAGING RESULTS
E.5.1 SUBJECT-LEVEL PARCELLATION
As pointed out in section 3.5 the MS-HBM model aims at representing the functional connectivity of the brain at different levels, allowing for estimates of population characteristics and individual variability (Kong et al. 2019 It is of experimental value to compare the parcellation for given subject, that is to say the soft label we give to a vertex X and how this subject parcellation can deviate from the population parcellation. Those differences underline how individual brain have a unique local organization. Similarly, we can obtain the subject networks jı and observe how those can deviate from the population networks μ Those results underline how given subject can have his own connectivity, or, very roughly, his own "wiring" between different areas of the brain. Results can be seen Fig. E.6
E.5.2 COMPARISON OF LABELLING WITH THE MF- VI RESULTS
We can compare the subject-level parcellation resulting from the latent networks recovered using ADAVI vs MF-VI method. The result, for the same subjects as previous section, can be seen in Fig. F.7. where the difference in ELBO presented in our main text translates into marginal differences for our downstream task of interest.
F BASELINES FOR EXPERIMENTS
F. BASELINE CHOICE
In this section we justify further the choice of architectures presented as baselines the our experi- ments:
32

---

113215
122622
113822
Connectiityy Networks
Networt 9
Netook
Netwok
Soft
Figure E.6: Subject-level parcellations and networks. For 3 different HCP subjects, we display the individual parcellation the bottom) and the individual μ² networks (on the top). Note how the individual parcellations, though showing the same general split between the pars opercularis and pars triangularis, slightly differ from each other and from the population parcellation (Fig, Similarly, networks μz differ from each other and from the population networks μ (Fig. 4 but keep their general association semantic/phonologic processing (0, in blue) and language production (1, in red) (Heim et al.) 2009 Zhang et al. 2020). To be able model and display such variability is one of the interests of models like the HB- (Kong et al. 2019).
Soft Chastrrig Difference ADAVI vs
Soft Clusteting Difference ADAVI us Subject
Soft Clustering Difference ADAVI HE
635
025
0.2
425
65
Figure E.7: Gap in labelling between ADAVI and MF-VI results. Following eq. EE.14 the difference latent space between the odds for the ADAVI and MF-VI techniques, before ap- plying sigmoid function. Differences marginal, and interestingly located the edges between networks, where the labelling is less certain.
compute
33

---

2017). This methods stands a common -aantoonn-rrciceee non- Mean Field VI (MF-VI) (Blei et al. amortized baseline, is fast to compute, and due to our choice of conjugate examples (see section 3.2) be considered as a proxy to the ground truth posterior. We implemented MF-VI in its usual setup, fitting to the posterior a distribution of the prior's parametric form;
(Sequential) Neural Posterior Estimation (SNPE-C) architecture (Greenberg et al. 2019). NPE-C is an typical example from the SBI literature (Cranmer et al. 20200) and functions as likelihood-free, black box method. Indeed, NPE-C trained using forward KL (samples from the latent parameters), and is not made "aware" of any structure in the problem. NPE- C fits single normalizing flow over the entirety of the latent parameter space, and its number of weights scales quadratically with the parameter space size. When ran over several simulation rounds, the method becomes sequential (SNPE-C), specializing for certain parameter regime to improve performance, but loosing amortization in the process;
Total Latent Space Flow (TLSF) architecture (Rezende & Mohamed 2016). Following the original normalizing flow implementation from Rezende Mohamed (2016), we posit TLSF as counterpoint over the entirety of the latent parameter space, and is not made "aware" of the structure of the model. But contrary to SNPE-C, TLSF is trained using reverse KL and benefits from the presence of a likelihood function. We can use TLSF in a non-amortized setup (TLSF- NA), in amortized setup (TLSF-A) trough observed data encoder conditioning the single normalizing flow;
SNPE-C. Like SNPE-C, TLSF TLSF fits single normalizing flow
Cascading Flows (CF) (Ambrogioni et al. 2021b). CF is an example of a structure-awar prior-aware VI method, trained using reverse KL. By design, its number of weights scales linearly with the plate's cardinalities. CF can be ran both in a non-amortized (CF-NA) and amortized (CF-A) setup, with the introduction of amortization through observed data encoders in the auxiliary graph. As structure-aware amortized architecture, CF-A is our main point of comparison in this section;
F.2 IMPLEMENTATION DETAILS
In this section, we describe with precision and per experiment the implementation details for the ar- chitectures described in appendix EI We implemented algorithms in Python, using the Tensorflow probability (TFP, Dillon et al. 2017) and Simulation Based Inference (SBI, Tejero-Cantero et al. 2020) libraries. For all experiments in TFP, used the Adam optimizer (Kingma & Bal (015)) For normalizing flows, we leveraged Masked Autoregressive Flow (MAF, Papamakarios & Murray 2018).
For all experiments:
Mean Field VI (MF-VI) (Blei et al. 2017). We implemented MF-VI in TFP. The precise form of the variational family is described below for each experiment.
Sequential Neural Posterior Estimation (SNPE-C) architecture (Greenberg al. 2019) We implemented SNPE-C with the SBI library, using the default parameters proposed by the API. Simulations were ran over 5 rounds, to ensure maximal performance. We acknowl- edge that this choice probably results in an overestimate of the runtime for the algorithm. To condition the density estimation based on the observed data, we designed an encoder that is a variation of our Hierarchical Encoder (see section 2.2p. Its architecture is the same as HE -the hierarchical stacking of Set Transformers (Lee al. 2019)- but the encoder's output is the concatenation of the G per-group encodings with the population encodings. This encoder is therefore parsimoniously parametrized, and adapted to the structure of the problem.
Neural Posterior Estimation (NPE-C). Though we acknowledge that NPE-C can be imple- mented easily using the SBI library, we preferred to use our own implementation, notably to have more control over the runtime of the algorithm. We implemented the algorithm using TFP. We used the same encoder architecture as for SNPE-C.
Total Latent Space Flow (TLSF). We implemented TLSF using TEP. Our API is actually the same for NPE-C, since TLSF- and NPE-C only differ by their training loss.
34

---

Cascading Flows (CF) (Ambrogioni et al. 2021b). We implemented Cascading Flows, using TFP, and having consulted with the authors. An important im- plementation detail that is not specified explicitly in Ambrogioni et al. (022b) (whose notations we follow here) is the implementation the target distribution over the auxiliary variables notably in the amortized setup. Following the authors specifications during our discussion, we implemented r as the Mean Field distribution = II, P5(e))
version of
ADAVI (ours). We implemented ADAVI using TFP.
Regarding the training data:
All amortized methods were trained over a dataset of 20, 000 samples
All non-amortized methods except SNPE-C were trained on a single data point (separately for 20 different data points)
SNPE-C was trained over 5 rounds of simulations, with 1000 samples per round, for an effective dataset size of 5000
For the non conjugate experiment (see section 3.2
MF-VI: variational distribution is
q - -ammmala; consentration= rate= hottssooos.e
We used the Adam optimizer with learning rate of 10,2, The optimization was ran for 20, 000 steps, with a sample size of 32.
CF: auxiliary size 8, observed data encoders with 8 hidden units. Minibatch size 32, 32 theta draws per X point (see appendix C.5), Adam (10-2), 40 epochs using a reverse KI. loss.
ADAVI: NF with Affine block with triangular scale, followed by MAF with [32, 32] units. HE with embedding size 8, 2 modules with 2 ISABs (2 heads, 8 inducing points), PMA (seed size 1), SAB and linear unit each. Minibatch size 32, 32 theta draws per X point (see appendix C.5). Adam (10-3) 40 epochs using reverse KL loss.
For the Gaussian random effects experiment (see section 3.3
MF-VI: variational distribution is
q (R; (tuummmmmn//ooppssttttfffoottlop
N([p "t.l= ..... 11 .....
mean=V std= Softplus(V
We used the Adam optimizer with learning rate of 10 10-2 The optimization was for 10, 000 steps, with a sample size of 32.
SNPE-C: 5 MAF blocks with 50 units each. Encoder with embedding size 8, 2 modules with 2 SABs (4 heads) and 1 PMA (seed size 1) each, and 1 linear unit. See SBI for optimization details.
NPE-C: 1 MAF with 32, 32] units. Encoder with embedding size 8, 2 modules with ISABs (2 heads, 8 inducing points), 1 PMA (seed size 1), 1 SAB and linear unit each. Minibatch size 32, 15 epochs with Adam (10 -3 using a forward KL loss.
TLSF: same architecture as NPE-C. TLSF-A: minibatch size 32, 32 theta draws per X point (see appendix C.5). 15 epochs with Adam (10-3) using reverse KL loss. TLSF. NA: minibatch size I, 32 theta draws per X point (see appendix C.5 1500 epochs with Adam (10-3) using a reverse KL loss.
CF: auxiliary size 16, observed data encoders with 16 hidden units. CF-A: minibatch size 32, 32 theta draws per X point (see appendix C.55), Adam (10 (10-3), 40 epochs using reverse KL loss. CF-NA: minibatch size 1, 32 theta draws per X point (see appendix CC.). Adam (10-2) 1500 epochs using a reverse KT. loss.
35

---

ADAVI: NF with Affine block with triangular scale, followed by MAF with [32, 32, 32] units. HE with embedding size 8, 2 modules with 2 ISABs (2 heads, inducing points), PMA (seed size 1). 1 SAB and linear unit each. Minibatch size 32, 32 theta draws per X point (see appendix C.5). Adam (10-9), 10 epochs using an unregularized ELBO loss, followed by 10 epochs using reverse KL. loss.
For the Gaussian mixture with random effects experiment (see section 3.4
MF-VI: variational distribution is
q ([p.. ME]: mean=Viu, -otp:
X W[[le...... wil, .... [ ...., #i]l mean= htt.
[of
X SirrcooncettaaionnSSffpppppplpp
We used the Adam optimizer with learning rate of 10-2 The optimization was ran for 10, 000 steps, with sample size of 32.
SNPE-C: 5 MAF blocks with 50 units each. Encoder with embedding size 8, modules with 2 SABs (4 heads) and 1 PMA (seed size 1) each, and 1 linear unit. See SBI for optimization details.
NPE-C:
MAF with (32, 32, units. Encoder with embedding size 8, 2 modules with 2 ISABs (2 heads, 8 inducing points), 1 PMA (seed size 1), 1 SAB and linear unit each. Minibatch size 32, 20 epochs with Adam (10-3) using forward KL loss.
TLSF-A: same architecture as NPE-C, minibatch size 32, 32 theta draws per X point (see appendix C.5). 250 epochs with Adam (10-3) using reverse KL loss.
TLSF-NA: same NF architecture as NPE-C. Encoder with embedding size 16, 2 modules with ISABS (2 heads, 8 inducing points), PMA (seed size 1), 1 SAB and linear unit each. Minibatch size 1, 32 theta draws per X point (see appendix C.5 1000 epochs with Adam (10-3 using a reverse KL loss.
CF: auxiliary size 8, observed data encoders with 8 hidden units. CF-A: minibatch size 32, 32 theta draws per X point (see appendix C.5) Adam (10-3), 200 epochs using a reverse KL loss. CF-NA: minibatch size 1, 32 theta draws per X point (see appendix C.5) Adam (10 1500 epochs using a reverse KL loss.
ADAVI: NF with 1 Affine block with diagonal scale, followed by 1 MAF with [32] units. HE with embedding size 16, 2 modules with 2 ISABs (4 heads, 8 inducing points). PMA (seed size 1), 1 SAB and 1 linear unit each. Minibatch size 32, 32 theta draws per X point (see appendix C.5). Adam (10-3), 50 epochs using a MAP loss on the affine blocks, followed by 2 epochs using an unregularized ELBO loss on the affine blocks, followed by 50 epochs of reverse KL loss (see appendix E.4 for the training strategy, total 102 epochs).
For the MSHBM example in toy dimensions (see appendix E.3
ADAVI: NF with Affine block with diagonal scale, followed by 1 MAF with [32] units.
HE with embedding size 32, 2 modules with 2 SABs (4 heads), PMA (seed size 1), SAB and linear unit each. Minibatch size 32, 32 theta draws per point (see appendix C.5). Adam (10-3), 5 epochs using a MAP loss on the affine blocks, followed by epoch using an unregularized BLBO loss on the affine blocks, followed by epochs of reverse KI. loss (see appendix E.4 for the training strategy, total 11 epochs).
For the MSHBM example in (see section 3.5
36

---

MF-VI: variational distribution is
LoNN( C ba (L.))
x Exp oN EL
o ((ri
Exp (oo ØL): meunn-Vt....stdd-EpPPV/
std= x o (N ; - std= Exp(V(L))
x ExpoN (I; Rupp
Dirichlet(II, concentratioon
We used the Adam optimizer with learning rate of 10 10-2. The optimization was ran for 10, 000 steps, with sample size of 32.
ADAVI: NF with Affine block with diagonal scale, followed by MAF with [1024] units. HE with embedding size 1024, 3 modules with 1 SABs (4 heads), PMA (seed size 1), SAB and linear unit each. Minibatch size 1, theta draws per point (see appendix C.5). Adam (10-3 3), 1000 epochs using MAP loss on the affine blocks, followed by 1000 epochs using an unregularized ELBO loss on the affine blocks, followed by 1000 epochs of reverse KL. loss (see appendix E.4 for the training strategy, total 3000 epochs).
37


Reviews:
[{"review_id": "fKDtBF0TvCk", "reviewer": "Reviewer_bbMw", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "# Strengths\nThe motivation for the work is clear and it's likely to be of broad interest. Factorizing the approximation network to match the structure of the model is interesting and novel to my knowledge, and the experiments build up from synthetic toy examples to a real-world problem. \n\n# Opportunities for improvement\n\n## Overall framing and relation to broader approximate inference\nThe overall framing of the paper starts from the perspective of improving the tractability of normalizing flows for posterior approximation rather than improving approximate inference, and takes its aim primarily at flow-based methods (notably, cascading flows) rather than approximate inference more broadly. I think broader engagement with the VI literature may improve the paper, including particularly some mention of the approximation vs amortization gap. This then situates the MF VI baseline as not just a \"common practice\" method but (as far as I can tell) the upper bound on performance of the remaining methods, which share its approximation gap but add an amortization performance gap. I'm not sure I understand the non-conjugate results in this context, though -- is it because the link function is applied to the flow approximation but not in the MF context? Relatedly, I think the discussion from supplement B2 is worth including in the main paper -- the HVM and SSVI approaches among others seem like a natural complement to the present work. Making these connections seems more important than engaging with SBI and likelihood-free inference -- I don't think likelihood-free approaches are ever competitive with likelihood-based approaches if likelihoods are available (nor is this their intent), so I'm not sure why they are an important baseline or particularly relevant to the current work. Finally, as a minor point, I wonder if the more typical reader here would come from the VI literature, for whom leading with the ELBO (rather than Reverse KL) framing would be more intuitive. \n\n## Pushing the synthetic examples to regimes where ADAVI wins\nNone of the examples actually show that ADAVI is a superior approach: in the NC case it is roughly equivalent in performance to CF. In GM ADAVI may achieve the best ELBO out of the amortized methods, but takes an order of magnitude more time than just running MF-VI, so the benefit of amortization is not shown. In the gaussian random effects example (fig3) it's clear that ADAVI has favorable scaling, but even in the largest example MF-VI has fewer parameters and achieves higher accuracy in less time. Furthermore, I'm not sure if the selected metrics truly show the benefit of amortization: rather than showing optimization time for all the models, shouldn't the paper show the predictive inference time (and performance) for held out data for all the models? Then the amortized models should be dramatically faster by not requiring any optimization. Finally, I'm not sure why timings are not done on consistent hardware -- it seems like everything in the paper should run on both CPU and GPU, so picking one (or showing timing for both) would be better than trying to compare timings between two different sets of hardware. \n\n## Including quantitative evaluation on the real problem\nWhile the real-world results match qualitative patterns expected from prior work, the paper needs to do more to quantitatively show how ADAVI is indeed superior for these kinds of models. The best way to do this would be reporting cross-validated predictive loss, e.g. by holding out subjects, times, or connectivity measures, and comparing both the loss and runtimes against other models. \n\n## Minor comments and typos\n\nSection 2.2. could be clarified further. For example: \n- Is the set transformer really elementwise? Or is it \"elementwise\" w.r.t. the leaf random variables in the graph? \n- If the \"contraction\" operator isn't to tensor contraction, could a different word be used, considering that tensor notation is already used and \"contraction\" has a common meaning in that setting? \n- $\\mathcal{B}_h$ is defined twice above and below expr. 1. \n- Both upper and lowercase $x^{i,j}$ is used -- is it the same or different? \n- The meaning of the superscripts on $x$ is not explained (I think one is the plate index and the other is the datapoint index, but not sure). \n\nIn addition: \n- The paper should clarify the full algorithm: is everything trained end-to-end? Or are the STs trained first, and then the model parameters? \n- Section 3.2 could potentially pick a more practically-relevant non-conjugate example from the literature (there are plenty of examples in neuroscience, though I'm not sure if any have non-discrete latents). \n- Fig2 and 3 are missing error bars/ribbons. \n- $x \\sim \\mathrm{Mix}(\\ldots)$ is not defined and not standard to my knowledge -- I'm assuming this is multinomial over mixture components given $\\pi$ and each component is Gaussian? \n- Expressions 9a-9b are not arranged in an intuitive way (definitions and priors are intermixed, everything is in one giant block). An annotated figure (even a plate diagram) would be better, and the specifics can be left to the supplement, especially since the paper doesn't give enough info to understand the model beyond the Kong et al. citation. \n\n", "summary_of_the_paper": "The paper tackles approximate inference for hierarchical Bayesian models with fully nested structure. The specific approach taken is variational inference with a q-distribution that iteratively applies conditional normalizing flows to derive a hierarchical representation, and factorizes in a manner parallel to the generative model by reusing flow parameters, thus having a number of parameters that does not grow with cardinality of each plate. The benefits of the model are illustrated in a few synthetic experiments as well as in application to a human neuroimaging dataset. ", "summary_of_the_review": "Hierarchical graphical models of the sorts considered here are in wide usage by scientific practitioners, and automatic approximate inference methods that are efficient and accurate are of longstanding interest -- thus the motivation of the work is clear. The nature of the contribution is likewise clear. My primary concern is that the evaluation doesn't fully demonstrate the benefits of the approach (e.g. using predictive out-of-sample log-likelihoods rather than training ELBOs, including a quantitative evaluation of the real-world problem, providing consistent timings, pushing the scale until ADAVI dominates). Secondarily, I think the overall framing, clarity, and writing could be improved. I think the work is below the bar now and am rating accordingly, but I think my concerns should be sufficiently addressable in rebuttal for the work to be appropriate for the conference. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "F9hS8gBL4uf", "reviewer": "Reviewer_bt4m", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strength:\n\n- A good motivation: Extending the applications of simulation-based inference and structured variational inference to high-dimensional data settings (such as neuroimaging data) is interesting.\n\nWeaknesses:\n\n- The proposed ADAVI method only applies to a pyramidal class of Bayesian networks in which dependency structure follows a pyramidal graph. The authors need to motivate in the text that this class of problems covers most cases in the target applications in neuroimaging studies. \n- The proposed method in its ultimate performance becomes similar to a simple mean-field approximation. The authors claim the proposed method instead provides more expressivity while failing to show this point in the experiments. \n- The performance improvements (in terms of quality of inference and inference time) compared to other alternatives remain marginal (results in Table 1,2 and Figure 3).\n- Despite the promising abstract and introduction setting a high expectation on the high-dimensionality of target problems (thousands of brain locations), the experiment on neuroimaging data is merely conducted on a relatively small dataset with 30 subjects and 1483 measures. This barely fits with the requirements in the field. These days we need methods that can deal with considerably larger datasets consisting of thousands of subjects with thousands of measures (for example UKBiobank). Furthermore, it would be nice to also compare quantitatively the time complexity and inference quality of ADAVI with other alternatives when applied to neuroimaging data.\n\nMinor suggestions, comments, and questions:\n- It may be nice to motivate the importance of HBR in the neuroimaging context. Why it is important to be able to handle the hierarchically organized data in the neuroscience context? \n- Section 2.1: Direct Acyclic Graphs => Directed Acyclic Graphs\n- Throughout the text: fig. => Fig.\n- Section 2.1: A RV's hierarchy => An RV's hierarchy\n- What do you mean by \"symmetry\" in \"... exploiting the symmetry induced by the plates ...\"?\n- Section 3: \"a a hierarchy\" => \"a hierarchy\"\n- Is the curves in Fig. 2a derived empirically? Then why the wall time for ADAVI remains fixed until $10^5$ examples? Please explain.\n- While the main arguments are around the lower computational complexity of the proposed method, it is very difficult to judge the time improvements when diverse hardware are used in a heterogeneous cluster of computers (for example two types of CPU). Would be nice to fix the hardware setup to ensure fairness and reliability of comparisons.\n- I suggest moving sections B.2 and B.3 to the main text (possibly to the final discussion) as they contain important information about the presented method.\n- Would be great also to see some results in section E.5 in the main text (on how the individual network can deviate from the population network). That would be of high interest to the audience from the neuroscience community. \n", "summary_of_the_paper": "This manuscript proposes an amortized variational inference to produce a dual variational family for hierarchical Bayesian models (HBM) that can be represented as pyramidal Bayesian models. The presented method exploits the exchangeability of parameters in HBM to reduce its parametrization for faster inference on high-dimensional data such as neuroimaging. The authors compare empirically the proposed method with several amortized and non-amortized alternatives and on several experimental data in terms of the size of parametrization, inference time, and quality of inferences. ", "summary_of_the_review": "Despite a nice motivation, the empirical results show a minor improvement over the already existing methods. The results are also weakly presented for the target problem in the neuroimaging context (small data, no quantitative comparison, lack of discussion). The method is not tested on very high dimensional settings as they are in the neuroimaging context.\n ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ADsjSoJRNqG", "reviewer": "Reviewer_KUup", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Overall, I am not an expert in HBMs and Neuroscience, but based on the introduction in this work, the proposed method is meaningful to the \nfield of Bayesian modelling for the neuroimaging, as it can estimate posterior distributions for a generative model of pyramidal structure, which is difficult to achieve. Authors provide sufficient experiments to support their claim in terms of both the inference quality and the speed. So I hold a positive attitude toward this paper. The paper is not very reader-friendly, so it would be better that authors can provide the introduction of some basic concepts and more related works in their final version to improve the readability of their paper.  ", "summary_of_the_paper": "This paper proposes to derive an automatic methodology that takes as input a generative HBM and generates a dual variational family able to perform amortized parameter inference. The proposed method can be used to the context of  pyramidally structured data with good inference quality and a favorable training time. ", "summary_of_the_review": "Please refer to the Main Review section.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "3h6_AwDm-c", "reviewer": "Reviewer_4eiB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strengths:\n- The method deviates from black-box simulators and instead exploits the hierarchical nature of many forward processes to scale to high-dimensional parameter spaces.\n- Experiments are detailed, and results are discussed carefully. ADAVI shows good performance, in particular against (S)NPE, but non-amortized methods are shown to achieve higher ELBO. \n- The supplementary materials are thorough and include helpful experimental details as well as further discussions.\n\nWeaknesses:\n- Population studies are taken as an example of large HBMs with millions of parameters. I believe it would be fair to mention that oftentimes only a few of those many parameters are actual of interest for scientific inquiry. In this case, most of the parameters are latent variables for which no explicit estimation is necessary. Would you argue that ADAVI is still relevant in this case? [I would say it may still be relevant, but I would be curious in having your opinion.] Could you also better motivate in which scientific cases inference over millions of parameters is strictly necessary?\n- In Section 2.3, the variational distribution is defined as a mean-field approximation. Can you comment on the constraints that result from this assumption?\n- Results are discussed in terms of the ELBO only. The quality of the approximate posteriors is never diagnosed explicitly (at least as far as I can see). \n- I would appreciate some comments on the impact of the size of the encoding space produced by the SetTransformers. How shall one set this size? Should it be the same across all levels?\n\nSmall remarks:\n- Avoid footnotes if you can.\n- Replace \"fig. X\" with \"Fig. X\". ", "summary_of_the_paper": "This work introduces ADAVI, an approximate inference algorithm for hierarchical Bayesian models (HBMs). The approach is similar to NPE from simulation-based inference but exploits the hierarchical structure of the forward model to generate an efficient variational family automatically. Experiments demonstrate the applicability of the method on HBMs of increasing complexity, including a challenging neuroimaging model. Results indicate good performance against other amortized methods. ", "summary_of_the_review": "This paper addresses a common problem of simulation-based inference and proposes a sound and efficient solution to enable inference in high-dimensional parameter spaces. Experiments show convincing results.\n\nMy main issue is the lack of quality checks of the approximate posteriors produced by ADAVI. This should matter the most, in my opinion, well above efficiency, wall-clock times, and the number of parameters. If the approximate posteriors are wrong, none of those matter. For now, I do not recommend this paper for acceptance, but I will be pleased to change and increase my evaluation if the authors can present diagnostics of the resulting approximate posteriors. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a unique contribution to the scalability of Bayesian inference to Pyramidal Bayesian Models with application to neuroimaging. The major point of concern by the reviewers is around how close is the inference approach to the more classical Mean-Field VI. However, in my opinion, the authors have addressed these concerns in the rebuttal. Therefore, I recommend Accept.
Decision: Accept (Poster)