PAVI: Plate-Amortized Variational Inference
Authors: Louis Rouillard
Thomas Moreau
Demian Wassermann
Venue: Submitted to ICLR 2023 (2023)
Keywords: structured Variational Inference
Bayesian inference
Hierarchical Bayesian Models
Inference amortization
Neuroimaging

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

In this work, we design structured VI families that can 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 plates. We name this concept plate amortization, and illustrate the powerful synergies it entitles, resulting in expressive, parsimoniously parameterized and orders of magnitude faster to train large scale hierarchical variational distributions.

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

Under review as a conference paper at ICLR 2023
PAVI: PLATE-AMORTIZED VARIATIONAL INFERENCE
Anonymous authors
Paper under double-blind review
ABSTRACT
Given observed data and probabilistic generative model, Bayesian inference aims at obtaining the distribution of a model's ' latent parameters that could have yielded the data. This task is challenging for large population studies where thou- sands of measurements are performed over a cohort of hundreds of subjects, re- sulting in massive latent parameter space. This large cardinality renders off- the-shelf Variational Inference (VT) computationally impractical. In this work, we design structured VI families that can efficiently tackle large population studies. Our main idea is to share the parameterization and learning across the differentt i...d variables in a generative model -symbolized by the model's plates. We name this concept plate amortization, and illustrate the powerful synergies it entitles, resulting in expressive, parsimoniously parameterized and orders of magnitude faster train large scale hierarchical variational distributions. We illustrate the practical utility of PAVI through a challenging Neuroimaging example featuring a million latent parameters, demonstrating significant step towards scalable and expressive Variational Inference.
INTRODUCTION
Population studies analyse measurements over large cohorts of human subjects. These studies are ubiquitous in health care (Fayaz al., 2016; Towsley al., 2011), and can typically involve hun- dreds of subjects and measurements per subject. For instance in Neuroimaging (Kong et al., 2019), measurements X can correspond signals in hundreds of locations in the brain for a thousand subjects. Given this observed data X, and a generative model that can produce data 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 all subjects -such as the connectivity corresponding to each label. We want to recover the distribution of the θ that could have produced X. Following the Bayesian formalism (Gelman et al., 2004), we cast both @ and as Random Variables (RVs) and our goal is to recover the posterior distribution p()]X). Due to the nested structure of our applications we focus on the case where corresponds to a Hicrarchical Bayesian Model (HBM) (Gelman al., 2004). In the context of population studies, the multitude of subjects and measurements per subject implies a large dimensionality for both θ and X. This large dimensionality in turn creates computational hurdles that we tackle through our method.
Several inference methods exist in the literature. Earliest works resorted to Markov Chain Monte Carlo (Koller & Friedman, 2009), which tend be slow in high dimensional settings (Blei al., 2017). Recent approaches, coined Variational Inference (VI), cast the inference as an optimization problem (Blei et al., 2017: Zhang et al., 2019; Ruiz & Titsias, 2019). Inference reduces to finding the distribution g Φ) E Q closest to the unknown posterior p ()] in a variational family Q chosen by the experimenter. Historically, VI required to manually derive and optimize over Q. which remains an effective method where applicable (Dao et al., 2021). In contrast, we follow the idea of automatic VI: deriving an efficient family @ directly from the HBM p (Kucukelbir et al., 2016; Ambrogioni et al., 2021a;b). Our method also relies on the pervasive idea of amortization in VI: deriving posteriors usable across multiple data points, which can be linked to meta- meta-nen (Zhang et al., 2019; Ravi & Beatson, 2019; lakovleva al., 2020; Yao et al., 2019). In particular, Neural Processes share with our method the conditioning a density estimator by the output of permutation- invariant encoder (Garnelo et al., 2018; Dubois al., 2020; Zaheer al., 2018). Meta-learning

---

Under review as a conference paper at ICLR 2023
studies problems with few hierarchies, similar to many VI methods (Ravi & Beatson, 2019; Tran et al., 2017). Agrawal & Domke (2021) notably studied the case of 2-1evel HBMs, providing with theoretical guarantees. In contrast, our focus is rather computational, and we study generic HBMs with an arbitrary number of hierarchies, aiming at tackling large population studies efficiently.
Modern VI is effective in low-dimensional settings, but does not scale up to large population studies -that involve millions of random variables. In this work we identify and tackle 2 challenges to enable this scale-up. first challenge with scalability is a detrimental trade-off between expressivity and high-dimensioonality (Rouillard & Wassermann, 2022). To reduce the inference gap. VI requires the variational family Q to contain distributions closely approximating p()]X) (Blei et al., 2017). Yet the form of p( 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 parameterization, and normalizing flows scale poorly with the dimensionality of O. In large populations studies, as this dimensionality grows, the paramcterization of normalizing flows becomes prohibitively large. To tackle this challenge, Rouillard & Wassermann (2022) recently proposed -in the ADAVI architecture- to partially share the parameterization of normalizing flows across the hierarchies of a generative model. While ADAVI tackled the over-parante wver-parameterization of VI in population studies, it still could not perform inference in very large data regimes. This is due second challenge with scalability: as the size of θ increases, the evaluation of a single gradient over the entirety of an architecture's weights quickly requires too much memory and compute. This second challenge can be overcome using stochastic VI (Hoffman et al., 2013), which subsamples the parameters é inferred for at each optimization step. However, using SVI, the weights for the posterior for a local parameter are only updated when # is visited by the algorithm. In the presence of hundreds of thousands of such local parameters, stochastic VI can become prohibitively slow.
θ
In this work, we introduce the concept of plate amortization (PAVI) 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 cquivalently across a model's plates. We first propose an algorithm to automatically derive an ex- pressive yet pursimoniously-ppaaeeeeiie parameterized variational family from a plate-enriched HBM. We then propose a hierarchical stochastic optimization scheme to train this architecture efficiently. obtaining orders of magnitude faster convergence. We detail two variants of our method, with different trade-offs between parameterization and inference quality. PAVI leverages the repeated structure of plate-enriched HBMs via an original combination of amortization and stochastic training. Through this combination, our main claim is to enable inference over arbitrarily large population studies uup to million RVs- with reduced parameterization and training time as the cardinality of the problem augments. We illustrate this by applying PAVI to a challenging human brain cortex parcellation, featuring inference over a cohort of 1000 subjects with tens of thousands of measurements per subject, demonstrating a significant step towards scalable, expressive and fast VI.
2 METHODS
2.1 HIERARCHICAL BAYESIAN MODELS (HBMS), TEMPLATES AND PLATES
Our objective is to perform inference in large population studies. As an example of how inference becomes impractical in this context, consider M in fig. I as a model for the height distribution in a population. 02.0 denotes the mean height across the population. 01,0, 01,1, 01.2 denote the mean heights for 3 groups of subjects, distributed around the population mean. Xo. represent the observed heights of 2 subjects from group 0, distributed around the group mean. Given the observed subject heights X, the goal is to determine the posterior distributions of the group means p(01.n X)) and population mean p(02.o)X)
As the number of groups and subjects per group augments, the parameterization and time required to infer the posteriors of a growing number of RVs become prohibitively large. Our goal is to keep this inference computationally tractable.
In the rest of this section, we define the inference problem and the associated notations. Population studies can be modelled using Hierarchical Bayesian Models (HBMs). Those HBMs feature samples from common conditional distributions multiple levels. For instance subject heights are ii.d. given group heights, which are i.i.d. given population height. Due to data being collected across hundreds of subjects (Kong al., 2019), HBMs representing population studies feature
2

---

