PAVI: Plate-Amortized Variational Inference
Authors: Louis Rouillard
Alexandre Le Bris
Thomas Moreau
Demian Wassermann
Venue: Accepted by TMLR ()

Abstract: Given observed data and a probabilistic generative model, Bayesian inference searches for the distribution of the model's parameters that could have yielded the data. Inference is challenging for large population studies where millions of measurements are performed over a cohort of hundreds of subjects, resulting in a massive parameter space. This large cardinality renders off-the-shelf Variational Inference (VI) computationally impractical.

In this work, we design structured VI families that efficiently tackle large population studies. Our main idea is to share the parameterization and learning across the different i.i.d. variables in a generative model -symbolized by the model's $\textit{plates}$.
We name this concept $\textit{plate amortization}$. Contrary to off-the-shelf stochastic VI --which slows down inference-- plate amortization results in orders of magnitude faster to train variational distributions. Applied to large-scale hierarchical problems, PAVI yields expressive, parsimoniously parameterized VI with an affordable training time --effectively unlocking inference in those regimes.

We illustrate the practical utility of PAVI through a challenging Neuroimaging example featuring 400 million latent parameters, demonstrating a significant step towards scalable and expressive Variational Inference.

Published in Transactions on Machine Learning Research (11/2023)
PAVI:
Plate-Amortized Variational Inference
Louis Rouillard Université Paris-Saclay, Inria, CEA Paloiseau, 91120, France
Alexandre Le Bris
Meczandre.le-brii@@irrii
Université Paris-Saclay, Inria, CEA Palaiseav, 91 91120, France
Thomas Moreau Université Paris-Suclay, Inria, CEA Palaiseau, 91 120, France
Nhomas.morauuBarri..a
Demian Wassermann Université Paris-Sackee Inria, CEA Palaiseau, 91120, France
demian, www.eeeennnnnniiiiiaaai..o
Reviewed on OpenReview: https: // openreview. net/ forum? id=u1Y96000A6
Abstract
Given obscrved data and a probabilistic generative model, Bayesian inference searches for the distribution of the model's parameters that could have yielded the data. Inference is challenging for large population studies where millions of measurements are performed over a cohort of hundreds of subjects, resulting a massive parameter space. This large cardi- nality renders off-the-shelf Variational Inference (VI) computationally impractical. In this work, we design structured VT families that efficiently tackle large population studies. Our main idea is to share the parameterization and learning across the different i.idd. variables in a generative model, symbolized by the model's plates. We name this concept plate amortiza- lion. Contrary to off-the-shelf stochastic VI, which slows down inference, plate amorlization results in orders of magnitude faster to train variational distributions. Applied to large-scale hierarchical problems, PAVI yields expressive, parsimoniously parameterized with an fordable training time. This faster convergence effectively unlocks inference those large regimes. We illustrate the practical utility of PAVT through challenging Neuroimaging example featuring 400 million latent parameters, demonstrating significant step towards scalable and expressive Variational Inference.
1 Introduction
Inference graphical models allows explain data in an interpretable and uncortainty-aware manner (Koller & Friedman, 2009; Gelman et al., 2004; Zhang et al., 2019). Such graphical models are leveraged in Neu- roimaging to capture complex phenomena, for instance, the localization of cognitive function across the human brain cortex (Thomas Yeo et al., 2011; Kong et al., 2019). Yet, current inference methods struggle with the high dimensionality of Neuroimaging applications, which can feature hundreds of millions of latent parameters. This work unlocks inference for those large-scale applications via novel scalable VI met.odolo- gies.
Our target applications are population studies. Population studies analyze measurements over large cohorts of subjects. They allow model population parameters -for instance, the mean weight in population

---

Published in Transactions on Machine Learning Research (11/2023)
and different groups of subjects- along with individual variabilityy the variation of each subject's weight compared to the group weight. These studies are ubiquitous in health care (Fayaz et al., 2016; Towsley et al., 2011), and can typically involve hundreds of subjects and individual measurements per subject. For instance, in Neuroimaging (o.g. Kong et al., 2019), mcasurements X can correspond to signals in thousands of locations in the brain for a thousand subjects. Given this observed data X, and a generative model that can produce X given model parameters O we want to recover the parameters é that could have yielded the observed X. In our Neuroimaging example, O can be local labels for each brain location and subject, together with global parameters common to all subjects, such as the connectivity fingerprint corresponding to each label. In many real-world applications, the link between latent parameters and observed data is not deterministic. For instance, the observed signal can be noisy, or the modeled process in itself may be stochastic (Koller & Friedman, 2009). As such, several parameters @ instances could yield the same data X. We want to model that parameter uncertainty by recovering the distribution of the that could have produced X. Following the Bayesian formalism (Gelman et al., 2004), we cast both θ and X as random variables (RVs), and our goal is to recover the posterior distribution p()))). Due to the nested structure of our applications, we focus on the case where p corresponds to hierarchical Bayesian model (HBM) (Gelman et al., 2004). In population studies, the multitude of subjects and measurements per subject implies large dimensionality for both θ and X. This large dimensionality in turn creates computational hurdles that we tackle through our method.
Several inference methods for HBMs exist in the literature. Earliest works resorted to Markov chain Monte Carlo (Koller & Friedman, 2009), which become slower as dimensionality increases (Blei et al., 2017). Recent approaches coined Variational Inference (VT) cast. the inference an optimization problem (Blei et al., 2017; Zhang et al., 2019; Ruiz & Titsias, 2019). In VI, inference reduces to choosing a variational family Q and finding inside that family the distribution ((O;)) E Q closest to the unknown posterior p(O X). Historically, VI required to manually derive and optimize over 2, which remains an effective method where applicable (Dao et al., 2021). This manual VT requires technical mastery. The experimenter has to choose an appropriate family @ to approximate the unknown posterior closely. At the same time, the experimenter must rely on properties such as conjugacy to derive an optimizing routine (Dao et al., 2021; Kucukelbir et al., 2016). This means that time and effort must be spent not only on the modeling part laying out hypothesis in generative model- but also on the inference part- -inferring parameters given the generative model. For many experimenters, manual VI thus features strong barriers to eniry, both in time and technical skill (Kucukelbir et al., 2016). In contrast, we follow the idea of automatic VI: deriving an efficient family Q directly from the HBM p. Automatic VI reduces inference to the modeling part only, simplifying the experimenter's ' research cycle (Kucukelbir et al., 2016; Ambrogioni et al., 2021a;b). We propose an automatic VT method computationally applicable to large population studies.
Our method relies on the pervasive idea of amortization in VI. Amortization can be broadly defined as finding posteriors usable across multiple data points (Zhang et al., 2019). As such, amortization can be interpreted as the meta-learning of the inference problem (Ravi & Beatson, 2019; Takovleva et al., 2020; Yao et al., 2019). A particular example of meta-learning is Neural Processes, which share with our method the conditioning of a density estimator by the output of a permutation-invariall encoder (Garnelo el al., 2018; Dubois et al., 2020; Zaheer et al., 2018). Though close in spirit to our work, meta-learning studies problems with a single hicrarchy. Ono-hicrarchy cases correspond models distinguishing local paramcters, representing the subjects, from global parameters, representing the population (Ravi & Beatson, 2019; Tran et al., 2017). In contrast, we study generic HBMs with an arbitrary number of hierarchies to tackle population studies. Another interesting focus is the two-hierarchy case poopulatio, group and subject for which Agrawal & Domke (2021) provide theoretical guarantees and explore in length the Gaussian approximation. In contrast, our focus is rather computational to enable the training of generic density approximators such as normalizing flows (Papamakarios et al., 2019). Furthermore, we propose a general method to treat any plate-enriched HBM, as opposed the particular factorization considered by Agrawal & Domke (2021).
Modern VT is effective in low-dimensional settings but does not scale up large population studies, which can involve millions of random variables (c.g. Kong et al., 2019). In this work, we identify and tackle two challenges to enable this scale-up. A first challenge with scalability is a detrimental trade-off between expressivity and -igediimensonnality (Rouillard & Wassermann, 2022). To reduce the inference gap, VI
2

---

Published in Transactions on Machine Learning Research (11/2023)
requires the variational family Q to contain distributions closely approximating p(o]X) (Blei et al., 2017). Yet. the form of x)] is usually unknown to the experimenter. Instead of a lengthy search for a valid family, one can resort to universal density approximators: normalizing flows (Papamakarios et al., 2019). But the cost for this generality is a heavy paramcterization, and normalizing flows scale poorly with the dimensionality of O. As result, in large population studies, the paramederization of normalizing flows becomes prohibitive. To tackle this challenge, Rouillard & Wassermann (2022) recently proposed, via the ADAVI architecture, partially share the parameterization of normalizing flows across the hierarchies of a generative model. We build upon this shared parameterization idea while improving the ADAVI architecture several aspects: removing the mean-field approximation; extending from pyramidal to arbitrary graphs; generalizing to con-sample-amortizee and stochastically trained schemes. Critically, while ADAVI tackled the sereparrameteriattion of VI in population studies, it still could not perform inference in large data regimes, This is due to a second challenge with scalability: as the size of & increases, evaluating a single gradient over the entircty of an architecture's weights quickly requires too much memory and computation. This second challenge can be overcome using stochastic VI (SVI, Hoffman et al., 2013), which subsamples the parameters θ inferred for at each optimization step. However, using SVI, the weights for the posterior for local parameter 0 E θ are only updated when the algorithm visits 0. In the presence of hundreds of thousands of such local parameters, stochastic VI can become prohibitively slow.
This work introduces the concept of plate amortization (PAVT) for fast inference in large-scale HBMs. Instead of considering the inference over local parameters 0 as separate problems, our main idea is to share both the parameterization and learning across those local parameters, or equivalently across a model's plates. We first propose algorithm to automatically derive an expressive yet patiiiooniosslyppaameteeettee family from a plato-cnriched HBM. We then propose hierarchical stochastic optimization scheme to train this architecture efficiently. PAVI leverages the repeated structure of plate-enriched HBMs via novel combination of amortization and stochastic training. Through this combination, our main claim is to enable inference over arbitrarily large population studies with reduced parameterization and training time as the cardinality of the problem augments. Critically, while traditional Stochastic VI unlocked large-scale inference at the cost of slower convergence, PAVT does not. feature such a detrimental trade-off. On the contrary, PAVT converges orders of magnitude faster than non-stochastic VI in large regimes. In practice, this quicker training unlocks inference in applications previously unattainable, bringing inference time down from weeks to hours. We illustrate this by applying PAVI to challenging human brain cortex parcellation, featuring inference over a cohort of 1000 subjects with tens of thousands of measurements per subject, for a total of half a billion RVs. This demonstrates a significant step towards scalable, expressive and fast VI.
variational
2 Problem statement: inference in large population studies
Here we introduce a motivating example for our work that we'll use as an illustration in the rest of this section. The concepts introduced in this section will be more formally defined below in Section 3.
Our objective is to perform inference in large population studies. As an example of how inference becomes impractical in this context, consider M in Figure 1 (top) as a model for the weight distribution in a pop- ulation. 02,0 denotes the mean weight across the population. 01,0, 01,1, 01,2 denote the mean weights for 3 groups of subjects, distributed around the population mean. Xo, X, represent the observed weights of 2 subjects from group 0, distributed around the group mean. Given the observed subject weights X, the goal is to determine the posterior distributions of the group and population means p(01,0, f1,1, (2.
To infer using the VI framework, we choose a variational family @ and search inside this family for the distribution closest to our unknown distribution of interest: q()) ~ p(O]X). Applying automatic VI to the example in Figure 1, the variational family Q will oftentimes factorize to (Ambrogioni et al., 2021a;b):
q(; 2) 1000-gg00000 g(02,00 22.) II
12,00 21,n)
n=U
where T - (02.0. $1,0.... $1.3} represent the weights associated to the variational family Q and each factor in q will approximate the corresponding posterior distribution: as an example q(02,0: (22,)) p(02.0 X).
3

---

Published in Transactions on Machine Learning Research (11/2023)
Compared to the mean-field factorization (Blei et al., 2017), in Equation (1) the posterior for the RVs 81.n is conditional to the RV 000 for greater expressivity. We will follow this dependency scheme in Section 4.1. During training, we will search for the optimal weights & to best approximate the posterior distribution p(O|X).
Yet, as the number of groups and subjects per group augments, this inference problem becomes computa- tionally intractable. On the parameterization side, each additional group in the population study requires additional weights 61... This constitutes first computational hurdle: the total number of weights & in the variational family becomes prohibitively large. On the computing side, optimizing over an increasing mum- ber of weights 61,n also requires a growing amount of calculation. This constitutes second computational hurdle. To circunvent this issue, one must resort to Stochastic VI and only optimize over a subset of the weights & at a given training step. In practice, this amounts to inferring the group weight of only a subset of the subject groups at time. Yet, stochastic training means that the weights Φ1,n corresponding to a given subject group will only be optimizod for a fraction of the training time. Consequently, the larger the population study becomes, the slower the inference. Our goal is to keep inference computationally tractable as the cardinality of the problem a augments.
To keep the inference tractable, we will build upon the natural factorization of the problem into subjects and groups of subjects. This factorization is symbolized by a model's plates as detailed in Section 3. First, we will harness this natural factorization in the parameterization of our variational family to reduce its number of weights. This strategy, coined plate umortization, will be detailed in Section 4.1. Second, we will reflect this factorization in the design of a stochastic training scheme -to control the computing required during the training. This will be detailed in Section 4.2. Critically, our contribution does not lie only in adding those two items. Combining shared parametrization and stochastic training, PAVI yields orders- of mf-maiite speedups in inference. Consequently, contrary to Stochastic VI, inference does not become prohibitively slow as the problem's 's cardinality augments. In practice, this unlocks inference in very large population studics. We illustrate this speed-up in our experiments in Scction and the practical utility of our method a challenging Neuroimaging setup in Section 5.4.
3
Generative modeling
3.1 Hierarchical Bayesian models (HBMs), templates and plates
Here we define more formally the inference problem and the associated notations.
Population studies can be compactly represented via plate-enriched directed acyelic graph (DAG) templates T (plates are defined in chapter 8.1 in Bishop, 2006; Gilks et al., 1994; Koller & Friedman, 2009, describes plates and template variables in chapter 6.3). T feature RV templates that symbolise multiple similar ground RVs. As an example, the RV template 0, Figure 1 (left) represents a generic group of subjects, of which there would be as many instances as the cardinality of the plate PI. The plate structure in T also denotes that the multiple ground RVs corresponding to the same RV template are conditionally i.i.d. As an example, two different group weights could be two Gaussian perturbations of the same population weight. The two groups weight RVs thus have the same conditional distribution, and are independent given the population weight. The plate structure in T thus symbolizes a strong symmetry in the problem that we exploit in our method.
We denote the template T's vertices, corresponding to RV templates, as X and θ= (0: X denotes the RVs observed during inference, and θ the parameters we infer: our goal is to approximate the posterior distribution p(O We denote T's plates as (Pp)po..p, and the plates 6, belongs to Plates(0:)) I and P respectively denote the number of latent RV templates and plates in T, which are in general not equal. In the toy example from Figure 1, there are two latent RV templates: O1 and fgy respectively the group and population mean weights. T also features two plates Pi, Po. which respectively denote groups in the population and the subjects in each group. Graphically, we can see that Plates(02) = 0, whereas Plates(0,) = (Pi) and Plates(X) = (Po. P.). To understand how we could have P I, one could keep the plates Po. but add two more RV templates to symbolize the population and group heights, in addition to their weights. In this case, we would have P = 2 and 1 == 4.
4

---

Published in Transactions on Machine Learning Research (11/2023)
M
grounding
,,,0
f1,1
L,22
1-0
X
X
X₃
X3
X
X5
Template T
8
stochastic training
Reduced distribution g²
Uss
automatic
VI
Jz
Po
Reduced model
optimization stap
e = (4,
2₂0
$2,A
we aptimize over:
Using as observed RVs:
X = (D,E}
8,,
8,,3
11,0
grounding
X
X;
F
0,,
Figure 1: PAVI working principle: The template T (left) can be grounded into the full model M (top). We aim to perform inference over M. Yet, M can feature large cardinalities in population studies. As an example, instead of 01.0...., 01,2. M can feature 01,0....,ff 10,000 -corresponding to a thousand different subject groups. This can make inference over M computationally intractable. To circumvent this issuc, we will train over M stochastically. To this end, instantiate M's template into smaller replica: the reduced model M (bottom left). Following the automatic VT framework, derive the reduced distribution q (bottom right) directly from M The reduced distribution q² features conditional normalizing flows Fi and F2 respectively associated to the RV templates 01 and 62. During the stochastic training (right), g² is instantiated over different. branchings the full model M -highlighted
blue. The branchings have M cardinalities and change at each stochastic training step The branching determine the encodings E conditioning the flows F as symbolised by the letters A, B, C and the observed data slice as symbolised by the letters D, E.

---

Published in Transactions on Machine Learning Research (11/2023)
By instantiating the repeated structures symbolized by the plates P in T, we obtain heavier graph repre- sentation: the hierarchical Bayesian model (HBM) M. This instantiation is visible in Figure 1, where we go from the template T (left) to the model M (top). In the context of Neuroimaging (Kong et al., 2019), M could feature millions of RVs, corresponding to thousands of subjects and mcasurements per subject, making it a much less compact representation. We build upon the interplay between the template T and the HBM M. To go from one representation to the other, T can be grounded into M given some plate cardinalities (Card(Pp)) (Koller & Friedman, 2009). Card(P) represents the number of elements in the plate P, for instance the number of groups in the study, Going from T to M, RV template 8, is instantiated into multiple ground RVs (0:,n) (00.), with the same parametric form, where N; = II =ITpePiatesen, Card Note that when RV templates belong to multiple plates, they symbolize as many ground RVs as the product of the cardinalities of those plates. In Figure 1, the RV template 0 is grounded into the ground RVs 01.0, , 01.2. There are as many ground RVs X, as the product of the number groups on the study Card (Pi) times the number of subjects per group Card(Po). We denote as x(f....) the (potentially cmpty) set of parents of the RV 0g,n in the ground graph corresponding to the model M. =(O..)) are the RVs whose value condition the distribution of fn For instance, in Figure subject's weight the child RV- is perturbation of the group's weight the parent RV, conditioning the child RV's distribution. This is denoted as x x(Xo) {01,0).
The benefit of the template T and the plates P is thus to symbolize the natural factorization of the graphical model M into i.i.d RVs. This factorization de-clutters the representation of hierarchical models. PAVT exploits this simplifying factorization: we consider the inference over the ground RVs 0g,n as symmetrical problems, over which we share parameterization and learning, as described in Section
3.2
Full model
The full model M is associated with the density p. In p, the plate structure indicates that a RV template 0, associated a conditional distribution P. shared across all ground RVs 0z
log X)) = =logp(X(0)) log
Nx
(2)
(5. + ΣΣ log (00.
where T(,..) ) is the (potentially empty) set of parents of the RV 0s,ms which condition its distribution. We denote with a x index all variables related to the observed RVs X. Exploiting the factorization visible
in Equation (2), our goal is to obtain a variational distribution (()) usable to approximate the unknown posterior X)) for the target model M.
4 Methods
4.1 PAVI architecture
Full variational family design Here we define the variational distribution corresponding to the full model M. To derive q, we push forward the prior r(O) using trainable normalizing flows, denoted as F (Pa- pamakarios et al., 2019). A normalizing flow is an invertible and differentiable neural network, also known as a parametric diffcomorphism. samples from an arbitrarily complex larget distribution. To every ground RV 0g,, we associate the learnable flow FiJ, to approximate its posterior distribution:
normalizing flow transforms samples from a simple base distribution into
log 9()) = Σ
log Ri, (0.,nI
i= n=0
(3)
0n, (ui,)) Unn (oluu,n//((00))
where H.n is the fush-forward of the prior ps through the flow Fen. This fush-forward is illustrated in Figure 1, where flows F push RVs into 0. Since it replicates the prior conditional dependencies, g()) in

---

Published in Transactions on Machine Learning Research (11/2023)
Equation (3) mimics the haaht term p(0) in Equation (2). This cascading scheme was first introduced by Ambrogioni et. al. (2021b). In the 2 plates case, this dependency scheme is equivalent to the branch approximation coined by Agrawal & Domke (2021). In this shallow case, Agrawal & Domke (2021) show that modeling only the prior dependencies does not limit the expressivity of the variational family compared to a fully joint factorization. However, this result does not generalize to cases featuring more plates, as detailed in Appendix C.5.
Focus on normalizing flows Expressivity is essential to make VI experimentally viable: the variational ((0) should be good proxy for the true p()]X) and not artifact due to g())'s limited expressivity (Blei et al., 2017). Though PAVI could be applicd to simpler density estimators -such as the Gaussian approxi- mations explored in depth by Agrawal & Domke (2021)- we focus on normalizing flows which can model a large range of distributions. As detailed in the next paragraph, PAVI allows the experimenter to always de- fault to the most expressive flow architecture possible -scch as Masked Autoregressive Flows (Papamakarios et al., 2018)- without sacrificing the variational family's scalability. This helps the experimenter minimize the inference's approximation gap (Cremer et al., 2018). Nonetheless, in cases where the experimenter has good sense of the form the terms Iin should take in the posterior due for instance to conjugacy the usage of very expressive flows Finn can be relaxed. Optimizing lower number of weights in Finn can then further improve the inference speed.
Though maximizing q's expressivity, using fully parameterized flows would not scale to the context of large hierarchical models featuring millions of RVs. This justifies the need for plate amortization, as detailed in the next paragraph.
Plate amortization Here, we introduce plate amortization: sharing the parameterization of density es- timators across a model's plates. Plate amortization reduces the number of weights in a variational family as the cardinality of the inference problem augments. In Section 4.2, we show that plate amortization also results in faster inference.
VT searches for a distribution (0;;)) Φ) that best approximates the posterior of Q given a value X, for X: (();6)) = p(O)xx Xo). When presented with new data point X1, optimization has to be performed again to search for the weights d1, such that q(;; 01) = p(e)X X1). Instead, sample amortized (sa) inference (Zhang et al., 2019; Cremer et al., 2018) infers in the general case, regressing the weights & using an encoder f of the observed data X: ((O; ø = f(X:)) = p(O|X = X,). The cost of learning the encoder weights is amorlized since inference over any new sample X requires no additional optimization. In Section 5.3, we compare PAVI to several sample amortized baselines, which will be denoted by the suffix (sa), and should not be confusod with plate amortized mothods. A more in-depth discussion about amortization can be found in Appendix C.1. We propose to exploit the concept of amortization but to apply it at a different granularity, leading to our notion of plate amortizalion.
Similar to amortizing across the different data samples X, we amortize across the different ground RVs In.... corresponding to the same RV template 6, Instead of casting every flow Fi,,, defined in Equation (3), as a separate, fully-parameteriized flow, we will share some parameters across the Fi.. To the template 0N, we associate a conditional flow Fi( ; ; o) with weights o, shared across all the [0,m))=..M The flow Fi,n associated with given ground RV 6,,. will be an instance of this conditional flow, conditioned by an encoding E,,.
;0 Ei... yielding Ii, = (0.. (0)..)) os, E...)
(4)
The distributions ai,n thus have 2 sets of weights, oi and E creating a parameterization trade-off. Con- centrating all of qnn's 's parameterization into Φ; results in all the ground RVs 0., having the same posterior distribution. On the contrary, concentrating all of qi, 's parameterization into E,,, allows the 6i,, to have completely different posterior distributions. But in a large cardinality setting, this freedom can result in a massive number of weights, proportional to the number of ground RVs times the encoding size. This double parameterization therefore efficient when the majority of the weights of qú is concentrated into di. Using normalizing flows F,, the burden of approximating the correct parametric form for the posterior is placed onto oi, while the E,,. encode lightweight. summary statisties specific to each 0z.. For instance,
7

---

Published in Transactions on Machine Learning Research (11/2023)
F, could learn to model Gaussian mixture distribution, while the E would encode the location and variance of each mode for each ground RV. Encodings Ei, allow to individualize for 0gn only the strictly necessary information necessary to approximate 0,, 's posterior. A more in-depth analysis of the impact of plate amortization on the expressivity of the variational family Q can be found in Appendix A.4.
Intermediate summary This section defined q, the variational distribution to approximate full model M's posterior. q features plate amortization, which helps maintain a tractable number of weights as the cardinality of M augments. The next. section introduces stochastic scheme train q. Critically, combining shared parameterization and stochastic training does not simply result in combining the advantages of both, but features synergies resulting in faster training, as further explained in Section 4.2.
4.2 PAVI stochastic training
Reduced model Our goal is to train the variational distribution ((),, defined in Equation (3), that approximates the posterior p(O)X). q corresponds to the full model M. M typically features large plate cardinalities Card(P) -with many subject groups and subjects per group and thus many ground RVs, making it computationally intractable. We will therefore train q stochastically, over smaller subsets of RVs at a time. In this section, we interpret stochastic training as the training over a reduced model M..
Instead of inferring directly over M, we will train over smaller replica of M. To this end, we instantiate the template T into second HBM Mr, the reduced model, of tractable plate cardinalities Card < Card(P). M has the same template as M, meaning the same dependency structure and the same parametric form for its distributions. The only difference lies in M smaller cardinalities, resulting in fewer ground RVs, as visible in Figure 1.
Reduced distribution and loss Here we define the distribution q used in stochastic training. q' features the smaller cardinalities of the reduced model M making it computationally tractable. At each optimization step t, we randomly choose inside M paths of reduced cardinality, as visible in Figure 1. Selecting paths is equivalent to selecting from X subset X of size N), and from θ a subset ('()). For a given 0,, we denote [.[t] the batch of selected ground RVs, of size Nj. Bx [t] equivalently denotes the batch of selected observed RVs. Inferring over (1), we will simulate training over the full distribution q:
N N
log 9 (t
log (0 ( (8e..)) nEB.10
(5)
where the factor N:/NF emulates the observation of as many ground RVs as in M by repeating the RVs from M (Hoffman et al., 2013). Similarly, the loss used at step t is the reduced ELBO constructed using X*(t] as observed RVs:
[t = Eggint [logp' (x'(t), (t) - log (e"(t))
logp" (+(+)) = Nx
Σ (e. nEBx((]
1 N N
Σ (( BEB
(6)
NN
This scheme can be viewed as the instantiation of M over batches of M's ground RVs. In Figure 1, we see that q² has the cardinalities of M and replicates its conditional dependencies. This training is analogous to stochastic VI (Hoffman et al., 2013), generalized with multiple hicrarchics, dependencies in the posterior, and mini-batches of RVs.
Sharing learning across plates Here we detail how our shared parameterization, detailed in Section 4.1, combined with our stochastic training scheme, results in faster inference. In stochastic VI (SVI, Hoffman et al., 2013), every 0t corresponding to the same template 0g is associated with individual weights. Those weights are trained only when the algorithm visits fn,, that is to say, at step when E B, As plates become larger, this event becomes rare. If 6.,, is furthermore associated with a highly- paralley-aammttiizeedd density estimator such normalizing flow- many optimization steps required for U. to converge. The
8

---

Published in Transactions on Machine Learning Research (11/2023)
combination of those two items leads to slow training and makes inference impractical in contexts such Neuroimaging, which can feature millions of RVs. With plate amortization, we aim to unlock inference in those large regimes by reducing the training time.
Instead of treating the ground RVs 8.n independently, we share the learning across plates. Due to the problem's plate structure, we consider the inference over the 0i, as different. instances of common density estimation task. In PAVI, large part of the parameterization of the estimators qi,n (0i.n//(:; Φi, Ei,)) is mutualized via the plate- hide-wide- weights Φj. This means that most of the weights of the flows Fun: concentrated in 6i, are trained at every optimization step across all the selected batches [t This results in drastically faster convergence than SVI, as seen in experiment 5.1.
4.3 Encoding schemes
PAVI-F and PAVI-E schemes PAVI shares the parameterization and learning of density estimators across an HBM's plates. In practice the distributions 4i.n (Oi,n (O0,n)) Φi, Ei,)) from Equation (4) with different n only differ through the value of the encodings E,,-- We detail two schemes to derive those encodings:
Free plate encodings (PAVI-F) In our core implementation, E are free weights. We define encoding arrays with the cardinality of the full model M, one array E, [E,,n per template 0,. This means that additional ground RV for instance, adding subject in a population study requires an additional encoding vector. The associated increment in the total number of weights is much lighter than adding a fully parameterized normalizing flow, as would be the case in the non-plate-amortizee regime. The PAVI-F scheme cannot be sample amortized: when presented with an unseen X, though o, can be kept as an efficient warm start, the optimal values for the encodings E,,, have to be searchod again.
During training, the encodings Ei,, corresponding to [. are sliced from the arrays E, and are optimized for along with Φi In the toy example from Figure 1, at t = 0, B, [0] (1,2} and the trained encodings are (1,),, and at t = 1 B, [1] = {0,1) and we train {E1,0, E1,
Deep set encoder (PAVI-E) The parameterization of PAVI-F scales lightly but linearly with Card Though lighter than the non-platc-amorttied case, this scaling could still become unaffordable in large population studies. We thus propose an alternate scheme, PAVI-E, with a parameterization independent of cardinalities. In this more experimental scheme, free encodings are replaced by an encoder f with weights y applied to the observed data: E f(X; n). As encoder Appendix 1.3 (Zaheer et al., 2018; Lee et al., 2019). Due to the plate structure, the observed X features multiple permutation invariances across data points corresponding to i.i.d. RVs. Deep sets are attention architectures that can model generic permutation invariant functions, As such, they constitute a natural design choice to incorporate the problem's invariances. The PAVI-E scheme allows for sample amortization across different data samples X,, XI,..,, ..., as described in Section 4.1. Note that an encoder will be used to generate the encodings whether the inference is sample amortized or not. Our encoder architecture is similar to the one from Agrawal & Domke (2021), considering only the case of fixed-cardinality plates, and generalizing to an arbitrary number of plates.
use deep-set architecture, detailed in
During training, shared learning is further amplified as all the architecture's weights -6; and p7 are trained at every step f. To collect the encodings to plug into q", we build up on a property of f: set size generulization (Zaheer et al., 2018). Instead of encoding the full-sized data X, f is applied to the slice (*(t). This amounts to aggregating summary statistics across subset of the observed data instead of the full data (Lee et al., 2019; Agrawal & Domke, 2021). The PAVI-E scheme has an even greater potential in the sample amortized context: we train sample amortized family over the lightweight model M and use it 'for free" to infer over the heavyweight model M.
Stochastic training and bias A key consideration is a potential bias introduced by our stochastic scheme. Namely, training stochastically over a variational family Q, we want to converge to the same solution q" as if we trained over the entirety of Q. In this section, we show that the PAVI-F scheme is unbiasod. In contrast, the PAVI-E scheme theoretically biased though we seldom noticed any negative impact of that bias in practice. A more thorough analysis can be found in Appendix A.2. Note that in this manuscript the term
9

---

Published in Transactions on Machine Learning Research (11/2023)
bias systematically refers to the stochastic training scheme. A different form of bias consists in the limited expressivity of the variational family Q which may not contain the true posterior p(O[X). We refer to this other bias as the variational family's gap, as further detailed in Appendix A.4.
To show that our training scheme is unbiased, we need to prove that the expectation of our stochastic loss is equal to the non-stochastic loss. This amounts to showing that:
Epaths [ELBO*[4] = ELBO
(7)
Eery [lo Q) - (())
where Epaths loosely refers to the paths that are sampled stochastically into the full model M's graph, defined in Section 4.2. Equation (7) means that the expectation of the reduced ELBO over all the stochastic paths that can be sampled inside the full model's graph is equal to the full ELBO, To prove Equation (7), it is sufficient to prove that:
Epaths (loga (+*()) log (()
(8)
meaning that the expectation of the reduced distribution q" over the stochastic paths equals the full distri- bution q. In Appendix A.2 we show that:
Epaths [loga (O"(t)
ΣΣ
patths [log qi (1. E.
(9)
1=1 ==
an expression that mimics the definition of in Equation (3) -albeit the expectations over paths. The rest
of the bias analysis depends on the encoding scheme.
Free plate encodings (PAVI-F) In the PAVI-F scheme, the encodings Ei are free weights. As consequence, their value does not depend the paths that are selected by the stochastic algorithm. This means that paths [logs,n ft, (0, (0 E..n)] log (0,,. E,,), proving Equation (8) and the unbiasedness of the PAVI- schome. With the PAVI-F scheme, training over M we converge to the same distribution as if training directly over M.
Deep set encoder (PAVI-E) In the PAVI-E scheme, encodings are obtained by applying the encoder f to a slice of the observed data. With the implementation presented in this work, the value of E, will depend which children of Oz, the stochastic algorithun selects. For instance, the encoding for a group will depend which subjects are selected inside the group to compute the group's summary statistics. Though computationally efficient, our implementation is thus thcoretically biased. The negative impact of this bias on PAVI-E's performance was seldom noticeable and always marginal throughout our experiments.
Technical summary In Section 4.1, we derived an architecture sharing its parameterization across a model's plates. This allows performing inference in large cardinality regimes without incurring an exploding number of weights. In Section 4.2 we derived a stochastic scheme to train this architecture over batches of data. Stochastic training allows large-scale inference without incurring an exploding memory and computing. Off-the-shelf stochastic VI would however result in siguificantly slower inference. Our novelty lies in our original combination of amort ization and stochastic training, which avoids this slow-down, as demonstrated in the following experiments.
5 Results and discussion
In this section, we show how PAVI unlocks hicrarchical Bayesian model inference for large-scale problems by matching the inference quality of SOTA methods while providing faster convergence and lighter param- eterization. Our experiments also highlight the differences between our two encoding schemes PAVI-E and PAVI-F. In summary:
1. Section 5.1 shows how plate amortization results in faster convergence compared to non-plate- amortized SVI;
2. Section 5.2 illustrates the role of the encodings E summary statistics in inference;
10

---

Published in Transactions on Machine Learning Research (11/2023)
3. Section 5.3 exemplifies -against baselines- PAVT's favorable scaling as problem's cardinality aug- ments;
4. Section 5.4 showcases the practical utility of PAVT by applying the method to a challenging popu- lation study.
Supplemental experiments In Appendix B.3, we evaluate the impact of the reduced model cardinalities performance. In Appendix B.4, we compare method against basclines over a mixture model; over model featuring the aggregation of higher-order summary statistics; and over smaller version of our Neuroimaging model used in Section 5.4. Those models complement the analysis of our main text, which is mostly concentrated on Gaussian models. Gaussian models constitute standard model featuring an approximate closed-form posterior which we use for validation.
ELBO metric Throughout this section, we use the ELBO as a proxy for the KL divergence between the variational posterior and the unknown true posterior (Blei al., 2017). ELBO is measured across 20 samples X, with repelitions per sample. The ELBO allows to compare the relative performance of different architectures given inference problem. In Appendix B.2, we provide sanity checks to assess the quality of the results.
Gaussian random effects (GRE) model In our experiments Section 5.1, 1,1,5,. 5.2, 5.3, we focus on a standard hierarchical model: a Gaussian random effects model (Digle et al., 2013; Gelman et al., 2004). The GRE model constitutes an inluitive inference problem inferring group and population means given observations- and features an approximate closed-form solution, as described in the next paragraph. This closed-form solution helps us display thcoretical bounds for the ELBO, for instance in Figure 2.
The GRE model can be described using the following set of equations:
X mg ( (01, oi)
vn,=1.. Card() =.. Card(Po)
(10)
0 0 02 NN2,0 of
1ml=1. Card(P.)
02,0 ~ ~N(On, o3)
where D represents the data X's feature size, with group means 0 and population means 02 as D-dimensional Gaussians. The GRE model features two nested plates: the group plate Pı and the sample plate Po as in Figure Taking our introductory example from Section 3, X represents the weight for subject no in group Th. 01,m, represents the mean weight in the group ni. 02 represents the mean weight in the population. Inferring over the GRE model, the objective is to retrieve the posterior distribution of the group and population means given the observed sample.
As baseline for comparison, we compare the ELBO of varions methods to an approximate closed-form baseline's ELBO. Though a closed-form posterior cannot be derived in the 3-level case, a good approximation can be constructed using Gaussian distributions centered on the empirical group and population means. This asymptolic ELBO represented using dashed lines Figures 2 and 3.
Asymptotic closed-form ELBO
5.1 Plate amortization and convergence speed
In this experiment, we illustrate how plate amortization results in faster training.
We use the GRE model, described in Equation (10), setting D 8, Card = 100 and Card"(Pi) = In this experiment, we set Card (Pi) & Card(Pi). This emulates regime in which SVI is slow because only small fraction of size Card"((P)- of large parameter space of size Card(Pi) gets optimized at given stochastic training step. We compare our PAVI architecture to bascline with the same architecture, trained stochastically with SVI (Hoffman et al., 2013), but without plate amortization. The only difference is that ground RVs 0z,, are associated in the baseline to individual fully-parameterized flows Fi, instead of sharing the same conditional flow F,, as further described in Section 4.1.
Figure 2 (left) displays the evolution of the ELBO across training steps for the bascline and PAVI with free encoding (PAVI-F) and deep-set encoders (PAVI-E). We see that both plate-amortized methods reach
11

---

Published in Transactions on Machine Learning Research (11/2023)
5000
ELBO thigher is better)
-5000
ELBO Ihighete better
-5000
-10000
-10000
-1500
encoding size:
-15000
SVI baseline
encoding
(not plate amortized)
encoding s2e:
-20000
PAVI-E (plate amortized)
(plate amortized)
-20000
encoding
encoding size:
encoding size: 16
Asymptotic ELBO
Asymptotic ELBO
25000
10a
10a
10a
500
1000
1500
2000
2500
3000
Steps
Steps
Figure 2: Left panel: Plate amortization increases convergence speed Plot of the ELBO (higher is better) as a function of the optimization steps (log-scale) for our methods PAVI-F (in green) and PAVI-E (in blue) versus a non-plate-amortized bascline (in purple). Standard deviation across repetitions is displayed as a shaded area. A dashed line denotes the asymptotic closed-form performance. Due to plate amortization, our method converges ten to a hundred times faster to the same asymptotic ELBO as its non-plate-amortine ized counterpart.; Right panel: Encodings as ground RVs summary statistics Plot of the ELBO (higher is better) as function of the optimization steps for the PAVI-F architccture with increasing encoding sizes. Standard deviation across repetitions is displayed as a shaded area. A dashed line denotes the asymplolic closed-form performance. As the encoding size augments, does the asymptotic performance until reaching the dimensionality of the posterior's 's sufficient statistics (D= 8), after which performance plateaus. Encoding size allows for a clear trade-off between memory and inference quality.
asymptotic ELBO equal the non-plate-amortized baseline's but with orders of magnitudes faster conver- gence and more numerical stability. This stems from the individual flows Fi.. in the baseline only being trained when the stochastic algorithm visits the corresponding 0. In contrast, our shared flow F; is updated at every optimization step in PAVI. Intuitively, the PAVI-E scheme should converge faster than PAVI-F by sharing the training not only of the conditional flows but also of the encoder across the different optimization steps. However, the computation required to derive the encodings from the observed data results in longer optimization steps and in slower inference, as illustrated in Section 5.3. The usage of an encoder nonetheless allows for sample amorlization with the PAVI-F scheme, which is impossible with the PAVI-F scheme.
5.2 Impact of encoding size
Here we illustrate the role of encodings as ground RV posterior's summary statistics, as further described in Section 4.1. We use the GRE HBM detailed in Equation (10), using D 8, Card(Pi) = 20 and CardF(Pi) = 2. We use single PAVI-F architecture, varying the size of the encodings Ei,,, which are defined in Section 4.3.
Due to plate amortization, encodings determine how much individual information each RV 0z.''s 's posterior is associated with. The encoding size varying from 2 to 16 is to be compared with the problem's dimensionality, D = 8. In GRE, D = 8 corresponds to the size of the sufficient statistics needed to reconstruct the posterior of a group mean since all other statistics, such as the variance, are shared ween the group mcans.
Figure 2 (right) shows how the asymptotic performance steadily increases when the encoding size augments and plateaus once reaching the sufficient summary statistic size D = 8. Interestingly, increasing the encoding size also leads to faster convergence: redundancy can likely be exploited in the optimization. Incrcasing the encoding size also leads experimentally lo diminishing returns in terms of performance. This properly can be exploited in large settings to drastically reduce the memory footprint of inference while maintaining acceptable performance by choosing the encoding size approximately equal the expected size of the sufficient statistics. Encoding size appears as an unequivocal hyperparameter allowing to trade inference quality for computational efficiency.
12

---

Published in Transactions on Machine Learning Research (11/2023)
Inference quality
Parameterization
Time
Delimin aatimization time (lower is better)
10*
Card
10a
Card(P))==2
ELBO (higher better)
Weights (lower is better)
Card/P))==00
Heavy GPU
10a
memory
10
requirement
Asymptotic ELBO (first pannel)
10a
10a
Crica) ADANN(aa)
Cr
unn
PAUs-R
Plyt-Elsa) Elsa)
PAINEE
ADAVilsa)
Chisa)
Cr
Uivi
Poyn (ass))
PAVI-E
ADAVRisal
(risa)
Cr
Unn PAVi-E
PAyS Etsa)
PAI-EE
Figure 3: PAVI provides favorable parameterization and training time as the cardinality of the target model augments Basclines are compared in each panel, with the suffix (sa) indicating sample amortization. Our architecture PAVI is displayed on the right of each panel. We augment the cardinality (P) of the GRE model, which is described in Equation (10). While doing so, we compare three different metrics. In the first ponel: inference quality, as measured by the ELBO. An asymptotic closed-form form ELBO is displayed using a dark blue dash. None of the presented SOTA architecture's performance dogrades the cardinality of the problem augments. In the second panel: parameterization, comparing the number of trainable weights of each architecture. PAVI similar ADAVI- displays constant number of weights the cardinality of the problem increases -or almost constant for PAVI-F. Third panel: GPU training time. Benefiting from learning across plates, PAVI has a short and almost constant training time as the cardinality of the problem augments. At Card(Pi) = 200, CF, UTVI, and ADAVT required large GPU memory, a constraint absent from PAVI due to its stochastic training. For convenience, the results used to plot this figure are reproduced in supplemental Table B.1.
5.3 Scaling with plate cardinalities
Here put in perspoctive the gains from plate amortization when scaling up inference problem's car- dinality. We consider the GRE model in Equation (10) with D 2 and augment the plate cardinalities (Pi), Card (2, 1) > (20,5) - (200, 20). In doing we augment the number of parameters θ: 6 + 42 + 402.
Baselines We compare our PAVI architecture against three state-of automatic VI baselines.
Cascading Flows (CF) (Ambrogioni et al., 2021b) is non-plate-amortizee structured VI archi- tecture improving on the baseline presented in Section 5.1. CF pushes the prior into the posterior using Highway Flows. CF follows a cascading dependency structure complemented by a backward auxiliary coupling. CF thus consists a structured baseline that does not pay particular attention to scalability to large cardinalities;
Automatic Dual Amortized VI (ADAVI) (Rouillard & Wassermann, 2022) is a structured VI architecture with constant parameterization with respect to problem's cardinality but large training times and memory. ADAVI has several limitations compared to PAVI. First, ADAVI implements a Mean Field approximation (Blei et al., 2017) while PAVI implements a cascading flow. This means that, contrary to PAVI, ADAVI can be biased in cases where the posterior features strong dependencies. Second, ADAVI is limited pyramidal HBMs, while PAVI tackles generic plate- enriched HBMs, making it more general. Third, ADAVT is limited to full-model sample-amortizel variant, which ultimately caps the cardinality of problems it can be used on. ADAVI thus consists in a structured bascline that tackles the parameterization hurdle of large cardinalities, but not the computational efficiency hurdle;
Unbiased Implicit VI (UIVI) an unstructured implicit VI architecture. UIVI infers over the full parameter space θ without any SVT-amenable factorization contrary to CF, ADAVT, and
13

---

Published in Transactions on Machine Learning Research (11/2023)
PAVI. To do so, UIVI reparameterizes a base distribution with stochastic transform. UIVI does not explicitly define density q and relies on optimization steps intertwined with MCMC runs. UTV thus consists in a non-structured VI baseline that does not pay particular attention to scalability to a large paramcter space This means that UIVI could not be applied above a certain cardinality due to its impossibility to be stochastically trained.
For all architectures, we indicate with the suffix (sa) sample amortization, as defined in Section 4.1. More implementation details can be found in Appendix B.5.
As the cardinality of the problem augments, Figure 3 shows how PAVI maintains a state-of-the-art inference quality while being more computationally attractive.
Parameterization In terms of parameterization, both ADAVI and PAVI-E provide a heavyweight but constant parameterization as the cardinality Card(P;) of the problem augments. This is due to both methods' usage of an encoder of the observed data, which makes their parameterization independent of the problem's cardinality. Comparatively, both CF and PAVI-F's parameterization scale linearly with Card(Pi), but with a drastically lighter augmentation for PAVI-F. The difference can be explained by the additional weights required by each architecture for an additional ground RV. CF requires an additional fully parameterized normalizing flow, whereas PAVI-F only requires an additional lightweight encoding vector. In detail, PAVI- F's parameterization due to the shdde- Φ1 represents a constant Re 2k weights, while the part due to the encodings Ei,x grows lincarly from 16 to 160 to 1.6k weights. UIVT's paramcterization scales quadratically with the size of the parameter space O This is due to UIVI's usage of a neural network to regress the weights of a transform applied to a base distribution with the size of e, As the cardinality augments, UIVT's quadratic weight scaling would be limiting before CF and PAVI-F's linear scaling, which would be limiting before ADAVI and PAVI-E's constant scaling.
Memory Regarding computational budget, PAVI's stochastic training allows for controlled GPU memory during optimization. This removes the need for a larger memory as the cardinality of the problem aug- ments, a hardware constraint that can become unaffordable at very large cardinalities. CF could be trained stochastically to remove this memory constraint but, without plate amortization, would suffer from slower inference, as illustrated in Section 5.1. In contrast, UIVI could not be trained stochastically, as infers over the full parameter space at once instead of factorizing it. As result, UIVI would be ultimately limited by memory to infer over larger problems.
Speed Regarding convergence speed, PAVI benefits from plate amortization to have orders of magnitude faster convergence compared lo structures ured VI baselines CF and ADAVI. This means that a stochast ically- trained architecture (PAVI) trains faster than non-stochastic baselines (CF, ADAVI) This result is opposite to the result we would have obtained without plate amortization since SVI slows down inferenco. For CTVI, as the cardinality augments, training amounts to evaluating a neural network with increasingly larger layers. GPU training time is thus constant, but this property would not translate to larger problems as the GPU memory would become insufficient. Plate amortization is particularly significant for the PAVI-E(sa) scheme, in which a sample-amortized variational family is trained over dataset of reduced cardinality yet. performs "for free" inference over an HBM of large cardinality. Maintaining Card"(P:) constant while Card(P;) augments allows for constant parameterization and training time as the cardinality of the problem augments This property is not limited to any maximum cardinality, contrary to UIVI. This is a novel result with strong future potential. The effect of plate amortization is particularly noticeable at Card = 200 between the PAVI-E((sa) and (F(sa) architectures, where PAVT performs sample-amortized inference with 10x fewer weights and 100x lower training time.
Scaling even higher the cardinality of the problem (Pr) 2000 for instance renders ADAVI, CF and UTVI computationally intractable In contrast, PAVI maintains a light memory footprint and a short training time, as exemplified in the next experiment.
Using PAVI on small scale problems By design, PAVI is meant to tackle large cardinality problems effectively. The use of PAVI on smaller problems is thus an open question. As Figure 3 underlines, in the Card(Pi) 2 case, PAVI-F matches the ELBO and training speed of CF but features a heavier parameter- ization. In addition, plate amortizing variational family theoretically reduces its expressivity, further
14

---

Published in Transactions on Machine Learning Research (11/2023)
explained in Appendix A.4. In theory, simpler, non-plate-amortizee baseline such as CF could thus be preferred for small problems. However, in practice, VT amounts to an optimization problem, and the most expressive variational family will not necessarily yield the best performance on a complex problem (Bottou & Bousquet, 2007). Our comparative experiments in Appendix B.4 exemplify this. In Appendix B.4, PAVT's repeated paramelerization oftentimes eases the optimization and yields a better ELBO. We leave for future work the analysis of the impact of PAVT's repeated parameterization wwhich could be considered as an inductive bias on properties such as inference calibration (Talts et al., 2020). In conclusion, when its parameterization is affordable, PAVI remains relevant choice for small plate-enriched problems.
5.4 Neuroimaging application: full cortex probabilistic parcellation for cohort of 1,000 subjects
To illustrate its usefulness, we apply PAVI to a challenging population study for full-cortex functional parcellation.
A parcellation clusters brain vertices into different connectivity networks: fingerprints describing co-activation with the rest of the brain, as measured using functional Magnetic Resonance Imaging (IMRI). parcellation is essentially a clustering of a subject's 's cortex vertices based on their connectivity with the rest of the brain. Different subjects exhibit a strong variability, as visible in Figure 4. However, fMRI is costly acquire: few noisy measurements usually made for given subject. It is thus essential to combine information across the different measurements for given subject and across subjects and to display uncertainty in the results. Those 2 points motivate hicrarchical Bayesian modelling and VI in Neuroimaging (Kong et al., 2019): we search the posteriors of connectivity networks and vertex labels, measuring [MRI over a large cohort of subjects.
We use the HCP dataset (Van Essen et al., 2012): 2 acquisitions from a 1000 subjects, with millions of measures per acquisition, and over 400 million parometers O to infer. We use a model with three plates: subjects, sessions and brain vertices. We cluster the brain vertices into 7 clusters, following Thomas Yeo et al. (2011). None of the baselines presented in Section 5.3 -CF, ADAVI, UIVI- can computationally tackle this high cardinality problem. We nevertheless show superior performance over those baselines over a tractable problem size with 2,000 parameters in Appendix B.4.3.
Despite the massive dimensionality of the problem, thanks to plate amortization, PAVI converges in less than hours of GPU time. Results are visible in Figure 4. Qualitatively, the recovered networks match existing expert knowledge motor control or vision (Heim et al., 2009; Zhang et al., 2020). Quantitatively, using the individual subject's parcellation as feature for cognitive score regression task, we obtain better scores compared to state-of- the-art methods (Thomas Yeo et al., 2011; Gordon et al., 2017; Wang et al., 2015; Kong et al., 2019). This mcans that a subject's brain spatial organization -where each function is "located" in the brain- can be used to predict the subject's coguitive ability, such as memory, reading, or spatial orientation.
the brain's functional organization. For instance, we recover networks responsible for
Using PAVI, we obtain state-of-thee the--rrt functional parcellations for a thousand subjects, capturing the uncer- tainty in this challenging task via VI while maintaining a manageable training time despite a 400-million-sieed parameter space.
5.5 Conclusion
In this work we present the novel PAVI architecture, combining a structured variational family and a stochas- tic training scheme. PAVI is based the on concept of plate amortization, allowing to share parameterization and learning across a model's plates. We demonstrated the positive impact of plate amortization on training speed and scaling to large plate cardinality regimes, making a significant step towards scalable, expressive Variational Inference.
Reproducibility statement
All experiments were performed in Python using the Tensorflow Probability library (Dillon et al., 2017). All experiments were conducted on computational cluster nodes equipped with a Tesla V100-16Gb GPU and 4 AMD EPYC 7742 64-Core processors. VRAM intensive experiments in Figure 3 were performed
15

---

Published in Transactions on Machine Learning Research (11/2023)
motor
attention
Accuracy for
13 cognitive measures (Higher is better)
Method
vision
MS-HBM
0.1321 (+0.0053)
from Kong et al. (2019)
Population parcellation
YeoBeckProject
0.1057 (+0.0060)
from Calhoun Adali (2012)
Gordon2017
0.0545 (#0.0062)
from Gordon al. (2017)
Wang2015
0.1202 (#0.0054)
from Wang et al. (2015)
Ours (PAVI)
0.1645 (=0.0047)
Subject parcellation
Subject 1000 parcellation
Figure 4: Probabilistic full cortex parcellation PAVI can be applied in the challenging context of Neuroimaging population studies. For a cohort of 1000 subjects, of which are represented here (in the bottom 2 items) we cluster 60, 000 cortex vertices according to their connectivity with the rest of the brain. On the left, we show the obtained probabilistic parcellations. Each color in the parcellation corresponds one of functional network (Thomas Yeo et al., 2011). Networks represent groups of neurons that co- activate in the brain and can be associated with certain cognitive functions, such vision motor control. Through our method, recover each subject's 's parcellation (al the bottom), which are i.i.d. perturbations of the population's parcellation (at the top). Our method also models uncertainty: coloring represents the dominant label for cach vertex, and the lovel of white increases with the uncertainty in the labeling. Wo can use the subject parcellations as features for a cognitive score prediction task, as visible in the table on the right. The table shows the mean predictive accuracy across cognitive measures, including memory, pronunciation, processing speed or spatial orientation. The baseline methods scores are reproduced from Kong et al. (2019)'s implementation. Our method produces individual maps that predictive of the subject's cognitive ability.
16

---

Published in Transactions on Machine Learning Research (11/2023)
Ampere 100 PCIE-40Gb GPU. Appendix B.5 lists implementation details for all our syuthetic experiments. Appendix B.6 is related to our Neuroinaging experiment 5.4, and details both our data pre-processing steps and our implementation. As part of our submission we furthermore packaged and release the code associated to our experiments.
Broader Impact Statement
Our Neuroimaging data come from the Human Connectome Project dataset (Van Essen et al., 2012). All data in the HCP is strongly anonymized, as per the HCP protocols. We used in this paper only Open Access imaging data data, following the HCP data use terms.
Acknowledgments
This work was supported by the ERC-StG NeuroLang ID:757672. This work was performed using HPC resources from GENCI-IDRIS (Grant 2022-102043).
References
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Cor- rado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg. Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sulskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogencous systems, 2015. URL https://wwz tensorflow....... org/. Software available from tensorflow.org.
Alexandre Abraham, Fabian Podregosa, Michael Eickenberg. Philippe Gervais, Andreas Mueller, Jean Kos- saifi, Alexandre Gramfort, Bertrand Thirion, and Gael Varoquaux. Machine learning for neuroimaging with seikit-learn. Frontiers in Newroinformatics, 8, 2014. ISSN 1662-5196. doi: 1..389/ffinff.. manff201440000.. URL https: / /ttp///wwwww//ooo www.tterroinnnnnommmmr wwwp/rrtigeee//...//o/o /tte//rrtiiiee/...38////iittt 2014.00014.
Abhinav Agrawal and Justin Domke. Amortized Variational Inference for Simple Hicrarchical Models. M. Ranzato, Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan (eds.), Advances Neural Information Processing Systems, volume 34, pp. 21388 21399. Curran Associates, Inc., 2021. URL https: Pacold--aprrpppaee .af..r..ff...
//proceedings. www.eeeddinnnnnunis......
Luca Ambrogioni, Kate Lin, Emily Fertig, Sharad Vikram, Max Hinne, Dave Moore, and Marcel van Gerven. Automatic structured variational inference. In Arindam Banerjec and Kenji Fukumizu (cds.), Proceedings of The 24th International Conference on Artificial Intelligence and Stalislics, volume 130 of Proceedings of Machine Learning Research, pp. 676 684. PMLR, 13 15 Apr 2021a. URL press/v130/amie ptres//v30//mmrrogiomm22 awtp/aaarrooiiiii1. html.
mlr.
Luca Ambrogioni, Gianluigi Silvestri, and Marcel van Gerven. Automatic variational inference with cascading flows. In Marina Meila and Tong Zhang (ods.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 254 263. PMLR, 18 Jul 2021b. URL httpp://tttttttc.t
htwp:/ns........./.eeeeeeesse ambrogioni21a. htmml
Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science and Statistics). Springer-Veriee Berlin, Heidelberg, 2006. ISBN 0387310738.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational Inference: A Review for Statisticians. Journal of the American Statistical Association, 112(518):859 April 2017. ISSN 0162-1459, 1537-274X. doi: 10.1080 10.1080001683 /01 459.2017..22 15952717.1177777777 URT htt.:ttt../aaaav..orr/b//////// 00670. ar Xiv: 1601.00670.
17

---

Published in Transactions on Machine Learning Research (11/2023)
Léon Bottou and Olivier Bousquet. The Tradeoffs of Large Scale Learning. In J. Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Neural Information Processing Systems, vol- une 20. Curran Associates, Inc., 2007. URL https://proceeddnng wttp:sesww..ppppppee////e/ee cc/paper/2007/iile Paper. påf.
Vince D. Calhoun and Tülay Adali. Multisubject Independent Component Analysis of fMRI: A Decade of Intrinsic Networks, Default Mode, and Neurodiagnostic Discovery. IEEE Reviews in Biomedical
Engineering, 5:60-73 73, 2012. ISSN 1937-3333, 1941-1189.
doi: 241109/RBME.2011220
URL
http:///eeexploree ieee. eeree.rr/eeeeeeeeeeeeeeeeeeeeeee wtt.:i......oommatt/eeeee//3/3 /oc./mmntt/26632///
Chris Cremer, Xuechen Li, and David Duvenaud. Inference Suboptimality in Variational Autoencoders. arXiv:1801.0353 [es, stat], May 2018. URL http://arraiiv.orggab////e 03555. arXiv: 1801.03558.
Kamalaker Dadi, Gael Varoquaux, Antonia Machlouzarides-Shalii Krzysatof J. Gorgolewski, Demian Wasserman, Bertrand Thirion, and Arthur Mensch. Fine-grain atlases of functional modes for fMRI
http://o...rr.......///eerrraaaa analysis. Neurolmage, 221:117126, 2020. ISSN 1053-8119. doi: htt.////eeeettteee/////////asas0
2020.117126. URL http::///www.sccnedeeeettttt htw.eddireettttteeeeeeeeeeeeeeee
Viet-Hung Dao, David Gunawan, Minh-Ngoc Tran, Robert Kohn, Guy E. Hawkins, and Scott D. Brown. Efficient selection between hicrarchical cognitive models: Cross-validation with variational bayes, 2021. URL http://wwww...a///////........
Peter Diggle, Patrick Heagerty, Kung-Yee Liang, and Scott Zeger. Analysis of Longitudinal Data. OUP Oxford, March 2013. ISBN 978-0-19-967675-0. Google-Books-DD ID: urOBIXPuOukC.
Joshua 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.1060) [cs, stat/, November 2017. URL http:// arxiv. org/abs//711. 10604. arXiv: 1711.10604.
Yann Dubois, Jonathan Gordon, and Andrew YK Foong. Neural process family. wtweeerral-Poocsss--ee
gi github.
September 2020.
A Fayaz, P Croft, RM Langford, LJ Donaldson, and GT Jones. Prevalence of chronic pain in the uk: a systematic review and meta-analysis of population studies. BMJ open, 6(6):0010364, 2016.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Cham- bon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python opti- mal transport. Journal of Machine Learning Research, 22(78):1 8, 2021. URL http://jmlr, http:jillrrrggppee org. org/papers/ v22/20-455 html.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J. Rezende, S. M. Ali Eslami, and Yee Whye Teh. Neural processes, 2018. URL https://arxiv..rgg
01622.
Andrew Gelman, John B. Carlin, Hal S. Stern, and Donald B. Rubin. Bayesian Dala Analysis. Chapman and Hall/CRC, 2nd ed. edition, 2004.
W. R. Gilks, A. Thomas, and D. J. Spiegelhalter. A Language and Program for Complex Bayesian Modelling, The Statistician, 43(1):169, 1994. ISSN 00390526. doi: (0.2307/2348941) URL https://www http://www..oorrr/oo//////o///.t jstor. org/ stable/10 2307 http:////9911rriinnrrroorae
Evan M. Gordon, Timothy O. Laumann, Adrian W. Gilmore, Dillan J. Newbold, Deanna J. Greene, Jef- frey J. Berg, Mario Ortega, Catherine Hoyt-Drazen, Caterina Gratton, Haoxin Sun, Jacqueline M. Hamp- lon, Rebecca S. Coalson, Annie L. Nguyen, Kathleen B. McDermott, Joshua S. Shimony, Abraham Z. Snyder, Bradley L. Schlaggar, Steven E. Petersen, Steven M. Nelson, and Nico U. F. Dosenbach. Precision Functional Mapping of Individual Human Brains. Newron, 95(4):791-807.c7 2017. ISSN 0896-6273. doi: https:///oo.org/100006/// 200.0...007.07.011.. article/aai/S00667773733333
URL https: //ww::://www... wew.ncedireect.................. com/science/
18

---

Published in Transactions on Machine Learning Research (11/2023)
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 20512.
Matthew D. Hoffman and David M. Blci. Structured Stochastic Variational Inference. arXiv:140444114 [es). November 2014. URL http: /awxiv.xvvvvv//////////////v/v org/abs//404..4 4114. arXiv: 1404.4114.
Matthew D. Hoffman, David M. Blei, Chong Wang. and John Paisley. Stochastic variational inference.
Journal of Machine Learning Research, 14(4):1303-133 1347, 2013. URL htttp:/jjmr..rgge ortpppprss///11/ hoffman13a.het html.
Ekaterina lakovleva, Jakob Verbeek, and Karteek Alahari. Meta-nga with shared amortized variational inference. In Hal Daumé III and Aarti Singh (cds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 4572-4582 PMLR, 13-18 Jul 2020. URL https: http::s//proceeding mlr ppress/v119/aakklleeeeee hhmml.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In In- ternational Conference on Learning Representations, 2017. URL https://openreeviek htwp:pppwrrreeeeeeeeeeererrrraea net/forumiid rkE3y85ee.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL
Daphne Koller and Nir Friedman. Probabilistic gruphical models: principles and techniques. Adaptive computation and machine learning. MIT Press, Cambridge, MA, 2009. ISBN 978-0-262-01319-2.
Ru Kong, Jingwei Li, Csaba Orban, Mert Rory Sabuncu, Hesheng Liu, Alexander Schaefer, Nanbo Sun, Xi-Nian Zuo, Avram J. Holmes, Simon B. Eickhoff, and B. T. Thomas Yeo. Spatial topography of indivicuull-specific cortical networks predicts human cognition, personality, and emotion. Cerebral cortex, 29 6:2533-25555 2551, 2019.
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M. Blei. Automatic Differ- entiation Variational Inference. arXiv: 160.000888888 fes, statf, March 2016. URL http: .rgg abs/ 1603. 00788.
arXiv: 1603.00788.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Soungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-basod cermutation-invariae neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3744 3753. PMLR, 09 15 Jun 2019. URL http://
proceedings mlr. wwt..preessv//// /ete//el91e html.
Chris J. Maddison, Andriy Mnih, and Yee whye Tch. The concrete distribution: A continuous relaxation of discrete random variables. CoRR, abs/161 hhs/1611.00712 2016. URL http://arxiit wwwp:rr.g//eeoooeee http//aa.../1/11/01//1/1/11
George Papamakarios and lain Murray. Fast SkepsilonS-fre Inference Simulation Models with Bayesian Conditional Density Estimation. arXiv:1605, 06376 fcs, stat), April 2018. URL httt:://axxxt org/ abs/ 06376. arXiv: 1605.06376.
George Papamakarios, Theo Pavlakou, and lain Murray, Masked Autoregressive Flow for Density Estimation. arXiv: (5l::1705000577 fes, stat], June 2018. URL http:///axxw..or//ab////////// wwtp//////////////////// arXiv: 1705.07057.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshmi- narayanan. Normalizing Flows for Probabilistic Modeling and Inference. arXi::1912.0222 02762 [es, stat), December 2019. URL http:// http://arxiv..tt wwwaaaavvvv.o. http:p//aaxii.ooogaa/////2222.2
arXiv: 1912.02762.
Sachin Ravi and Alex Beatson. Amortized bayesian meta-learning. In International Conference on Learning Representations, 2019. URL http:: htttps://openreviee. http://www....eeeoooommm/d//rrgs
19

---

Published in Transactions on Machine Learning Research (11/2023)
Louis Rouillard and Demian Wassermann. ADAVI: Automatic Dual Amortized Variational Inference Applied To Pyramidal Bayesian Models. In ICLR 2022, Virtual, France, April 2022. URL https: //hal archives wwweerve-oouvertssssssee.... ft/haa-03267966
Francisco Ruiz and Michalis Titsias. A contrastive divergence for combining variational inference and MCMC. In Kamalika Chaudhuri and Ruslan Salakhutdinov (ods.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 5537 5545. PMLR, html.
09-15 Jun 2019. URL https:
mlr. ppwss/v97/ruuiie
Hojjat Salehinejad, Julianne Baarbe, Sharan Sankar, Joseph Barfett, Errol Colak, and Shahrokh Valaee. Recent advances in recurrent neural networks. CoRR, abs/1801 0011100010007,, 2018. URL http://arxiv. org/ 0tts/800..01077
Stephen M Smith, Christian F Beckmann, Jesper Andersson, Edward J Auerbach, Janine Bijsterbosch, Gwenaälle Donaud, Eugene Duff, David A Feinberg, Ludovica Griffanti, Michael P Harms, et al. Resting- state fmri in the human connectome project. Neuroimage, 80:144 168, 2013.
Sean Talts, Michael Betancourt, Daniel Simpson, Aki Vehtari, and Andrew Gelman. Validating Bayesian Inference Algorithms with Simulation-Bar Based Calibration, October 2020. URL hotp://arxiv.oot abs/ 1804.06788. arXiv:1804.06788 [stat].
B. T. Thomas Yeo, Fenna M. Krienen, Jorge Sepulcre, Mert R. Sabuncu, Danial Lashkari, Marisa Hollinshead, Joshua L. Roffman, Jordan W. Smoller, Lilla Zöllei, Jonathan R. Polimeni, Bruce Fischl, Heshong Liu, and Randy L. Buckner. The organization of the human cerebral cortex estimated by intrinsic functional connectivity. Journal of Newrophysiology, 106(3):1125-1165, 2011. doi: PMID: 21653723.
/tt////..003..21111111 2011.
URL ttttps://don http://do..or///.///// 1152/jn jn. 00338 2011.
Michalis K. Titsias and Francisco Ruiz. Unbiased implicit variational inference. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 167-176. PMLR, 16 18 Apr 2019. URL http:::/ https://proceeddngg waapiinnn..m.....mmm............ mlr wwl.sppres//aee wtt////ttiisili.eeeeeeeeeaaaaatt html.
Kayle Towsley Michael Shevell, Lynn Dagenais, Repacq Consortium, et al. Population-based study of neuroimaging findings in children with cerebral palsy. European journal of paediatric neurology, 15()): 39-35, 2011.
Dustin Tran, Rajesh Ranganath, and David M. Blei. Hierarchical Implicit Models and Likelihood Variational Inference. arXiv: 08896 fes, stat), November 2017. URL http:://rrtt org/ abs/ 1702. 08896. arXiv: 1702.08896.
D. C. Van Fssen, K. Ugurbil, F. 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 F. Yacoub. The Human Connectome Project: a data acquisition perspective. Neuroimage, 62(4):2222 2231, Oct 2012.
Danhong Wang, Randy L Buckner, Michael D Fox, Daphne J Holt, Avram J Holmes, Sophia Stoecklein, Georg Langs, Ruiqi Pan, Tianyi Qian, Kuncheng Li, Justin T Baker, Steven M Stufflcbeam, Kai Wang, Xiaomin Wang, Bo Hong, and Hesheng Liu. Parcellating cortical functional networks in individuals. Nature Neuroscience, 18(12):1853 1866,16068, December 2015. ISSN 1097-6256, 1546-1726, http://www.natune com/ ammaarticles/mm 4164.
doi: 10. 10.1038/mn.41 URL
Stefan Webb, Adam Golinski, Robert Zinkov, N. Siddharth, Tom Rainforth, Yee Whye Teh, and Frank Wood. Faithful Inversion of Generative Models for Effective Amortizod Inference. 1712 00287 fes, stat/, November 2018. URL http://arxiv. wwtp:///////...////////./o 002877...000....... arXiv: 1712.00287.
Karl Weiss, Taghi M. Khoshgoftaar, and DingDing Wang- A survey of transfer learning. Journal of Big Data, 3(1):9, May 2016. ISSN 2196-1115. doi: 01.1186/440537-01--008333 URL attps://dot org/10. 10. 1 -186//44533-001-0004
20

---

Published in Transactions on Machine Learning Research (11/2023)
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and Philip S. Yu. A Comprehensive Survey Graph Neural Networks. JEEE Transactions on Neural Networks and Learning Systems, pp. 1 1-21, 2020. ISSN 2162-237X, 2162-2388. doi: 10.110/1NNN.S.202227733333 URL http://arxiv.orr 00596. arXiv: 1901.00596.
Huaxiu Yao, Ying Wei, Junzhou Huang, and Zhenhui Li. Hicrarchically structured meta-learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 7045- 1055. PMLR, 09-15 Jun 2019. URL https: https::/proceeddngg mlr. wwt..prressvvvv wawl///albb.... html.
Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1860-5669. PMLR, 10-15 Jul 2018. URL https:// mlr.
press/v80/ limili.nnnhhiie...mm
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexan- der Smola. Deep Sets. Xrxiv: [es, stat), April 2018. URL wttp://wwww.rg///////o........
arXiv: 1703.06114.
Cheng Zhang, Judith Butepage, Hedvig Kjellstrom, and Stephan Mandt. Advances in Variational Inference. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(8):2008-200 2026, August 2019. ISSN 0162-8828, 2160-9292, 1939-3539. doi: 1LL10/TPPAM22018..0 28897777 URL https://tteexploreeeeee (tfoccmmmnt/88833999/
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.1038 -414467-000-580444
21

---

Published in Transactions on Machine Learning Research (11/2023)
Appendix
A Supplemental methods
A.1 PAVI implementation details
A.1.1 Plate branchings and stochastic training
As exposed in Section 4.2, at each optimization step t we randomly select branchings inside the full model M, branchings over which we instantiate the reduced model M In doing so, we define batches [t for the RV templates 0g. Those batches have to be coherent with one another: they have to respect the conditional dependencies of the original model M. As example, if ground RV is selected part of M then its parent RV needs to be selected well. To ensure this, during the stochastic training we do not sample RVs directly but plates:
For every plate Pp. we we sample without replacement Card' (Pp) indices amongst the Card possible indices.
Then, for every RV template 0,, we select the ground RVs 0n,, corresponding to the sampled indices for the plates Plates(0,).
3. The selected ground RVs 0₂2, constitute the set O O*(t] of parameters appearing in Equation (5). The same procedure yields the observed RV subset X [t] and the data slice x (=(t).
For instance, in the toy example from Figure 1, X2 will be chosen iff the index is selected as part of sub-sampling P1 and the index selected as part of sub-sampling Po- Less formally, this is equivalent to going middle, then left in the full graph representing M. This stochastic choice is illustrated in Figure where X2 corresponds to the node E. This stochastic strategy also applies to the selected encoding scheme described in Section 4.3&4.2- detailed in the next sections.
A.1.2 PAVI-F details
In Section 4.3 we refer to encodings E, = [E..n corresponding RV templates 0,. In practice, we have some amount of sharing for those encodings: instead of defining separate encodings for every RV template, define encodings for every plate level. A plate level is combination of plates with at least one parameter RV template 0; belonging to it:
PlateLevels ((Pp..Pi) (l.)es(0.))a.cc
(A.1)
For every plate level, we construct large encoding array with the cardinalities of the full model M:
Encodings P())..P)) R xCon(PP)x.xCaddP)) (raeaaaadddaaaaaaaaaaaaaaooooo PPeattsiell
(A.2)
E, Encodings (Picooddng((Patess(()))) (Plates(8.))
Where D is encoding size that we kept constant to de-clutter the notation but vary between plate levels. The encodings for given ground RV 0zn then correspond to an element from the encoding array E,
A.1.3 PAVI-E details
Tn the PAVI-E scheme, encodings not free weights but the output of encoder S( (,))) applied to the observed data X. In this section we detail the design of this encoder.
As in the previous section, the role of the encoder will be to produce one encoding per plate level. We start from a dependency structure for the plate levels:
V(Pa..Pb) E PlateLevels
V(P..PPa PlateLevels
(A.3)
(Pa.Po) E (((...P)) + / (lates(0,)((PP--P)) T(.))
20,/ (1)tes(0.)=(P..-33)
22

---

Published in Transactions on Machine Learning Research (11/2023)
note that this dependency structure is in the backward direction: plate level will be the parent of another plate level, if the former contains a RV who has a child in the latter. We therefore obtain a plate level dependency structure that reverts the conditional dependency structure of the graph template T. To avoid redundant paths in this dependency structure, we take the maximum branching of the obtained graph.
Given the plate level dependency structure, we will recursively construct the encodings, starting from the observed data:
Vx X: Encodings(Plates/) p(x)
(A.4)
where x is the observed data for the RV 2, and p simple encoder that processes every observed ground RV's value independently through identical multi-layer perceptron. Then, until have exhausted all plate levels, we process existing encodings to produce new encodings:
ViPa.PPi E PlateLevels/ / Ax E X, Plates(x) = (Pk..Pi
(A.5)
Encodings( (Pp...nn = (Prroddingg(((P...PP)
where g is the composition of attention deep-set networks called Set Transformers (Lee et al., 2019; Zaheer et al., 2018). For every plate Pp present in the parent plate level but absent in the child plate level, 9 will compute summary statistics across that plate, effectively contracting the corresponding batch dimensionality in the parent encoding (Rouillard & Wassermann, 2022).
In the case of multiple observed RVs, we run this 'backward pass" independently for each observed data -with one encoder per observed RV. We then concatenate the resulting encodings corresponding to the same plate level.
For more precise implementation details, we invite the reader to consult the codebase released with this supplemental material.
A.2 Analysis of bias in the stochastic training
A key concern in our stochastic training scheme is its unbiasedness: we want stochastic optimization to converge to the same variational posterior as if we trained over the full model directly -without any stochasticity. In this section we first show that the PAVI-F is unbiased. Second, we identify strategies to obtain an unbiased PAVI-E scheme, while showing how the approximations we do in practice can theoretically result in biased training. As an important note, the negative impact of this bias on the performance PAVI-E remained limited throughout our experiments -as seen in experiments 5.3, B.3, B.2 and B.4.
A.2.1 General derivation (applicable to both the PAVI-F and PAVI-E schemes)
We first formalize the plate sampling strategy described in Appondix A.1.1. To every plate P we associate the RV Ip corresponding to the CP)si()) set of indices sampled without replacement from the Card(P) possible index values. As an example, with a plate Po with Car) = 4 and Card = 2, {0,2) or {2 can be 2 different samples from Ip.- At a given optimization step 4, we sample independently from the RVs {Ip,)p=..PP This defines the batches [t] and the distribution q² in Equation (5).
To check the unbiasedness of our stochastic training, we need to show that:
Eipp ... ...EE [EBOO ELBO
(A.6)
Where:
ELBO = Eeng Dogg () log (())
(A.7)
And ELBO*(I] is defined Equation (6). In that expression, 9 and p have symmetrical roles. As the ELBO amounts the difference between the logarithms of distributions p and 9, we can prove the equality in Equation (A.6) if we prove that the expectation of cach reduced distribution is equal to the corresponding full distribution. To prove the equality in Equation (A.6), sufficient condition is therefore to prove that:
Eir, ... Eirp (log'((0(())) = Er [log (o"g""(0" ((*(t)) log (e)
(A.8)
23

---

Published in Transactions on Machine Learning Research (11/2023)
where to de-clutter the notations we denote the expectation over the collection of RVs [Ipp)p=o.. as Eip.
Consider a given ground RV (g.m corresponding to the RV template I; and to the plates Plates((,). At a given stochastic step 6, Oz will be chosen if and only if its corresponding branching is chosen. Recall that when sampling equiprobably without replacement a set of & elements from a population of n elements, a given element. will be present in the set with probability k/n. We can apply this reasoning to the choice of given ground RV. For instance, in Figure 1, X2 will be chosen iff the index 1 is selected as part of sub-sampling Pi and the index 0 is selected as part of sub-sampling Po- As Card' indices are chosen inside the plate Pi of full cardinality and Card' (Po) 1 indices chosen inside the plate Po of full cardinality 2, X2 is therefore chosen with probability 2/3 3 x 1 1/2. More formally, for ground RV Oi we have:
branching corresponding
2
Card' (P)
Vn = 0..N; P(O., E BB((t))=
II
Card(P)
PPPllaes(0..)
(A.9)
N
N
Applying this reasoning to every RV template 0;, we have that:
N
Eip [log 2ogg"((' (()())
Eip
g0
MEB.
Ni
Bip
Σ B (0) log vi,n (0 (0.n (0a.))
(A.10)
N
N:
Eip
LLEB, log qr
E.,))]
where we exploited the fact that the expectation of the sum of RVs is the sum of the expectations, even in the case of dependent RVs. The term nEB,1() x log qi,n (0i,a (0),.); di, Ei,)) is the product of 2 RVs related to the stochastic choice of plate indices:
the RV IngB, is an indicator that 0i branching has been chosen via the stochastic sampling of plate indices. By construction, this RV depends only the indices of the plates P Plates(0:).
the RV log 9.m (O.n scheme:
E,,) depends on Ei whose construction depends on the encoding
In the PAVI-F scheme, Ei, is constant.
In the PAVI-E scheme, Ei,, results of the application of an encoder to the observed data of a subset 0,,''s descendants. By construction, this subset will only depend on the indices of plates containing 0,'s descendants, but not containing 0, The value of E,,, therefore only depends on the indices of plates P g Plates(O:)
As an example of this reasoning, consider the model M ratudt in Figure 1. We can evaluate both terms for the ground RV 0,,2 in the PAVI-E scheme:
[2 depends on whether the index 2 chosen as part of sub-sampling the plate Pi, and thercfore only depends on the RV Ir.- In this case, the associated probability is 2/3;
to evaluate log f (0,,2 (01,2 02 Φ1, E122), the value of E1,2 will result from the application of the encoder over the value of either X, or Xş. This choice depends on whether the index 0 or chosen as part of sub-sampling the plate Po. Therefore, the value of the term log 91 only depends on the RV Tro
24

---

Published in Transactions on Machine Learning Research (11/2023)
As result, in both PAVI-F and PAVI-E, the terms 1ng., and log qt,n (0s.n (T(,.)); (O. 6s, Ei,.) depend on the sampled indices of disjoint sets of plates, and are therefore independent. This means that the expectation of their product can be rewritten as the product of their expectations:
B [log & (0
Σ
BI
[Tne Br [log qi,n (i
(n (O...))
N
N
No
Eip log qs (0.
(A.11)
N
N
N E [log tsn an (Osml
n=0
This equality can be further simplified in the PAVI-F case PAVI-E case, as detailed in the sections below.
proving its unbiasedness- but not in the
A.2.2 Unbiasedness of the PAVI-F scheme
In the PAVI-F scheme, detailed in Section 4.3, the encodings E are constants with respect to the branching choice, therefore we have:
N
Eip olog do (()(
Esp [log qi,n qi,n (0.n(...1 Er,))
1
N
(A.12)
Σ
log qún (6 (0L,)): Er..)
n=0
log (())
which proves Equation (A.8) and Equation (A.6). In the above example of 01,2 in M, in the PAVI- scheme the expression Erp [log gi, (0.m(( Φb, E,,n]] can be evaluated into log (0.2(0002)002000000 Φ1, E1,2). This demonstrates that the PAVI-F scheme is unbiased: training over stochastically chosen mub-graphs for g² is in expectation equal to training over the full graph of
A.2.3 Approximations the PAVI-E scheme
In the PAVI-E scheme, detailed in Section 4.2, the encodings Ei,, are computed from the observed data X. Specifically, considering the ground RV 0g we have E,,, = f(x(,,,tt) where [f,,,(t] corresponds to the observed data of a subset of fi,n² 's descendants. Depending on the chosen branching downstream of 0. the value Ei can therefore vary. This mcans we cannot further simplify Equation (A.11): the terms log ni (0i,n (O are not constants with respect to the RVs Ip. In the above example of 01,2 in M, in the PAVI-E scheme the expression Erp [log ai.g (00..): de, Ei,))] can be evaluated into:
2 (log 2. 02 $1 ((x)) log 21 (01. oi ((x)))
How could the PAVI-E scheme be made unbiased? Specifically, by making the value of E inde- pendent of the choice of downstream branching. A possibility would be to parameterize Ein as an average expectation- over all the possible sub-branching, downstream of 6 Yet, in practical cases, the cardinalities of the reduced model are much inferior to the ones of the full model: Card" < Card(P). This mcans that numerous Card' (P))'(PPsszed subsets can be chosen inside the Card(P) possible descendants. In order to average over all those subset choices to compute Ei,,, numerous encoding calculations would be required at each stochastic training step. For large-scale cases, we deemed this possibility impractical. Other possibilitics could exist, all revolving around the problem of aggregating collections of stochastic estimators into one general estimator -in an unbiased and efficient manner. To our knowledge, this is complex and still open research question, whose advancement could much benefit our applications.
25

---

Published in Transactions on Machine Learning Research (11/2023)
Practical approximation for the PAVI-E scheme In practice, compute the encoding E based the single downstream branching corresponding to the sampling of the RVs 7p. Compared to the previous paragraph, this amounts to estimating the expectation of E -ver all downstream branchings using a single one of those branchings. Note that, oven this encoding cstimate was unbiased, log ft. would remain a highly non-linear function of E,,- a consequence, we need rely on the approximation:
Eip [log (0. ((E = log (0in(O../TT00. ((X;,))
(A.13)
which can theoretically introduce some bias in our gradients. The approximation Equation (A.13) can interpreted as follow: 'the expectation of the density of 0g,, when collecting summary statistics over a stochastic subset of Ozn 's descendants approximately equal the densily of 0,,n when collecting summary statistics the entirety of Oz,n descendants'. Another interpretation is that the distribution associated with the summary of the full data be approximated by annealing the distributions associated with sum- maries of subsets of this data. In practice, this approximation did not yield significantly worse performance for the PAVI-E scheme over the generative models we tested. At the same time, computing the encodings over a single branching allows the computation of all the Ei,, encodings in a single lightweight pass over the data [t]. This simple solution thus provided a substantial increase in training speed with seldom noticeable bias. Yet, we do not bar the existence of pathological generative HBMs where this approximation would become coarse. Experimenters should bear in mind this possibility when using the PAVT-E scheme. In practice, using the PAVI-F scheme sanity check over synthetic, toy-dimension implementation of the considered generative models is good way to validale the PAVI-E scheme before moving on to the real problem instantiating the same generative model with a larger dimensionality.
A.3 PAVI algorithms
More technical details can be found the codebase provided with this supplemental material.
A.3.1 Architecture build
Algorithm 1: PAVI architecturo build
Input: Graph template T, plate cardinalities {(Card(pp), (Prd encoding scheme Output: 9 distribution
for 1..7 do
Construct conditional flow Fú:
Define conditional posterior distributions d.n as the push-forward of the prior via F,, following Equation (3);
Combinc the 9i.n distributions following the cascading flows scheme, as Section 4.1 (Ambrogioni el. al., 2021b);
if PAVI-F encoding scheme then Construct encoding arrays {E, (E as in Appendix A.1.2;
else if PAVT-E encoding scheme then Construct encoder f as in Appendix A.1.3;
26

---

Published in Transactions on Machine Learning Research (11/2023)
A.3.2 Stochastic training
Algorithm 2: PAVI stochastic training
Input: Untrained architecture 9, observed data X, encoding scheme, number of steps T
Output: trained architecture q
for 1= 0..7 do
Sample plate indices to define the batches B,[t], the latent @²(t] and the observed [x(t) and [t))),
following Appendix ;
Define reduced distribution p'
if PAVI-F encoding scheme then
Collect encodings E by slicing from the arrays E, the elements corresponding to the batches
[
elsc if PAVI-E encoding scheme then
L Compute encodings as = f(x
Feed obtained encodings into q²
Compute reduced ELBO as in Equation (6), back-propagate its gradient Update conditional flow weights foih=11..;
if PAVI-F encoding scheme then Update encodings (E.,)eet.. clsc if PAVI-E encoding scheme then
L
Update encoder weights 77;
A.3.3 Inference
Algorithm 3: PAVI inference
Input: trained architecture q, observed data X, encoding scheme Output: approximate posterior distribution if PAVI-F encoding scheme then
Collect full encoding arrays E;
else if PAVI-E encoding scheme then
Compute encodings as E f(X) using set size generalization
Feed obtained encodings into
A.4 Inference gaps
In terms of inference quality, the impact of our architecture be formalized following the gaps terminology (Cremer et al., 2018). Consider joint distribution X), and value X for the RV template X. We pick variational family Q, and in this family look for the parametric distribution ((O; Φ) that best approximates p(e|X = x). Specifically, we want to minimize the Kulback- divergence (Blei et al., 2017) between our variational posterior and the true posterior, that Cremer et al. (2018) refer to as the gap G:
G KL(g(O; x) ( (p(e)))
(A.14)
log p( (X) ELBO(g;
We denote g"(O; (+) the optimal distribution inside Q that minimizes the KL divergence with the true posterior:
Gapprox (2; (²) (() -LBOOL"BB""""00
0
(A.15)
Ivanilla VI 9upprox
The approrimation gop Gupprox depends on the expressivity the variational family Q, specifically whether Q contains distributions arbitrarily close to the posterior -in the KL sense.
27

---

Published in Transactions on Machine Learning Research (11/2023)
Note: Gapprox is property of the variational family Q. Gapprox is an asymptotic bound for the KL divergence between any distribution Q and the true posterior. This gap is therefore a form of bias but is not to be mistaken with the stochasticity-indues whether q" can be found by training stochastically over 2, whereas Gapprox relates to the the bias between q² and the true posterior.
bias studied in Appendix A.2. The bias in Appendix A.2 relates to
Cremer et al. (2018) demonstrate that, in the case of sample amortized inference, when the weights no longer are free but the output of an encoder e F, inference cannot be better than in the non-sample- amortized case, and positive amortization gap is introduced:
Gapprox (2;s(X, Gapprox (2: (*)
0
(A.16)
Gsample amortized VI Gapprox + Gsa
Where denote 17" the optimal weights for the encoder f inside the function family F. The gap terminology can be interpreted as follow: "thcoretically, sample amortization cannot be beneficial in terms of KL divergence for the inference over given sample X.'
Using the same gap terminology, we can define gaps implied by our PAVI architecture. Instead of picking the distribution inside the family Q, consider picking q from the plate-amortized family Qpa corresponding Q. Distributions in Qps are distributions from with the additional constraints that some weights have to be equal. Consequently, 2pa is subset of Q:
Ppa C Q
(A.17)
As such, looking for the optimal distribution inside 2pa instead of inside Q cannot result in beller perfor- mance, leading to plate amortization gap:
9pa 22, Spai v/,,, ,0") = Cupprox( (Spai 15*) Gapprox (2 (²)
20
(A.18)
GPAVI-F Iupprox + Ipa
Where we denote as z0" the optimal weights for variational distribution q inside Qpa -n the KL sense. The equation A.18 is valid for the PAVI-F scheme see Section 4.3. We can interpret it as follow: "theoretically, plate amortization cannot be beneficial in terms of KL divergence for the inference over a given sample X"
Now consider that encodings are no longer free parameters but the output of an encoder f. Similar the case presented in Equation (A.16), using an encoder cannot result in better performance, leading to an encoder gap:
Gencoder (eriodder(@p... (pa. Fi = eapprox( f(X, )) Gapprox (Spai ())
2 0
(A.19)
GPAVIEE = Gapprox + 9pa Gencoder
The equation Equation (A.19) is valid for the PAVI-E scheme see Section 4.3.
The most complex case is the PAVI-E(sa) scheme, where we combine both plate and sample amortization. Our argument cannot account. for the resulting (PAVI-E(sa) gap: both the PAVI-E and PAVI-F(sa)) schemes rely upon the same encoder f. In the PAVI-E scheme, f overfitting over a dataset composed of the slices of given data sample X. In the PAVI-E(sa) scheme, the encoder trained over the whole distribution of the samples of the reduced model M Intuitively, it is likely that the performance of PAVI-E(sa)) will always be dominated by the performance of PAVI-E, but -as far as we understand it- the gap terminology cannot account for this discrepancy.
Comparing previous equations, we therefore have:
Gvanilla S GPAVLF S GPAVI-E
(A.20)
Note that those are theoretical results, that do not necessarily pertain optimization in practice. In particular, in Section 5. 5.185.3, this theoretical performance loss is not observed empirically over the studied
28

---

Published in Transactions on Machine Learning Research (11/2023)
examples. On the contrary, in practice, results actually be better than non-amortized basclines, is the case for the PAVI-F scheme in Figure 3 or experiments B.4. We interpret this as result of simplified optimization problem due to plate amortization -with fewer parameters to optimize for, and mini- biini-bacching effects across different ground RVs. A better framework to explain those discrepancies could be the one from Bottou & Bousquet. (2007): performance error but also of an optimization error. A less expressive architecture -usig plate amortization may practice yield better performance. Furthermore, for the experimenter, the theoretical gaps Gpa; Gencoder are likely to be well 'compensated for" by the lighter parameterization and faster convergence entitled by plate amortization.
practice is not only the reflection an upprorimation
B Supplemental results
In this section, we present various supplemental experiments that further compare the PAVI-E and -F between themselves and against baselines. Overall, we found the PAVI-F scheme faster train and yield better inference quality than the PAVI-E scheme. When its parameterization is affordable, PAVI-F should be preferred. PAVI-E nevertheless opens up promising research directions, with the potential for carameterizationnccnssee
time-constant sample-amortized inference as the problem cardinality augments. Though degraded compared to PAVI-F's, PAVI-E's performance is still on par or beats baselines in a variety of inference tasks scc experiments 5.3&
B.1 Tabular results for experiment 5.3
For convenience, we reproduce the results of our comparative experiment from Figure 3 in Table B.1.
B.2
GRE results sanity check
As exposed in the introduction of Section 5, in this work we focused on the usage of the ELBO as an inference performance metric (Blei et al., 2017):
ELBO(g) = ((X) - KL((())|lp())x))
(B.21)
Given that the likelihood term log (X) does not depend on the variational family q, differences in ELBOS directly transfer to differences in KL divergence, and provide straightforward metric to compare different variational posteriors. Nonetheless, the ELBO doesn't provide an absolute metric of quality. As sanity check, we want to assert the quality of the results presented in Section 5.3 -that are transferable to Sec- tion 5.12., based on the same model. In Figure B.1 we plot the posterior samples of various methods against approximate closed-form ground truths, using the Card(Pi) = 20 case. All the method's results are aligned with the closed-form ground truth, with differences in ELBO translating meaningful qualitative differences in terms of inference quality.
B.3 Effect of the reduced model cardinalities on the training efficiency
In Figure B.2 we show the impact of the augmentation of Card' (P:) on the efficiency of the variational posterior's optimization.
In practice, we noticed that the most efficient choice the case of the PAVI-F scheme was to maximize the cardinalities of the reduced model given the memory constraints of the GPU. Indeed, training over larger cardinalities does make each optimization step slightly slower but also makes the ELBO gradient estimates less noisy and allows to train more encoding vectors E,,, at given optimization step.
In the case of the PAVI-E scheme, the training speed is constant with respect to Card' This is due to the computing of the encodings being vectorized across plates in our deep-set encoder see Appendix A.1.3 (Lee et al., 2019). We observe a slight if barely noticcable reduction of the inference bias when augmenting Card Similar to the PAVI-F scheme, Card"(Pi) should be maximized with respect to the GPU memory constraint. The intuition behind this choice is that the generalization of the learning of summary statistics
29

---

Published in Transactions on Machine Learning Research (11/2023)
Architecture
ELBO (10')
Parameterization
Optimization time (s)
Card(Pi)=
ADAVI (sa)
1.88 (= 0.21)
12,280
503 (# 50)
CF
(sa)
1.48 (= 0.05)
2,275
228 (+ 23)
CF
1.98 (+ 0.02)
355
63 ( 6)
UIVI
2.11 (+ 0.05)
2,764
220 (+ 22)
PAVI-F
2.14 (= 0.02)
4,058
65 (+ 6)
PAVI-E (sa)
2.03 (= 0.03)
27,394
140 (# 14)
PAVI-E
1.74 (+ 0.18)
27,394
37 (+ 4)
Asymptotic ELBO
2.20
Card(Pt) 20
ADAVI (sa)
21.6 (= 1.36)
12,280
800 (#: 80)
CF (sa)
17.3 (= 0.60)
21,751
1,612 (# 161)
CF
22.6 (= 0.01)
2,551
404 (# 40)
UIVT
23.9 (+ 0.02)
14,676
230 (+ 23)
PAVI-F
23.8 (= 0.02)
4,202
72 (# 7)
PAVI-E (sa)
22.3 (= 0.30)
27,394
145 (# 14)
PAVI-E
18.4
(+ 2.63)
27,394
50 (+ 5)
Asymptotic ELBO
24.1
Card(P;) 200
ADAVI (sa)
221.7 (# 10.93)
12,280
1,600 (+ 160)
CF (sa)
175.7 (1 0.52)
216,511
12,000 (# 1,200)
CF
228.6 (# 0.52)
24,511
2,600 (= 260)
UIVI
241.2 (+ 0.14)
139,300
240 (+ 24)
PAVI-F
237.7 (# 0.95)
5,642
81 (# 8)
PAVI-E (sa)
227.0 (# 4.76)
27,394
150 (# 15)
PAVI-E
224.2 (# 15.34)
27,394
85 (# 9)
Asymptotic ELBO
243.4
Table B.I: Scaling with plate cardinalities Reproduction of the results from Figure 3
30

---

Published in Transactions on Machine Learning Research (11/2023)
Obtrread 2000
Obeaned sample
otseved semple
sample
2000
A03
402
0.5
-0.3
4DM
1.0
0.
-0.5
-15
-10
85
WSH0.00
200 01--
WO-2IN
60-1aw
AR ginu211
-14
-15 -15
-55
85
1.3
13
26
25
85
10
002 COAAA
RO-2061
KD- 5.00
posterer
weggt
fran
00
-15
frincl
0.5
1011
-0.3
-15 -10 -25
15
Wo-cor
100 HANN
N0-2090
60 SIN
0.5
A# 27o22277
-10
22
-15 -10 -25 50
15
1.5
13
26
25
05 24
Figure B.1: GRE Sanity check Inference methods present qualitatively correct results, making ELBO comparisons relevant in our experiments. On the topmost line, we represent different: X samples for the GRE model described in Equation (10) with Card(Pi) 20. Each set of colored points represents the X points belonging to one of the 20 groups. Bottom lines represent the posterior samples for the methods used in Section 5.3. Colored points sampled from the posterior of the groups means 01, whereas black points are samples from the population mean 02. We represent as black circles closed-form ground truth, centered on the correct posterior mean, and with a radius equal to 2 times the closed-form posterior's standard deviation. Correct posterior samples should be centered on the same point as the corresponding black circle, and 95% of the points should fall within the black circle. Associated with each postcrior sample is a Wasserstein Distance (WD) computed with sample from an approximato closed-form posterior (Flamary et al., 2021). PAVI is represented on the last lines. Some minor bias can be observed for the PAVI-E scheme, but this approximation error is marginal compared to the optimization error that can be observed for unbiased methods, such as CF (sa) (Bottou & Bousquet, 2007). We can observe superior quality for the PAVI-F scheme, rivaling ADAVI and CF's performance with orders of magnitude fewer parameters and training time, as visible in Figure 3.
31

---

Published in Transactions on Machine Learning Research (11/2023)
240
240
220
220
EL9O
200
PAVI-F Card redu=1 redu-2
PAVI-E Card redu=1 Card, redu-2
PAVI-F Card redu-8
PAVI-E Card redu=g
redu-16
Card PAVI-E no stochasticity
160
PAVI-F stochasticity
160
10a
10a
10a
Time (4)
Times 16]
Figure B.2: Both panels: Effect of Card' (P1) on the training efficiency Experiment performed on the GRE model (see Equation (10)). We increase the cardinality of the reduced model Card (P1) from 1 to 20 while keeping Card(Pi) 20 fixed. At Card train directly M. This experiment is interesting to evaluate the bias introduced in the stochastic training, as detailed in Appendix A.2. Left pancl: PAVI-F scheme As Card" (P1) augments, the training gets faster and less noisy, likely due to less stochasticity in the gradient estimates and more encodings vectors E,,, being trained at once. This increase in speed quickly caps, and the speed approximately and Card"(Pi) = 16. This experiment also illustrates the unbiasedness
20, there is stochasticity in the training, meaning
the same between Card of the stochastic training: Card"(P:) = 20 there is no stochasticity in the training, and the asymptotic performance is the same for the stochastic training (see Appendix A.2). Right panel: PAVI-E scheme is pathological: the encoder "learns' to collect summary statistics across set
The case of Card" Card of element, and not surprisingly the learned function doesn't generalize well on sets of size 20. In all of the other cases, the ELBO converges approximately to the same values as with Card = 20, that is to say when there no stochasticity in the training (black curve). We furthermore observe slight reduction of the inference bias as (a)) augments. This illustrates how the theoretical bias of the PAVI- E scheme identified in Appendix A.2 does not translate in significantly worse empirical results albeit in the pathological case of Card'(Pi) all values of Card' Card This fcature is essential in the PAVI-E scheme: of our model witt light memory footprint and apply the obtained architecture to the full model. Note that this feature wouldn't necessarily be present when training CPU: computing summary statistics over larger sets would make the training slower when (Pi) augments.
1. Interestingly, the training speed is approximately constant across
train reduced version
32

---

Published in Transactions on Machine Learning Research (11/2023)
Architecture
ELBO (10*)
Optimization time (s)
CF
16.2 (+ 2.9)
3,100
ADAVI
200.0 (# 2.3)
3,000
UTVI
26.33 (+ 1.2)
500
PAVI-E (ours)
15.2 (+2.8)
470
PAVI-F (ours) - 11.1 (+1.1)
300
Table B.2: Analytical performance over Gaussian mixture HBM PAVI shows superior performance. PAVT converges in a fraction of CF's optimization time the accond-bestt-eeffrming architecture.
across sets of data points is easier the closer the reduced set size is to the full set size. This constant training speed allows for a controlled memory footprint of the stochastic training: contrary to the PAVI-F scheme, Card' (P1) could be maintained at value manageable by the GPU without reducing the training speed. This example also displays a pathological choice for Card' (P1): ): the encoder fails learn compule the correct summary statistics over sets of size However trivial, this example underlines that Card should be chosen a value logical with respect to the posterior's sufficient statistics. When generalizing to moments of higher order -such as the variance- there is therefore theoretical lower bound to consider when fixating the value of Card Our intuition is that the more complex to estimate the statistic is, the larger (ard"(P;) should be.
Note that in practice non-stochastic VI is intractable for large-scale models, because of its memory require- ment. large-scale experiments such as the one presented Figure 4, stochastic training necessary. Our claim is be faster than non-plate-amortized stochastic VI in those largo-scale contexts bbut not necessarily to be faster than non-stochastic VI in the small-scale regime of this experiment. For PAVI-F we nonetheless obtain similar training speed compared to non-stochastic VI with Card" Card"(Pi) small as 8. For PAVI-E, the stochastic training is as fast as the non-stochastic one.
B.4
Additional comparative experiments
B.4.1 Gaussian Mixture model
In this experiment, we test out various baselines over a challenging model: Gaussian mixture with random effects.
D, Card(Pi), Card(Po) 2, 20, 10
Vnu m. Card(P1) Vng=1. Card(Po)
Xni.nol0 Thi ~ Mixture( (2) am))
Vn = 1.. Card(P)) The Dirichlot(1 Tz)
(B.22)
Vai=1. Card(P)
o',m 02 -N oi) = 02 NNOD.
W=1.1b
where Mixture([Di,.... D2], R) denotes a mixture between the distributions [Di...., Dz] with mixture weights x. Results are visible in Table B.2, where PAVT displays the best asymptotic ELBO, well as the shortest optimization time. On that note, we underline that stochastic training over a mixture distri- bution is challenging, as due to the sub-sampling of points only a fraction of the mixture components could be expressed at a given step, requiring the architecture to dynamically cluster the data points across time.
B.4.2 Hierarchical Variance model
In this experiment, we test out our architectures over a non-canonical model, which parent RVs play the role of variance for the distribution of their children. Our goal in particular to evaluate potential empirical bias for the PAVI-E scheme studied in Appendix A.2. We will refer to the following HBM
33

---

Published in Transactions on Machine Learning Research (11/2023)
Architecture
ELBO (10²)
Optimization time (s)
CF
11.2 (# 2.9)
500
UIVI
-6.7 (# 0.8)
100
PAVI-E (ours)
-6.7 (+1.0)1
90
PAVI-F (ours)
- 6.7 (+0.8)
20
Table B.3: Analytical performance over Hicrarchical Variance HBM In this setting with hard- to- oo-estimate summary statistics, PAVI-E doesn't show any empirical bias. : PAVI-E suffers from numerical instability in this example, with runs degenerating into NaN results.
Architocture
ELBO (10*)
Optimization time (s)
CF
139.7 (= 21)
3,200
ADAVI
UIVI
PAVI-F (ours) -.2 (# 0.1)
1,600
Table B.4: Analytical performance over our Neuroimaging HBM This HBM is not pyramidal and consequently camnot be processed by the ADAVT architecture. Due to the large parameler space (approx- imately 2k parameters) and the complex density involved, despite intensive efforts we did not manage to obtain numerically stable UIVI optimization. PAVI yiclds higher ELBO than CF, in half of CF's opti- mization time.
the Hierarchical Variance model:
D, Card(Pi), Card(Po) 2, 15,
Ve,=1.. Card(P,)
log ((m) 10,m ~N(0, 0 0m,))
vmg=1., Card(Po)
(B.23)
Vn, = 1.. Card(P1)
log(l,(m1)))0 12.0 - NNO, 0o.0)
1og02,0 -N(OD,
In this model, the encoder used in the PAVI-E scheme- has collect non-trivial summary statistics: pirical variances across variable subsets of the observed data. This is the context in which would expect to observe the most bias due to the approximation in Equation (A.13). The performance of PAVI-E re- mains nonetheless competitive. This illustrates how PAVI-E's theoretical bias inttroduced the stochastic training- does not result in significantly worse inference compared to state-of-tho-art architoctures. PAVI also displays the best ELBO, in conjunction with UIVI, but a times shorter optimization time.
B.4.3 Small dimension version of our Neuroimaging model
In this experiment, we validate the performance of our architecture on synthetic data generated using a small dimension version of the model presented in Equation (B.24). To allow the use of the comparative baselines, we fixate S = = 5,3, T.T = 3, NT 12,3, D = 4, L = 7. Results PAVI can deal with HBMs unavailable to the ADAVI architecture (Rouillard & Wassermann, 2022), and so so with a performance superior to the CF bascline (Ambrogioni et al., 2021b). Despite intensive efforts, we did not manage to obtain a numerically stable UTVI optimization. This is likely due to the complexily of this inference problem: a large dimensionality (2k parameters) combined with complex dependencies between RVs in the posterior. In this context, automatic structured VI baselines such as CF PAVI exploit the parametric form of the prior p to help with inference. PAVI's plate amortization likely facilitates the inference, resulting in reduced oplimization error (Bottou & Bousquet, 2007).
visible in Table B.4:
34

---

Published in Transactions on Machine Learning Research (11/2023)
B.5 Experimental details analytical examples
All experiments were performed in Python, using the Tensorflow Probubilily library (Dillon et al., 2017). Through this section, we refer to Masked Autoregressive Flows (Papamakarios et al., 2018) as MAF. All experiments are performed using the Adam optimizer (Kingma & Ba, 2015). At training, the ELBO was estimated using a Monte Carlo procedure with samples. All architectures were evaluated over a fixed set of samples X, with seeds per sample. Non-sample-anortiedd architectures were trained and evaluated on each of those points. Sample amortized architectures were trained over dataset of 20, 000 samples separate from the 20 validation samples, then evaluated over the 20 validation samples.
Termination condition Our termination condition amounts to plateau detection across the statistics of multiple runs. In details, we ran the inference with a large, conservative number of epochs for all methods and collected the evolution of the ELBO for each run. We then look as runtime the time to reach an ELBO close to the one at convergence for the method. We averaged over 100 repetitions to avoid stochastic noise due to Monte Carlo estimation of the ELBO. The order of magnitude difference between the results is clearly above the potential noise effects due our procedure.
B.5.1 Plate amortization and convergence speed (5.1)
All 3 architectures (bascline, PAVI-F, PAVI-E) used:
for the flows Fi, MAF with [32, 32] hidden units;
as encoding size, 128
For the encoder f in the PAVI-E scheme, used multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 64 inducing points.
B.5.2 Impact of encoding size (5.2)
All architectures used:
for the flows Fú, MAF with [32, hidden units, after an affine block with triangular scaling matrix.
as encoding size, value varying from to 16
B.5.3 Scaling with plate cardinalities (5.3)
ADAVI (Rouillard & Wassermann, 2022) used:
for the flows Fs MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.
for the encoder, encoding size of 8 with a multi-head architecture with 2 heads of 4 units each, 2 ISAB blocks with 32 inducing points.
Cascading Flows (Ambrogioni al,, 2021b) we used:
a mean-field distribution over the auxiliary variables r
auxiliary size, fixed value of
as flows, Highway Flows as designed by the Cascading Flows authors
PAVI-F used:
35

---

Published in Transactions on Machine Learning Research (11/2023)
for the flows Fr, MAF with [32 hidden units, after affine block with triangular scaling matrix.
an encoding size of 8
PAVI-E we used: for the flows Fú, MAF with [32,32] hidden units, after affine block with triangular scaling matrix.
for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads of 8 units each, 2 ISAB blocks with inducing points.
UIVI we used as Card = 2 + 20 + 200:
base distribution, standard Gaussian with dimensionality - 42 402
as transform h, an affine transform with diagonal scale
embedding distribution, standard Gaussian with dimensionality 3 3 6 +
as transform weights regressor, an MLP with hidden units [32, - [64, 64 + [128, 128]
lo sample uncorrelated samples e, HMC run with 5 burn-in steps and 5 samples
an Adam optimizer with exponential learning rate decay, starting al le - 2, x0.9 every 300 steps
B.5.4 Gaussian mixture (B.4.1)
ADAVI (Rouillard & Wassermann, 2022) we used:
for the flows F,, a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.
for the encoder, an encoding size of 16 with multi-head architecture with 2 heads of 8 units each, 2 TSAB blocks with 8 inducing points.
Cascading Flows (Ambrogioni et al., 2021b) we used:
mean-field distribution over the auxiliary variables r
as auxiliary size, fixed value of 16
as flows, Highway Flows as designed by the Cascading Flows authors
PAVI-F used:
for the flows F,, a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.
encoding size of 16
Card Card 5
PAVI-E we used: for the flows F,, MAF with [128, 128] hidden units, after an affine block with triangular scaling matrix.
36

---

Published in Transactions on Machine Learning Research (11/2023)
for the encoder, encoding size of 128 with multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 128 inducing points.
Card 5
UIVI we used:
as base distribution, standard Gaussian with dimensionality 82
transform h, affine transform with diagonal scale
embedding distribution, standard Gaussian with dimensionality 6
as transform weights regressor, MLP with hidden units [64, 64]
to sample uncorrelated samples am HMC run with 5 burn-in steps and 5 samples
an Adam optimizer with exponential learning rate decay, starting at 1e- 2, x0.9 every 300 steps
B.5.5 Hierarchical Variances (B.4.2)
Cascading Flows (Ambrogioni et al., 2021b) used:
a mean-field distribution over the auxiliary variables r
as auxiliary size, fixed value of 32
as flows, Highway Flows as designed by the Cascading Flows authors
PAVI-F we used:
for the flows F,, a MAF with [32, hidden units, after an affine block with diagonal scaling matrix.
an encoding size of 32
Card' Card =
Card' (Po)
PAVI-E we used:
for the flows F,, MAF with [128, 128] hidden units, after an affine block with diagonal scaling matrix.
for the encoder, an encoding size of 128 with a multi-hean architecture with 4 heads of 32 units each, 2 ISAB blocks with 128 inducing points.
Card' (P1) 3
Card' CPr) 3
UIVI we used:
as base distribution, a standard Gaussian with dimensionality 32
as transform h, an affine transform with diagonal scale
as embedding distribution, a standard Gaussian with dimensionality
as transform weights regressor, MLP with hidden units [64, 64]
to sample uncorrelated samples e, HMC with 5 burn-in steps and 5 samples
an Adam optimizer with exponential learning rate decay, starting at 1c 2, x0.9 every 300 steps
37

---

Published in Transactions on Machine Learning Research (11/2023)
B.5.6 Small Neuroimaging example (B.4.3)
Cascading Flows (Ambrogioni et al., 2021b) we used:
mean-field distribution over the auxiliary variables r
as auxiliary size, fixed value of 32
as flows, Highway Flows as designed by the Cascading Flows authors
PAVI-F used:
for the flows F,, MAF with [32, hidden units, after an affine block with diagonal scaling matrix.
an combination encoding sizes of 32 and 8
UIVI we used:
as base distribution, a standard Gaussian with dimensionality 1, 802
as transform h, an affine transform with diagonal scale
embedding distribution, standard Gaussian with dimensionality 16
as transform weights regressor, an MLP with hidden units [128, 256, 512, 1024]
to sample uncorrelated samples an HMC run with 5 burn-in steps and 5 samples
an Adam optimizer with exponential learning rate decay, starting at le - x0.9 every 300 steps
Optimization systematically degenerated into NaN results after around 200 optimization steps.
B.6
Details about our Neuroimaging experiment (5.4)
B.6.1
Data description
In this experiment, we use data from the Human Conneclome Project (HCP) dataset (Van Essen el al., 2012). We randomly select cohort of S 1,000 subjects from this dataset, each subject is associated with T = 2 resting-state fMRI sessions (Smith et al., 2013). We minimally pre-process the signal using the nilearn python library (Abraham et al., 2014):
removing high variance confounds
2. detrending the data
3. band-filterngg the data (0.01 to 0.1 Hz), with a repetition time of 0.74 seconds spatially smoothing the data with a 4mm Full-Width at Half Maximum
For every subject, extract the surface Blood Oxygenation Level Dependent (BOLD) signal of N 59, 412 vertices across the whole cortex. We compare this signal with the extracted signal of D = 64 DiFuMo components: a dictionary of brain spatial maps allowing for an effective [MRI dimensionality reduction (Dadi et al., 2020). Specifically, we compute the one-to-one Pearson's correlation coefficient every vertex with every DiFuMo component. The resulting connectome, with subjects, sessions, vertices and a connectivity signal with dimensions, is of shape (S x T N D). We project this data correlation coefficients lying in ]-1;1[- in unbounded space using inverse sigmoid function.
38

---

Published in Transactions on Machine Learning Research (11/2023)
B.6.2 Model description
We use a model inspired by the work of Kong et al. (2019). We hypothesize that every vertex in the cortex belongs to either one of L. 7 functional networks. This number inspired by the work of Thomas Yeo et al. (2011) and Kong al. (2019), in which authors identify 7 general networks on resting-state fMRI data.
Each network corresponds to pattern of connectivity with the brain, represented the correlation of the BOLD signal with the signal from the D = 64 DiFuMo components. We define I, = functional networks at the population level. Each network / corresponds to the connectivity fingerprint
At the population level, for a vertex n we denote as logits,, E R the logits of its probability to belong to each network 1. Each subject & is associated with the logits logits, which are considered as perturbation of the population logits. This creates a regularization across subjects: the same vertex is encouraged to have the same label across all subjects. The variable "/ controls the inter-subject spatial variability across all subjects, for the network I.
For every subject 8, session t, and vertex n, we denote as Xattnn the observed connectivity. This connectivity is modeled via a mixture model: Xs,t.n is perturbation of either one of the pr's connectivity fingerprints. label,,n denotes the label in the mixture (that depends on the subject-specific logits). Ki controls the variability between X,, and the mixture component plabel,...
The resulting model be described as:
D,T,N,D.L- 1000, 2, 59412, 64,
W=1.L: Pr N'(0 x ID. 6)
W= log ki ~ N'O x I,, 1)
V==I.N: logits, ~ NNO X Iz, 6)
(B.24)
1=1..: log 71- AN(O x I2, 1)
V==1.. labels logits,,, Categoricalllogits,, v==I.N
logits s,n logits,, M (logits, 1/1))
X.4mm//4) [m.) labely,n N(Hlabel,,. Klabel,,..)
The model contains 4 plates: the nelwork plate of full cardinality L (that we did not exploit in our imple- mentation), the subject plate of full cardinality S, the session plate of full cardinality T and the vertez plate of full cardinality N.
Our goal is lo recover the posterior distribution of the networks μ and the labels label represented as the parcellation in Figure 4 given the observed connectome described in Appendix B.6.1.
B.6.3 PAVI implementation
We used in this experiment the PAVI-F scheme, using:
for the RVs pr:
for the flows Fs, MAF with [64, 128, 64] hidden units, following an affine block with diagonal
scale
for the encoding size: 16
for the RVs m
MAP estimator using a regressor with [64, 128, 64] hidden units
for the encoding size: 8
for the RV' logits,, logits,.n
for the flows F,, affine block with diagonal scale regressed using 164, 64) ML.P
39

---

Published in Transactions on Machine Learning Research (11/2023)
for the encoding size: 8
for the reduced model, we used 5 = 128, 7 and N = 128.
To allow for the optimization over the discrete labels, RV, we used the Gumbell-Softmax trick, using fixed temperature of 1.0 (Jang et al., 2017; Maddison al., 2016).
B.6.4 Cognitive scores prediction
For the cognitive scores prediction in Figure 4, we reproduced the standard methodology from Kong et al. (2019).
We perform a 20-fold cross-validation across 1,000 subjects. We start from the logits,.n s.n associated with each subject. We use PCA to project the features from the 19 training folds to their 100 first components (with 33% explained variance). We then train linear regression predict each of the 13 cognitive scores from the training-fold PCA features. We compute the test performance the test fold using the training set PCA and linear regression averaged across the 13 cognitive measures. This process is reproduced 100 times. The reported scores are the triple average across folds, measures, and repetition, while the standard deviation is computed across the 100 repetitions only.
C Supplemental discussion
C.1 Plate amortization as a generalization of sample amortization
In Section 4.1, we introduced plate amortization as applying the generic concept of amortization to the granularity of plates. There is actually an even stronger connection between sample amortization and plate amortization.
An HBM p models the distribution of a given obscrved RV X -jointly with the parameters O. Different samples Xg, X1,... of the model p are i.i.d. draws from the distribution p(X). p can thus be considered the model for 'one sample". Instead of p, consider "macro" model for the whole population of samples one could draw from p. The observed RV of that macro model would be the infinite collection of samples drawn from the same distribution p(X). In that light, the i..d. sampling of different X values from p could be interpreted as a plate of the macro model. Thus, we could consider sample amortization as an instance of plate amortization for the 'sample plate". Or equivalently: plate amortization can be soen as the generalization of amortization beyond the particular case of sample amortization.
C.2 Alternate formalism for SVI PAVI-E(sa) scheme
In this work, we propose different formalism for SVI, based around the concept of full HBM M versus reduced HBM Mredu sharing the same template T. This formalism is helpful to set up GPU-accelerated stochastic VI (Dillon et al., 2017), as it entitles a fixed computation graph wiith the cardinality of the reduced model Mredu_ which encodings are 'plugged in" -either sliced from larger encoding arrays or as the output of an encoder applied to a data slice, see Section 4.18. Particularly, our formalism docsn't entitle a control flow over models and distributions, which can be hurtful in the context of compiled computation graphs such as in Tensorflow (Abadi et al., 2015).
The reduced model formalism is also mcaningful in the PAVI-E(sa), where we train and amortized variational posterior over Mredu and obtain 'for free" a vaniational posterior for the full model Mffll see Section 4.2. In this context, our scheme is no longer a different take on hierarchical, batched SVI: the cardinality of the full model is truly independent of the cardinality of the training and is only simulated as scaling factor the stochastic training sec Section 4.2. We have the intuition that fruitful research directions could stem from this concept.
40

---

Published in Transactions on Machine Learning Research (11/2023)
C.3 Benefiting from structure in inference
Our contributions can be abstracted through the concept of plate amorlization -see Section 4.1. Plate amortization is particularly useful in the context of heavily parameterized density approximators such as normalizing flows, but not tied to it: plato-amortized Mcan Field (Blei et al., 2017), ASVI (Ambrogioni et al., 2021a), or implicit (Yin & Zhou, 2018; Titsias & Ruiz, 2019) schemes are also possible to use. Plate amortization can be viewed as the amortization of common density approximators across different sub- structures of problem. This general concept could have applications in other highly-structured problem classes such graphs or sequences (Wu et al., 2020; Salchinejad et al., 2018).
Central to our design, PAVT adjoins an encoding structure to an inference problem. The encodings E,, embed all the 'individualized! information in the problem, while the shared conditional density estimators F; translate those encodings into posterior distributions. We believe there potential in externalizing part of the inference problem into an embedding. For instance, embeddings could be learned separately from density estimators, opening up the possibility for transfer lcarning (Weiss et al., 2016). Or the encodings could integrate some known symmelry of the problem, such as convolution-based encodings in the case of Random Fields (Bishop, 2006).
C.4 Connection with meta-learning
In Section 4.1, introduced plate amortization: sharing the parameterization and learning across model's plates. This section discusses the connection between plate amortization and meta-learning (Ravi & Beatson, 2019; Iakovleva et al., 2020; Yao et al., 2019).
Supervised learning can be seen as the mapping from a given context set C {(x, y}} to a predictive function f (Bishop, 2006) such that f(x) = y. Meta-learning or "learning to learn" instead recovers this mapping C CHS in the general case. Once the mela-training is completed, a predictive function I conditioned by an unseen context can be obtained in a single forward pass without any training done on C. As an instance of meta-lcarning, the Neural Process Family encodes the context C via deep set encoder (Garnelo et al., 2018; Dubois et al., 2020; Zaheer al., 2018). The encoded context, along with the data point x is then used to condition an estimator for the density C).
This framework similar to the PAVI-E scheme, where a combination of a deep set encodor and a normnllizingglow-beese density estimator output the posterior probability of a ground RV 0z.- This encoder- estimator pair repeatedly used across plate. This is analogous meta-learning to solve the inference problem across the different elements of plate such as the different subjects of population study, In PAVI-E, the encoding E lower hierarchy plays role similar to the context C in meta-leaning. A few differences, however, exist between the two frameworks:
meta-learning is typically concatenated to the 1-plate regime, whereas PAVI is designed for the multi-hierarchy scenario;
meta-learning typically considers set of i.i.d. tasks, whereas in PAVI the inference of different ground RVs Onn are conditionally dependent through the hierarchical model
meta-learning is trained using the forward KL loss, that is to say in the sample-amortized regime, maximizing the probability q(y) of samples from the underlying generative process. In contrast, PAVI though possible to train using the forward KL loss is trained via the reverse KL, which requires to evaluate the density of the generative process explicitly.
We suspect there would be interesting applications for the PAVI architecture in hicrarchical mota-lcarning scenarios.
C.5 Conditional dependencies modeled in the variational family
Consider an inference problem with three latent parameters 01, 92 03. and a model p(O, X) that factorizes p(0,, O2 0s) = ((0))) 02)))))))0)))((XXX1 0s). To approximate the posterior X)) different conditional
41

---

Published in Transactions on Machine Learning Research (11/2023)
dependencies can be modeled in the chosen variational family Q. The mean-field (MF) approximation fully factorizes the variational distribution as q(01,02. = 4(0) ((2))))))))(())) MF was originally rntrodee in VI to facilitate computation, allowing dedicated optimization schemes (Blei et al., 2017). Nonetheless, this approximation assumes independence between RVs in the posterior and ultimately limits the expressivity of the variational family. To remove this approximation, modeling complex conditional dependencies in the variational family is an open research subject (Ambrogioni et al., 2021b; Webb et al., 2018). In the PAVI design, we inherit our statistical dependency structure from the prior distribution p, as detailed in Equation (3). This amounts to factorizing q(01, 0x) q ((()))) 02). This choice of dependencies follows the line of research from Hoffman & Blei (2014) and Ambrogioni et al. (2021a). As pointed out by Ambrogioni et al. (2021b), when modeling only those forward dependencies, the modeling of colliders can be an issue More generally, one would like to model arbitrary dependencies in the variational posterior. As an example, the factorization q(01, 02, 0a) = ( ((00))0)((()1)))))) 0a) would more faithfully represent the conditional dependencics that can arise in the postcrior. Modeling arbitrary dependencics reduces the inference gap but can become computationally intractable as the number of RVs increases. In PAVI, we strike particular balance between expressivity and computational tractability, as further motivated in this section.
In the case of single plate- -the 2-1evel case- Agrawal & Domke (2021) demonstrate that modeling only the forward dependencies does not result in reduced expressivity compared to the modeling of the full dependencies. Yet this result does not hold in the n-level case, as Webb et al. (2018) shows that faithful inversion features conditional dependencies in the posterior between ground RVs of the same template. We can dub those dependencies as horizontal dependencies across RVs in the same plate. Like Structured SVT (Hoffman & Blei, 2014) or Automatic SSVI (Ambrogioni et al., 2021a), the PAVI design thus results in reduced expressivity when stacking multiple plates in the model p. This issue can be partially alleviated with the usage of backward encoding scheme gooin in the reverse direction compared to the prior's dependencios- as in Cascading Flows (Ambrogioni et al., 2021b).
In practice, though limiting our expressivity, horizontal dependencies are difficult to inject back inlo our architecture. Critically, in the PAVI design, the use of a common density estimator across the ground RVs of the same template (sce Section 4.1) and the stochastic training over batches of those RVs (sce Section 4.2) prevent the direct modeling horizontal dependencies. Put differently, the fact that we consider the inference over different ground RVs as conditionally independent inference problems is central to our design and adverse to the modeling of horizontal dependencies.
Injecting horizontal dependencies back into our variational family is therefore a non-trivial research direction, that is not at the core of this paper. This opens up promising research directions: how could arbitrary conditional dependencies be modeled in the variational posterior in the context of stochastic training?
C.6 Towards user-friendly Variational Inference
By re-purposing the concept of amortization at the plate level, we aim to propose clear computation versus precision trade-offs in VI. dtyper-paramedd such as the encoding size -as illustrated in Figure 2 (right)- allow to clearly trade inference quality in exchange for a reduced memory footprint. On the contrary, in clas- sical VI, changing Q's parametric form for instance, switching from Gaussian to Student distributions can have a strong and complex impact on the number of weights and inference quality (Blei et al., 2017). By allowing the usage of normalizing flows in very large cardinality regimes, our contribution aims at dis- entangling approximation power and computational fcasibility. In particular, having access to expressive density approximators for the posterior can help experimenters diversify the proposed HBMs, removing the need for properties such conjugacy to obtain meaningful inference (Gelman et al., 2004). Combining clear pyper-parameters and scalable yet universal density approximators, we tend towards user-friendly methodology in the context of large population studies VI.
In its current implementation, PAVI requires to set up a rather large number of syper-parametery. This includes choosing a normalizing flow architecture, the size of the encodings at different plate levels, and the reduced model cardinalities. Those hyper- parametesss are in addition to generic ones such as the number of samples from the variational family used to estimate the FLBOO, or the choice of an optimizer. However,
42

---

Published in Transactions on Machine Learning Research (11/2023)
experiments suggest that those hyper-parameters could default to conservative values. As shown in Section 5.2, the encoding size should be maximized with respect to the available memory, as should the reduced cardinalities, as detailed in Appendix B.3. As for the normalizing flow architecture, since the latter is shared across plates, the user can always default to the most expressive architecture available without incurring an exploding number of weights. As an example, architecture (Papamakarios & Murray, 2018). Since all hyper-parameters can be set conservatively, a user- friendly automatic VI API can be designed, where the user only provides the system with the generative model p. and some observed value X. The variational family q can then automatically be derived from p. and directly optimized over X. This is the design principle we implemented in the codebase adjoined to this publication.
our experiments, we stuck the MAF
43


Reviews:
[{"review_id": "Ba2kyCu1je", "reviewer": "Reviewer_nUtr", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n\n-Bayesian inference in hierarchical problems with large numbers of parameters -- such as the ones found in large population studies with high numbers of hierarchies, subjects and measurements -- is a challenging problem in general. This study introduces the concept of plate amortization, which efficiently tackles this issue;\n\n-the study develops an approach for VI that outperforms previous methods, such as stochastic VI methods, in terms of computational cost;\n\n-the developed method is shown to be of practical utility in a Neuroimaging task with 400 million parameters. This suggests that the developed approach is a solid step towards applications in real-world and challenging problems.\n\n\n## Weaknesses:\n\n-my main concern is the fact that this study only validates its performance on tasks with Gaussian models. While these results seem quite promising, in order to properly assess the extent to which this approach works more generally, it would be important to also test the approach on other models;\n\n-this is more a comment than a specific request: writing style and notation are not always easy to follow. In particular, while the formalism with plates is intended to declutter the notation, it ends up making the reading quite difficult. That said, I do not have specific recommendations on how to improve this aspect;\n\n-while the study shows the competitive advantages of PAVI for large problems, it would be beneficial for the reader to have an explicit discussion and associated guidelines for lower-dimensional problems, where PAVI might not be the method of choice, given the potential comparatively high computational cost. Indeed, there is a section devoted to the analysis of performance for different parameter dimensionalities, but the take-home messages are left to the reader rather than explicitly written;\n\n-it would be good if the authors more explicitly discuss how easy it would be to adapt such approach to new problems, in particular regarding the code and for someone not involved in the study;\n\n-how calibrated is PAVI (https://arxiv.org/abs/1804.06788)? It would be good to provide such validation, and in case this is impractical to assess, given the high dimensional parameter spaces, this should be discussed;\n\n-in figure 3, I assume the shaded regions correspond to standard deviations around the mean (or interquartile ranges around the median) over 100 traces. In any case, it would be good to clarify this in the figure caption.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This study proposes a new method for variational inference (VI) in the context of hierarchical problems with large numbers of parameters. In a few tasks, the method is shown to perform on par with other methods with respect to ELBO, and substantially better in terms of computational efficiency.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I believe my main concern regarding the types of tasks used to validate the approach (see first point in subsection \"Weaknesses\" above) is critical to address. All other changes proposed (in subsection \"Weaknesses\") are minor, but would allow the reader to better assess the generality and limitations of the approach and codebase.", "broader_impact_concerns": "Given the exploratory nature of this study, I do not anticipate major ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "IIVaxeuTVA", "reviewer": "Reviewer_DxTo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strength\n- The new approximations are applied to HBMs of huge sizes, which have not been seen before, and are a successful and practical attempt to apply  VI to the real huge dataset.\n- Although several approximations, such as amortization, are included in applying the model to huge sizes,  the almost similar samples from the approximate posterior distribution are obtained as well as those from ordinary VI, as shown in Fig B.2. Therefore, it is considered to be of great utility.\n\n# Weakness\n- Although there is some ingenuity in the structure of the normalizing flow, it is just a combination of commonly used techniques and is not very novel.\n- The presentation of the paper is not good. For example, there are the following problems\n  1. The order of the sub-sections in Sec 2 seems strange. Although the formal definition of problem settings and notations are introduced in Sec 2.2, the authors discuss the problem settings in Sec 2.1 with this notation. This makes the paper very difficult to read.\n  2. As I wrote in \"Requested changes\", the paper has too many typos and unusual terminologies. Due to this, I felt that this paper was not ready to be published aside from the technical contributions.\n- Some of the definitions are unclear. \n   1. What is $\\pi (x_i)$ in Eq.(2) and how is it related to $\\Theta?\n   1. What is the relation between $P$ and $I$ in the last paragraph on page 4 ? I thought that in Fig.1, $P=I$ holds, but in general, does $P=I$ holds?\n- What do you mean by unbiased in Sec 2.5？ In statistical inference, our ultimate goal is inferring the unbiased estimate of $\\theta$. The author only showed an unbiasedness of the stochastic ELBO during the training. The author should discuss the bias introduced by the VI, amortization, and plate binding technique, for example, by showing Fig B.2", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a technique for efficiently performing variational Bayesian inference in hierarchical models even when the local latent variables are huge (in the order of millions). In order to reduce such a huge number of parameters, the amortization technique is utilized. \nThe new technique proposed in this paper uses summary statistics combined with the Normalizing flow to reduce the parameter size and improve the flexibility of the approximation. The authors also suggest encoding the summary statistics, making the approximation more flexible and compact.\nThe training is performed efficiently by utilizing SVI techniques. Unlike SVI, the proposed method has the advantage of constantly updating the shared parameter during training, resulting in faster learning convergence.\nNumerical experiments have confirmed these advantages, and the performance on actual brain data is much better than existing methods.\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "- Fig B.2. in the Appendix should be moved into the main paper and discussed in relation to the parameter estimation bias caused by approximations. It would be better to study this bias not only qualitatively but also quantitatively, for example, by studying the Wasserstein distance between the obtained posterior samples. \n- The definition of $\\pi (\\theta_i,n)$ should be modified since if there are no parents for $\\theta_i$, this cannot be defined.\n- The presentation should be improved, as I wrote in Weakness.\n- I think there are too many typos, so I strongly recommend conducting English proofreading.\n# Typos\nSec 1\n- Random Variables (RVs) -> random variables (RVs)\n-  Hierarchical Bayesian Model ->  hierarchical Bayesian model\n- Markov Chain Monte Carlo -> Markov chain Monte Carlo\n-  Mean-Field -> mean-field\nSec 2\n- in fig. 1 -> in Figure. 1\n- in section ?? -> in Section. ??\n\nThere are too many typos, so I gave up listing all of them.", "broader_impact_concerns": "The section about the broader impact was not included in the text and needed to be added. In particular, this study was used to analyze brain data, and it is necessary to discuss whether the approximate posterior distribution obtained by the proposed method really represents the complexity of the posterior distribution that should be captured in such high-dimensional data. In general, statistical inference suffers from modeling bias and variation due to parameter estimation. In addition to these, in this study, the situation is further complicated by bias due to variational inference and bias due to the amortization technique.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ruaJW8Ql4Y", "reviewer": "Reviewer_ojb4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\nThe proposed approach can potentially scale variational inference to an unprecedented scale. Moreover, the empirical performance is clear and I think this work has great merit. \n\n### Weaknesses\n\nI think the overall writing of the paper is a bit cluttery. The overuse of ``em dash'' results in needlessly long and complicated sentences (I am not sure if the authors used em dash, it seems more like an en dash.) There is also a repetition of the abbreviations and concepts across the paper (for instance, the HBM definition, the goal of the paper and the overall idea of the approach is repeated in throughout the different sections.) I also think the section 2 is too large and covers several different concepts that may read better if made individual sections.  Moreover, the considerable distance between Figure 1 and Figure 2 disrupts the reading flow. I would recommend placing these figures closer together or perhaps merging them into a single figure with a comprehensive caption. I also think some of the terms need a more thorough explanation. Certain terms could also benefit from more detailed explanations, especially the concepts of templates and template variables. Section 2.2 merely refers readers to a 1200-page book. A brief introduction to templates, template variables, and ground variables would undoubtedly enhance the readability of the paper. Overall, I think a thorough reread and rewrite can improve the paper immensely.\n\nFurthermore, I suggest adding a reference to the model description in Figure 3. I'm also curious as to why the authors did not use the same model (the same caridnality of plates) for both the left and right sub-figures (i.e., to demonstrate the convergence speed and the embedding dimensionality). It would be more informative if this was done using the same model. Moreover, I suspect the marginal $\\log p(X)$ should be available in closed-form in the GRE model. Could the authors plot this as a reference in the ELBO figures? An ELBO of 4000 is intriguing as it suggests the posterior is highly concentrated. Could the authors elaborate on why this is the case?\n\nConcerning Figure 4, could the authors specify what the termination criteria were for the different approaches? Reading the ELBO values from the bar plots is quite challenging, especially given the scale of the ELBO is in the thousands of nats. Might the authors also provide a tabular representation of the ELBO results for this figure?\n\nOverall, I remain somewhat uncertain about how this approach constitutes an instance of automatic VI. It appears that multiple aspects of the training procedure need careful design before someone can implement the proposed approach. For example, you must select the normalizing flow architecture, the dimensions of the embeddings, the path sampling procedure, and the number of samples from q to approximate the gradient, among other design decisions. While I can envision this as an instance of a Black-Box VI algorithm, I struggle to see how it represents an automated approach.\n\n### Small issues\n- I could not find a reference to $\\mathcal{B}_{X}$ term used in eq. 6.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors introduce a technique to scale variational inference to hierarchical models with a large number of latent variables. The proposed technique involves employing the same factorization as the model for the variational approximation and a doubly stochastic scheme to train the parameters. The paper's centeral contribution is to amortize the parameters and the learning cost across the variables of the same hierarchy level. This amortized stochastic approach has shown to result in faster convergence in large-scale problems.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the section on strengths and weaknesses. ", "broader_impact_concerns": "None. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Thank you for your patience around the review process for this manuscript. 

The reviewers are aligned in leaning towards accepting this paper with two minor revisions. 1) It would be helpful to provide additional clear context around the relationship of this work to "Amortized Variational Inference in Simple Hierarchical Models" [Agrawal and Domke, 2021]. 2) The readers would benefit from additional explanation around the sensitivity of the normalizing flows component of your proposal.

I expect these to be minor revisions and believe that the authors will be able to easily address them. Congratulations on the acceptance of your manuscript.
Decision: Accept (poster)