Amortized Variational Inference for Simple Hierarchical Models
Authors: Abhinav Agrawal
Justin Domke
Venue: NeurIPS 2021 Poster (2021)
Keywords: Variational Inference
Black Box Variational Inference
Amortized Variational Inference
Scalable Variational Inference
Hierarchical Models
Structured Variational Inference
Gaussian VI

Abstract: It is difficult to use subsampling with variational inference in hierarchical models since the number of local latent variables scales with the dataset. Thus, inference in hierarchical models remains a challenge at a large scale. It is helpful to use a variational family with a structure matching the posterior, but optimization is still slow due to the huge number of local distributions. Instead, this paper suggests an amortized approach where shared parameters simultaneously represent all local distributions. This approach is similarly accurate as using a given joint distribution (e.g., a full-rank Gaussian) but is feasible on datasets that are several orders of magnitude larger. It is also dramatically faster than using a structured variational distribution.

Amortized Variational Inference in Simple Hierarchical Models
Abhinav Agrawal
Justin Domke
College of Information and Computer Science Univeristy Of Massachusetts Amherst aagrawal@cs, umass. edu
College of Information and Computer Science Univeristy Of Massachusetts Amherst domke@cs. umass. edu
Abstract
It is difficult to use subsampling with variational inference in hierarchical models since the number of local latent variables scales with the dataset. Thus, inference in hierarchical models remains a challenge at large scale. It is helpful to use a variational family with structure matching the posterior, but optimization is still slow due to the huge number of local distributions. Instead, this paper suggests an amortized approach where shared parameters simultancously represent all local distributions. This approach is similarly accurate as using given joint distribution (e.g., a full-rank Gaussian) but is feasible on datasets that are several orders magnitude larger. It is also dramatically faster than using structured variational distribution.
Introduction
Hierarchical Bayesian models are general framework where parameters of "groups" drawn from some shared distribution, and then observed data is drawn from distribution specified by each group's parameters. After data observed, the inference problem is to infer both the parameters for each group and the shared parameters. These models have proven useful in various domains [13] including hierarchical regression amd classification [12), topic models [4, 22, 3], polling [11, 24), epidemiology [23], ecology (8), psychology [37], matrix-factorizzation [35], and collaborative filtering [26, 33].
proven technique for scaling variational inference (VT) to large datasets is subsampling. The idea is that if the target model has the form p(z, y) p(z) II. p()l=)) then unbiased gradient can be estimated while only evaluating p(2) and p(ya 2) at a few [29, 16, 20, 31, 30, 36, 15].
This paper addresses hierarchical models of the form p(0, y) p(0) II. p(zz, 31 (0), where only observed. There are two challenges. First, the number of local latent variables z; increases with the dataset, meaning the posterior distribution increases in dimensionality. Second, there is often dependence between 2; and which must captured to get strong results [15, 18].
The aim of this paper is to develop a black-box variational inference scheme that can scale to large hierarchical models without losing benefits of a joint approximation. Our solution takes three steps. First, in the true posterior, the different latent variables z; are conditionally independent given 0 which suggests using a variational family of the same form. We confirm this intuition by showing that for any joint variational family q(8, 2), one can define a corresponding "branch" family q(0) II. (2x1(0) such that inference will be equally accurate (theorem 2). We call inference using such family the "branch" approach.
Second, observe that if using the branch approach, the optimal local variational parameters be computed only from 0 and local data (eq. (12)). Thus, we propose to amortize the computation of the local variational parameters by learning network approximately solve that optimization. We
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

show that when the target distribution is symmetric over latent variables, this will be as accurate as the original joint family, assuming sufficiently capable amortization network (claim 5).
Third, we note that many real hierarchical models, there are many i.i.d. data generated from each local latent variable. This presents challenge for learning amortization network, since the full network should deal with different numbers of data points and naturally reflect the symmetry between the imputs (that is, without having relearn the symmetry.) We propose an approach where a preliminary "feature" network processes each datum, after which they are combined with a pooling operation which forms the input for standard network (section 6). This is closely related to the "deep sets" [39] strategy for permutation invariance.
We validate these methods preference model for the Moviel dataset with 162K users who make 25M ratings of different movies. At small scale (2.5K ratings), show similar accuracy using dense joint Gaussian, branch distribution, or our amortized approach. At moderate scale (180K ratings). joint inference intractable. Branch distributions gives a meaningful answer, and the amortized approach is comparable or better. At large scale (18M ratings) the amortized approach is thousands of nats better on test-likelihoods even after branch distributions were trained for almost times as long as the amortized approach took to converge (fig. 6).
synthetic model where exact inference is possible, and
2 Hierarchical Branched Distributions
86
9211
SE mh
2E (L.....)
Figure I: The graphical model for the HBDs. On the left, we have plate notation for the generic HBD from eq. (3). Note, we can have an edge from 9 to Yıy (we skip it for clarity.) On the right, we have an example model with N 3.
We focus two-level hierarchical distributions. A generic model of this type is given by
N
= p()) II p(a, (0P =1,
where 0 and [ are latent variables, y { are observations, and x {x; M are covariates. As the visual representations of these models resemble branches, we refer them as hierarchical branch distributions (HBDs).
Symmetric. We call an HBD symmetric if the conditionals are symmetric, i.e., if 2, 2j. = Zj. and y = Vj, implies that
10= p(2; [0), and p(u: 0 21, ..) P(y, 0, 25,
Locally i.i.d. Often local observations yi (and z,)) are collection of conditionally i.i.d observations. Then, an HBD takes the form of
N
p(0. 2, 2, = p(0) II p(a () II 10, x,,
(3)
where y: {11 }= and x; ) are collections of conditionally i.i.d observations and covariates; 2 1 is the number of observations for branch

---

No local covariates, a simplified form of
Some applications do not involve the covariates 24. In such cases, HBDs have
p(0, 2) p(0) II 1(2
(4)
In this paper, we will be using eq. (1) and eq. (3) to refer HBDs results extend easily to case where there local covariates. (For instance, in section 5, we amortize using (Tz, Vi) as inputs. When there are no covariates, we can amortize with just yz-)
2.1 Related Work
Bayesian inference in hierarchical models is an old problem. The most common solutions are Markov chain Monte Carlo (MCMC) and VI. A key advantage VI that gradients can sometimes be estimated using only a subsample of data [29]. Hoffman et al. [16] observe that inference in hierarchical models still slow at large scale, since the number of parameters scales with the dataset. Instead, they assume that 0 and 2; from eq. (1) are in conjugate exponential families, and observe that for a mean-field variational distribution ((0) II. ((2)), ). the optimal ((z;) can be calculated in closed form for fixed (0). This is highly scalable, though it is limited to factorized approximations and requires a conditionally conjugate target model.
A structured variational approximation like g(0) II, (zz;00) can be used which reflects the dependence of 2, on [15, 34, 2, 18]. However, this still has scalability problems in general since the number of parameters grows in the size of the data (section 7). To the best of our knowledge, the only approach that avoids this is the framework of structured stochastic VI [15, 18]. which assumes the target is conditionally conjugate, and that for fixed 0 an optimal "local" distribution q(=;) ()) can be calculated from local data. Hoffman and Blei [15] address matrix factorization models and latent Dirichlet allocation, using Gibbs sampling to compute the local distributions. Johnson et al. [18] use amortization for conjugate models but do not consider the setting where local observations are a collection of i.i.d observations. Our approach is not strictly instance of either of these frameworks, as we do not assume conjugacy or that amortization can exactly recover optimal local distributions [15, Eq. 7J. Still the spirit is the same, and our approach should be seen part of this line of research.
Amortized variational approximations have been used to learn models with local variables [20, 10, 17, 5]. A particularly related instance of model learning is the Neural Statistician approach [10]. where "statistics network" learns representations of closely related datasets-the construction of this network is similar to our "pooling network" in section 6. However, in the Neural Statistician model there are no global variables f, and is not obvious how to generalize their approach for HBDs. In contrast, our "pooling network" results from the analysis in section 5, reducing the architectural space when 0 present while retaining accuracy. Moreover, learning a model is strictly different from our "black-box" setting where we want to approximate the posterior of given model, and the inference only has access to log p or Vo. log (or their parts) [30, 21, 2, 1].
3 Joint approximations for HBD
When dealing with HBDs, a non-structured distribution does not scale. To see this, consider the naive VI Drrecttve--videenee Lower Bound (ELBO, L). Let goiint be a joint variational approximation over O and z; we use sans-serif font for random variables. Then,
p(0, z, y/ log
C (git [p) (0g0tt((022)
(5)
4 toint (0,z)
where á are variational parameters. Usually the above expectation is not tractable and one uses a Monte-Carlo estimator. A single sample estimator is given by
C = log p(0. 2()x)
(6)
0 go (0,z)
where (0,z) ~ gggint and C is unbiased. Since guoint need not factorize, even a single estimate requires sampling the latent variables at each step. This problematic when there are a large number of local latent variables. One encounters the same scaling problem when taking the gradient of the ELBOO We can estimate the gradient using any of the several available estimators [30, 20, 31, 32]; however, none of them scale if goiu does not factorize [16].
3

---

Branch approximations for HBD
It is easy to see that the posterior distribution of the HBD eq. (1) takes the form
N
p(0, =/y, p(0ly, x) II 31,
1=1
As such, it is natural to consider variational distributions that factorize the same way (see fig. 2.) In this section, we confirm this intuition-we start with any joint variational family which can have any dependence between 0 and 21: N are conditionally independent given 0. We show that inference using the branch family will be at least as accurate as using the joint family. We formalize the idea of branch distribution in next definition.
Then, we define a corresponding "branch" family where
Definition 1. Let goint be any variational family with parameters ø. We define B be corresponding branch family if, for all Φ, there exists (v, {we ) such that,
N
Branch qu,1
(0, 2) qu (0) II qua (z/10) go
Joint (0)
II 9
Joint (2)
(7)
Given joint distribution, a branch distribution can always be defined by choosing and as the components of that influence q ggoint(()) and To go (2, respectively, and choosing (w)(6) and qu. (2 (0) correspond- ingly. However, the choice is not unique (for instance, the parameteri- zation can require transformations. different transformations can create different variants.)
Figure eq. (7).
Branch as in
The idea to use qu,w Branch is natural [15, 2]. However, one might question if the branch variational family is as good as the original d goiin; in theorem 2, we establish that this is indeed true.
Theorem 2.
Let be a HBD, and gloint (0,z) be joint approximation family parameterized by ø. Choose corresponding branch variational family g.ranch (0, z) as in definition 1. Then,
min KL (g 8,00
Branch [[p) < min K L (g (goint((p)) [lp).
We stress that g.rancch is new variational family derived from but not identical to q gointt Theorem 2 implies that we can optimize a branched variational family ge,andh without compromising the quality of approximation (see appendix D for proof.) In the following corollary, apply theorem 2 to joint Gaussian to show that using a branch Gaussian will be equally accurate.
Corrolary 3. Let p be a HBD, and let gimint (0, 2) N ((0, 2), Σ) be a joint Gaussian approxima- tion (with Φ = (μ, E)). Choose variational family
N Branch (0, N Eo) II ( + A,0, E,) Lu
with = (po, Σo) and wi = (m, Er, A.). Then, min KL (g.ranch Ip) S min K L (g [[p).
In the above corollary, the structured family Branch is chosen such that it can represent any branched Gaussian distribution. Notice, the mean of the conditional distribution is an affine function of This affine relationship appears naturally when you factorize the joint Gaussian over (0, z). For more details see appendix
4.1 Subsampling in branch distributions
In this section, we show that if p is an HBD and garanch is as in definition 1, we can estimate ELBO using local observations and scale better. Consider the ELBO
log p (0, z, y/ xz)
(g
Branch (pp)
Branct dv.as
(0.z)
Branch (0, z)

---

JointELBO @., y.
BranchELBO 11, w.
0, 4, /uint (8,2)
0 ~~ 4.(
E log pl0 2. y02) go
2, ~ 9w, (2: for [1, J.
roint
E < log p(0)
log (as, Me 0, x.))
(0)
Dwo (2/
(a) Estimation with goint as in eq. (6)
(b) Estimation with Te.u
Branch eq. (9)
SubSampledBraacchELB00: v.
Amort: waa.riiesiuuuapleeeeeaaalllLLLL 0 ~ q.(
0 ~ (.(
B ~ Minibatch(B)
Minibatch(B)
Pwn (z, (8) for B
ut net,, (3r 91 for E B
2 log
9u (8)
p
KEB
log
p (2 21)
Ow (2, ()
2, qu, (z,
log
go
(())
B
Σ
log p (2 v1)
Aws (2)
(c) Estimation with Qu,w Branch in eq. (10)
(d) Estimation with ge,u for as in eq. (2)
Figure 3: Pscudo codes for ELBO estimation with different variational methods;
{wo Ji=1
) (xr is defined similar to y.) (a) Estimates ELBO for joint approximation; (b) to (d) estimate ELBO for branch approximations: (c, d) use subsampling to estimate ELBO; (d) uses amortized conditionals; (a) to (c) work for any HBD, and (d) assumes is symmetric HBD in eq. (2). For models where ni > 1, we use the net, as in fig. 4. Minibatch
y = (y: J and y/
some distribution over the set of possible minibatches and B denotes the number of samples in minibatch B.
log (0)
N
log p (Zz, Mi 0,
E
(e) 4u (0)
(B) Ow, (2:10)
Guu; (zi 0)
Without assuming special structure (e.g. conjugacy) the above expectations will not be available in
closed form. To estimate the ELBO, let (0, (z, (m)),m)) ~ g Then, an unbiased estimator is
p (0)
N
log p Da 0.
Č log
Σ
(0)
qw, (z)
Unlike the joint estimator of eq. (6), one can subsample the terms in eq. (9) to create new unbiased estimator. Let B be randomly selected minibatch of indices from (1.2....,NV N). Then,
E log (00)
N
P (z;, Yi 0 x;:)
log
(10)
Iu (0)
B
Jw,
(z)
EEB
is another unbiased estimator of ELBO. In figs. 3b and 3c, we present the complete pseudocodes for ELBO estimation with and without subsampling in branch distributions.
Unsurprisingly. the same summation structure appears for gradients estimators of branch ELBO, allowing for efficient gradient estimation. With subsampled evaluation and training, branch distri- butions are immensely computationally efficient-in our experiments, we scale to models with 3 times more latent variables by switching branch approximations (see fig. 6).
While branch distributions are immensely more scalable than joint approximations, the number of parameters still scales as (N). In the next section, demonstrate that for symmetric HBDs, can share parameters for the local conditionals (amortize) to allow further scalability.
5 Amortized branch approximations
In this section, we discuss how one can amortize the local conditionals of a branch approximations when the target HBD is symmetric (see eq. (2).) first formally introduce the amortized branch distributions in the next definition and then justify the amortization for symmetric HBD.
5

---

Definition 4. Let 9 g be a joint approximation and let qu be as in definition 1. Suppose (T;, Vi) is some parameterized map (with parameters u) from local observations (2,, Mi) to space of ue Then,
Amort (0,z) (0)
II Qee (2.,) (2.,0))
(11)
a corresponding amortized branch distribution.
The idea to amortize is natural once you examine the optimization for symmetric HBDs. Consider the optimization for objective eq. (8).
P (0)
N
p (Zz, Mi x44)
max & (ge,wnc ([p) = max 0,80
log
E
E
log
(v(0)
2n (0)
T. Qw, (7)
Qeu; (z:
log (0)
N
(Zi, 0,
max
EI gu(
Σ
max E
E
log
4o (0)
ws
(u(0)
(2:00)
Tw; (z)
E1
The crucial observation in the above equation is that for any given v, the optimal solution of inner optimization depends only on local data points (æ,, ye), i.e.
p (Zi, U4 06 x+)
_ = argmax E
E at; (8( Qi, (2)0)
log
(12)
qw, (z+ 0)
Now, notice that if p and f0,10 grranch have symmetric conditionals, then, for each i, solve the same optimization over Wi, just with different parameters and d;. Thus, one could replace the optimization over wy with optimization over parameterized function from (æi, yt) the space of wy- Formally, when the network net,, is sufficiently capable, we make the following claim.
Claim 5. Let p symmetric HBD and let gg be some joint approximation. Let T be as in definition 1. Suppose that for all v, there exists u, such that,
log p (Zi U/ mg/00,0 (0, qw: (z:
net (xt 91) argmax E
E
(13)
wy (e) qw (2.
Then,
min KL (g
Amort
(p) S min KL (9% 1oint (pp)
(14)
Note, we only amortize the conditional distribution qw. (2, ()) and leave zu (0) unchanged. In practice, of course, we do not have perfect amortization functions. The quality the amortization depends on our ability to parameterize and optimize powerful neural network. In other words, we make the following approximation
P (z, y/ v,)0, log
net, (Z,, Vi) FO argmax
E
E (z
(15)
Iw, (z
W/
2-
In our experiments, we found the amortized approaches work well even with moderately sized networks. Due to parameter sharing, the amortized approaches converge much faster than other alternatives (especially, true for larger models; see fig. 6 and table 2).
6 Amortized branch approximations for i.i.d. observations
In the previous section, we discussed how we could amortize branch approximations for symmetric HBDs. However, in some applications, the construction of amortization network net, is not as straightforward. Consider the case when we have a varying number of local i.i.d observations for each local latent variable. In this section, we highlight the problem with naive amortization for locally ii.d HBDs, and present simple solution alleviate them.

---

Table All variational families used in our experiments. Σ denotes generic covariance matrix and x² denotes a diagonal covariance.
Gaussian Family
in eg. (5)
goranch us (7)
Altsort us in definition 4
Dense
N/ No) IT [M 1,0,5))
Mo Ma) II N + A,0.24)
(v.2)
(no, Eol. (M,, E.)
(no, Eo). (p,, Ldj (etu
N M0 Dal I Ma Ea)
Block Diagonal
N M0 (z 21)
po Fol T M E4) fml u. (No. Eol. (jt Ee) (a
5m1 (po Fol, ay (Ml Ee)
= (po: No Fa)
Diagonal
N 100,=1,,00 of)
N M(00po og)
(0 II N==lu
= (n.o")
(po,00), un (H,.of)
(po, oi). (Mi, of) cet.(x..
Mathematically, for locally i.i.d HBDs have yi {yi ) and oi [(ra))]=1 such that the conditional over factorizes as
Te II p (uig (ayg. 24.
P (M 21,
Now, if 2; and are directly input to the amortization network net.. the input size to the network would change for different (notice have n, observations for 7h local variable.) Another problem is that the optimal variational parameters are invariant to the order in which the i.i.d. observations are presented. For instance, consider two data points: (J;; y/) = [(ra1 :u.)..... (Fim. Vin.)) and (a', v/) = [[inn Vin, )......... (#,1. 3.)]] A naive amortization scheme will evaluate very different conditionals for these two data points because net,, (24, y:) and net. (24, v/) will be different.
To deal with both issues: variable length input and permutation invariance, we suggest learning "feature network" and "pooling function" based amortization network; this is reminiscent of "deep sets" [39] albeit here intended not just to enforce permutation invariance but also to deal with inputs of different sizes. Firstly, a feature network feat feat_net _nt takes each (Jaj- Mey) pair and returns a vector of features ej. Secondly, pooling function pool takes the collection {e, ) and achieves the two aims. First, collapses n, feature vectors into single fixed-sized feature (with the same dimensions as e,). Second, pooling invariant by construction to the order of observations (for example, pooling function would take a dimension- -iie mean or sum across j.) Finally, this pooled feature vector e is input to another network param perm__eetttaaaa that returns the final parameters wy. The pscudocode for net, with feature networks is available in fig. 4. In table 4, in appendix, we summarize the applicability of proposed variational methods to different HBD variants.
net, (X:;, (01, for in {1, 21, ni) + feat fett_nett, (1is,, =tpooolfee ) W/ f- pappaammmee return wi
Figure 4: Psuedocode for net, for locally i.i.d symmetric HBD.
7 Experiments
We conduct experiments on synthetic and a real-world prob- lem. For each, we consider three inference methods: using joint distribution, using a branch variational approximation, and using our amortized approach. For each method, we consider three variational approximations a completely diagonal Gaussian, block-diagonal Gaussian (with blocks for 0 and 2) and dense Gaussian; see fig. 5 for a visual. (For each choice of a joint distribu- tion, the corresponding nranch is used for the branch variational approximation and the corresponding 90,u approach; see table for details.)
8
Figure 5: Visualization of dense, block-diagonal, and diagonal co- variances. For each, we exper- iment with Joint. Branch and Ie,ar Ammrtt methods.
Amort for the amortized
We use reparameterized gradients [20] and optimize with Adam [19] (see appendix G for complete experimental details.) In appendix A, discuss some of the observations had during experimen-
7

---

tation involving parameterization, feat-net architecture, batch-size selection, initialization, and gradient estimators.
7.1 Synthetic problem
The aim of the synthetic experiment is to work with models where we have access to closed- form posterior (and the marginal likelihood.) If variational family contains posterior, one can expect the methods to perform close to ideal (provided we can optimize well.)
For experiments, we the following hierarchical regression model
nl II
- Nooo, 1) II (=
(16)
i=1
1=1
where N denotes a Gaussian distribution, and 7 is an identity matrix (see appendix G.2 for posterior and the marginal closed-form expressions.)
To demonstrate the performance of our methods, we experiment with three different problems scale (correspond to three different models with N 10. IK, and 100K.) Synthetic data is created using forward sampling for each of the scale variants independently. We avoid any test data and metrics for the synthetic problem as the log-manginal known in closed form. The inference results are present in fig. 7 (in appendix.) In all cases, amortized distributions perform favorably when compared to branch and joint distributions.
7.2 MovieLens
Next, we test our method on the MovieLens25M [14], dataset of 25 million movie ratings for over 62,000 movies, rated by 162,000 users, along with set of features (tag relevance scores [38]) for each movie.
Purely, make experiments we pre-process the data to drop users with more than 1,000 ratings-leaving around 20M ratings. Also, for the sake of efficiency, we PCA the movie features to reduce their dimen- sionality to 10. We used a train-tes split such that, for each user, one-tenth of the ratings are in the test set. This gives us 18M ratings for training (and 2M ratings for testing.)
efficient GPU hardware,
Table 3: Details for different scales of the MovieLens 25M problem.
Scale
of ratings
# of users
raiing:
Frd mnn
Small
2.5K
156.3
Moderate
180K
1600 159978
112.5
Large
18M
112.5
We use the hierarchical model
N p(0, 2,))) No I) II (())p((). E(0) II
Wy
Bly sigmoid(x.f..)).
(17)
j=1
Table 2: Inference results for the MovieLens25M problem. For both metrics, we draw fresh batch of 10,000 samples from the final posterior. All values are in nats (higher is better).
Metric
Final ELBO 180K
Test likelihood
# train ratings
2.5K
18M
2.5K
180K
18M
Methods (see table
Dense
goint
-1572.31
-166.37
Branch
-1572.39
-1.0368e+05
-1+1413e+07
-666.66
-11054.43
-1.3046e+00
gaimors
-1572.45
-1.0352e+05
-1.0665e+07
-166.64
-1096..38
-1.1476e+06
Block
Iiitt
-1579.04
-167.36
Diagonal
-1579.05
-1 -1.0350e+05
-1.1078c+07
-166.97
-10987.17
-1.2538c400
To.u Amors
-1579.06
-1.0353c405
-1.0665e+07
-166.96
-10975.96
-1.1484c+06
Diagonal
Tolat
-1592.59
-167.39
g Alnort
Branch
-1592.64
-1.0428c405
-1.1325c+07
-167.31
-10977.95
-1.2713c+006
-1592.64
-1.0430c+05
-1.0736c+07
-167.29
-10980.75
-1.1497c+06

---

# ratings 2.5K
le5
# ratings
180K
le7 # ratings so 18M
-1570
-1.0
-1.04
1106
-1.2
(Dense)
Joint (g(**)
-1590
Branch (
-1.08
-1.4
-1.10
-1600
100000
200000
100000
200000
100000
200000
1e5
1e7
-1.0
( (Block Diagonal)
-1580
-106
-1.2
-1390
1.08
-1.4
-1.10
11500
100000
200000
100000
200000
100000
200000
le5
1e7
-1.0
(Diagonal)
ELBO
-158
-1.06
-1.2%
-159
1.08
-1.4
1110
100000
200000
100000
200000
100000
200000
Figure Training ELBO trace for the MovieLens25M problem. Top bottom: dense, block diagonal, and diagonal Gaussian (for each, we have gointt Branch and qe method.) Left to right: small, moderate, and large scale of the MovieLens25M problem. For clarity, we plot the exponential moving average of the training ELBO trace with smoothing value of 0.001. These traces correspond to the values reported in table
where 9 represents distribution over user preferences; for instance, might represent that users who like action films tend to also like thrillers but tend to dislike musicals; 2; determine the user specific preference; Tay are the features of the jth movie rated by user i; Yby is the binary movie ratings; n; is the number of movies rated by user and B denotes Bernoulli distribution. Here, f and E are functions of 0, such that, for [8p, p. (z), we have
(0) 0μ and
(0) tril(Os) (rri(()))))
where tril is a function that transforms unconstrained vector into a lower-triangular, positive Cholesky factor. As movie features æij E 10 we have 8,2 E I 10 0g E 1255, and 2; E 11 Note that V depends on 0, and the likelihood is Bernoulli, the model is non-conjugate.
For inference, we use the methods as described in table Note, we hold the amortization network architecture constant across the scales theles-thehiee number of parameters remains fixed for q.mort (for all Gaussian variants) while the number of parameters scale as O(N) for ftranch (see appendix F for more details.)
In fig. 6, we plot the training time ELBO trace, and in table 2, we present the final training ELBO and test likelihood values. We approximate the test likelihood p(y x. y) with E. [ptytest [z 9)], and draw a fresh batch of 10,000 samples to approximate the expectation (see appendix G for complete details.) For the smaller model, the amortized and branch approaches perform similar to the joint approach for all three variational approximations: this supports theorem and claim 5. For the moderate size model, the branch and amortize approaches are very comparable to each other, while joint approaches fail to scale. For the large model (18M ratings), amortized approaches are significantly better than branch methods. We conjecture this is because parameter sharing in amortized approaches improves convergence for models where batch size is smaller com- pared to total iterates-trus only for large model. Interestingly, the performance of amortized Dense and Block Gaussian approximation very similar in the large and moderate setting (see q results in tables 2, 7 and 8.) We conjecture this is because the posterior over the global parameters is

---

very concentrated for this problem. As behaves like single fixed value, Block Gaussian performs just as good as the Dense approach (see appendix B for more discussion.)
Discussion
In this paper, present structured amortized variational inference scheme that can scale to large hierarchical models without losing benefits of joint approximations. Such models are ubiquitous in social sciences, ecology, epidemiology, and other fields. Our ideas not only inspire further research in inference but also provide a formidable baseline for applications.
Acknowledgements and Disclosure of Funding
This material is based upon work supported in part by the National Science Foundation under Grant No. 1908577.
References
[1] Abhinav Agrawal, Daniel R. Sheldon, and Justin Domke. Advances in black-box VI: normaliz- ing flows, importance weighting, and optimization. NeurIPS 2020, 2020.
[2] Luca Ambrogioni, Kate Lin, Emily Fertig, Sharad Vikram, Max Hinne, Dave Moore, and Marcel Gerven. Automatic structured variational inference. In International Conference Artificial Intelligence and Statistics, pages 676-684. PMLR, 2021.
[3] David M Blei. Probabilistic topic models. Communications of the ACM, 55(4))77-84, 2012.
[4] David M Blei, Andrew Y Ng, and Michael Jordan. Latent dirichlet allocation. the Journal of machine Learning research, 3:223-1022, 2003.
[5] Diane Bouchacourt, Ryota Tomioka, and Sebastian Nowozin. Multi-level variational autoen- coder: Learning disentangled representations from grouped observations. In Secood AAAI Conference on Artificial Intelligence, 2018.
[6] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Mille and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://giie www.ttuu.cooooeeeeommmm google/jax.
[7] Thomas M. Cover and Joy A. Thomas. Elements of information theory (2. ed.). Wiley, 2006.
[8] Noel Cressie, Catherine A Calder, James S Clark, Jay M Ver Hoef, and Christopher K Wikle. Accounting for uncertainty in ecological analysis: the strengths and limitations of hierarchical statistical modeling. Ecological Applications, 19(3):553-570, 2009.
[9] Justin Domke. Provable smoothness guarantees for black-box variational inference. In Interna- tional Conference on Machine Learning, pages 2587-2596. PMLR, 2020.
[10] Harrison Edwards and Amos Storkey. Towards a neural statistician. arXiv preprint 1rXiv:1606.02185, 2016.
[11] Andrew Gelman. Red state, blue state, rich state, poor state: Why Americans vote the way they do-expanded edition. Princeton University Press, 2009.
[12] Andrew Gelman and Jennifer Hill. Data analysis using regression and multileve//hieracchical models. Cambridge university press, 2006.
[13] Andrew Gelman and Aki Vehtari. What are the most important statistical ideas of the past 50 years? arXiv preprint arXiv:2012.00177 2020.
[14] Maxwell Harper and Joseph Konstan. The movielens datasets: History and context, Acm transactions interactive intelligent systems (tiis, 5(4):1-19, 2015.
10

---

[15] Matthew D. Hoffman and David M. Blei. Stochastic structured variational inference. In Guy Lebanon and S. V. N. Vishwanathan, editors. Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics, AISTATS 2015, San Diego, California, USA, May 9-12, 2015, volume 38 of JMLR Workshop and Conference Proceedings. JMI..orgg 2015. URL. httpttppppp//eeeeeeeeeeep
mlr. pwwss//esvseeeeeee v38/ 2//////ffft fman 15.tmml.
[16] Matthew D Hoffman. David M Blei, Chong Wang. and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 14(5), 2013.
[17] Maximilian llse, Jakub M Tomczak, Christos Louizos, and Max Welling- Diva: Domain invariant variational autoencoders. arxiv e-prints, article. arXiv preprint arXiv: 10427,
2019.
[18] Matthew Johnson, David Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta. Composing graphical models with neural networks for structured representations and fast inference. Advances in neural information processing systems, 29:2946-2954, 2016.
[19] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 2015.
[20] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 13X2v:1312.61 14, 2013.
arXiv preprint
[21] Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M Blei. Automatic differentiation variational inference. The Journal of Machine Learning Research, 18(1):430-44) 14430-474, 2017.
[22] John D Lafferty and David M Blei. Correlated topic models. Advances in neural information processing systems, 18:147-,54, 2006.
[23] Andrew Lawson. Bayesian disease mapping: hierarchical modeling in spatial epidemiology. Chapman and Hal//CRC, 2008.
[24] Jeffrey R Lax and Justin H Phillips. The democratic deficit in the states, American Journal of Political Science, 56( 14):))):1616... 2012.
[25] Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klbert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pages 9-48. Springer, 2012.
[26] Yew Jin Lim and Yee Whye Teh. Variational bayesian approach movie rating prediction. In Proceedings of KDD cup and workshop, volume 7, pages 15-21. Citeseer, 2007.
[27] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. arXiv preprint 1rXiv:191..02762. 2019.
[28] K. B. Petersen and M. S. Pedersen. The matrix cookbook, nov 2012. URL http: http://tw.........t/////epeete/e
/
ftmmmmm.l. Version 20121115.
[29] John Platt, Emre Kiciman, and Dave Maltz. Fast variational inference for large-scale internet diagnosis. In NIPS, 2007.
[30] Rajesh Ranganath, Sean Gerrish, and David Blei. Black Box Variational Inference. AISTATS. 2014.
[31] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In ICML., 2014.
[32] Geoffrey Roeder, Yuhuai Wu, and David Duvenaud. Sticking the landing: Simple, lower- variance gradient estimators for variational inference. In NIPS, 2017.
[33] Ruslan Salakhutdinov and Andriy Mnih. Bayesian probabilistic matrix factorization using markov chain monte carlo. In Proceedings of the 25th international conference on Machine learning, pages 880-887, 2008

---

[34] Rishit Sheth and Roni Khardon. Monte carlo structured svi for two-level non-conjugate models. arXiv preprint arXiv: 1612.0397, 2016.
[35] Michael E Tipping and Christopher M Bishop. Probabilistic principal component analysis.
Journal of the Royal Statistical Society: Series B (Statistical Methodology), 61(3):611-622
1999.
[36] Michalis Titsias and Miguel Lázaro-Gredilla. Doubly stochastic variational bayes for non- conjugate inference. In International conference machine learning, pages 1971-1979. PMLR, 2014.
[37] Robert Vallerand. Toward hierarchical model of intrinsic and extrinsic motivation. Advances in experimental social psychology, 29:271-360, 1997.
[38] Jesse Vig, Shilad Sen, and John Riedl. The tag genome: Encoding community knowledge to support novel interaction. ACM Transactions on Interactive Intelligent Systems (TiSS), 2(3): 1 1-44, 2012.
[39] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. In I. Guyon, U. V. Luxburg, S. Bengio. H. Wallach. R. Fergus, Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: wwwrrceedngg...uuis......./// 20.7://ll//2aa17777/2/7226666662
12


Reviews:
[{"review_id": "xKK_XsUe4jl", "reviewer": "Reviewer_qcmy", "summary": "This work proposes a Gaussian naive bayes approximate posterior for a 2-level hierarchical naive bayes model. Various parameterizations for the variational posterior are explored, including a fully coupled joint posterior, a factored branch posterior, and an amortized version of the latter. Results and theory indicate that having an approximate posterior that makes the same independence assumptions as the true posterior does not hurt, and that parameter sharing improves performance at scale.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Originality*\n* It is unclear whether the contribution of this paper, the exploration of the parameterization of the Gaussian naive bayes approximate posterior, is novel. To remedy this, the related work must cover similar models and inference methods. Currently the related work is too high-level, and offers little context. Similar generative models must have been used elsewhere, such as hierarchical naive bayes [1]. Please find and cite the relevant work in order to convince us that the model is useful and that amortized variational inference has not been applied in this exact setting.\n* The preference for the branch posterior rather than the fully-connected joint can be motivated by Webb et. al. [2].\n\n*Quality*\n* The exploration of the approximate posterior is thorough. However, not enough context is given for experimental results. How do the models perform relative to other baseline models? How is evaluation normally conducted on MovieLens?\n\n*Clarity*\n* The paper is very clearly written, and seems technically sound.\n\n*Significance*\n* Providing more context will help show the significance of the work. The methods and experiments section are complete, but giving more context for the background and results will complete the paper.\n\n*Comments*\n* The title is a bit too high-level and vague as well. This paper proposes amortized variational inference for a very particular model, I suggest including that in the title.\n\n[1] Langseth, H. and Thomas D. Nielsen. “Classification using Hierarchical Naïve Bayes models.” Machine Learning 63 (2006): 135-159.\n\n[2] Stefan Webb, Adam Golinski, Robert Zinkov, N Siddharth, Tom Rainforth, Yee Whye Teh, and Frank Wood. Faithful inversion of generative models for effective amortized inference. In Advances in Neural Information Processing Systems, 2018.\n\nEdited score: After calibrating across other reviews, the issues I raised are either not high priority (experimental) or can easily be fixed (more context), and have lowered confidence accordingly. I would be willing to raise the score higher if those are addressed.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "* The appendix includes some discussion of the limitations of the different posterior parameterizations.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "okIH9ckIPl", "reviewer": "Reviewer_k4wz", "summary": "The paper deals with scalable approximate inference for hierarchical models and proposes a framework for hierarchical approximations that exhibit the same factorization as the exact posterior of the hierarchical model. This efficiently enables subsampling the data when evaluating the evidence lower bound and its gradients. The authors further propose a pooling-based strategy for amortized inference to avoid the usual linear scaling in the number of approximation parameters.  The paper is concluded with a set of numerical experiments with both synthetic and real data demonstrating the merits of the proposed framework.\n\n\n", "questions": "", "limitations": "", "rating": 7, "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": "The paper proposes a novel (to the best of my knowledge) and elegant solution to an important problem and I think the proposed methods will be of general interest to the approximate inference community. \n\nAll technical details in the paper appear to be correct and the claims in the paper are theoretically motivated and reasonably justified using empirical experiments. However, the empirical experiments could be further improved by\n\n- The results for the synthetic experiment are only compared using the lower bounds. Since the target distribution is Gaussian, why not 1) compute and plot the true log marginal likelihood in Figure 7 in the appendix for comparison 2) compute and evaluate the error on the estimated moments/parameters for completeness since the true quantities are available.\n\n- Use replications and add error bars to all experiments\n\nOverall, the paper is well-written and well-structured and I really enjoyed reading it. However, I do think the notation can be slightly improved. For example,\n\n- I am a little bit puzzled by the exact meaning of the definition in Eq. (2). Perhaps, this is caused by the usual casual notation in probability, where there is no distinction in the notation for a random variable and the values it takes. Can this be rephrased or made more precise?\n\n- Along the same lines, in eq. (6), (8)-(10) it would be a nice service to the reader if the authors would distinguish between the ranom variables and samples of the random variables. For example, use \\theta^s or \\hat{theta} to represent a sample of the random variable \\theta and the same for the other variables.\n\n- I think there may be missing a minimization operator on the left hand side of eq. (14) since currently, the left hand side is a function of v (and the existence of a uniform bound seems unlikely). I noticed that \\min_v is also present in the proof in the appendix, so it is most likely just a typo.\n\n\nOther comments:\n\n- the xlabels in fig 6 are missing.\n\n - extra parenthesis in eq. (16)\n\n- It's nice that the authors included a section describing some of the practical challenges and considerations.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Not applicable", "review_text": "", "extra_scores": {"ethical_concerns": "Not applicable", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "3YCR_EtxkoP", "reviewer": "Reviewer_F7hF", "summary": "This paper makes an attempt to address the problem of large-scale inference for hierarchical probabilistic graphical models that contain both global and local latent variables by proposing to use Amortized VI instead of VI in [1] for more efficient and scalable inference. The authors first define a general class of hierarchical models which they call Hierarchical Branch Distributions (HBDs). They then propose a variational family matching this class of models for performing amortized inference. The authors then perform a series of experiments on synthetic datasets and the MoveLens dataset where they evaluate different inference models in terms of the final ELBO and convergence.  ", "questions": "", "limitations": "", "rating": 5, "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": "[General Comment]\n\nIn my opinion, the main weakness of the current manuscript is the lack of novelty or at least the lack of transparency of the contributions. While I very appreciate the importance of accurate/scalable inference for hierarchical graphical models and therefore am a big fan of this line of research, I found it very difficult to pin-point the exact contributions of this paper. Amortized variational inference has been applied not only to models with solely local variables, but also models with global and local variables as well, so this is not a new idea by any means. Here are just a few examples: [2-6] that all use AVI and most contain local and global variables. In fact, [2] also uses pooling to extract group content from local features. So if the main contribution can be described as Figure 3 (‘d’ in particular), it’s hard for me to see the novelty here as I would classify that as basically choosing an appropriate conditional variational family for your model. Ofcourse if this was an application paper, where an extensive set of experiments had been done  with the specified model to a particular task to achieve good results, it would have been a different story but this is not the case here. \n\n\n\n[Clarity]\n\nI did not find the paper easy to read, which I believe also partially contributed to the contributions not being very clear. I would encourage the authors to proofread the paper once more. I also do not find all the mathematical statements clear or necessary. For example, the “symmetric“ definition in page 2 is redundant as far as I can see (a=b ⇒ p(a|c) = p(b|c)). Theorem 2 (& paragraph 4 in page 1) is also obvious and redundant I would argue. Choosing a variational family that factorizes just your target would of course yield a more accurate inference. \n\n-There is a font switch for theta on page 3 (compare equation 4 & 5)\n-Typo in equation (16)\n\n\n\n[Novelty & Prior work]\n\nI don't think the current manuscript did a good job of discussing related work. See my general comment for more details. \nAnother missing reference is [7]\n\n\n[Minor comments]\n\nIn order to properly test the inference, I also would advise the authors to test the model containing no ‘x’. While the authors state that “results extend easily to case where there are no local covariates”, this is not necessarily the case in practice depending on  how much information the covariates provide about ‘y’. In such cases, the model does not have to rely on the latent variables so accurate inference might not matter as much to achieve a good likelihood. \n\n\n\n\n[References]\n\n[1] Hoffman, Matthew D., et al. \"Stochastic variational inference.\" Journal of Machine Learning Research 14.5 (2013).\n\n[2] Edwards, Harrison, and Amos Storkey. \"Towards a neural statistician.\" arXiv preprint arXiv:1606.02185 (2016).\n\n[3] Lee, Dong Bok, et al. \"Meta-GMVAE: Mixture of Gaussian VAE for Unsupervised Meta-Learning.\" International Conference on Learning Representations. 2020.\n\n[4] Esmaeili, Babak, et al. \"Structured neural topic models for reviews.\" The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.\n\n[5] Ilse, Maximilian, et al. \"Diva: Domain invariant variational autoencoders.\" Medical Imaging with Deep Learning. PMLR, 2020.\n\n[6] Bouchacourt, Diane, Ryota Tomioka, and Sebastian Nowozin. \"Multi-level variational autoencoder: Learning disentangled representations from grouped observations.\" Thirty-Second AAAI Conference on Artificial Intelligence. 2018.\n\n[7] Johnson, Matthew J., et al. \"Composing graphical models with neural networks for structured representations and fast inference.\" Advances in neural information processing systems 29 (2016): 2946-2954.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "To the best of my knowledge, this work does not raise any major social impact.  ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 + 2 (rebuttal)"}, "extra_text": {}}, {"review_id": "-KIrO4qCxM", "reviewer": "Reviewer_NXpx", "summary": "An approach vor variational inference in large-scale hierarchical models is presented that is more efficient than the base lines. This is achieved using shared parameters (amortization) and a feature pooling network. Experiments confirm the efficiency of the approach on synthetic data as well as the MovieLens dataset.", "questions": "", "limitations": "", "rating": 7, "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": "The paper is original and clearly written (all terms are defined, although some only later after the first mention).\n\nRegarding the significance my thoughts are as follows:\nThis paper has a focus on the theory and thus does not spend too much space on the experiments. The theory regarding the branch approximations is interesting and seems to be relevant. What is not entirely clear to me from the paper is how widely applicable the proposed methods are. The cited application papers in the introduction are all older works before the boom of neural network based approaches. In my view the main limitation of the paper is thus an appropriate comparison to a state-of-the-art model (such as VAE maybe, or a NN approach for learning hierarchical models) to show, either theoretically or practically, how the proposed method compares to more recent methods. This would make the results, which are more or less just comparing variants of the same method, much more convincing. In general, though, I think this is a good paper and the above reservations should maybe not be weighted too heavily.\n\nCorrolary -> Corollary", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"ethical_concerns": "no issues", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors consider variational inference for a hierarchical model where each observation is generated using a local/observation-specific latent variable, which in turn depends on the global latent variable. They prove that using a variational posterior with the same dependence structure as the model results in inference accuracy comparable to a general joint posterior approximation at a much lower computational cost and allows observation subsampling for an additional speedup. Experiments confirm the theoretical results and show that amortizing local inference enables training on much larger datasets without affecting accuracy.

The paper describes a simple but effective idea clearly and shows that it works in practice. The discussion of related work is insufficient however, and should be extended to include the Neural Statistician and its descendants and their relationship to this work. Also, as pointed out by a reviewer, the title of the paper is too general and should really be made more specific, as the method applies only to hierarchical models with a very specific, if common structure.
Decision: Accept (Poster)