Under review as a conference paper at ICLR 2023
thousands of RVs, most of which correspond to symmetrical distributions. For instance, the distribution subject's height given the group height is the same across all subjects. As a result, those large HBMs M can be compactly represented via plate-enriched Directed Acyclic Graph (DAG) templates T (Gilks et al., 1994; Koller & Friedman, 2009). The plates in the template T translate graphically the hierarchical i.i.d. structure in the model M.. We exploit this plate-induced repeated structure to design scalable inference methods.
Full Grounding,
Full model M
We denote T's vertices as X and θ {0.)i=1..7- X denotes the RVs observed during inference. and e the parameters we infer: our goal posterior distribution X)) We denote 7 plates as
approximate the
Template
GL
OLL
(Pp)p=o.Ps and the plates 0, belongs to as Plates(O,).. (O,). In the toy example from fig. 1, there are 2 latent RV templates: 01 and 2z. two plates Po. P1 and we want to approximate p(01, (x()).
X₁
XL
X2
X;
X₁
X5
Reduced model J-l
Reduced
distribution
Inference
6xt
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 plate cardinalities (Card(Pp)) p=... p=
Reduoed Grounding
(Koller & Friedman, 2009). Card(P) represents the number of elements in the plate P, for in- stance the number of groups in the study. The grounding operation instantiates the repeated structures symbolized by the plates: a given RV template 0; corresponds to mul- tiple ground RVs (D.m)n=u...v with the same parametric form, where N₂ [[PEPtatee(o.) Card Template grounding is illustrated in fig. where 'T is instantiated into models M. In PAVI, we consider the inference over the ground RVs 0z,n as symmetrical problems, over which we share parameterization and learning.
Figure I: Plate Amortized Variational Inference (PAVI) working principle. Starting on the left, the graph template T is grounded into 2 separate HBMs: M (top) and .M' (down) of respec- tive plate cardinalities (3, 2) -large and (2, 1) -small. Based on M the re- duced distribution q is constructed. We train qº over data slices of small cardi- nality, before performing inference over the full model of large cardinality.
Our goal is to perform inference over the full model
M, which is associated to the density p. In p, the plate structure indicates that a RV template 0 is associated to a conditional distribution p3 shared across all ground RVs ii
Nx
N Σ Σ
log (O)
Σ (2, +
log p. (0.
where T00...) are the parents the RV 0z,, -onnditioninig its distribution. We denote with ex index all variables related to the observed RVs X. .M typically features large plate cardinalities Card(P). and thus many ground RVs, making it computationally intractable. Instead of inferring di- rectly over M, will train over smaller replica of M. To this end, instantiate the template T into a second HBM M the reduced model- of tractable plate cardinalities Card' (P) < 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, visible in fig. We train over the tractable reduced model M obtain variational distribution usable to approximate the posterior for the intractable target model M.
2.2
PAVI ARCHITECTURE
2.2.1 FULL VARIATIONAL FAMILY DESIGN
Here we define the variational distribution g corresponding to the full model M- introduced the previous section. q is the distribution used for inference, but that -due to M's large cardinalities- cannot directly train over. To derive q, we push forward the prior p(O) using trainable normalizing flows, denoted as J (Papamakarios et al., 2019). To every ground RV 0gn, associate learnable
3

---

Under review as a conference paper at ICLR 2023
flow Fi,, to approximate its posterior distribution:
N: log zi.n (Oi,n (oll)(((nn))))
log (()) Σ
i=1 ==0
(2)
log qi,n (0 (Bt,n)) log |det J (usmn)) + log (o
Wi, = -F,,((8e..))
where qt, is the push-forward of the prior Pt through the flow Fi..- This fusw-forward is illustrated in fig. 2, where flows FF push RVs u into 8. This cascading scheme was first introduced by Ambrogioni et al. (2021b), and makes q inherit the conditional dependencies of the prior p. More details about the dependencies modelled in the variational distribution can be found in appendix C.5.
2.2.2 PLATE AMORTIZATION
Here we introduce plate amortization: sharing the parameterization of density estimators across model's plates. Traditional VI searches for a distribution y( (O:; Φ) that best approximates the posterior of given a value Xo for X: q(9; Φo) p X = Xo). When presented with a new value X1, optimization has to be performed again to search for the weights Φ1, such that (();))) ~ p(O)X = X1). Sample amortized inference (Zhang et al., 2019; Cremer et al., 2018) in- stead infers in the general case, regressing the weights & using encoder f of the observed data X: g
Φ = f(X)) = p(e = X). The cost of learning the encoder weights is amortized since infer- ence over any new sample X requires additional optimization. We propose to exploit the concept of amortization, but to apply it at a different granularity, leading to our notion of plate amorlization.
Similar to amortizing across the different data samples X, amortize across the different ground RVs (0..n) > corresponding to the same RV template 0. Instead of casting every flow Fin. deeinee in eq. (2)-as as an separate, palllyaarrmetteieeedd flow, we will share some parameters across the Fin- To the template 0;, we associate a conditional flow Fi ;d, with weights Ø; shared across all the { In The flow Fs,n associated given ground RV 8s.. will be an instance of this conditional flow, conditioned by an encoding Eg.
Fr, Fi( 0i, Ei,.) yielding ii, (o,n (Oo..:: 6z, E,,..
(3)
The distributions oi thus have 2 sets of weights, @; and E..., creating a parameterization trade-off. Concentrating all of qi.n's 's parameterization into di results in all the ground RVs 0ü.n having the same posterior distribution. On the contrary, concentrating all of qi,n's 's parameterization into E,, allows the or to have completely different posterior distributions. But in large cardinality setting, this freedom can result a massive number of weights, proportional to the number of ground RVs times the encoding size. This double parameterization is therefore efficient when the majority of the weights of gi.n is concentrated into o; Using normalizing flows F,, the burden of approximating the correct parametric form for the posterior is placed onto di, while the E encode lightweight summary statistics specific to each 0,,. In section 2.3.2, we also see that this shared parameterization has synergies with stochastic training.
2.3
PAVI STOCHASTIC TRAINING
2.3.1 REDUCED DISTRIBUTION AND LOSS
Because of the large number of RV's in M, optimizing over the full distribution 9 -defined in eq. (2)- is computationally intractable. Instead, we optimize over distribution with the smaller cardinalities of the reduced model. At each optimization step I, we randomly choose inside M paths of reduced cardinality, as visible in fig. 2. Selecting paths is equivalent to selecting from X a subset X'[(] of size NY. and from a subset (+(t). For a given 0,, we denote as B, [t] the batch of selected ground RVs, of size NF. Inferring over [ we will simulate training over the full distribution q:
1
log ('(()))
Σ
N
log qi,m (0i.. | (0i..))
(4)
N
nEB((1
4

---

Under review as a conference paper at ICLR 2023
where the factor N//Nę simulates that observe as many ground RVs 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:
(x'( = Nx Σ lxn +
N NY
Nx nEBx((
(5)
1=
ELBO"[t] = Egr~o- (X"
(T())
log q (+*(t))
This scheme can be viewed as the instantiation of .M over batches of M's ground RVs. In fig. 2 we see that q' has the cardinalities of MM' and replicates its conditional dependencies. This training is analogous to stochastic VI (Hoffman et al., 2013), generalized with multiple hierarchies and minibatches of RVs. Our novelty lies in the interaction of this stochastic scheme with plate amortization, as explained in the next section.
Reduced distribution g"
t=0
B
F2
B
$2.A
2.3.2 SHARING LEARNING ACROSS PLATES
81,C
Here we detail how our shared parameterization detailed in section 2.2.2- combined with our stochastic training scheme, results in faster in- ference. In traditional stochastic VI, every 0n, corresponding to the same template 0; is associated to individual weights. Those weights are trained only when ii, is visited by the algorithm, that is to say at step when n E [4]. As plates become larger, this event becomes rare. If Oi.n. furthermore associated to pashbly-paammterizeed density estima- tor -such normalizing flow- many optimization steps are required for itn to converge. The combi- nation of those two items leads slow training.
At optimization we infer over:
Fj
[t] = {A, B,C)
Using as observed RVs:
0,D
01C
X [1] [D,
Figure 2: PAVI stochastic training scheme The reduced distribution g¹ features conditional normalizing flows Fi and F2 respectively associated to the RV templates 01 and 0z.. During the stochastic training, g¹ is instantiated over different branchings of the full model M -highlighted in blue on the left. The branchings have the cardinalities of M and change at each stochastic training step f. The branching determine the encod- ings E conditioning the flows F -as symbol- ised by the letters A, B, C- and the observed data slice as symbolised by the letters D. E.
Instead, our idea is to share the learning across the ground RVs 0.... Due to the problem's plate structure, we consider the inference over the 0zn. as different instances of a common density estimation task. In PAVI, large part of the parameterization o., E.,.. is
of the estimators q. ( mutualized via the plate-wide-shared weights Φi This means that most of the weights of the flows Fi.m -concentrated in 65- are trained at every optimization step, across all the selected batches [. This results in drastically faster convergence compared to SVI-a- seen in exp. 3.1.
2.4 ENCODING SCIIEMES
PAVI shares the parameterization and learning of density estimators across an HBM's plates. In practice the distributions qi (0,,, (88..)); 6r. E...) -defined in eq. (3)- with different n only differ through the value of the encodings E We detail one core and one alternate scheme 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, = |Ei,m]==o... per template 0, This means that an additional ground RV -for instance adding a subject requires an additional encoding vector. The associated increment in the total number of weights is much lighter than the addition of a fully parameterized normalizing flow -as would be the case in the non- non-plate-amonn amortizd regime. The PAVI-F scheme cannot be sample amortized: when presented with an unseen X, though 6; can be kept as an efficient warm start, the optimal values for the encodings E,,, have to be searched again. During training, the encodings E,,.. corresponding to
population study-
5

---

Under review as a conference paper at ICLR 2023
E B,[t] are sliced from the arrays E, and are optimized for along with di. In the toy example
from fig. 2, at & = O, [0] = {1.2} and the trained encodings are {E1.1. E1,2}. and at =
B = {0, 1] and we train {E,,o, }. In appendix A.2, show that the training of PAVI-F is
unbiased: training over M we converge to the same distribution as if training directly over M.
Deep set encoder (PAVI-E) The parameterization of PAVI-F scales lightly but linearly with Card This scaling could 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 applied to the observed data: E = f(X: 17). As encoder f we use deep-set architecture exploiting the data's plate-induced permutation invariance -detailed in appendix A.1.3 (Zaheer et al., 2018: Lee et al., 2019). This scheme allows for sample amortization across different data samples Xo, XI.... Note that an encoder will be used to generate the encodings whether inference is sample amortized or not. During training, shared learning is further amplified as all the architecture's 's weights -Ø; and 17- are trained every step f. To collect the encodings to plug into q², we build up on a property of f: ser size generalization (Zaheer et al., 2018). Instead of encoding the full-sized data X, f is applied to the slice X" [). This amounts to aggregating summary statistics across a subset of the observed data (Lee et al., 2019; Agrawal & Domke, 2021). This property is intensified in the sample amortized context: we train a sample amortized family over the lightweight model _M, and use it "for free" over the heavyweight model M. As detailed in appendix A.2, we rely on computationally efficient but theoretically biased training scheme for PAVI-E. The negative impact of this bias on PAVI-E's performance was seldom noticeable and always marginal throughout our experiments.
Summary In section 2.2 we derived an architecture sharing its parameterization across a model's plates. In section 2.3 derived stochastic scheme to train this architecture over batches of data. Our novelty lies in this original combination of amortization and stochastic training, which results in significantly faster inference, demonstrated in the following experiments.
RESULTS AND DISCUSSION
Throughout this section we use the ELBO as a proxy to the KL divergence between the variational posterior and the unknown true posterior. ELBO is measured across 20 samples X, 5 random seeds per sample. The ELBO allows to compare the relative performance of different architectures on a given inference problem. In appendix B.1 we provide with sanity checks to assess the quality of the results. In appendix B.2 we evaluate the impact of the reduced model cardinalities on performance. In appendix B.3 we compare numerically our method against basclines over a Gaussian mixture model; a model featuring the aggregation of higher order summary statistics; and a smaller version of our Neuroimaging model used in section 3.4.
Experimentally, we found the PAVI-F scheme overall faster to train and yielding 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 purameterizztioonconssse
time-constant sample-amortized inference as cardinality augments. Though degraded compared to PAVI-F's, PAVI-E's performance is still on par beats baselines in a variety of inference tasks -see exp. 3.3&..3.
3.1 PLATE AMORTIZATION AND CONVERGENCE SPEED
In this experiment, we illustrate how plate amortization results in faster training. We consider the following Gaussian Random Effects model (GRE):
X 0o 1m.m ~ NN 2
Vny=1.. CP))(P;) V.g=l.. (Po
(6)
1m 02,0 ~ N(22,00 of)
(ard(P))
8,00 ~ ~N
where D represents the data X's feature size, with group means 0, and population means 02 as D-dimensional Gaussians. We opt in this equation for a double indexing scheme instead of a simple indexing as in methods. The GRE model features two nested plates: the group plate Pı and the sample plate Po as in fig. 1. Inferring over the GRE model, the objective is to retrieve the posterior distribution of the group and population means given the observed sample.

---

Under review as a conference paper at ICLR 2023
0-
4000
ELRD thigher is better)
ELSO Shigher Is (retter)
-50.0
-10000
-2000
-4000
115000
enceding sire:
encoding sice:
SVI baseline (not plate amortized)
PNVI-F (plate amortized)
-20000
encoding
-9000
PAVIEE E (plate amortized)
cncoding Sice: 36
encoding size:
110000
10P
10*
10a
-25000
500
1000
2000
2000
3000
Steps
Staps
Figure 3: Left panel: Plate amortization increases convergence speed Plot of the ELBO (higher 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 baseline (in purple). Due to plate amortization, method converges ten to a hundred times faster to the same asymptotic ELBO as its non-plate- amortized counterpart.; Right panel: Encodings
ground RVs summary statistics Plot of the ELBO (higher better) as a function of the optimization steps for the PAVI-F architecture with increasing encoding sizes. As the encoding size augments, does the asymptotic performance, until reaching the dimensionality of the posterior's sufficient statistics (D = 8), after which performance plateaus. Encoding size allows for clear trade-off between memory and inference quality.
Here we set D 8, (Prd 100 and Card = We compare our PAVI architecture to a stochastic non-plate-amortizer baseline with the same architecture as PAVI (Hoffman et al., 2013). The only difference that ground RVs 0n, are associated in the baseline to individual uully-parameterized flows Fr,, instead of sharing the same conditional flow F, -as described in section 2.2.2. Figure (left) displays the evolution the ELBO for the baseline and PAVI with free encoding (PAVI-F) and deep encoders (PAVI-E). We that both plate amortized methods reach asymptotic ELBO equal to the non- non-plate-ann Illate-amotizzr baseline's, s, but with orders magnitudes faster convergence, and more numerical stability. This stems from the individual flows F.,. in the baseline only being trained when the corresponding 0i,, is visited by the stochastic training, while the shared flow F, updated at every optimization step in PAVI. We also note that the PAVI-E scheme con- verges faster in theory than PAVI-F -in terms of number of steps- by sharing not only the training of the flows, but also of the encoder through the optimization. In practice however, computing the encodings results in longer steps, and in slower convergence, as illustrated in section 3.3.
3.2 IMPACT OF ENCODING SIZE
Here we illustrate the role of encodings as ground RV posterior's summary statistics -as described in section 2.2.2. We use the GRE HBM detailed in eq. (6), using D = 8, Card(P)) 20 and Card' Card'(P)) = 2. We use single PAVI-F architecture. varying the size of the encodings E.,. -see section 2.4. Due to plate amortization, encodings determines how much individual information each RV 0n,n is associated to. The encoding size -varying from 2 to 16- is to be compared 8. In GRE, D = corresponds to the size of the
with the dimensionality of the problem, sufficient statistics needed to reconstruct the posterior of group mean -all other statistics such as the variance being shared between the group means. Figure (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 appears as unequivocal hyperparameter allowing to trade inference quality for computational efficiency. Increasing the encoding size also leads experimentally to diminishing returns in terms of performance. This property can be exploited in large settings to drastically reduce the memory footprint of inference while maintaining acceptable performance -choosing the encoding size approximately equal to the expected size of the sufficient statistics.
likely be exploited in the optimization. Encoding size
3.3 SCALING WITH PLATE CARDINALITIES
Here we put in perspective the gains from plate amortization when scaling up an inference problem's cardinality. We consider the GRE model in eq. (6) with D and augment the plate cardinalities
7

---

Under review as a conference paper at ICLR 2023
Inference quality
Parameterization
Time
Polimizaauu time [lower is better)
EUSO (higher better)
107
Weights (lower is better)
Cand1P
Heavy GPU memory requirement
30
207.
10"
tonvioo
Oriaa)
Ung PAUt
ACAVilla))
(144)
Uny Pive
A04//19))
Orica)
Din PAv,
Myn tin toee
Figure 4: PAVI provides with favorable parameterization and training time as the cardinality of the target model augments Our architecture PAVI is displayed on the right of each panel. We augment the cardinality Card(Pi) of the GRE model -described in eq. (6). While doing so, we compare 3 different metrics: In the first panel: inference quality, as measured by the ELBO. None of the presented SOTA architecture's s performance degrades as the cardinality of the problem aug- ments. In the second pannel: parameterization, comparing the number of trainable weights of each architecture. PAVI similar to ADAVI displays constant number of weights as 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(P1) 200, CF, UIVI and ADAVI required large GPU memory, constraint absent from PAVI due to its stochastic training.
(C)rd(Pi), Card (2 + 5) + (200, 20). In doing so, we augment the number of 6 - 42 + 402.
parameters θ
Baselines We compare our PAVI architecture against 3 state-of-the-art baselines. Cascading Flows (CF) (Ambrogioni et al., 2021b) is non-plate-amortizee structured VI architecture improving on the baseline presented in section 3.1. CF push the prior into the posterior q using Highway Flows. CF follows a cascading dependency structure, complemented with backward auxiliary coupling. ADAVI (Rouillard & Wassermann, 2022) is a structured VI architecture similar to PAVI, with constant parameterization with respect to problem's cardinality, but large training times and memory. ADAVI has several limitations compared to PAVI: ADAVI implements a Mean Field approximation (Blei et al., 2017) while PAVI implements a cascading flow; ADAVI is limited pyramidal HBMs while PAVI tackles generic saate-enriched HBMs; ADAVI is limited to full-model sample-amortee amortized variant. Unbiased Implicit VI (UIVI) is an unstructured implicit VI chitecture. UIVI infers over the full parameter space wwithout any SV SVI-amenable factorization- by reparameterizing a base distribution with a stochastic transform. For all architectures, we indicate with the suffix (sa) sample amortization, corresponding to the classical meaning of amortization, as detailed in section 2.2.2. More implementation details can be found in appendix B.4.
As the cardinality of the problem augments, fig. 4 shows how PAVI maintains state- of-the-art inference quality, while being more computationally attractive. Specifically, in terms of parameterization, both ADAVI and PAVI-E provide with heavyweight but con- stant parameterization both CF and PAVI-F's parameterization scale linearly with Card(Pi). but with a drastically lighter augmentation for PAVI-F. For an additional ground RV, CF requires an additional fully parameterized normalizing flow, whereas PAVI-F only requires an additional lightweight encoding vector. UIVI's parameterization scales quadratically with the size of the parameter space O due to a neural network regressing a transform applied to a base distribution with the size of O In detail, PAVI-F's parameterization due to the plate-wide-sharre while the part due to the encodings Ei,n grows linearly from 16 to 160 to 1.6k weights. Note that PAVI's stochastic training also allows for a controlled GPU memory during optimization, removing the need for a larger memory as the cardinality of the problem augments -a hardware constraint that can become unaffordable at very large cardinalities. To remove this memory constraint, CF could be trained stochastically, but wwithout plate amortization- would suffer from slower inference, as illustrated in section 3.1. In contrast, UIVI could not be trained stochastically, as it infers over the full θ at once instead of factorizing it. As result UIVT would be ultimately limited by memory to
the cardinality Card(P:) of the problem augments. Comparatively,
Φ1 represents a constant R 2k weights,
8

---

Under review as a conference paper at ICLR 2023
Network Network 1 Network 0&1
population
Network 0
Networks
Uncertain
Paroelletions
Natwork
subject 103212
subject
Figure 5: Probabilistic parcellation of Brocas's area PAVI can be applied in the challenging con- text Neuroimaging population studies. For cohort of 1000 subjects, 2 of which are represented here -in the bottom 2 items- we present 2 results. First, connectivity networks with the brain's left hemisphere -left purple items: those represent the brain regions vertices with each label are "wired" to. Second, Broca's 's area probabilistic parcellation -rightmost orange items: cluster the brain's vertices, associating them to connectivity networks. Our method models uncertainty: coloring tran- sitions from red an uncertain white to blue, representing the probability of vertex to belong to one network or the other.
infer over larger problems. In terms of convergence speed, PAVI benefits from plate amortization to have orders of magnitude faster convergence compared to structured VI baselines CF and ADAVI. For UIVI, the cardinality augments, training amounts evaluating neural network with increasingly larger layers. GPU training time is therefore constant, but this property would not translate larger problems the GPU memory would become insufficient. Plate amortization is particularly significant for the PAVI-E(sa) scheme, in which a sample-amortized variational family trained over a dataset of reduced cardinality, yet performs "for free" inference over HBM of large cardinality. Maintaining Card' (Pi) constant while Card(P1) augments allows for a constant parameterization and training time as the cardinality of the problem augments -without any maximum cardinality, contrary to UIVI. The effect of plate amortization is particularly noticeable at (P1) = 200 between the PAVI(sa) and CF(sa) architectures, where PAVI performs amortized inference with 1 fewer weights and x lower training time. Scaling even higher the cardinality of the problem -Card(P:) = 2000 for instance- renders ADAVI, CF and UIVI computationally intractable, while PAVI maintains light memory footprint, and short training time, exemplified in the next experiment.
3.4 NEUROIMAGING APPLICATION: BROCA'S S AREA PARCELLATION
To illustrate its usefulness, we apply PAVI to challenging population study for Broca's s area's functional parcellation. A parcellation clusters brain vertices into different connectivity networks: labels describing co-activation with the rest of the brain as measured using functional Magnetic Resonance Imaging (IMRI). Different subjects exhibit strong variability, (MRI costly to acquire: few noisy data usually gathered for a given subject. is thus essential to combine information 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 fMRI subjects. We use the HCP dataset (Van Essen et al., 2012): acquisitions from a 1000 subjects, with thousands of measures per acquisition, for over a million parameters O We use model with 3 plates; subjects, sessions and brain vertices. None of the baselines presented in section 3.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 thousand parameters in appendix B.3.3. Despite the massive dimensionality of the problem, thanks to plate amortization PAVI converges in dozen epochs, under an hour of GPU time. Results are visible in fig. 5, supporting the hypothesis of a functional bi-partition of Broca's area into posterior part involved in phonology and an anterior part in lexical/semantic processing ffollowing the anatomy of pars opercularis and triangularis (Heim et al., 2009; Zhang et al., 2020).
in fig. 5. However,
large cohort of
Conclusion In this work we present the novel PAVI architecture, combining structured variational family and a stochastic training scheme. PAVI based the concept of plate amortization, allowing to share parameterization and learning across 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.

---

Under review as a conference paper at ICLR 2023
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 Tesla V100- 16Gb GPU and AMD EPYC 7742 64-Core processors. VRAM intensive experiments in fig. performed on an Ampere 100 4CIE-40Gb GPU. Appendix 4 lists implementation details for all our synthetic experiments. Appendix B.5 related to our Neuroimaging experiment 3,4, and details both our data pre- veeprocessing steps and implementation. As part of our submission furthermore packaged and release the code associated to our experiments.
ETHICS STATEMENT
Our Neuroimaging data come from the Human Connectome Project dataset (Van Essen al., 201 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.
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard, Yangqing Jía, 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 Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wat- tenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learn- ing on heterogeneous systems, 2015. URL https:/ /www. tensorflow org/. Software available from tensorftow.org.
Alexandre Abraham, Fabian Pedregosa, Michael Eickenberg, Philippe Gervais, Andreas Mueller, Jean Kossaifi, Alexandre Gramfort, Bertrand Thirion, and Gael Varoquaux. Machine learning for neuroimaging with scikit-learn. Frontiers in Neuroinformatics, 8, 2014. ISSN 1662-5196. doi: 10.3389//minf.22444.000 URL http:///www..o...ieeee.ooooo/eea 3389 /ninnf.inf.../. 2014 4. 00014.
Abhinay Agrawal and Justin Domke. Amortized Variational Inference for Simple Hierarchical Mod- els. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 21388-21399 Curran As- sociates, Inc., 2021. URL http:: pttttp::/pproeetttt http::///ccceeeeeo/seooeoo wtwp:rooceeeinnnnneeeiesseeee.ee neurips....... ca ccc/paaee http:/apprrr/222 faaper/2021/flee/ Pwwer.
Luca Ambrogioni, Kate Lin, Emily Fertig, Sharad Vikram, Max Hinne, Dave Moore, and Marcel van Gerven. Automatic structured variational inference. Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 676 684. PMLR, 13-15 Apr 2021a. URL https: /ttpas://proceeee pttppppecceennnneeeee mlr press / v130 aaaabmaboooooooooooooo illaalllll..m html.
Luca Ambrogioni, Gianluigi Silvestri, and Marcel van Gerven. Automatic variational inference with cascading flows. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp- 254-263. PMLR, 18-24 Jul 2021b. URL https: / httpass//proceeee ppwpceeeiigg...pcp/ppppp mlr...mrr.... wtt...peess/1339 cabbrogionill...om html.
Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science and Statis- tics). Springer-Veriee Berlin, Heidelberg, 2006. ISBN 0387310738.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational Inference: A Review for Statis- ticians. Journal of the American Statistical Association, (2(518):859-877 April 2017. ISSN 0162-1459, 1537-274X. doi: 10. (08000162144992277.22228 URL attp: :/ /ttp/v..xg./g///////........ abs/1601 60670.
arXiv: 1601.00670.
10

---

Under review as a conference paper at ICLR 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- ume 20. Curran Associates, Inc., 2007. URL https:/ btttss:/pprcceeeiggss..o paper/2007/fileg 0d311 2ate
Chris Cremer, Xuechen Li, and David Duvenaud. Inference Suboptimality in Variational Autoen- coders. vrXiv:/8010035. [cs, stat], May 2018. URL 03558. arXiv: 1801 003558858
Kamalaker Dadi, Gaël Varoquaux, Antonia Machlouzarides-Shalle Krzysztof Gorgolewski, Demian Wassermann, Bertrand Thirion. and Arthur Mensch. Fine-grain atlases of functional modes for IMRI analysis. Neurolmage, 221:117126,, 2020. ISSN 1053-8119, doi: https:
URL https:::: /wwww//wwww/////////////../. www.w.scieneeieee.....o com/
Viet-Hung Dao, David Gunawan, Minh-Ngoe Tran. Robert Kohn, Guy E. Hawkins, and Scott D. Brown. Efficient selection between hierarchical cognitive models: Cross-validation with varia- tional bayes, 2021. URL https: /arpiv//aaivvvv//aa////i .ogg abs/ 2202 06814.
Joshua Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A. Saurous, TensorFlow Distributions.
1111v:1711.1004 [cs, stat), November 2017. URL 10604. arXiv: 1711.10604.
Yann Dubois, Jonathan Gordon, and Andrew YK Foong. Neural process family. http://
yanndubs, www.ggiitthooooo iot.oorrNNeeaaaa Proccesrressssse Pareesss-Famlll Family/ September 2020.
Fayaz, P Croft, RM Langford, LJ Donaldson, and GT Jones. Prevalence of chronic pain the uk: systematic review and meta-analysis of population studies. BMJ open, 6(6):e010364, 2016.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J. Rezende, S. M. Ali Eslami, and Yee Whye Teh. Neural processes, 2018. URL
org/ abs/
0107.016220
Andrew Gelman, John B. Carlin, Hal S. Stern, and Donald B. Rubin. Bayesian Data 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: 10.2307/2348941, URL https:/ //ww///.w////....///// jistissto....... htt.://aaa...........2..222.err.
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. Blei. Structured Stochastic Variational Inference. 1401::4404.411 [cs], November 2014. URL arXiv: 1404.4114.
argg/abs//400/ 4114.
Matthew D. Hoffman, David M. Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 14(4):1303-1347 2013. URL http: //jmmr. wwwmmlo.rrpppper///ooo /ho fman1 133.. html.
Ekaterina lakovleva, Jakob Verbeek, and Karteek Alahari. Meta-learning with shared amortized vari- ational inference. Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 4572-4582. PMLR, 13-18 Jul 2020. URL https: http///p//eeeee//eeeee ngs iakovleva20a. html.
wog..ml..... mlr. prwss / 9/
Eric Jang, Shixiang Gu, and Ben Poole. GUMBEL-SOFIMAX pp. 12, 2017.
CATEGORICAL REPARAMETERIZATION WITH
11

---

Under review as a conference paper at ICLR 2023
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 http: //arxivv.org abs /142 6980.
Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. Adap- tive computation and machine learning. MIT Press, Cambridge, MA, 2009. ISBN 978-0-262- 01319-2.
Ru Kong, Jingwei Li, Csaba Orban, Mert Rory Sabuncu, Hesheng Liu, Alexander Schaefer, Nanbo Sun, Xi-Nian Zuo, Avram J. Holmes, Simon B. Eickhoff, and B. T. Thomas Yeo. Spatial topogra- phy of -ndivicual-specific cortical networks predicts human cognition, personality, and emotion. Cerebral cortex, 29 6:2533-223 2551, 2019.
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M. Blei. Automatic Differentiation Variational Inference. arXiv: 103.008803 [cs, stat), March 2016. URL http: http://aab/11000078888 arXiv: 1603.00788.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-basee based permutation-invatie invariant neural networks. In Kama- lika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Confer- ence on Machine Learning, volume 97 Proceedings of Machine Learning Research, pp. 3744- 3753. PMLR, 09-15 Jun 2019, URL http: http:://pcccciiiieeeccmoo html.
wwt.press/v9//eess
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. CoRR, abs//611..07711 2016. URL http://:://axxx oxg/abs/1611 00712.
George Papamakarios, Theo Pavlakou, and Iain Murray. Masked Autoregressive Flow for Density Estimation. wwXivv/705.07000 [cs, star], June 2018. htt.http//axxiv.orggg////// 1705. 07057. arXiv: 1705.07057.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. arXiv:1912.0211 [cs, stat]. December 2019. URL http:/ /ttp///axivvv/vvv org/ abs/ 19t1. 02762. 1912.02762.
arXiv:
Sachin Ravi and Alex Beatson. Amortized bayesian meta-learning. In International Confer- ence on Learning Representations, 2019. URL sttps:://operreview net.forummfdd rkgpy3C5tX.
Louis Rouillard and Demian Wassermann. ADAVI: Automatic Dual Amortized Variational Infer- ence Applied To Pyramidal Bayesian Models. In ICLR 2022, Virtual, France, April 2022. URL
https:/ http:/haaaaaaaacsseeeeeee
archives www.s-uuveeeeeeeeeeeeeeeeee -t/oee/-6.///111-007777756 03267956.
Francisco Ruiz and Michalis Titsias. contrastive divergence for combining variational inference and MCMC. 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. 5537-5545. PMLR, 09-15 Jun 2019. URL https:/t///c. /ttsss///roceeeee ngs m.rr.
press htsss//////uuiil..hhh html.
Hojjat Salehinejad, Julianne Baarbe, Sharan Sankar, Joseph Barfett, Errol Colak, and Shahrokh Valaee. Recent advances in recurrent neural networks. CoRR, abs/1801.01078, 2018. URL
http:// /ttt://axiiv..//ii arg/abs/180/ /tt1//11111111....11111111111111
Stephen M Smith, Christian Beckmann, Jesper Andersson, Edward Auerbach, Janine Bijster- bosch, Gwenaëlle Douaud, Eugene Duff, David A Feinberg, Ludovica Griffanti. Michael P Harms, et al. Resting-state fmri in the human connectome project. Neuroimage, 80: 10.18,4168, 2013.
12

---

Under review as a conference paper at ICLR 2023
Michalis K. Titsias and Francisco Ruiz. Unbiased implicit variational inference. In Kamalika Chaud- huri and Masashi Sugiyama (eds.), Proceedings of the Twenty-Second International Conference Artificial Intelligence and Statisties, volume 89 of Proceedings of Machine Learning Research, pp. 167-176. PMLR, 16-18 Apr 2019. URL https: / httpc//prooceeeeeee ngs ngs. aag.... ml wwwps/eessBB. titsiasl .htmml.
Kayle Towsley, Michael I Shevell, Lynn Dagenais, Repacq Consortium, et al. Population-base study of neuroimaging findings in children with cerebral palsy, European journal of paediatric neurology, 15(1):99-35, 2011.
Dustin Tran, Rajesh Ranganath, and David M. Blei. Hierarchical Implicit Models and Likelihood- Free Variational Inference. 172. 702.08896 [cs. stat], November 2017. URL http:// arxiv. org/ abs /1702 08896.
arXiv: 1702.08896.
D. C. Van Essen, K. Ugurbil, E. Auerbach, D. Barch, T. E. Behrens, R. Bucholz, A. Chang, L. Chen, M. Corbetta, W. Curtiss, S. Della Penna, D. Feinberg, M. F. Glasser, N. Harel, A. C. Heath, L. Larson-Prior, D. Marcus, G. Michalareas, S. Moeller, R. Oostenveld, S. E. Petersen, F. Prior, B. L. Schlaggar, S. M. Smith, A. Z. Snyder, J. Xu, and E. Yacoub, The Human Connectome Project: data acquisition perspective. Neuroimage, 62(4):2222-2D) 2231, Oct 2012.
Stefan Webb, Adam Golinski, Robert Zinkov, N. Siddharth, Tom Rainforth, Yee Whye Teh, and Frank Wood. Faithful Inversion of Generative Models for Effective Amortized Infer- ence. 1.12v:1712.002888 [cs, stat], November 2018. URL 00287. arXiv: 1712.00287.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and Philip S. Yu. A Comprehensive Survey on Graph Neural Networks. IEEE Transactions on Neural Networks and Learning Systems, pp. 1-21, 2020. ISSN 2162-237X, 2162-2388. doi: 10.1109/TNLS arXiv: 1901.00596.
2978386. URL http: arxiv. org. 1901. 00596.
Huaxiu Yao, Ying Wei, Junzhou Huang, and Zhenhui Li. Hierarchically 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-7054. PMLR, 09-15 Jun 2019. URL https: http://ppcceeeeeeeeeeeeee/. wtt..mmw............mmmmmm yaol 9b. html.
Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. In Jennifer Dy and An- dreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 5660-5669. PMLR, 10-15 Jul 2018. URL https: http::s//proceeddeee
mlr www.mmlr....ss.....m https//v0//y/// inn 8b html.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep Sets. arXiv:1703.066 14 [cs, stat], April 2018. URL htlhtt:p/aaxii org/abs /aaab///03/ 11703 06114. arXiv: 1703.061 14.
Cheng Zhang, Judith Butepage, Hedvig Kjellstrom, and Stephan Mandt. Advances in Variational Inference. IEEF Transactions Pattern Analysis and Machine Intelligence, 41 (8):2008-2026 August 2019. ISSN 0162-8828, 2160-9292, 1939-3539. doi: 1LLL10//PAAM.2211888877777 URL wttph:t///// wttss///eeeexppeeeeoee / wtw.eeeee.rg//ccmeeereee/e dowument.tetntcoeooo/ httg///occmeeee//8/33/// 8588899/88
org
Yizhen Zhang, Kuan Han, Robert Worth, and Zhongming Liu. Connecting concepts the brain by mapping cortical representations of semantic relations. Nature Communications, 11777
April 2020. ISSN 2041-1723 doi: 10. (073//41467-000-1111 467-020-1180444
13

---

Under review as a conference paper at ICLR 2023
SUPPLEMENTAL MATERIAL
A SUPPLEMENTAL METHODS
A.1 PAVI IMPLEMENTATION DETAILS
A.1.1 PLATE BRANCHINGS AND STOCHASTIC TRAINING
As exposed in section 2.3.1, at each optimization step & 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 0,, Those batches have to be coherent with another: they have to respect the conditional dependencies
the original model M. As an example, if ground RV is selected as part of M", then its parent RV needs be selected as well. To ensure this, during the stochastic training we do not sample RVs directly but plates:
1. For every plate Pp. sample without replacement Card'(Pp) indices amongst the Card(P,) possible indices.
2. Then, for every RV template 8,, we select the ground RVs 8s,n corresponding to the sampled indices for the plates Plates(O:).
3. The selected ground RVs fj.n constitute the set (*(t] of parameters appearing eq. (4). The same procedure yields the observed RV subset [t] and the data slice X" [t].
For instance, in the toy example from fig. I, X2 will be chosen iff the index is selected as part of sub-sampling Pi and the index 0 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 fig. 2 t 1 where X2 corresponds to the node E. This stochastic strategy also applies to the selected encoding scheme -described in section 2.48223.2- as detailed in the next sections.
A.1.2 PAVI-F DETAILS
In section 2.4 refer encodings E, -E,,n] == corresponding RV templates 0,. In prac- tice, we have some amount of sharing for those encodings: instead of defining separate encodings for every RV template, we define encodings for every plate level. A plate level combination of plates with at least one parameter RV template 0; belonging to it:
PlateLevels ((PR..Pt) = Plates(0;) (0,) ba
(A.1)
For every plate level, we construct large encoding array with the cardinalities of the full model M:: Encodings = (PP..P)) - R (ar XCar(()))
((PE.. P P.Pe..P))PPlsteeeve
(A.2)
E, Encodings CeroodinggPPaatessiti (0,))
Where D is encoding size that we kept constant to de-clutter the notation but vary between plate levels. The encodings for a given ground RV 0i, then correspond to an element from the encoding array E.
A.1.3 PAVI-E DETAILS
In the PAVI-E scheme, encodings not free weights but the output of en encoder f(. applied to the observed data X. 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 dependency structure for the plate levels:
((Pa..P)) E PlateLevels
V(P..Pa) PlateLevels
(A.3)
(P..Po) E T (PP... Pa)) +
38./ Plates(B.)=(P...))) Pb) E T(,))
30,/ ((P..)(
note that this dependency structure is in the backward direction: a plate level will be the parent of another plate level, if the former contains RV who has a child in the lattr. We therefore obtain
14

---

Under review as a conference paper at ICLR 2023
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 dependeney structure, we will recursively construct the encodings, starting from the observed data:
Vr E X: Encodings(Plates())) = p(x)
(A.4)
where x is the observed data for the RV 2, and is a simple encoder that processes every observed ground RV's value independently through an identical multi-layer perceptron. Then, until we have exhausted all plate levels, we process existing encodings to produce new encodings:
V(Pk.. Pi) E PlateLevels / Ar E X, Plates(x)) = (Pk..P)): g (ETcoddngs((PPa.PP)))
(A.5)
(eecodings((px-...
where g is the composition of attention-baset deep-set networks called Set Transformers (Lee al., 2019; Zaheer et al., 2018). For every plate P, present in the parent plate level but absent in the child plate level, 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 wwith 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 BIAS IN THE STOCHASTIC TRAINING
A key concern in our stochastic training scheme is its unbiasedness: we want our stochastic opti- mization to converge to the same variational posterior as 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, yet show how the approximations we do in practice theoretically result in biased training. As important note, the negative impact of this bias on the performance PAVI-E remained limited throughout our experiments as seen in experiments 3.3, B.2, B.1 and B.3.
A.2.1 GENERAL DERIVATION (APPLICABLE TO BOTH THE PAVI-F AND PAVI-E SCHEMES)
We first formalize the plate sampling strategy described in appendix A.I.I. To every plate P associate the RV Ip corresponding to the Card"(P) (si-d set of indices sampled without replace- ment from the Card(P) possible index values. As example, with plate Po with Card(P,) and Card' (Po) = 2, {0.2) or 3) can be 2 different samples from P.- At a given optimization step sample independently from the RVs (Ir,)p=0.. O..P. This defines the batches B:((] and the distribution g² in eq. (4).
To check the unbiasedness of our stochastic training, we need to show that:
E
EIp [ELBO*[]] = ELBO
(A.6)
Where:
ELBO = Ee [logp(X, θ) log q())]
(A.7)
And ELBO*((] defined eq. (5). In that expression, q and p have symmetrical roles. As the ELBO amounts to the difference between the logarithms of distributions p and g. we can prove the equality in eq. (A.6) we prove that the expectation of each reduced distribution equal to the corresponding full distribution. To prove the equality in eq. (A.6), sufficient condition is therefore to prove that:
Etg. Etpp [log (t)) E, [log q" (+T(t) log (())
(A.8)
where to de-clutter the notations we denote the expectation over the collection of RVs (Ip, Ip-o..P E
15

---

Under review as a conference paper at ICLR 2023
Consider given ground RV 8i,n corresponding the RV template 8, and to the plates Plates(0.) At a given given stochastic step t, f,, will be chosen and only if its corresponding branching is chosen. Recall that when sampling equiprobably without replacement set of & elements from population of n elements, a given element will be present the set with probability k/n. We can apply this reasoning to the choice of branching corresponding will be chosen iff the index is selected as part of sub-sampling P1 and the index O is selected as part of sub-sampling Po- As Card" = 2 indices are chosen inside the plate Pi of full cardinality 3, and Card' Card = indices are chosen inside the plate Po of full cardinality 2, X2 is therefore chosen with probability 2/3 1/2. More formally, for ground RV fr,n we have:
given ground RV. For instance, in fig. 1, X2
Card Ca)
=n= = O..N, P(0 B.((])
II
PEPlates(8.)
(A.9)
N
N
Applying this reasoning to every RV template 0,, we have that:
N
Ep ( =
EiP
Σ log m. (O..
N
i=1
nEB.I4
N₂
Bip N
L log (0
(A.10)
N
N;
E Eip [1nes,[] log qt,n qt.n ((.))) E..
where we exploited the fact that the expectation of the sum of RVs is the sum of the expectations, in the case of dependent RVs. The term 1l (eB,(() [t] x log gt,m (9.,. (00..)))))))))))) E,,.) is the product of 2 RVs -related to the stochastic choice of plate indices:
the RV nEB, [t] is an indicator that Oi,n's branching has been chosen via the stochastic sampling of plate indices. By construction, this RV depends only on the indices of the plates Plates(0.).
the RV log 9t (0.,. ((0.../((0.n)::0)) E.,.. depends on E whose construction depends on the encoding scheme:
In the PAVI-F scheme, E,,, is constant.
In the PAVI-E scheme, E,, results of the application of encoder the observed data of subset of 0;,," 's descendants. By construction, this subset will only depend on the indices of plates containing 9, descendants, but not containing 0,. The value of E,n therefore only depends on the indices of plates P Plates(O:)
As an example of this reasoning, consider the model M illustrated fig. 2. We can evaluate both terms for the ground RV ,,2 in the PAVI-E scheme:
1 228, [t] depends on whether the index chosen as part of sub-sampling the plate Pi. and therefore only depends on the RV Lp, In this case the associated probability is 2/3;
to evaluate log q1.2 (0..2000000000000 Φ1, E1,2), the value of E1 will result from the application of the encoder f over the value of either X4 Xş. This choice depends whether the index or is chosen as part of sub-sampling the plate Po. Therefore, the value of the term log 41 only depends on the RV Ip..
In both PAVI-F and PAVI-E- the terms [ngB,(t) and log 9:,n (0. (00..)) Φ.. E...) are therefore independent, meaning that the expectation of their product can be rewritten the product of their
16

---

Under review as a conference paper at ICLR 2023
expectations:
N
E, [log (0
Bip
[AnEB,4]]
Etp [log gi, (i, (0.
N
N
N;
N
Eip 9i 1 ((
(A.11)
N
N
n=
N; Ei Dlg Yi (0.(OOn/TT(0ina) Es.n)
Am
This equality can be further simplified in the PAVI-F case -proving its unbiasedness- but not in the PAVI-E case, detailed in the sections below.
A.2.2 UNBIASEDNESS OF THE PAVI-F SCHEME
In the PAVI-F scheme, detailed in section 2.4, the encodings E are constants with respect the branching choice, therefore we have:
N2
E log (e
Σ
Σ
Es [log Ri, (0s,n 1 (0i,n);;
i=1
n=0
1
N;
(A.12)
Σ
log Pin (0s,a T (0inn);; Ein))
1=1
log (())
which proves eq. (A.8) and eq. (A.6). In the above example of 01.2 in M. in the PAVI-F scheme the expression E [log qa.n (0;,, (r(( ():)):: E. can be evaluated into log 91,((0122 02 Φ1, E1.2). This demonstrates that the PAVI-F scheme unbiased: training over stochastically chosen sub- graphs for 9º is in expectation equal to training over the full graph of q.
A.2.3 APPROXIMATIONS IN THE PAVI-E SCHEME
In the PAVI-E scheme, detailed in section 2.3.2, the encodings E.,, are computed from the ob- served data X. Specifically, considering the ground RV 6,,,, we have E,,, f(( where [4] corresponds to the observed data of a subset of 0,,n's descendants. Depending on the chosen branching downstream of 0z,., the value of Ei,n can therefore vary. This means we can- not further simplify eq. (A.11): the terms log 9s.n (f..( (#...|((O..n))) E.,.) are not constants with re- spect the RVs Ip. In the above example of 81.2 in M, in the PAVI-E scheme the expression E, [log n (0 7 (O;..): Φi, Ei,))) can be evaluated into:
2 (log q1
di, + log 2111 6i, ((Xs))
How could the PAVI-E scheme be made unbiased? Specifically, by making the value of E,, independent of the choice of downstream branching. A possibility would be to parameterize E as an average -an expectation- over all the possible sub- bub-branchings downstream of 6 Yet, in practical cases, the cardinalities of the reduced model are much inferior to the ones of the full model: Card (P) < Card(P). This means that numerous Card' (P))((zzd subsets can be chosen inside the Card(P) possible descendants. order to average over those subset choices to compute E,,., numerous encoding calculations would be required at each stochastic training step. For large-scale cases, we deemed this possibility impractical. Other possibilities could exist, all revolving around the problematic of aggregating collections of stochastic estimators into general estimator -in an unbiased and efficient manner. To our knowledge, this a complex and still open research question, whose advancement could much benefit our applications.
Practical approximation for the PAVI-E scheme practice, we compute the encoding E.,, based on the single downstream branching corresponding to the sampling of the RVs Jp. Compared to
17

---

Under review as a conference paper at ICLR 2023
the previous paragraph, this amounts estimating the expectation of E,,, -over all downstream branchings- using a single one of those branchings. Note that, even this encoding estimate was unbiased, log Ri, would remain highly non-linear function of E Ei,.. As consequence, need to rely on the approximation:
Eip [log qi, (0,, ((.. = log (0,,
(A.13)
which can theoretically introduce some bias in gradients. The approximation eq. (A.13) can be interpreted as follow: "the expectation the density of 8,,, when collecting summary statistics over stochastic subset of 0,.n s descendants is approximately equal to the density of 8s,n when collecting summary statistics over the entirety of Oz, descendants".. Another interpretation that the distribution associated the summary of the full data be approximated by annealing of the distributions associated to summaries subsets of this data. In practice, this approximation did yield significantly worse performance for the PAVI-E scheme over the generative models tested. At the same time, computing the encodings over a single branching allows to compute all the E.,, encodings in a single lightweight pass over the data ("(t). This simple solution therefore provided with 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 PAVI-E scheme. In practice, using the PAVI-F scheme as a samity check over synthetic, toy-dimension implementations of the considered generative models is good way to validate the PAVI-E scheme -before moving onto the real problem instantiating the same generative model with a larger dimensionality.
A.3 PAVI ALGORITHMS
More technical details can be found in the codebase provided with this supplemental material.
A.3.1 ARCHITECTURE BUILD
Algorithm 1: PAVI architecture build
Input:
Graph template T, plate cardinalities ((Card(Pp), Card (Ppdd"(pp)) encoding scheme
Output: q distribution
for
1../ do
Construct conditional flow F;;
Define conditional posterior distributions qi,a as the push-forward of the prior via F., following eq. (2);
Combine the qi distributions following the cascading flows scheme, as section 2.2.1
(Ambrogioni et al., 2021b);
PAVI-F encoding scheme then
Construct encoding arrays (E, (E., hei.2 as in appendix A.1.2;
else if PAVI-E encoding scheme then Construct encoder o as in appendix A.1.3;
18

---

Under review as a conference paper at ICLR 2023
A.3.2 STOCHASTIC TRAINING
Algorithm 2: PAVI stochastic training
Input: Untrained architecture y, observed data X, encoding scheme, number of steps T Output: trained architecture q
for t = 0..I' do
Sample plate indices to define the batches B,(4), the latent @*((] and the observed ['((] and X'It], following appendix A.1.1;
Define reduced distribution p²
if PAVI-F encoding scheme then
Collect encodings Ei,, by slicing from the arrays E; the elements corresponding to the batches [t];
else if PAVI-E encoding scheme then
Compute encodings as f(X" [t):
Feed obtained encodings into q²
Compute reduced ELBO as in eq. (5), back propagat its gradient:
Update conditional flow weights {
PAVI-F encoding scheme then
Update encodings I..BB,....
else if PAVI-E encoding scheme then
L
Update encoder weights
A.3.3 INFERENCE
Algorithm 3: PAVI inference
Input: trained architecture y, observed data X, encoding scheme Output: approximate posterior distribution
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 q
A.4 INFERENCE GAPS
In terms of inference quality, the impact of our architecture can be formalized following the gaps terminology (Cremer et al., 2018). Consider joint distribution p(O, X). and value X for the RV template X. We pick variational family Q, and in this family look for the parametric distribution g(O; Φ) that best approximates p(o]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. (201 8) refer to the gap G:
g = (I(((()))))
(A.14)
log p(X) ELBO(T;))
We denote q" (O; the optimal distribution inside that minimizes the KI. divergence with the true posterior:
Gappnox Q: x²) logg((X :L"O(
2 0
(A.15)
Gvanilla VI = Gapprox
The approximation gap Gapprox depends the expressivity of the variational family Q, specifically whether Q contains distributions arbitrarily close to the posterior -in the KL sense.
Gapprox is property of the variational family Q. Gapprox is asymptotic bound for the KL divergence between any distribution g E Q and the true posterior. This gap therefore a form of bias, but is not to be mistaken with the itdchast bias studied in appendix A.2. The bias
Note:
19

---

Under review as a conference paper at ICLR 2023
in appendix A.2 relates to whether q" can be found by training stochastically over Q, whereas Sappexx relates to the the bias between q² and the true posterior.
Cremer et al. (2018) demonstrate that, in the of sample amortized inference, when the weights 6 no longer are free but the output of an encoder f E F, inference cannot be better than in the non-sample-amorrtied case, and positive amortization gap is introduced:
Gaa (u, F;"n" = Gapprox (2; (XX, 1j") - 9approx (2 2"
0
(A.16)
Ssample amortized =
Gapprex + 9mu
Where we denote as 17" the optimal weights for the encoder f inside the function family F. The gap terminology can be interpreted as follow: "theoretically, 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 q inside the family Q, consider picking q from the plate-mmrriizer family 2px corresponding to Q. Distributions in 2m are distributions from Q with the additional constraints that some weights have be equal. Consequently, 2pa is a subset Q:
Op Q
(A.17)
As such, looking for the optimal distribution inside 2pu instead inside Q cannot result in better performance, leading to plate amortization gap:
gpa Q, Quai 46" xx) = Guppaux Opuiz/") - Gapprox (2; Φ²)
0
(A.18)
GEAVI-F =
Gupprox + 9pu
Where we denote vi the optimal weights for variational distribution q inside 2p -in the KL sense. The equation A.18 is valid for the PAVI-F scheme -see section 2.4. We can interpret it as follow: "theoretically, plate amortization cannot be beneficial in terms of KL divergence for the inference over given sample X"
Now consider that encodings are no longer free parameters but the output of encoder f. Similar to the case presented in eq. (A, using encoder cannot result in better performance, leading to an encoder gap: Gencoter Spa, Fi 1p m²) = Gupprox ( (Sps: ((X, Gapprox (Qpi w!*))
2 0
(A.19)
GEWIE = Gupprox + 9pa + Gencuder
The equation eq. (A.19) is valid for the PAVI-E scheme -see section 2.4.
The most complex case is the PAVI-E(sa) scheme, where we combine both plate and sample amor- tization. Our argument cannot account for the resulting GPAVI-E/sa) gap: both the PAVI-E and PAVI- E(sa) schemes rely upon the same encoder f. In the PAVI-E scheme, is overfit over dataset composed of the slices of a given data sample X. the PAVI-E(sa) scheme, the encoder is 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 understand it- the gap terminology cannot account for this discrepancy.
Comparing previous equations, we therefore have:
Gvanilla VI S GPAVI-F S GPAVI-E
(A.20)
Note that those are theoretical results, that do not necessarily pertain to optimization in practice. In particular, in section 3. 3.3., this theoretical performance loss not observed empirically over the studied examples. On the contrary, in practice our results can actually be better than non -mooriztt baselines, as is the case for the PAVI-F scheme in fig. or experiments B.3. We interpret this as a result of simplified optimization problem due to plate amortization -with fewer parameters to optimize for, and mini-batching effects across different ground RVs. better framework to explain those discrepancies could be the one from Bottou & Bousquet (2007): performance in practice is not only the reflection of an approximation error, but also of an optimization error. A less expressive architecture -using plate amortization- may in practice yield better performance. Furthermore, for the experimenter, the theoretical gaps Gpu Gensnder are likely to be well "compensated for" by the lighter parameterization and faster convergence entitled by plate amortization.
20

---

Under review as a conference paper at ICLR 2023
B SUPPLEMENTAL RESULTS
B.1 GRE RESULTS SANITY CHECK
As exposed in the introduction of section 3, in this work we focused on the usage of the ELBO as an inference performance metric (Blei et al., 2017):
ELBO(q) log ((X) (I(((())l((Ix))
(B.21)
Given that the likelihood term log ((X) does not depend on the variational family g. differences in ELBOs directly transfer differences in KL divergence, and provide with a straightforward metric to compare different variational posteriors. Nonetheless, the ELBO doesn't provide with an absolute metric quality. As a sanity check, we want to assert the quality of the results presented in section 3.3 -that transferable section 3.1&3.2, based on the same model. In fig. B.1 plot the posterior samples of various methods against analytical ground truths, using the Card = 20 case. All the method's results are aligned with the analytical ground truth, with differences in ELBO translating meaningful qualitative differences in terms of inference quality.
B.2 EFFECT OF THE REDUCED MODEL CARDINALITIES ON THE TRAINING EFFICIENCY
In fig. B.2 we show the impact of the augmentation of Card' (Pi) on the efficiency of the variational posterior's 's optimization.
In practice, we noticed that the most efficient choice in 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 make the BLOO 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'(P:), This due to the compute of the encodings being vectorized across plates in our deep-set encoder see appendix A.1.3 (Lee et al., 2019). We observe slight if barely noticeable reduction of the inference bias when augmenting Card Similar to the PAVI-F scheme, Card' should be maximized with respect to the GPU memory constraint. The intuition behind this choice is that the generalization of the learning summary statistics across sets of data points is casier the closer the reduced set size is to the full size. This constant training speed allows for controlled memory footprint of the stochastic training: contrary to the PAVI-F scheme, Card' Card'(P:)) could be maintained value manageable by the GPU without reducing the training speed. This example also displays a pathological choice for Card"(Pi): the encoder fails to learn to compute the correct summary statisties over sets of size However trivial, this example underlines that Card'(P,) should be chosen at a value logical with respect to the posterior's 's sufficient statistics. When generalizing moments of higher order -such the variance- there is therefore theoretical lower bound to consider when fixating the value of Card'(P,) Our intuition is that the more complex to estimate the statistic is, the larger Card' (Pi) should be.
Note that in practice non-stochastic VI is intractable for large-scale models, because of its memory requirement. In large scale experiments such as the one presented in fig. 5, stochastic training is necessary. Our claim is to be faster than non-plate-amortize stochastic VI in those large-scale contexts but not necessarily to be faster then non-stochastic VI in the small-scale regime of this experiment. For PAVI-F nonetheless obtain similar training speed compared VI with Card (Prr"(PP) as small as 8. For PAVI-E, the stochastic training is as fast as the non-stochastic one.
sto-stcccattic
21

---

Under review as a conference paper at ICLR 2023
Observed sample
Obseneed sample
Obserecc sample
Obearved sample
100
300
303
400
Data
03
-05
-03
-10
-18
45
0.0
63
1.0
15
20
63
10
15
1.4
03
0.0
43
14
sample
1.
25
64
3.0
03
valgrack:
-05
-0.3
-10
-10
-95
00
60
10
0.5
0.0
95
1.4
05
Sempres
10
2.5
84
3.0
-10
-15
-23
0.0
30
0.5
-5
-10
-55 00
Figure B.l: GRE Sanity check Inference methods present qualitatively correct results, making ELBO comparisons relevant in experiments. On the topmost line, represent 4 different X samples for the GRE model described in eq. (6) with Card(Pi) = 20. Each set of colored points of the 20 groups. Bottom lines represent the posterior
represent the X.,. points belonging samples for the methods used in section 3.3. Colored points are sampled from the posterior of the groups means 0,, whereas black points samples from the population black circles an analytical ground truth, centered on the correct posterior mean, and with radius equal 2 times the analytical posterior's '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. PAVI is represented on the last lines. Some minor bias can observed for the PAVI-E scheme, but this approximation error is marginal compared to the optimization error that can observed for unbiased methods, such as CF(sa) (Bottou & Bousquet, 2007). We can observed a superior quality for the PAVI-F scheme, rivaling ADAVI and CF's performance with orders of magnitude less parameters and training time, as visible fig. 4.
0z. We represent as
22

---

Under review as a conference paper at ICLR 2023
240
240
220
220
1130
200
RAWAT Card redum1
PAU-E Card, redum1
PAVI Cans_redu-2
PAVI Cardurddu
PAVIT Card redure
Card
Card
rudu-16
stccheciiityy
10
Time 10
Times 102
Figure B.2: Both panels: Effect of Card' Card'(P:) on the training efficiency Experiment performed on the GRE model (see eq. (6)). We increase the cardinality of the reduced model Card (Pi) from 20 while keeping Card ing, meaning we train directly on M. This experiment is interesting to evaluate the bias introduced in the stochastic training. as detailed in appendix A.2. Left panel: PAVI-F scheme As Card (P1) augments, the training gets faster and less noisy. likely due less stochasticity in the gradient esti- mates and more encodings vectors E being trained at once. This increase in speed quickly caps, and the speed is approximately the same between Card (P1) 8 and Card' (Prd"((1)) periment also illustrates the unbiasedness of the stochastic training: at Card' (Pi ) 20 there is no for the stochastic train-
20 fixed. Card' (Pi) = 20, there is no stochasticity in the train-
16. This ex-
stochasticity in the training, and the asymptotic performance is the ing (see appendix A.2). Right panel: PAVI-E scheme The case of Card Card'(Pi)) = 1 is pathological: the encoder "learns" to collect summary statistics across the learnt 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 (PP:)((PP))))))) = 20, that is say when there is no stochasticity in the training (black curve). We furthermore observe slight reduction of the in- ference bias as Card' (Pt) augments. This illustrates how the theoretical bias the PAVI-E scheme identified in appendix A.2 does not translate in significantly worse empirical results -albeit in the pathological case Card (P:) = 1. Interestingly, the training speed is approximately constant train
of element, and not surprisingly-
across all values of Card" (P1). This feature is essential in the PAVI-E scheme: reduced version of our model -with a light memory footprint- and apply the obtained architecture 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 Carrd((P) augments.
23

---

Under review as a conference paper at ICLR 2023
Architecture
ELBO (10")
Optimization time (s)
CF
166.2 2.9)
3,100
ADAVI
-20.0 (+: 2.3)
3,000
UIVI
-26.3 (# 1.2)
500
PAVI-E (ours)
-15.2 (42.8)
470
PAVI-F (ours) -1.1 (+1.1)
300
Table I: Analytical performance Gaussian mixture HBM PAVI shows superior perfor- mance. PAVI converges fraction of CF's optimization time the second best performing archi- tecture.
B.3
ADDITIONAL COMPARATIVE EXPERIMENTS
B.3.1 GAUSSIAN MIXTURE MODEL
this experiment we test out various baselines over a challenging model: random effects.
Gaussian mixture with
D. Card(P1), Card = 2, 20, Vng=l mo I.. CP) Vng=1., Card(Pa)
10
Xm 01,n1, Thi ~ Mixture( ([V(0),,.) .2),..,, . knn))
Vn = Card(( Thi Dirichlet(1 Iz)
(B.22)
(nr=(.Card(P),
8,,,, 1020 ~ -N VI=I.1
VI ===1. 1..L 82 ~N(0, o2)
where Mixture(Pn,, D2J, denotes mixture between the distributions [Di,.... Dr. with mixture weights w. Results are visible table 1, where PAVI displays the best asymptotic ELBO, well as the shortest optimization time. On that note, underline that stochastic training a mixture distribution challenging, as -due to the sub-sampling of points- only a fraction of the mixture components could be expressed given step, requiring the architecture dynamically cluster the data points across time.
B.3.2 HIERARCHICAL VARIANCE MODEL
In this experiment we test out our architectures over a non-canonical model, in which parent RVs play the role of variance for the distribution of their children. Our goal is in particular evaluate a potential empirical bias for the PAVI-E scheme -as studied in appendix A.2. We will refer the following HBM as the Hierarchical Variance model:
D, Card(Pi), Card(Pn) = 15, 15
Voy -1 Card(Pi) log X(mi.no) 10,m. No 0 0 Vmg=1.. Card(Po)
(B.23)
Vn, = 1.. Card(Pi) log log 12,0 N 6o.u)
log log02,0 N)
In this model, the encoder used the PAVI-E scheme- has collect non non-trivial summary statistics: empirical 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 eq. (A.13). The performance of
PAVI-E remains nonetheless competitive. This illustrates how PAVI-E's theoretical bias -introduced in the stochastic training- does not result in significantly worse inference compared to state-of-the- art architectures. PAVI also displays the best ELBO, in conjunction with UIVI, but in times shorter optimization time.
B.3.3 SMALL DIMENSION VERSION OF OUR NEUROIMAGING MODEL
In this experiment we validate the performance of our architecture on synthetic data generated using small dimension version of the model presented in eq. (B.24). To allow the of the comparative
24

---

Under review as a conference paper at ICLR 2023
Architecture
ELBO (109)
Optimization time (s)
CF
-11.2 2.9)
500
UIVI
-6.7 (#: 0.8)
100
PAVI-E (ours)
-6.7 (=)0))
90
PAVI-F (ours)
6.7 (++0.8)
20
Table 2: Analytical performance over a Hierarchical Variance HBM In this setting with hard- to-estimate summary statisties, PAVI-E doesn't show any empirical bias. 1: PAVI-E suffers from numerical instability on this example, with runs degenerating into NaN results.
Architecture
ELBO (10*)
Optimization time (s)
CF
139.7 (# 21)
3,200
ADAVI
UIVI
PAVI-F (ours)
8.2 (+ 0.1)
1,600
Table 3: Analytical performance over our Neuroimaging HBM This HBM is not pyramidal and consequently cannot be processed by the ADAVI architecture. Due to the large parameter space (approximately 2k parameters) and the complex density involved, despite intensive efforts we did not manage to obtain numerically stable UIVI optimization. PAVI shows superior compared to CF,, in half its optimization time.
baselines, we fixate S, table 3: PAVI can deal with HBMs unavailable to the ADAVI architecture (Rouillard & Wassermann, 2022), and so so with a performance superior to the CF baseline (Ambrogioni et al., 2021b). Despite intensive efforts did not manage obtain numerically stable UIVI optimization. This is likely due to the complexity 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 or PAVI exploit the parametric form the prior p to help with inference. PAVI's plate amortization likely facilitates the inference, resulting in reduced optimization error (Bottou & Bousquet, 2007).
12 -,3,
Results are visible in
B.4 EXPERIMENTAL DETAILS- ANALYTICAL EXAMPLES
All experiments were performed in Python, using the Tensorflow Probability library (Dillon al., 2017). Throughout this section we refer to Masked Autoregressive Flows (Papamakarios et al., 2018) as MAF. All experiments performed using the Adam optimizer (Kingma & Ba, 2015). At train- ing, the ELBO was estimated using a Monte Carlo procedure with samples. All architectures were evaluated over fixed of 20 samples X, with seeds per sample. Non-sample-amorre tectures were trained and evaluated on each of those points. Sample amortized architectures were trained over a dataset of 20, 000 samples separate from the 20 validation samples, then evaluated over the 20 validation samples.
archi-
B.4.1 PLATE AMORTIZATION AND CONVERGENCE SPEED (3.1)
All 3 architectures (baseline, PAVI-F, PAVI-E) used:
for the flows F,, a MAF with (32,2]] 32] hidden units;
as encoding size, 128
For the encoder f in the PAVI-E scheme, we used a multi-heat architecture with 4 heads of 32 units each, 2 ISAB blocks with 64 inducing points.
25

---

Under review as a conference paper at ICLR 2023
B.4.2
IMPACT OF ENCODING SIZE (3.2)
All architectures used:
for the flows F,, MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.
as encoding size, value varying from 2 to 16
B.4.3 SCALING WITH PLATE CARDINALITIES (3.3)
ADAVI (Rouillard Wassermann, 2022) used:
for the flows Fú, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.
for the encoder, an encoding size of 8 with multi-head architecture with 2 heads of units each, 2 ISAB blocks with 32 inducing points.
Cascading Flows (Ambrogioni et al., 2021b) we used:
a mean-field distribution over the auxiliary variables r
auxiliary size, fixed value of 8
flows, Highway Flows designed by the Cascading Flows authors
PAVI-F we used:
for the flows Fi, a MAF with 322,32] hidden units, after an affine block with triangular scaling matrix.
encoding size of 8
PAVI-E we used:
for the flows F. a MAF with [32,32] hidden units, after an affine block with triangular scaling matrix.
for the encoder, an encoding size of 16 with muat architecture with 2 heads of 8 units each, ISAB blocks with inducing points.
UIVI used as Card
20 - 200:
base distribution, standard Gaussian with dimensionality 6 - 42 - 402
as transform h, an affine transform with diagonal scale
as embedding distribution, a standard Gaussian with dimensionality - + 9
as transform weights regressor, MLP with hidden units [32,32] - [64, 64] - [128, 128] sample uncorrelated samples e, a HMC with 5 burn-in steps and ssamples
Adam optimizer with exponential learning rate decay, starting at 1e 2, x0.9 every 300
steps
B.4.4 GAUSSIAN MIXTURE (B.3.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 head architecture with 2 heads of 8 units each, ISAB blocks with inducing points.
Cascading Flows (Ambrogioni et al., 2021b) used:
26

---

Under review as a conference paper at ICLR 2023
mean-field distribution over the auxiliary variables r
as auxiliary size, a fixed value of 16
as flows, Highway Flows as designed by the Cascading Flows authors
PAVI-F we used:
for the flows F,, MAF with [32] hidden units, after affine block with diagonal scaling matrix.
an encoding size of Card' Card
PAVI-E used:
for the flows Fi, a MAF with [128, 128] hidden units, after affine block with traingular scaling matrix.
for the encoder, an encoding size of 128 with a multi- architecture with 4 heads of 32 units each, 2 ISAB blocks with 128 inducing points.
Card' Card"(Pi) =
UIVI we used:
base distribution, standard Gaussian with dimensionality 82
transform h, an 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 e, a HMC run with 5 burn-in steps and samples
Adam optimizer with exponential learning rate decay, starting at x0.9 every 300
steps
B.4.5 HIERARCHICAL VARIANCES (B.3.2)
Cascading Flows (Ambrogioni et al., 2021b) used:
mean-field distribution over the auxiliary variables r
as auxiliary size, a 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, 32] hidden units, after an affine block with diagonal scaling matrix.
an encoding size of
Card'
Card
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, encoding size of 128 with a multi-head architecture with 4 heads of 32 units each, ISAB blocks with 128 inducing points.
Card"(Pt)
Card' (Po)
27

---

Under review as a conference paper at ICLR 2023
UIVI 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 3
transform weights regressor, MLP with hidden units [64, 64]
sample uncorrelated samples HMC run with 5 burn-in steps and 5 samples
Adam optimizer with exponential learning rate decay, starting at 1e x0.9 every 300
steps
B.4.6 SMALL NEUROIMAGING EXAMPLE (B.3.3)
Cascading Flows (Ambrogioni et al., 2021b) we used:
mean-field distribution over the auxiliary variables r
auxiliary size, a fixed value of
flows, Highway Flows designed by the Cascading Flows authors
PAVI-F we used:
for the flows Fi, a MAF with [32, 32] hidden units, after an affine block with diagonal scaling matrix.
combination of encoding sizes of 32 and 8
UIVI we used:
as base distribution, a standard Gaussian with dimensionality 802
as transform h, an affine transform with diagonal scale
embedding distribution, standard Gaussian with dimensionality 16
transform weights regressor, MLP with hidden units [128, 256, 512, 1024] to sample uncorrelated samples a HMC run with 5 burn-in steps and samples
an Adam optimizer with exponential learning rate decay, starting at le - 3, x0.9 every 300 steps
Optimization systematically degenerated into NaN results after around 200 optimization steps.
B.5 DETAILS ABOUT OUR NEUROIMAGING EXPERIMENT (3.4)
B.5.1 DATA DESCRIPTION
In this experiment we use data from the Human Connectome Project (HCP) dataset ( (Van Essen et al., 2012). We randomly select a cohort S 1, 000 subjects from this dataset, each subject being 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):
1. removing high variance confounds
2. detrending the data
3. band-filtering the data (0.01 to 0.1 Hz), with repetition time of 0.74 seconds
4. spatially smoothing the data with a 4mm Full-Width at Half Maximum
For every subject, we extract the surface Blood Oxygenation Level Dependent (BOLD) signal of N 314 vertices corresponding this signal with the extracted signal of = 64 DiFuMo components: a dictionary of brain spatial maps allowing for effective fMRI dimensionality reduction (Dadi al., 2020). Specifically,
average Broca's area (Heim et al., 2009). We compare
28

---

Under review as a conference paper at ICLR 2023
compute the one-to-one Pearson's 's correlation coefficient of every vertex with every DiFuMo component. The resulting connectome, with S subjects, T sessions, N vertices and a connectivity signal with D dimensions, is of shape (8 x x N D). We project this data -corelation coefficients lying in 1-1; [- in an unbounded space using an inverse sigmoid function.
B.5.2 MODEL DESCRIPTION
We model inspired from the work of Kong al. (2019). We hypothesize that every vertex in Broca's 's area belongs to either one of L = 2 functional networks. This functional bi bi-partition would reflect the anatomical partition between pars opercularis and pars triangularis (Heim et al., 2009; Zhang et al., 2020).
Each network a pattern of connectivity with the brain cortex, represented as a the correlation of the BOLD signal with the signal from the D = 64 DiFuMo components. We define L 2 such functional networks at the population level, that correspond to some "average" across the cohort of subjects. Every subject has an individual connectivity, and therefore individual L 2 networks, that are considered as a Gaussian perturbation of the population networks, with variance c. The connectivity of given subject also evolves through time, giving rise session-specifie networks, that are a Gaussian perturbation of the subject networks with variance a. Finally, every vertex in Broca's ' area has its individual connectivity, and is perturbation of one network's connectivity the other's s. We model this last step as a Gaussian mixture distribution with variance k. We explicitly model the label label of given vertex, and we consider this label constant across sessions.
The resulting model can be described as:
S.T.N.D.L WWW.D,,1000,,.... 2,34, 64,
8 g -6,0
W=1.L: 14 Uniform(-44 ID. Ip)
W=I.L:: log e ~ Uniform(s x TD, + x Tn)
WEil:::: Mrel 4 ~ N(p,,e)
W=1.:: log Uniform(s x TD, x xTn)
(B.24)
Mtae (ut,,,
W=1.L: log Uniform(s x TD, x xT))
W=l..: probs Dirichlet(1 [z)
-= labela.n probs.... Categorical(probs..
Ka.. [st)e1.. label,, ((Maaabell.n.. Kllbbl....)
The model contains 4 plates: the network plate of full cardinality (that we did not exploit in our implementation). the subject plate of full cardinality S, the session plate of full cardinality T and the vertex plate of full cardinality N.
Our goal is to recover the posterior distribution of the networks μ -represented as networks in fig. 5- and the labels label -represented as the parcellation fig. 5- given the observed connectome described in appendix B.5.1.
B.5.3 PAVI IMPLEMENTATION
We used in this experiment the PAVI-F scheme, using:
for the RVs μt Ml,: Mi,ett
for the flows F,, MAF with [128, 128) hidden units, following an affine block with
diagonal scale
for the encoding size: 128
for the RVs 61, 01, Ki, probs labelss,
29

---

Under review as a conference paper at ICLR 2023
for the flows F, MAF with [8,8] hidden units, following affine block with diag-
onal scale
for the encoding size: 8
for the reduced model, we used S 30, T and N 332.
To allow for the optimization over the discrete labelg,n RV, we used the using a fixed temperature of 1.0 (Jang et al., 2017; Maddison et al., 2016).
trick,
C SUPPLEMENTAL DISCUSSION
C.I PLATE AMORTIZATION AS A GENERALIZATION OF SAMPLE AMORTIZATION
In section 2.2.2 we introduced plate amortization as the application of the generic concept of amor- tization to the granularity of plates. There is actually an even stronger connection between sample amortization and plate amortization.
A HBM p models the distribution of given observed RV X -jointly with the parameters O Dif- ferent samples Xo, X,... ... of the model p are i.idd draws from the distribution p(X). p can thus be considered as the model for "one sample". Consider, instead of "macro" model for the whole population of samples one could draw from p. The observed RV that macro model would be the infinite collection of samples drawn from the same distribution p(X). In that light, the ii.d. sampling of different values from p could be interpreted as a plate of the macro model. Thus, could consider sample amortization as instance of plate amortization for the "sample plate". Or equivalently: plate amortization can be seen 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 a different formalism for SVI, based around the concept of full HBM M versus reduced HBM Mredil sharing the same template T. This formalism helpful to set up GPU-accelerated stochastic VI (Dillon et al., 2017), as it entitles a fixed computation graph -with the cardinality of the reduced model M in which encodings are "plugged in" -either sliced from larger encoding arrays
the output of an encoder applied to data slice, section 2.2.1 12.3.2. Particularly, our formalism doesn't entitle a control flow over models and distributions, which can be hurtful in the context of compiled computation graphs such in Tensorflow (Abadi et al., 2015).
The reduced model formalism is also meaningful in the PAVI-E(sa), where we train and amortized variational posterior over n and obtain "for free" variational posterior for the full model M -see section 2.3.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 from the cardinality of the training, and is only simulated as scaling factor in the stochastic training -see section 2.3.1. We have the intuition that fruitful research directions could stem from this concept.
C.3 BENEFITING FROM STRUCTURE IN INFERENCE
Our contributions can be abstracted through the concept of plate amortization -see section 2.2.2. Plate amortization is particularly useful in the context of heavily parameterized density approxima- tors such as normalizing flows, but is not tied to it: plate-amortized Mean 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 a problem. This general concept could have ap- plications in other highly-structure, problem classes such graphs sequences (Wu et al., 2020; Salchinejad et al., 2018).
C.4 CONNECTION WITH META-LEARNING
In section 2.2.2 we introduced plate amortization: sharing the parameterization and learning across a model's plates. In this section we discuss the connection between plate amortization and meta- learning(Ravi & Beatson, 2019; lakovleva et al., 2020; Yao al., 2019).
30

---

Under review as a conference paper at ICLR 2023
Supervised learning can be seen as the mapping from given context set C = {(2, y)} } to a predictive function f (Bishop, 2006) such that f(z) = y. Meta-learning -or "learning to learn"- instead recovers this mapping C + in the general case. Once the meta-training is completed, a predictive function I conditioned by an unseen context C can be obtained in a single forward pass -without any training done on C. As instance of meta-learning. the Neural Process Family encodes the context C' via a deep set encoder (Garnelo et al., 2018; Dubois et al., 2020; Zaheer et al., 2018). The encoded context, along with the data point x are then used to condition an estimator for the density (3)22,C) C).
This framework is similar to the PAVI-E scheme, where a combination of a deep set encoder and a normalizing- saazzinggllow-baar flow-based density estimator output the posterior probability of a ground RV 0,n- This encoder-estimator pair is repeatedly used across plate. This is analogous to meta-learning to solve the inference problem across the different elements of a plate -such as the different subjects of a population study. In PAVI-E, the encoding E,, at lower hierarchy play 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 -ittrhhirrrrchy scenario;
meta-learning typically considers a set of i.i.d. tasks, whereas in PAVI the inference of different ground RVs fn,n are conditionally dependent through the hierarchical model p;
meta-learning is trained using the forward KL loss, that is say in the sample-amortized regime, maximizing the probability q(v) of samples from the underlying generative pro- cess. In contrast PAVI -though possible to train using the forward KI. loss- is trained via the reverse KL, needing to explicitly evaluate the density of the generative process.
We suspect there would be interesting applications for the PAVI architecture in hierarchical meta- learning scenarios.
C.5 CONDITIONAL DEPENDENCIES MODELLED IN THE VARIATIONAL FAMILY
The Mean-Field approximation was originally introduced in VI to facilitate computation, allowing dedicated optimization schemes (Blei et al., 2017). Nonetheless, this approximation assumes in- dependence between RVs in the posterior, and ultimately limits the expressivity of the variational family. To remove this approximation, modelling complex conditional dependencies in the varia- tional family is an open research subject (Ambrogioni et al., 202 b;; Webb al., 2018). In the PAVI design, we inherit our statistical dependency structure from the prior distribution p. as detailed in eq. (2). 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 modelling only those forward dependencies, the modelling of colliders can be an issue.
In the case of single plate -te 2-1evel case- Agrawal & Domke (2021) demonstrate that modelling only the forvard dependencies does not result in reduced expressivity compared to the modelling of the full dependencies. Yet this result does not hold in the n-level case, as Webb et al. (2018) show that faithful inversion features conditional dependencies in the posterior between ground RVs of the same template. We can dub those dependeneies as horizontal dependencies, across RVs in the same plate. Similar to Structured SVI (Hoffman & Blei, 2014) or Automatic SSVI (Ambrogioni et al., 2021a), the PAVI design therefore results in reduced expressivity when stacking multiple plates in the model This issue can be partially alleviated with the usage of a backward encoding scheme -going in the reverse direction compared to the prior's dependencies- as in the PAVI-E design (see section 2.3.2) or in Cascading Flows (Ambrogioni et al., 2021b).
In practice, though limiting our expressivity, horizontal dependencies are difficult to inject back into our architecture. Critically, in the PAVI design, the use of a common density estimator across the ground RVs of the same template (see section 2.2), and the stochastic training over batches of those RVs (see section 2.3.1) 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
31

---

Under review as a conference paper at ICLR 2023
could arbitrary conditional dependencies be modelled in the variational posterior in the context of stochastic training?
C.6 TOWARDS ESER-FRIENDLY VARIATIONAL INFERENCE
By re-purposing the concept of amortization the plate level, our goal is to propose clear compu- tation versus precision trade-offs in VI. Syyper-parameten such as the encoding size as illustrated in fig. 3 (right)- allow to clearly trade inference quality in exchanged for reduced memory foot- print. On the contrary, in classical VI, changing Q's parametric form -for instance switching from Gaussian to Student distributions- can have strong and complex impact both on number of weights and inference quality (Blei ct al., 2017). By allowing the usage of normalizing flows in very large cardinality regimes, our contribution aims at disentangling approximation power and computational feasibility. In particular, having access expressive density approximators for the posterior can help experimenters diversify the proposed HBMs, removing the need of properties such as conjugacy to obtain meaningful inference (Gelman et al., 2004). Combining clear hyper- arperpparameters and scalable yet universal density approximators, we tend towards a user-friendly methodology in the context of large population studies VI.
32


Reviews:
[{"review_id": "dVyBw_I03V", "reviewer": "Reviewer_1MMm", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents an approach for variational inference on large-scale hierarchical models which takes advantage of conditional independence in plate structures to factorise its variational approximations and better scale with problem dimensionality. The proposed framework is based on normalising flows taking shared parameters and data-encoding vectors. The resulting algorithm is equipped a stochastic optimisation scheme for training and is shown to converge faster than other algorithmic baselines in experiments.", "summary_of_the_review": "The paper's contributions are somewhat novel, but it lacks in experimental comparisons against the baselines in other challenging scenarios.", "strength_and_weaknesses": "### Strengths\n* The proposed method achieves significant (orders of magnitude) performance speed-ups when compared to baselines in experiments.\n* Although the idea of using factorised variational posteriors for hierarchical models is not new, this paper implements it in a more computationally efficient way, which is also amenable for per high-performance computing with hardware accelerators.\n* Diagrams help understand the concepts in the approach.\n\n### Weaknesses\n\n* Related work: There are other flexible variational inference approaches for hierarchical models which also factorise the variational posterior. For example, Tran et al. (2017) introduced hierarchical implicit models for this task, and other relatively simple approaches for random effects models (as in the example experiment) are also available (e.q., Dao et al., 2022).\n\n* Experiments: There was only one benchmarking model (Eq. 6) and variations thereof for the comparisons. The other experiment (Fig. 5) did not provide comparisons against baselines due to its complexity. Another representative set of experiments would help assessing the improvements over baselines in corner cases where other methods can be applied but PAVI should in principle perform better.\n\n#### References\nTran, D., Ranganath, R., & Blei, D. M. (2017). Hierarchical Implicit Models and Likelihood-Free Variational Inference. 31st Conference on Neural Information Processing Systems (NIPS).\n\nDao, V. H., Gunawan, D., Tran, M., Kohn, R., Hawkins, G. E., & Brown, S. D. (2022). Efficient Selection Between Hierarchical Cognitive Models: Cross-Validation With Variational Bayes. Psychological Methods, (April).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "EmlYZCl1I_", "reviewer": "Reviewer_d76j", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors consider approximate inference in Bayesian networks, specifically in the subset of Bayesian networks that can be represented as plate diagrams, which are common for large population studies, the population is composed of subgroups and each group is composed of many individuals. Given global the population parameters, each of the group parameters are i.i.d. and given group parameters, the individuals are i.i.d., a common hierarchical structure.\n\nAssuming a fully specified Bayesian generative model of the data, and only leaf nodes are observed, the authors propose a variational inference method to learn the full approximate posterior distribution to all of the unobserved nodes. The method combines normalising flows with the generative model priors to approximate the posterior. The method is applied to illustrative toy datasets and shows results to be expected and a high dimensional neurological time series dataset (but results are not presented in the main paper).\n", "summary_of_the_review": "- clarity of the main paper is very poor, verbose and possibly non-standard terminology for what appears to be standard Bayesian networks.\n- the technical contribution, parameterization of normalizing flows and mini batch training, seem novel but incremental.\n- the experiments are mostly toy, one non toy is only briefly considered.\n\nUPDATE: I have updated my score from 3 to 5 after reading through the responses.", "strength_and_weaknesses": "Overall, I struggled with the paper and I have given a low confidence score due to possible misunderstandings.\n\n**Strengths**\n- the Supplementary material seems very comprehensive and answered many of the questions I had after reading the main paper.\n- the method seems well designed and thought out\n- the encoding size experiment provided good confirmation that the degrees of freedom must match the sufficient statistics.\n\n**Weaknesses**\n\nI feel this paper is very hard to recommend for acceptance due to the writing and the somewhat basic numerical evaluation.\n\nI felt the writing and the structure of the paper was very poor, it took me a great deal of effort and multiple rounds of re-reading and re-evaluating my own understanding of the paper before finally appreciating the paper and its contributions (and I have experience in graphical models).\n  - I would make section 2.1 into it's own section that really makes the mathematical problem clear. E.g. a plate diagram, the generative model in  equation (1), a real world example with cardinalities and (non-linear) link functions and latent variables. Describe the practicalities of inferring parameters in such a model and why it is impractical scaling with cardinality.\n- I don't know if this is standard or not but I was personally thrown off by the terminology \"RV template $\\theta_i$\" and \"ground RVs $\\theta_{i,n}$\" and \"grounding a template with cardinalities\". Is is possible to reduce notation? E.g. $\\mathcal{T}$ is introduced but never used, plates $\\mathcal{P}$ is simply indicated by $i$, Card($\\mathcal{P}$) is just $N_i^{full}$. Perhaps $N^{full}$ and $N^{redu}$ could be replaced by $N$ and $\\tilde{N}$, i.e. use tilde \"~\" to denote reduced version of various quantities?\n- what does $\\pi(\\theta_{i,n})$ explicitly represent in equations (1) and (2)? I believe they should be precise, known, values if they are conditioning a distribution, but this is seems vaguely defined in the paper. In Ambrogioni 2021b they say \"$\\theta_j$ is the array of values of the parents of the parents of the $j^{th}$ variable, and $\\pi_j(theta_i)$ and $\\pi_j()$ is the link function\".\n- what is the prior for an RV, $p(\\theta_{i,n}| \\pi(\\theta_{i,n}))$? My intuition is that prior is the should be the generative model up to $\\theta_{i, n}$ with all of the parent variables marginalized out? Or is it the generative model term with conditioned on given parent values? In the latter case, having such a distribution as the input to a normalizing flow doesn't feel like it would lead to a good posterior approximation, parent values change, prior changes, normalizing flow function is constant but the input changes?\n- Section 2.3: please give dimensions of all **X**, and $\\theta$ and I assume one row/col of  **X** is a subgroup indicator variable?\n- (nitpick) I am not entirely sure I agree with the use of the word \"amortization\" in the PAVI-F, the fact that $\\phi_i$ is shared across all nodes in a layer does not makle inference faster at test time for new nodes? I still have to learn $E_{new}$ from scratch? (aAlthough amortization means \"spreading out\", in this community, I feel it refers to predicting an approximate posterior in one shot given observed variables.\n- the experiments are toy examples and the only non toy example is rather brief with no quantitative results in the main paper\n- Lack of Technical Novelty\n  - randomly subsampling nodes in the DAG seems like a very simple straightforward idea and I think the unbiasedness is significant and shows that this isn't \"just a hack\" but a real valid method, however it is rather incremental.\n  - claiming to have reduced memory footprint by \"using a low dimensional encoding\" seems ambitious claim of novelty. While it may be novel, forcing the method to use a low dimensional parameterization of a normalizing flow seems incremental.\n\n- I was surprised that the discussion about \"faithful posteriors\" (Webb et. el. 2019) was in the SM, to me this felt like very closely related work and provides a method for building approximate posterior Bayesian networks,", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "E5JUFb_BNU", "reviewer": "Reviewer_pbEf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes Plate-Amortized Variational Inference (PAVI), an inference framework for hierarchical Bayesian models with amortization over the population (dataset) level implemented with normalizing flows with shared parameters. The main idea is to posit an expressive variational family parameterized with normalizing flows, and yet reduce the number of parameters to be trained by sharing part of the variational parameters across populations. Another important concept is stochastic training, where for each step, instead of training all the variables present in the graphical model, subsample a small number of variables and compute stochastic gradient with them. Thanks to amortization, PAVI could scale to large-scale population studies involving high-dimensional parameters.", "summary_of_the_review": "Overall, the proposed method is an interesting contribution, but the paper could be strengthened by adding more discussions and hopefully more experimental results/interpretations.", "strength_and_weaknesses": "Strengths\n- The paper is well-written and easy to follow.\n- The description of amortization using the language of plate diagram is interesting.\n- The neuroimaging experiment seems interesting.\n\nWeaknesses\n- In my opinion, the notion of plate amortization is not entirely new, so I think the paper can be strengthened by discussing relevant existing works.\n- The experiments include a synthetic Gaussian model and neuroimaging experiments. The synthetic experiment looks promising, but it is of a small scale that does not particularly highlight the benefit of the proposed method. For the neuroimaging experiment, although it seems promising that only the proposed method was manageable to train, it is not clear whether the prediction given by the model is reasonable (at least for the non-experts like me).\n\nDetailed comments\n- The idea of plate amortization is indeed interesting. I guess the proposed framework virtually extends to hierarchical Bayesian models with an arbitrary level of hierarchies, but the main application would be the amortization across populations (datasets). Regarding amortization over datasets, there have been similar ideas (although motivated differently), especially in the context of meta-learning. Neural processes families, for instance, introduce an amortized inference network that can readily generalize to a new task (corresponding to the population) which is parameterized by the dataset encoding networks (deep sets or transformers); this would roughly correspond to PAVI-F, except that PAVI-F uses normalizing flows for the parameterization. Amortized Bayesian meta-learning (Ravi and Beatson, 2019) is also worth discussing in a similar sense. Depending on the application, neural processes can be considered as a baseline to be compared to the proposed method, since it can also scale to large population study due to its use of the amortized inference network.\n- The experiment section could be enhanced. I understand that the neuroimaging experiment can be compelling, but in the current form, I'm not entirely sure how well the proposed method is doing; the only result I could see from the paper is Figure 5, which I fail to interpret. It would be great to see the numbers, for instance, metrics could tell us how accurate PAVI is, or at least how well the predictions given by PAVI match the domain knowledge given by experts. Also, if the baselines are not trainable due to their limited scalability, one could consider comparing them on a smaller subsets of manageable sizes to see if PAVI would do better than baselines for real-world tasks.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-H3o6H47eU", "reviewer": "Reviewer_oFPk", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work proposes an amortized VI scheme that leverages the hierarchical structure of a Bayesian model, using the latter's plate representation. Plates are used to instantiate subsets of the variables of interest, and share weights in the encoding of sets of variables used on a conditional flow architecture. The result is a plate-amortized proposal distributionI. This idea seeks to enable VI in a large scale setup. The methodology is validated empirically against the state-of-the-art, with one application with real data in neuroimaging context.\n", "summary_of_the_review": "Given 1) the limited novelty, 2) short experimental campaign, 3) and possible issues with theory, I recommend rejection.", "strength_and_weaknesses": "## Strengths:\n* The methodology is generic and intuitive\n* Potential for large scale applications\n* The idea of embedding structure into the approximate inference scheme is potentially fruitful line of further research\n* The authors provide detailed discussion in the appendix, including some limitations regarding the approximation gap.\n* The empirical validation in a realistic large scale scenario.\n\n## Weakness:\n* Heavy notation makes the text hard to follow;\n* Limited novelty (e.g., the work is incremental when compared with ADAVI). It seems that the major difference between PAVI and ADAVI is the stochastic updates of variables;\n* I believe the gradient estimates for PAVI are biased;\n* Limited experimental campaing, limited to a small set of methods and models. Validation against multiple methods in smaller dataset could be beneficial.\n\n## Questions\n1. In the demonstration on Appendix A.2 it indicates that PAVI-F is unbiased by showing in A.12 that the ELBO term is equal the term with all the variables, however Section A.4 implies that there is a gap between vanilla VI, PAVI-F and PAVI-E: Gvanilla VI ≤ GPAVI-F ≤ GPAVI-E . This seems unintuitive and indicative that indeed PAVI-F might be biased, given that the term that differ in those approximation schemes is indeed the ELBO term. Are both of those results consistent?\n\n2. How does the proposed methodology would compared with MCMC/HMC? It could be beneficial given that a new approximation scheme is being proposed to compare it as well sampling methods. \n\n3. When it comes to the expressivity of the variational approximation. How does the predictive performance of the method compare with others variational inference proposals (empirically)? For example another family of expressive VI approximations are found in the works such as UIVI [1], SIVI [2] or VCD [3], which does not use flows and yet are generic and could be a better test bench and give the reader more confidence about the proposed method.\n\n4. In Eq. A.11, when two variables share the same descendants, hence leading to correlated E_{i,n} in the PAVI-E scheme, wouldn’t in this case lead also to dependency between the indicator random variables and log q_{i,n} (given the dependency of log q on the encoding)?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject