As a result, in both PAVI-F and PAVI-E, the terms $1_{n \in \mathcal{B}_i[t]}$ and $\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}); \phi_i, \mathbf{E}_{i,n})$ depend on the sampled indices of disjoint sets of plates, and are therefore independent. This means that the expectation of their product can be rewritten as the product of their expectations:

$$ 
\begin{aligned}
\mathbb{E}_{I_P}[\log q^{\mathbf{r}}(\Theta^{\mathbf{r}}[t])] &= \sum_{i=1}^{I} \frac{N_i}{N_i^{\mathbf{r}}} \sum_{n=0}^{N_i} \mathbb{E}_{I_P}[1_{n \in \mathcal{B}_i[t]}] \mathbb{E}_{I_P}[\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}))] \\
&= \sum_{i=1}^{I} \frac{N_i}{N_i^{\mathbf{r}}} \sum_{n=0}^{N_i} \frac{N_i^{\mathbf{r}}}{N_i} \mathbb{E}_{I_P}[\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}))] \qquad (A.11) \\
&= \sum_{i=1}^{I} \sum_{n=0}^{N_i} \mathbb{E}_{I_P}[\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}); \mathbf{E}_{i,n})]
\end{aligned}
$$

This equality can be further simplified in the PAVI-F case — proving its unbiasedness— but not in the PAVI-E case, as detailed in the sections below.

### A.2.2 Unbiasedness of the PAVI-F scheme

In the PAVI-F scheme, detailed in Section 4.3, the encodings $\mathbf{E}_{i,n}$ are constants with respect to the branching choice, therefore we have:

$$ 
\begin{aligned}
\mathbb{E}_{I_P}[\log q^{\mathbf{r}}(\Theta^{\mathbf{r}}[t])] &= \sum_{i=1}^{I} \sum_{n=0}^{N_i} \mathbb{E}_{I_P}[\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}); \mathbf{E}_{i,n})] \\
&= \sum_{i=1}^{I} \sum_{n=0}^{N_i} \log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}); \mathbf{E}_{i,n}) \\
&= \log q(\Theta)
\end{aligned}
\qquad (A.12)
$$

which proves Equation (A.8) and Equation (A.6). In the above example of $\theta_{1,2}$ in $\mathcal{M}$, in the PAVI-F scheme the expression $\mathbb{E}_{I_P}[\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}); \phi_i, \mathbf{E}_{i,n})]$ can be evaluated into $\log q_{1,2}(\theta_{1,2}|\theta_{2,0}; \phi_1, \mathbf{E}_{1,2})$. This demonstrates that the PAVI-F scheme is unbiased: training over stochastically chosen sub-graphs for $q^{\mathbf{r}}$ 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 4.2, the encodings $\mathbf{E}_{i,n}$ are computed from the observed data $X$. Specifically, considering the ground RV $\theta_{i,n}$, we have $\mathbf{E}_{i,n} = f(\mathbf{X}_{i,n}^{\mathbf{r}}[t])$ where $\mathbf{X}_{i,n}^{\mathbf{r}}[t]$ corresponds to the observed data of a subset of $\theta_{i,n}$'s descendants. Depending on the chosen branching downstream of $\theta_{i,n}$, the value of $\mathbf{E}_{i,n}$ can therefore vary. This means we cannot further simplify Equation (A.11): the terms $\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}); \mathbf{E}_{i,n})$ are not constants with respect to the RVs $I_P$. In the above example of $\theta_{1,2}$ in $\mathcal{M}$, in the PAVI-E scheme the expression $\mathbb{E}_{I_P}[\log q_{i,n}(\theta_{i,n}|\pi(\theta_{i,n}); \phi_i, \mathbf{E}_{i,n})]$ can be evaluated into:

$$ \frac{1}{2}(\log q_{1,2}(\theta_{1,2}|\theta_{2,0}; \phi_1, f(\mathbf{X}_4)) + \log q_{1,2}(\theta_{1,2}|\theta_{2,0}; \phi_1, f(\mathbf{X}_5))) $$

**How could the PAVI-E scheme be made unbiased?** Specifically, by making the value of $\mathbf{E}_{i,n}$ independent of the choice of downstream branching. A possibility would be to parameterize $\mathbf{E}_{i,n}$ as an average —an expectation— over all the possible sub-branchings downstream of $\theta_{i,n}$. Yet, in practical cases, the cardinalities of the reduced model are much inferior to the ones of the full model: $\mathrm{Card}^{\mathbf{r}}(\mathcal{P}) \ll \mathrm{Card}(\mathcal{P})$. This means that numerous $\mathrm{Card}^{\mathbf{r}}(\mathcal{P})$-sized subsets can be chosen inside the $\mathrm{Card}(\mathcal{P})$ possible descendants. In order to average over all those subset choices to compute $\mathbf{E}_{i,n}$, 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 problem of aggregating collections of stochastic estimators into one general estimator —in an unbiased and efficient manner. To our knowledge, this is a complex and still open research question, whose advancement could much benefit our applications.
**Practical approximation for the PAVI-E scheme** In practice, we compute the encoding $\mathbf{E}_{i,n}$ based on the single downstream branching corresponding to the sampling of the RVs $I_P$. Compared to the previous paragraph, this amounts to estimating the expectation of $\mathbf{E}_{i,n}$ —over all downstream branchings— using a single one of those branchings. Note that, even if this encoding estimate was unbiased, $\log q_{i,n}$ would remain a highly non-linear function of $\mathbf{E}_{i,n}$. As a consequence, we need to rely on the approximation:

$$ \mathbb{E}_{I_P} [\log q_{i,n}(\theta_{i,n} | \pi(\theta_{i,n}); \phi_i, f(\mathbf{X}_{i,n}^r[t]))] \approx \log q_{i,n}(\theta_{i,n} | \pi(\theta_{i,n}); \phi_i, f(\mathbf{X}_{i,n})) \quad (A.13) $$

which can theoretically introduce some bias in our gradients. The approximation Equation (A.13) can be interpreted as follow: "the expectation of the density of $\theta_{i,n}$ when collecting summary statistics over a stochastic subset of $\theta_{i,n}$'s descendants is approximately equal to the density of $\theta_{i,n}$ when collecting summary statistics over the entirety of $\theta_{i,n}$'s descendants". Another interpretation is that the distribution associated with the summary of the full data can be approximated by annealing the distributions associated with summaries of subsets of this data. In practice, this approximation did not yield significantly worse performance for the PAVI-E scheme over the generative models we tested. At the same time, computing the encodings over a single branching allows the computation of all the $\mathbf{E}_{i,n}$ encodings in a single lightweight pass over the data $\mathbf{X}^r[t]$. This simple solution thus provided a substantial increase in training speed with seldom noticeable bias. Yet, we do not bar the existence of pathological generative HBMs where this approximation would become coarse. Experimenters should bear in mind this possibility when using the PAVI-E scheme. In practice, using the PAVI-F scheme as a sanity check over synthetic, toy-dimension implementations of the considered generative models is a good way to validate the PAVI-E scheme —before moving on to the real problem instantiating the same generative model with a larger dimensionality.

## A.3 PAVI algorithms

More technical details can be found in the codebase provided with this supplemental material.

### A.3.1 Architecture build

**Algorithm 1: PAVI architecture build**

**Input:** Graph template $\mathcal{T}$, plate cardinalities $\{(\text{Card}(\mathcal{P}_p), \text{Card}^r(\mathcal{P}_p))\}_{p=0..P}$, encoding scheme

**Output:** $q$ distribution

for $i = 1..I$ do

Construct conditional flow $\mathcal{F}_i$;

Define conditional posterior distributions $q_{i,n}$ as the push-forward of the prior via $\mathcal{F}_i$, following Equation (3);

Combine the $q_{i,n}$ distributions following the cascading flows scheme, as in Section 4.1 (Ambrogioni et al., 2021b);

if PAVI-F encoding scheme then

Construct encoding arrays $\{\mathbf{E}_i = [\mathbf{E}_{i,n}]_{n=0..N_i}\}_{i=1..I}$ as in Appendix A.1.2;

else if PAVI-E encoding scheme then

Construct encoder $f$ as in Appendix A.1.3;
**Practical approximation for the PAVI-E scheme** In practice, we compute the encoding $\mathbf{E}_{i,n}$ based on the single downstream branching corresponding to the sampling of the RVs $I_P$. Compared to the previous paragraph, this amounts to estimating the expectation of $\mathbf{E}_{i,n}$ —over all downstream branchings— using a single one of those branchings. Note that, even if this encoding estimate was unbiased, $\log q_{i,n}$ would remain a highly non-linear function of $\mathbf{E}_{i,n}$. As a consequence, we need to rely on the approximation:

$$ \mathbb{E}_{I_P} [\log q_{i,n}(\theta_{i,n} | \pi(\theta_{i,n}); \phi_i, f(\mathbf{X}_{i,n}^r[t]))] \approx \log q_{i,n}(\theta_{i,n} | \pi(\theta_{i,n}); \phi_i, f(\mathbf{X}_{i,n})) \quad (A.13) $$

which can theoretically introduce some bias in our gradients. The approximation Equation (A.13) can be interpreted as follow: "the expectation of the density of $\theta_{i,n}$ when collecting summary statistics over a stochastic subset of $\theta_{i,n}$'s descendants is approximately equal to the density of $\theta_{i,n}$ when collecting summary statistics over the entirety of $\theta_{i,n}$'s descendants". Another interpretation is that the distribution associated with the summary of the full data can be approximated by annealing the distributions associated with summaries of subsets of this data. In practice, this approximation did not yield significantly worse performance for the PAVI-E scheme over the generative models we tested. At the same time, computing the encodings over a single branching allows the computation of all the $\mathbf{E}_{i,n}$ encodings in a single lightweight pass over the data $\mathbf{X}^r[t]$. This simple solution thus provided a substantial increase in training speed with seldom noticeable bias. Yet, we do not bar the existence of pathological generative HBMs where this approximation would become coarse. Experimenters should bear in mind this possibility when using the PAVI-E scheme. In practice, using the PAVI-F scheme as a sanity check over synthetic, toy-dimension implementations of the considered generative models is a good way to validate the PAVI-E scheme —before moving on to the real problem instantiating the same generative model with a larger dimensionality.

## A.3 PAVI algorithms

More technical details can be found in the codebase provided with this supplemental material.

### A.3.1 Architecture build

**Algorithm 1: PAVI architecture build**

**Input:** Graph template $\mathcal{T}$, plate cardinalities $\{(\text{Card}(\mathcal{P}_p), \text{Card}^r(\mathcal{P}_p))\}_{p=0..P}$, encoding scheme

**Output:** $q$ distribution

for $i = 1..I$ do

Construct conditional flow $\mathcal{F}_i$;

Define conditional posterior distributions $q_{i,n}$ as the push-forward of the prior via $\mathcal{F}_i$, following Equation (3);

Combine the $q_{i,n}$ distributions following the cascading flows scheme, as in Section 4.1 (Ambrogioni et al., 2021b);

if PAVI-F encoding scheme then

Construct encoding arrays $\{\mathbf{E}_i = [\mathbf{E}_{i,n}]_{n=0..N_i}\}_{i=1..I}$ as in Appendix A.1.2;

else if PAVI-E encoding scheme then

Construct encoder $f$ as in Appendix A.1.3;
### A.3.2 Stochastic training

**Algorithm 2:** PAVI stochastic training

**Input:** Untrained architecture $q$, observed data $\mathbf{X}$, encoding scheme, number of steps $T$

**Output:** trained architecture $q$

for $t = 0..T$ do

Sample plate indices to define the batches $\mathcal{B}_i[t]$, the latent $\Theta^{\mathbf{r}}[t]$ and the observed $X^{\mathbf{r}}[t]$ and $\mathbf{X}^{\mathbf{r}}[t]$, following Appendix A.1.1;

Define reduced distribution $p^{\mathbf{r}}$;

if PAVI-F encoding scheme then

Collect encodings $\mathbf{E}_{i,n}$ by slicing from the arrays $\mathbf{E}_i$ the elements corresponding to the batches $\mathcal{B}_i[t]$;

else if PAVI-E encoding scheme then

Compute encodings as $\mathbf{E} = f(\mathbf{X}^{\mathbf{r}}[t])$;

Feed obtained encodings into $q^{\mathbf{r}}$;

Compute reduced ELBO as in Equation (6), back-propagate its gradient;

Update conditional flow weights $\{\phi_i\}_{i=1..I}$;

if PAVI-F encoding scheme then

Update encodings $\{\mathbf{E}_{i,n}\}_{i=1..I, n \in \mathcal{B}_{i,t}}$;

else if PAVI-E encoding scheme then

Update encoder weights $\eta$;

### A.3.3 Inference

**Algorithm 3:** PAVI inference

**Input:** trained architecture $q$, observed data $\mathbf{X}$, encoding scheme

**Output:** approximate posterior distribution

if PAVI-F encoding scheme then

Collect full encoding arrays $\mathbf{E}_i$;

else if PAVI-E encoding scheme then

Compute encodings as $\mathbf{E} = f(\mathbf{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 a joint distribution $p(\Theta, X)$, and a value $\mathbf{X}$ for the RV template $X$. We pick a variational family $\mathcal{Q}$, and in this family look for the parametric distribution $q(\Theta; \phi)$ that best approximates $p(\Theta|\mathbf{X} = \mathbf{X})$. Specifically, we want to minimize the Kulback-Leibler divergence (Blei et al., 2017) between our variational posterior and the true posterior, that Cremer et al. (2018) refer to as the *gap* $\mathcal{G}$:

$$ \begin{aligned} \mathcal{G} &= \text{KL}(q(\Theta; \phi) || p(\Theta|X)) \\ &= \log p(X) - \text{ELBO}(q; \phi) \end{aligned} \qquad (A.14) $$

We denote $q^*(\Theta; \phi^*)$ the optimal distribution inside $\mathcal{Q}$ that minimizes the KL divergence with the true posterior:

$$ \begin{aligned} \mathcal{G}_{\text{approx}}(\mathcal{Q}; \phi^*) &= \log p(X) - \text{ELBO}(q^*; \phi^*) \\ &\geq 0 \end{aligned} \qquad (A.15) $$

$$ \mathcal{G}_{\text{vanilla VI}} = \mathcal{G}_{\text{approx}} $$

The approximation gap $\mathcal{G}_{\text{approx}}$ depends on the expressivity of the variational family $\mathcal{Q}$, specifically whether $\mathcal{Q}$ contains distributions arbitrarily close to the posterior —in the KL sense.
### A.3.2 Stochastic training

**Algorithm 2:** PAVI stochastic training

**Input:** Untrained architecture $q$, observed data $\mathbf{X}$, encoding scheme, number of steps $T$

**Output:** trained architecture $q$

for $t = 0..T$ do

Sample plate indices to define the batches $\mathcal{B}_i[t]$, the latent $\Theta^r[t]$ and the observed $\mathbf{X}^r[t]$ and $\mathbf{X}^r[t]$, following Appendix A.1.1;

Define reduced distribution $p^r$;

if PAVI-F encoding scheme then

Collect encodings $\mathbf{E}_{i,n}$ by slicing from the arrays $\mathbf{E}_i$ the elements corresponding to the batches $\mathcal{B}_i[t]$;

else if PAVI-E encoding scheme then

Compute encodings as $\mathbf{E} = f(\mathbf{X}^r[t])$;

Feed obtained encodings into $q^r$;

Compute reduced ELBO as in Equation (6), back-propagate its gradient;

Update conditional flow weights $\{\phi_i\}_{i=1..I}$;

if PAVI-F encoding scheme then

Update encodings $\{\mathbf{E}_{i,n}\}_{i=1..I, n \in \mathcal{B}_{i,t}}$;

else if PAVI-E encoding scheme then

Update encoder weights $\eta$;

### A.3.3 Inference

**Algorithm 3:** PAVI inference

**Input:** trained architecture $q$, observed data $\mathbf{X}$, encoding scheme

**Output:** approximate posterior distribution

if PAVI-F encoding scheme then

Collect full encoding arrays $\mathbf{E}_i$;

else if PAVI-E encoding scheme then

Compute encodings as $\mathbf{E} = f(\mathbf{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 a joint distribution $p(\Theta, X)$, and a value $\mathbf{X}$ for the RV template $X$. We pick a variational family $\mathcal{Q}$, and in this family look for the parametric distribution $q(\Theta; \phi)$ that best approximates $p(\Theta|\mathbf{X} = \mathbf{X})$. Specifically, we want to minimize the Kulback-Leibler divergence (Blei et al., 2017) between our variational posterior and the true posterior, that Cremer et al. (2018) refer to as the *gap* $\mathcal{G}$:

$$ \begin{aligned} \mathcal{G} &= \text{KL}(q(\Theta; \phi) || p(\Theta|X)) \\ &= \log p(X) - \text{ELBO}(q; \phi) \end{aligned} \qquad (A.14) $$

We denote $q^*(\Theta; \phi^*)$ the optimal distribution inside $\mathcal{Q}$ that minimizes the KL divergence with the true posterior:

$$ \begin{aligned} \mathcal{G}_{\text{approx}}(\mathcal{Q}; \phi^*) &= \log p(X) - \text{ELBO}(q^*; \phi^*) \\ &\geq 0 \end{aligned} \qquad (A.15) $$

$$ \mathcal{G}_{\text{vanilla VI}} = \mathcal{G}_{\text{approx}} $$

The approximation gap $\mathcal{G}_{\text{approx}}$ depends on the expressivity of the variational family $\mathcal{Q}$, specifically whether $\mathcal{Q}$ contains distributions arbitrarily close to the posterior —in the KL sense.
**Note:** $\mathcal{G}_{approx}$ is a property of the variational family $\mathcal{Q}$. $\mathcal{G}_{approx}$ is an asymptotic bound for the KL divergence between any distribution $q \in \mathcal{Q}$ and the true posterior. This gap is therefore a form of bias but is not to be mistaken with the stochasticity-induced bias studied in Appendix A.2. The bias in Appendix A.2 relates to whether $q^*$ can be found by training stochastically over $\mathcal{Q}$, whereas $\mathcal{G}_{approx}$ relates to the the bias between $q^*$ and the true posterior.

Cremer et al. (2018) demonstrate that, in the case of sample amortized inference, when the weights $\phi$ no longer are free but the output of an encoder $f \in \mathcal{F}$, inference cannot be better than in the non-sample-amortized case, and a positive amortization gap is introduced:

$$ \begin{aligned} \mathcal{G}_{\text{sa}}(\mathcal{Q}, \mathcal{F}; \eta^*) &= \mathcal{G}_{\text{approx}}(\mathcal{Q}; f(\mathbf{X}, \eta^*)) - \mathcal{G}_{\text{approx}}(\mathcal{Q}; \phi^*) \\ &\geq 0 \end{aligned} \qquad (\text{A.16}) \\ \mathcal{G}_{\text{sample amortized VI}} = \mathcal{G}_{\text{approx}} + \mathcal{G}_{\text{sa}} $$

Where we denote as $\eta^*$ the optimal weights for the encoder $f$ inside the function family $\mathcal{F}$. The gap terminology can be interpreted as follow: "theoretically, sample amortization cannot be beneficial in terms of KL divergence for the inference over a given sample $\mathbf{X}$."

Using the same gap terminology, we can define gaps implied by our PAVI architecture. Instead of picking the distribution $q$ inside the family $\mathcal{Q}$, consider picking $q$ from the plate-amortized family $\mathcal{Q}_{pa}$ corresponding to $\mathcal{Q}$. Distributions in $\mathcal{Q}_{pa}$ are distributions from $\mathcal{Q}$ with the additional constraints that some weights have to be equal. Consequently, $\mathcal{Q}_{pa}$ is a subset of $\mathcal{Q}$:

$$ \mathcal{Q}_{pa} \subset \mathcal{Q} \qquad (\text{A.17}) $$

As such, looking for the optimal distribution inside $\mathcal{Q}_{pa}$ instead of inside $\mathcal{Q}$ cannot result in better performance, leading to a plate amortization gap:

$$ \begin{aligned} \mathcal{G}_{\text{pa}}(\mathcal{Q}, \mathcal{Q}_{\text{pa}}; \psi^*, \phi^*) &= \mathcal{G}_{\text{approx}}(\mathcal{Q}_{\text{pa}}; \psi^*) - \mathcal{G}_{\text{approx}}(\mathcal{Q}; \phi^*) \\ &\geq 0 \end{aligned} \qquad (\text{A.18}) \\ \mathcal{G}_{\text{PAVI-F}} = \mathcal{G}_{\text{approx}} + \mathcal{G}_{\text{pa}} $$

Where we denote as $\psi^*$ the optimal weights for a variational distribution $q$ inside $\mathcal{Q}_{pa}$ —in the KL sense. The equation A.18 is valid for the PAVI-F scheme —see Section 4.3. We can interpret it as follow: "theoretically, plate amortization cannot be beneficial in terms of KL divergence for the inference over a given sample $\mathbf{X}$".

Now consider that encodings are no longer free parameters but the output of an encoder $f$. Similar to the case presented in Equation (A.16), using an encoder cannot result in better performance, leading to an encoder gap:

$$ \begin{aligned} \mathcal{G}_{\text{encoder}}(\mathcal{Q}_{\text{pa}}, \mathcal{F}; \psi^*, \eta^*) &= \mathcal{G}_{\text{approx}}(\mathcal{Q}_{\text{pa}}; f(\mathbf{X}, \eta^*)) - \mathcal{G}_{\text{approx}}(\mathcal{Q}_{\text{pa}}; \psi^*) \\ &\geq 0 \end{aligned} \qquad (\text{A.19}) \\ \mathcal{G}_{\text{PAVI-E}} = \mathcal{G}_{\text{approx}} + \mathcal{G}_{\text{pa}} + \mathcal{G}_{\text{encoder}} $$

The equation Equation (A.19) is valid for the PAVI-E scheme —see Section 4.3.

The most complex case is the PAVI-E(sa) scheme, where we combine both plate and sample amortization. Our argument cannot account for the resulting $\mathcal{G}_{PAVI-E(sa)}$ gap: both the PAVI-E and PAVI-E(sa) schemes rely upon the same encoder $f$. In the PAVI-E scheme, $f$ is overfitting over a dataset composed of the slices of a given data sample $\mathbf{X}$. In the PAVI-E(sa) scheme, the encoder is trained over the whole distribution of the samples of the reduced model $\mathcal{M}^r$. Intuitively, it is likely that the performance of PAVI-E(sa) will always be dominated by the performance of PAVI-E, but —as far as we understand it— the gap terminology cannot account for this discrepancy.

Comparing previous equations, we therefore have:

$$ \mathcal{G}_{\text{vanilla VI}} \leq \mathcal{G}_{\text{PAVI-F}} \leq \mathcal{G}_{\text{PAVI-E}} \qquad (\text{A.20}) $$

Note that those are theoretical results, that do not necessarily pertain to optimization in practice. In particular, in Section 5.1&5.3, this theoretical performance loss is not observed empirically over the studied
**Note:** $\mathcal{G}_{approx}$ is a property of the variational family $\mathcal{Q}$. $\mathcal{G}_{approx}$ is an asymptotic bound for the KL divergence between any distribution $q \in \mathcal{Q}$ and the true posterior. This gap is therefore a form of bias but is not to be mistaken with the stochasticity-induced bias studied in Appendix A.2. The bias in Appendix A.2 relates to whether $q^*$ can be found by training stochastically over $\mathcal{Q}$, whereas $\mathcal{G}_{approx}$ relates to the the bias between $q^*$ and the true posterior.

Cremer et al. (2018) demonstrate that, in the case of sample amortized inference, when the weights $\phi$ no longer are free but the output of an encoder $f \in \mathcal{F}$, inference cannot be better than in the non-sample-amortized case, and a positive amortization gap is introduced:

$$ 
\begin{aligned}
\mathcal{G}_{\text{sa}}(\mathcal{Q}, \mathcal{F}; \eta^*) &= \mathcal{G}_{\text{approx}}(\mathcal{Q}; f(\mathbf{X}, \eta^*)) - \mathcal{G}_{\text{approx}}(\mathcal{Q}; \phi^*) \\
&\geq 0
\end{aligned}
\qquad (\text{A.16}) \\
\mathcal{G}_{\text{sample amortized VI}} = \mathcal{G}_{\text{approx}} + \mathcal{G}_{\text{sa}}
$$

Where we denote as $\eta^*$ the optimal weights for the encoder $f$ inside the function family $\mathcal{F}$. The gap terminology can be interpreted as follow: "theoretically, sample amortization cannot be beneficial in terms of KL divergence for the inference over a given sample $\mathbf{X}$."

Using the same gap terminology, we can define gaps implied by our PAVI architecture. Instead of picking the distribution $q$ inside the family $\mathcal{Q}$, consider picking $q$ from the plate-amortized family $\mathcal{Q}_{pa}$ corresponding to $\mathcal{Q}$. Distributions in $\mathcal{Q}_{pa}$ are distributions from $\mathcal{Q}$ with the additional constraints that some weights have to be equal. Consequently, $\mathcal{Q}_{pa}$ is a subset of $\mathcal{Q}$:

$$ \mathcal{Q}_{pa} \subset \mathcal{Q} \qquad (\text{A.17}) $$

As such, looking for the optimal distribution inside $\mathcal{Q}_{pa}$ instead of inside $\mathcal{Q}$ cannot result in better performance, leading to a plate amortization gap:

$$ 
\begin{aligned}
\mathcal{G}_{\text{pa}}(\mathcal{Q}, \mathcal{Q}_{\text{pa}}; \psi^*, \phi^*) &= \mathcal{G}_{\text{approx}}(\mathcal{Q}_{\text{pa}}; \psi^*) - \mathcal{G}_{\text{approx}}(\mathcal{Q}; \phi^*) \\
&\geq 0
\end{aligned}
\qquad (\text{A.18}) \\
\mathcal{G}_{\text{PAVI-F}} = \mathcal{G}_{\text{approx}} + \mathcal{G}_{\text{pa}}
$$

Where we denote as $\psi^*$ the optimal weights for a variational distribution $q$ inside $\mathcal{Q}_{pa}$ —in the KL sense. The equation A.18 is valid for the PAVI-F scheme —see Section 4.3. We can interpret it as follow: "theoretically, plate amortization cannot be beneficial in terms of KL divergence for the inference over a given sample $\mathbf{X}$".

Now consider that encodings are no longer free parameters but the output of an encoder $f$. Similar to the case presented in Equation (A.16), using an encoder cannot result in better performance, leading to an encoder gap:

$$ 
\begin{aligned}
\mathcal{G}_{\text{encoder}}(\mathcal{Q}_{\text{pa}}, \mathcal{F}; \psi^*, \eta^*) &= \mathcal{G}_{\text{approx}}(\mathcal{Q}_{\text{pa}}; f(\mathbf{X}, \eta^*)) - \mathcal{G}_{\text{approx}}(\mathcal{Q}_{\text{pa}}; \psi^*) \\
&\geq 0
\end{aligned}
\qquad (\text{A.19}) \\
\mathcal{G}_{\text{PAVI-E}} = \mathcal{G}_{\text{approx}} + \mathcal{G}_{\text{pa}} + \mathcal{G}_{\text{encoder}}
$$

The equation Equation (A.19) is valid for the PAVI-E scheme —see Section 4.3.

The most complex case is the PAVI-E(sa) scheme, where we combine both plate and sample amortization. Our argument cannot account for the resulting $\mathcal{G}_{PAVI-E(sa)}$ gap: both the PAVI-E and PAVI-E(sa) schemes rely upon the same encoder $f$. In the PAVI-E scheme, $f$ is overfitting over a dataset composed of the slices of a given data sample $\mathbf{X}$. In the PAVI-E(sa) scheme, the encoder is trained over the whole distribution of the samples of the reduced model $\mathcal{M}^r$. Intuitively, it is likely that the performance of PAVI-E(sa) will always be dominated by the performance of PAVI-E, but —as far as we understand it— the gap terminology cannot account for this discrepancy.

Comparing previous equations, we therefore have:

$$ \mathcal{G}_{\text{vanilla VI}} \leq \mathcal{G}_{\text{PAVI-F}} \leq \mathcal{G}_{\text{PAVI-E}} \qquad (\text{A.20}) $$

Note that those are theoretical results, that do not necessarily pertain to optimization in practice. In particular, in Section 5.1&5.3, this theoretical performance loss is not observed empirically over the studied
examples. On the contrary, in practice, our results can actually be better than non-amortized baselines, as is the case for the PAVI-F scheme in Figure 3 or experiments B.4. We interpret this as a result of a simplified optimization problem due to plate amortization—with fewer parameters to optimize for, and mini-batching effects across different ground RVs. A 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 $G_{pa}, G_{encoder}$ are likely to be well "compensated for" by the lighter parameterization and faster convergence entitled by plate amortization.

# B Supplemental results

In this section, we present various supplemental experiments that further compare the PAVI-E and -F between themselves and against baselines. Overall, we found the PAVI-F scheme faster to train and to yield better inference quality than the PAVI-E scheme. When its parameterization is affordable, PAVI-F should be preferred. PAVI-E nevertheless opens up promising research directions, with the potential for parameterization-constant, time-constant sample-amortized inference as the problem cardinality augments. Though degraded compared to PAVI-F's, PAVI-E's performance is still on par or beats baselines in a variety of inference tasks —see experiments 5.3&B.4.

## B.1 Tabular results for experiment 5.3

For convenience, we reproduce the results of our comparative experiment from Figure 3 in Table B.1.

## B.2 GRE results sanity check

As exposed in the introduction of Section 5, in this work we focused on the usage of the ELBO as an inference performance metric (Blei et al., 2017):

$$ \mathrm{ELBO}(q) = \log p(X) - \mathrm{KL}(q(\Theta) \| p(\Theta|X)) \qquad (B.21) $$

Given that the likelihood term $\log p(X)$ does not depend on the variational family $q$, differences in ELBOs directly transfer to differences in KL divergence, and provide a straightforward metric to compare different variational posteriors. Nonetheless, the ELBO doesn't provide an absolute metric of quality. As a sanity check, we want to assert the quality of the results presented in Section 5.3—that are transferable to Section 5.1&5.2, based on the same model. In Figure B.1 we plot the posterior samples of various methods against approximate closed-form ground truths, using the $Card(P_1) = 20$ case. All the method's results are aligned with the closed-form ground truth, with differences in ELBO translating meaningful qualitative differences in terms of inference quality.

## B.3 Effect of the reduced model cardinalities on the training efficiency

In Figure B.2 we show the impact of the augmentation of $Card(P_1)$ on the efficiency of the variational posterior'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 makes the ELBO gradient estimates less noisy and allows to train more encoding vectors $E_{i,n}$ at a given optimization step.

In the case of the PAVI-E scheme, the training speed is constant with respect to $Card(P_1)$. This is due to the computing of the encodings being vectorized across plates in our deep-set encoder —see Appendix A.1.3 (Lee et al., 2019). We observe a slight if barely noticeable reduction of the inference bias when augmenting $Card(P_1)$. Similar to the PAVI-F scheme, $Card(P_1)$ should be maximized with respect to the GPU memory constraint. The intuition behind this choice is that the generalization of the learning of summary statistics
examples. On the contrary, in practice, our results can actually be better than non-amortized baselines, as is the case for the PAVI-F scheme in Figure 3 or experiments B.4. We interpret this as a result of a simplified optimization problem due to plate amortization—with fewer parameters to optimize for, and mini-batching effects across different ground RVs. A 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 $G_{pa}, G_{encoder}$ are likely to be well "compensated for" by the lighter parameterization and faster convergence entitled by plate amortization.

# B Supplemental results

In this section, we present various supplemental experiments that further compare the PAVI-E and -F between themselves and against baselines. Overall, we found the PAVI-F scheme faster to train and to yield better inference quality than the PAVI-E scheme. When its parameterization is affordable, PAVI-F should be preferred. PAVI-E nevertheless opens up promising research directions, with the potential for parameterization-constant, time-constant sample-amortized inference as the problem cardinality augments. Though degraded compared to PAVI-F's, PAVI-E's performance is still on par or beats baselines in a variety of inference tasks —see experiments 5.3&B.4.

## B.1 Tabular results for experiment 5.3

For convenience, we reproduce the results of our comparative experiment from Figure 3 in Table B.1.

## B.2 GRE results sanity check

As exposed in the introduction of Section 5, in this work we focused on the usage of the ELBO as an inference performance metric (Blei et al., 2017):

$$ \text{ELBO}(q) = \log p(X) - \text{KL}(q(\Theta) \| p(\Theta|X)) \quad (B.21) $$

Given that the likelihood term $\log p(X)$ does not depend on the variational family $q$, differences in ELBOs directly transfer to differences in KL divergence, and provide a straightforward metric to compare different variational posteriors. Nonetheless, the ELBO doesn't provide an absolute metric of quality. As a sanity check, we want to assert the quality of the results presented in Section 5.3—that are transferable to Section 5.1&5.2, based on the same model. In Figure B.1 we plot the posterior samples of various methods against approximate closed-form ground truths, using the $Card(P_1) = 20$ case. All the method's results are aligned with the closed-form ground truth, with differences in ELBO translating meaningful qualitative differences in terms of inference quality.

## B.3 Effect of the reduced model cardinalities on the training efficiency

In Figure B.2 we show the impact of the augmentation of $Card(P_1)$ on the efficiency of the variational posterior'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 makes the ELBO gradient estimates less noisy and allows to train more encoding vectors $E_{i,n}$ at a given optimization step.

In the case of the PAVI-E scheme, the training speed is constant with respect to $Card(P_1)$. This is due to the computing of the encodings being vectorized across plates in our deep-set encoder —see Appendix A.1.3 (Lee et al., 2019). We observe a slight if barely noticeable reduction of the inference bias when augmenting $Card(P_1)$. Similar to the PAVI-F scheme, $Card(P_1)$ should be maximized with respect to the GPU memory constraint. The intuition behind this choice is that the generalization of the learning of summary statistics
<table><caption>Table B.1: Scaling with plate cardinalities Reproduction of the results from Figure 3</caption><thead><tr><th></th><th>Architecture</th><th>ELBO (10<sup>1</sup>)</th><th>Parameterization</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td rowspan="7">Card(P<sub>1</sub>) = 2</td><td>ADAVI (sa)</td><td>1.88 (± 0.21)</td><td>12,280</td><td>503 (± 50)</td></tr><tr><td>CF (sa)</td><td>1.48 (± 0.05)</td><td>2,275</td><td>228 (± 23)</td></tr><tr><td>CF</td><td>1.98 (± 0.02)</td><td>355</td><td>63 (± 6)</td></tr><tr><td>UIVI</td><td>2.11 (± 0.05)</td><td>2,764</td><td>220 (± 22)</td></tr><tr><td>PAVI-F</td><td>2.14 (± 0.02)</td><td>4,058</td><td>65 (± 6)</td></tr><tr><td>PAVI-E (sa)</td><td>2.03 (± 0.03)</td><td>27,394</td><td>140 (± 14)</td></tr><tr><td>PAVI-E</td><td>1.74 (± 0.18)</td><td>27,394</td><td>37 (± 4)</td></tr><tr><td></td><td>Asymptotic ELBO</td><td>2.20</td><td></td><td></td></tr><tr><td rowspan="7">Card(P<sub>1</sub>) = 20</td><td>ADAVI (sa)</td><td>21.6 (± 1.36)</td><td>12,280</td><td>800 (± 80)</td></tr><tr><td>CF (sa)</td><td>17.3 (± 0.60)</td><td>21,751</td><td>1,612 (± 161)</td></tr><tr><td>CF</td><td>22.6 (± 0.01)</td><td>2,551</td><td>404 (± 40)</td></tr><tr><td>UIVI</td><td>23.9 (± 0.02)</td><td>14,676</td><td>230 (± 23)</td></tr><tr><td>PAVI-F</td><td>23.8 (± 0.02)</td><td>4,202</td><td>72 (± 7)</td></tr><tr><td>PAVI-E (sa)</td><td>22.3 (± 0.30)</td><td>27,394</td><td>145 (± 14)</td></tr><tr><td>PAVI-E</td><td>18.4 (± 2.63)</td><td>27,394</td><td>50 (± 5)</td></tr><tr><td></td><td>Asymptotic ELBO</td><td>24.1</td><td></td><td></td></tr><tr><td rowspan="7">Card(P<sub>1</sub>) = 200</td><td>ADAVI (sa)</td><td>221.7 (± 10.93)</td><td>12,280</td><td>1,600 (± 160)</td></tr><tr><td>CF (sa)</td><td>175.7 (± 0.52)</td><td>216,511</td><td>12,000 (± 1,200)</td></tr><tr><td>CF</td><td>228.6 (± 0.52)</td><td>24,511</td><td>2,600 (± 260)</td></tr><tr><td>UIVI</td><td>241.2 (± 0.14)</td><td>139,300</td><td>240 (± 24)</td></tr><tr><td>PAVI-F</td><td>237.7 (± 0.95)</td><td>5,642</td><td>81 (± 8)</td></tr><tr><td>PAVI-E (sa)</td><td>227.0 (± 4.76)</td><td>27,394</td><td>150 (± 15)</td></tr><tr><td>PAVI-E</td><td>224.2 (± 15.34)</td><td>27,394</td><td>85 (± 9)</td></tr><tr><td></td><td>Asymptotic ELBO</td><td>243.4</td><td></td><td></td></tr></tbody></table>
<table><caption>Table B.1: Scaling with plate cardinalities Reproduction of the results from Figure 3</caption><thead><tr><th></th><th>Architecture</th><th>ELBO (10<sup>1</sup>)</th><th>Parameterization</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td rowspan="7">Card(P<sub>1</sub>) = 2</td><td>ADAVI (sa)</td><td>1.88 (± 0.21)</td><td>12,280</td><td>503 (± 50)</td></tr><tr><td>CF (sa)</td><td>1.48 (± 0.05)</td><td>2,275</td><td>228 (± 23)</td></tr><tr><td>CF</td><td>1.98 (± 0.02)</td><td>355</td><td>63 (± 6)</td></tr><tr><td>UIVI</td><td>2.11 (± 0.05)</td><td>2,764</td><td>220 (± 22)</td></tr><tr><td>PAVI-F</td><td>2.14 (± 0.02)</td><td>4,058</td><td>65 (± 6)</td></tr><tr><td>PAVI-E (sa)</td><td>2.03 (± 0.03)</td><td>27,394</td><td>140 (± 14)</td></tr><tr><td>PAVI-E</td><td>1.74 (± 0.18)</td><td>27,394</td><td>37 (± 4)</td></tr><tr><td></td><td>Asymptotic ELBO</td><td>2.20</td><td></td><td></td></tr><tr><td rowspan="7">Card(P<sub>1</sub>) = 20</td><td>ADAVI (sa)</td><td>21.6 (± 1.36)</td><td>12,280</td><td>800 (± 80)</td></tr><tr><td>CF (sa)</td><td>17.3 (± 0.60)</td><td>21,751</td><td>1,612 (± 161)</td></tr><tr><td>CF</td><td>22.6 (± 0.01)</td><td>2,551</td><td>404 (± 40)</td></tr><tr><td>UIVI</td><td>23.9 (± 0.02)</td><td>14,676</td><td>230 (± 23)</td></tr><tr><td>PAVI-F</td><td>23.8 (± 0.02)</td><td>4,202</td><td>72 (± 7)</td></tr><tr><td>PAVI-E (sa)</td><td>22.3 (± 0.30)</td><td>27,394</td><td>145 (± 14)</td></tr><tr><td>PAVI-E</td><td>18.4 (± 2.63)</td><td>27,394</td><td>50 (± 5)</td></tr><tr><td></td><td>Asymptotic ELBO</td><td>24.1</td><td></td><td></td></tr><tr><td rowspan="7">Card(P<sub>1</sub>) = 200</td><td>ADAVI (sa)</td><td>221.7 (± 10.93)</td><td>12,280</td><td>1,600 (± 160)</td></tr><tr><td>CF (sa)</td><td>175.7 (± 0.52)</td><td>216,511</td><td>12,000 (± 1,200)</td></tr><tr><td>CF</td><td>228.6 (± 0.52)</td><td>24,511</td><td>2,600 (± 260)</td></tr><tr><td>UIVI</td><td>241.2 (± 0.14)</td><td>139,300</td><td>240 (± 24)</td></tr><tr><td>PAVI-F</td><td>237.7 (± 0.95)</td><td>5,642</td><td>81 (± 8)</td></tr><tr><td>PAVI-E (sa)</td><td>227.0 (± 4.76)</td><td>27,394</td><td>150 (± 15)</td></tr><tr><td>PAVI-E</td><td>224.2 (± 15.34)</td><td>27,394</td><td>85 (± 9)</td></tr><tr><td></td><td>Asymptotic ELBO</td><td>243.4</td><td></td><td></td></tr></tbody></table>
<figure><img src="image_6.png" /><figcaption>Figure B.1: **GRE Sanity check** Inference methods present qualitatively correct results, making ELBO comparisons relevant in our experiments. *On the topmost line*, we represent 4 different **X** samples for the GRE model described in Equation (10) with $Card(\mathcal{P}_1) = 20$. Each set of colored points represents the $\mathbf{X}_{n_1, \bullet}$ points belonging to one of the 20 groups. *Bottom lines* represent the posterior samples for the methods used in Section 5.3. Colored points are sampled from the posterior of the groups means $\theta_1$, whereas black points are samples from the population mean $\theta_2$. We represent as black circles a closed-form ground truth, centered on the correct posterior mean, and with a radius equal to 2 times the closed-form posterior's standard deviation. **Correct posterior samples should be centered on the same point as the corresponding black circle, and 95% of the points should fall within the black circle.** Associated with each posterior sample is a Wasserstein Distance (WD) computed with a sample from an approximate closed-form posterior (Flamary et al., 2021). PAVI is represented on the 3 last lines. Some minor bias can be observed for the PAVI-E scheme, but this approximation error is marginal compared to the optimization error that can be observed for unbiased methods, such as CF(sa) (Bottou & Bousquet, 2007). We can observe a superior quality for the PAVI-F scheme, rivaling ADAVI and CF's performance with orders of magnitude fewer parameters and training time, as visible in Figure 3.</figcaption></figure>
<figure><img src="image_7.png" /><figcaption>Figure B.1: **GRE Sanity check** Inference methods present qualitatively correct results, making ELBO comparisons relevant in our experiments. *On the topmost line*, we represent 4 different **X** samples for the GRE model described in Equation (10) with $Card(\mathcal{P}_1) = 20$. Each set of colored points represents the $\mathbf{X}_{n_1, \bullet}$ points belonging to one of the 20 groups. *Bottom lines* represent the posterior samples for the methods used in Section 5.3. Colored points are sampled from the posterior of the groups means $\theta_1$, whereas black points are samples from the population mean $\theta_2$. We represent as black circles a closed-form ground truth, centered on the correct posterior mean, and with a radius equal to 2 times the closed-form posterior's standard deviation. **Correct posterior samples should be centered on the same point as the corresponding black circle, and 95% of the points should fall within the black circle.** Associated with each posterior sample is a Wasserstein Distance (WD) computed with a sample from an approximate closed-form posterior (Flamary et al., 2021). PAVI is represented on the 3 last lines. Some minor bias can be observed for the PAVI-E scheme, but this approximation error is marginal compared to the optimization error that can be observed for unbiased methods, such as CF(sa) (Bottou & Bousquet, 2007). We can observe a superior quality for the PAVI-F scheme, rivaling ADAVI and CF's performance with orders of magnitude fewer parameters and training time, as visible in Figure 3.</figcaption></figure>
<figure><img src="image_7.png" /><figcaption>Figure B.2: **Both panels: Effect of $Card^r(P_1)$ on the training efficiency** Experiment performed on the GRE model (see Equation (10)). We increase the cardinality of the reduced model $Card^r(P_1)$ from 1 to 20 while keeping $Card(P_1) = 20$ fixed. At $Card^r(P_1) = 20$, there is no stochasticity in the training, meaning we train directly on $\mathcal{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^r(P_1)$ augments, the training gets faster and less noisy, likely due to less stochasticity in the gradient estimates and more encodings vectors $E_{i,n}$ being trained at once. This increase in speed quickly caps, and the speed is approximately the same between $Card^r(P_1) = 8$ and $Card^r(P_1) = 16$. This experiment also illustrates the unbiasedness of the stochastic training: at $Card^r(P_1) = 20$ there is no stochasticity in the training, and the asymptotic performance is the same as for the stochastic training (see Appendix A.2). **Right panel: PAVI-E scheme** The case of $Card^r(P_1) = 1$ is pathological: the encoder "learns" to collect summary statistics across a set of 1 element, and —not surprisingly— the learned function doesn’t generalize well on sets of size 20. In all of the other cases, the ELBO converges approximately to the same values as with $Card^r(P_1) = 20$, that is to say when there is no stochasticity in the training (black curve). We furthermore observe a slight reduction of the inference bias as $Card^r(P_1)$ augments. This illustrates how the theoretical bias of the PAVI-E scheme identified in Appendix A.2 does not translate in significantly worse empirical results —albeit in the pathological case of $Card^r(P_1) = 1$. Interestingly, the training speed is approximately constant across all values of $Card^r(P_1)$. This feature is essential in the PAVI-E scheme: we can train over a reduced version of our model—with a light memory footprint—and apply the obtained architecture to the full model. Note that this feature wouldn’t necessarily be present when training on CPU: computing summary statistics over larger sets would make the training *slower* when $Card^r(P_1)$ augments.</figcaption></figure>
<figure><img src="image_8.png" /><figcaption>Figure B.2: **Both panels: Effect of $Card^r(P_1)$ on the training efficiency** Experiment performed on the GRE model (see Equation (10)). We increase the cardinality of the reduced model $Card^r(P_1)$ from 1 to 20 while keeping $Card(P_1) = 20$ fixed. At $Card^r(P_1) = 20$, there is no stochasticity in the training, meaning we train directly on $\mathcal{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^r(P_1)$ augments, the training gets faster and less noisy, likely due to less stochasticity in the gradient estimates and more encodings vectors $E_{i,n}$ being trained at once. This increase in speed quickly caps, and the speed is approximately the same between $Card^r(P_1) = 8$ and $Card^r(P_1) = 16$. This experiment also illustrates the unbiasedness of the stochastic training: at $Card^r(P_1) = 20$ there is no stochasticity in the training, and the asymptotic performance is the same as for the stochastic training (see Appendix A.2). **Right panel: PAVI-E scheme** The case of $Card^r(P_1) = 1$ is pathological: the encoder "learns" to collect summary statistics across a set of 1 element, and —not surprisingly— the learned function doesn’t generalize well on sets of size 20. In all of the other cases, the ELBO converges approximately to the same values as with $Card^r(P_1) = 20$, that is to say when there is no stochasticity in the training (black curve). We furthermore observe a slight reduction of the inference bias as $Card^r(P_1)$ augments. This illustrates how the theoretical bias of the PAVI-E scheme identified in Appendix A.2 does not translate in significantly worse empirical results —albeit in the pathological case of $Card^r(P_1) = 1$. Interestingly, the training speed is approximately constant across all values of $Card^r(P_1)$. This feature is essential in the PAVI-E scheme: we can train over a reduced version of our model—with a light memory footprint—and apply the obtained architecture to the full model. Note that this feature wouldn’t necessarily be present when training on CPU: computing summary statistics over larger sets would make the training *slower* when $Card^r(P_1)$ augments.</figcaption></figure>
<table><caption>Table B.2: Analytical performance over a Gaussian mixture HBM PAVI shows superior performance. PAVI converges in a fraction of CF's optimization time —the second-best-performing architecture.</caption><thead><tr><th>Architecture</th><th>ELBO (10<sup>1</sup>)</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td>CF</td><td>- 16.2 (± 2.9)</td><td>3,100</td></tr><tr><td>ADAVI</td><td>- 20.0 (± 2.3)</td><td>3,000</td></tr><tr><td>UIVI</td><td>- 26.3 (± 1.2)</td><td>500</td></tr><tr><td>PAVI-E (ours)</td><td>- 15.2 (±2.8)</td><td>470</td></tr><tr><td>PAVI-F (ours)</td><td>- 11.1 (±1.1)</td><td>300</td></tr></tbody></table>

across sets of data points is easier the closer the reduced set size is to the full set size. This constant training speed allows for a controlled memory footprint of the stochastic training: contrary to the PAVI-F scheme, $Card^r(P_1)$ could be maintained at a value manageable by the GPU without reducing the training speed. This example also displays a pathological choice for $Card^r(P_1)$: the encoder fails to learn to compute the correct summary statistics over sets of size 1. However trivial, this example underlines that $Card^r(P_1)$ should be chosen at a value logical with respect to the posterior's sufficient statistics. When generalizing to moments of a higher order —such as the variance— there is therefore a theoretical lower bound to consider when fixating the value of $Card^r(P_1)$. Our intuition is that the more complex to estimate the statistic is, the larger $Card^r(P_1)$ 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 Figure 4, stochastic training is necessary. Our claim is to be faster than non-plate-amortized stochastic VI in those large-scale contexts —but not necessarily to be faster than non-stochastic VI in the small-scale regime of this experiment. For PAVI-F we nonetheless obtain similar training speed compared to non-stochastic VI with $Card^r(P_1)$ as small as 8. For PAVI-E, the stochastic training is as fast as the non-stochastic one.

## B.4 Additional comparative experiments

### B.4.1 Gaussian Mixture model

In this experiment, we test out various baselines over a challenging model: a Gaussian mixture with random effects.

$$ 
\begin{array}{l}
D, \ Card(P_1), \ Card(P_0) = 2, 20, 10 \\
\forall n_1=1.. \ Card(P_1) \quad X_{n_1,n_0} | \theta_{1,n_1}, \pi_{n_1} \sim \text{Mixture}([\mathcal{N}(\theta_{1,n_1}^1, \sigma_x^2), \dots, \mathcal{N}(\theta_{1,n_1}^L, \sigma_x^2)], \pi_{n_1}) \\
\forall n_1 = 1.. \ Card(P_1) \quad \pi_{n_1} \sim \text{Dirichlet}(1 \times \vec{1}_L) \qquad (B.22) \\
\forall n_1=1.. \ Card(P_1) \quad \theta_{1,n_1}^l | \theta_{2,0}^l \sim \mathcal{N}(\theta_{2,0}^l, \sigma_1^2) \\
\forall l = 1..L \quad \theta_{2,0}^l \sim \mathcal{N}(\vec{0}_D, \sigma_2^2),
\end{array}
$$

where Mixture($[D_1, \dots, D_L]$, $\pi$) denotes a mixture between the distributions $[D_1, \dots, D_L]$ with mixture weights $\pi$. Results are visible in Table B.2, where PAVI displays the best asymptotic ELBO, as well as the shortest optimization time. On that note, we underline that stochastic training over a mixture distribution is challenging, as —due to the sub-sampling of points— only a fraction of the mixture components could be expressed at a given step, requiring the architecture to dynamically cluster the data points across time.

### B.4.2 Hierarchical Variance model

In this experiment, we test out our architectures over a non-canonical model, in which parent RVs play the role of variance for the distribution of their children. Our goal is in particular to evaluate a potential empirical bias for the PAVI-E scheme —as studied in Appendix A.2. We will refer to the following HBM as
<table><caption>Table B.2: Analytical performance over a Gaussian mixture HBM PAVI shows superior performance. PAVI converges in a fraction of CF's optimization time —the second-best-performing architecture.</caption><thead><tr><th>Architecture</th><th>ELBO (10<sup>1</sup>)</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td>CF</td><td>- 16.2 (± 2.9)</td><td>3,100</td></tr><tr><td>ADAVI</td><td>- 20.0 (± 2.3)</td><td>3,000</td></tr><tr><td>UIVI</td><td>- 26.3 (± 1.2)</td><td>500</td></tr><tr><td>PAVI-E (ours)</td><td>- 15.2 (±2.8)</td><td>470</td></tr><tr><td>PAVI-F (ours)</td><td>- 11.1 (±1.1)</td><td>300</td></tr></tbody></table>

across sets of data points is easier the closer the reduced set size is to the full set size. This constant training speed allows for a controlled memory footprint of the stochastic training: contrary to the PAVI-F scheme, $Card^r(P_1)$ could be maintained at a value manageable by the GPU without reducing the training speed. This example also displays a pathological choice for $Card^r(P_1)$: the encoder fails to learn to compute the correct summary statistics over sets of size 1. However trivial, this example underlines that $Card^r(P_1)$ should be chosen at a value logical with respect to the posterior's sufficient statistics. When generalizing to moments of a higher order —such as the variance— there is therefore a theoretical lower bound to consider when fixating the value of $Card^r(P_1)$. Our intuition is that the more complex to estimate the statistic is, the larger $Card^r(P_1)$ 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 Figure 4, stochastic training is necessary. Our claim is to be faster than non-plate-amortized stochastic VI in those large-scale contexts —but not necessarily to be faster than non-stochastic VI in the small-scale regime of this experiment. For PAVI-F we nonetheless obtain similar training speed compared to non-stochastic VI with $Card^r(P_1)$ as small as 8. For PAVI-E, the stochastic training is as fast as the non-stochastic one.

## B.4 Additional comparative experiments

### B.4.1 Gaussian Mixture model

In this experiment, we test out various baselines over a challenging model: a Gaussian mixture with random effects.

$$ 
\begin{array}{l}
D, \ Card(P_1), \ Card(P_0) = 2, \ 20, \ 10 \\
\forall n_1=1.. \ Card(P_1) \quad X_{n_1,n_0} | \theta_{1,n_1}, \pi_{n_1} \sim \text{Mixture}([\mathcal{N}(\theta_{1,n_1}^1, \sigma_x^2), \dots, \mathcal{N}(\theta_{1,n_1}^L, \sigma_x^2)], \pi_{n_1}) \\
\forall n_1 = 1.. \ Card(P_1) \quad \pi_{n_1} \sim \text{Dirichlet}(1 \times \vec{1}_L) \qquad (B.22) \\
\forall n_1=1.. \ Card(P_1) \quad \theta_{1,n_1}^l | \theta_{2,0}^l \sim \mathcal{N}(\theta_{2,0}^l, \sigma_1^2) \\
\forall l = 1..L \quad \theta_{2,0}^l \sim \mathcal{N}(\vec{0}_D, \sigma_2^2),
\end{array}
$$

where Mixture($[D_1, \dots, D_L]$, $\pi$) denotes a mixture between the distributions $[D_1, \dots, D_L]$ with mixture weights $\pi$. Results are visible in Table B.2, where PAVI displays the best asymptotic ELBO, as well as the shortest optimization time. On that note, we underline that stochastic training over a mixture distribution is challenging, as —due to the sub-sampling of points— only a fraction of the mixture components could be expressed at a given step, requiring the architecture to dynamically cluster the data points across time.

### B.4.2 Hierarchical Variance model

In this experiment, we test out our architectures over a non-canonical model, in which parent RVs play the role of variance for the distribution of their children. Our goal is in particular to evaluate a potential empirical bias for the PAVI-E scheme —as studied in Appendix A.2. We will refer to the following HBM as
<table><caption>Table B.3: Analytical performance over a Hierarchical Variance HBM In this setting with hard-to-estimate summary statistics, PAVI-E doesn't show any empirical bias. ¹: PAVI-E suffers from numerical instability in this example, with runs degenerating into NaN results.</caption><thead><tr><th>Architecture</th><th>ELBO (10<sup>2</sup>)</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td>CF</td><td>- 11.2 (± 2.9)</td><td>500</td></tr><tr><td>UIVI</td><td>-6.7 (± 0.8)</td><td>100</td></tr><tr><td>PAVI-E (ours)</td><td>-6.7 (±1.0)<sup>1</sup></td><td>90</td></tr><tr><td>PAVI-F (ours)</td><td><b>- 6.7 (±0.8)</b></td><td><b>20</b></td></tr></tbody></table>

<table><caption>Table B.4: 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 a numerically stable UIVI optimization. PAVI yields a higher ELBO than CF, in half of CF's optimization time.</caption><thead><tr><th>Architecture</th><th>ELBO (10<sup>4</sup>)</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td>CF</td><td>- 139.7 (± 21)</td><td>3,200</td></tr><tr><td>ADAVI</td><td>—</td><td>—</td></tr><tr><td>UIVI</td><td>—</td><td>—</td></tr><tr><td>PAVI-F (ours)</td><td><b>- 8.2 (± 0.1)</b></td><td><b>1,600</b></td></tr></tbody></table>

the Hierarchical Variance model:

$$ \begin{array}{l} D, \mathrm{Card}(\mathcal{P}_1), \mathrm{Card}(\mathcal{P}_0) = 2, 15, 15 \\ \forall n_1 = 1.. \mathrm{Card}(\mathcal{P}_1) \quad \log X_{(n_1,n_0)} |\theta_{1,n_1} \sim \mathcal{N}(0, \theta_{1,n_1}) \\ \forall n_0 = 1.. \mathrm{Card}(\mathcal{P}_0) \quad \log \theta_{1,(n_1)} |\theta_{2,0} \sim \mathcal{N}(0, \theta_{0,0}) \\ \qquad \log \theta_{2,0} \sim \mathcal{N}(\vec{0}_D, 1) \end{array} \qquad (B.23) $$

In this model, the encoder—used in the PAVI-E scheme—has to collect non-trivial summary statistics: empirical variances across variable subsets of the observed data. This is the context in which we would expect to observe the most bias due to the approximation in Equation (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 a 5 times shorter optimization time.

### B.4.3 Small dimension version of our Neuroimaging model

In this experiment, we validate the performance of our architecture on synthetic data generated using a small dimension version of the model presented in Equation (B.24). To allow the use of the comparative baselines, we fixate $S, S^r = 5, 3, T, T^r = 6, 3, N, N^r = 12, 3, D = 4, L = 7$. Results are visible in Table B.4: 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, we did not manage to obtain a 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 of the prior $p$ to help with inference. PAVI's plate amortization likely facilitates the inference, resulting in a reduced optimization error (Bottou & Bousquet, 2007).
<table><caption>Table B.3: Analytical performance over a Hierarchical Variance HBM In this setting with hard-to-estimate summary statistics, PAVI-E doesn't show any empirical bias. ¹: PAVI-E suffers from numerical instability in this example, with runs degenerating into NaN results.</caption><thead><tr><th>Architecture</th><th>ELBO (10<sup>2</sup>)</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td>CF</td><td>- 11.2 (± 2.9)</td><td>500</td></tr><tr><td>UIVI</td><td>-6.7 (± 0.8)</td><td>100</td></tr><tr><td>PAVI-E (ours)</td><td>-6.7 (±1.0)<sup>1</sup></td><td>90</td></tr><tr><td>PAVI-F (ours)</td><td><b>- 6.7 (±0.8)</b></td><td><b>20</b></td></tr></tbody></table>

<table><caption>Table B.4: 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 a numerically stable UIVI optimization. PAVI yields a higher ELBO than CF, in half of CF's optimization time.</caption><thead><tr><th>Architecture</th><th>ELBO (10<sup>4</sup>)</th><th>Optimization time (s)</th></tr></thead><tbody><tr><td>CF</td><td>- 139.7 (± 21)</td><td>3,200</td></tr><tr><td>ADAVI</td><td>—</td><td>—</td></tr><tr><td>UIVI</td><td>—</td><td>—</td></tr><tr><td>PAVI-F (ours)</td><td><b>- 8.2 (± 0.1)</b></td><td><b>1,600</b></td></tr></tbody></table>

the Hierarchical Variance model:

$$ \begin{array}{l} D, \mathrm{Card}(\mathcal{P}_1), \mathrm{Card}(\mathcal{P}_0) = 2, 15, 15 \\ \forall n_1 = 1.. \mathrm{Card}(\mathcal{P}_1) \quad \log X_{(n_1,n_0)} |\theta_{1,n_1} \sim \mathcal{N}(0, \theta_{1,n_1}) \\ \forall n_0 = 1.. \mathrm{Card}(\mathcal{P}_0) \quad \log \theta_{1,(n_1)} |\theta_{2,0} \sim \mathcal{N}(0, \theta_{0,0}) \\ \qquad \log \theta_{2,0} \sim \mathcal{N}(\vec{0}_D, 1) \end{array} \qquad (B.23) $$

In this model, the encoder—used in the PAVI-E scheme—has to collect non-trivial summary statistics: empirical variances across variable subsets of the observed data. This is the context in which we would expect to observe the most bias due to the approximation in Equation (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 a 5 times shorter optimization time.

### B.4.3 Small dimension version of our Neuroimaging model

In this experiment, we validate the performance of our architecture on synthetic data generated using a small dimension version of the model presented in Equation (B.24). To allow the use of the comparative baselines, we fixate $S, S^r = 5, 3, T, T^r = 6, 3, N, N^r = 12, 3, D = 4, L = 7$. Results are visible in Table B.4: 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, we did not manage to obtain a 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 of the prior $p$ to help with inference. PAVI's plate amortization likely facilitates the inference, resulting in a reduced optimization error (Bottou & Bousquet, 2007).
## B.5 Experimental details - analytical examples

All experiments were performed in Python, using the *Tensorflow Probability* library (Dillon et al., 2017). Through this section, we refer to *Masked Autoregressive Flows* (Papamakarios et al., 2018) as MAF. All experiments are performed using the Adam optimizer (Kingma & Ba, 2015). At training, the ELBO was estimated using a Monte Carlo procedure with 8 samples. All architectures were evaluated over a fixed set of 20 samples X, with 5 seeds per sample. Non-sample-amortized architectures 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.

**Termination condition** Our termination condition amounts to plateau detection across the statistics of multiple runs. In details, we ran the inference with a large, conservative number of epochs for all methods and collected the evolution of the ELBO for each run. We then took as runtime the time to reach an ELBO close to the one at convergence for the method. We averaged over 100 repetitions to avoid stochastic noise due to Monte Carlo estimation of the ELBO. The order of magnitude difference between the results is clearly above the potential noise effects due to our procedure.

### B.5.1 Plate amortization and convergence speed (5.1)

All 3 architectures (baseline, PAVI-F, PAVI-E) used:

* for the flows $F_i$, a MAF with [32, 32] hidden units;

* as encoding size, 128

For the encoder $f$ in the PAVI-E scheme, we used a multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 64 inducing points.

### B.5.2 Impact of encoding size (5.2)

All architectures used:

* for the flows $F_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* as encoding size, a value varying from 2 to 16

### B.5.3 Scaling with plate cardinalities (5.3)

**ADAVI** (Rouillard & Wassermann, 2022) we used:

* for the flows $F_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* for the encoder, an encoding size of 8 with a multi-head architecture with 2 heads of 4 units each, 2 ISAB blocks with 32 inducing points.

**Cascading Flows** (Ambrogioni et al., 2021b) we used:

* a mean-field distribution over the auxiliary variables $r$

* as auxiliary size, a fixed value of 8

* as flows, *Highway Flows* as designed by the Cascading Flows authors

**PAVI-F** we used:
## B.5 Experimental details - analytical examples

All experiments were performed in Python, using the *Tensorflow Probability* library (Dillon et al., 2017). Through this section, we refer to *Masked Autoregressive Flows* (Papamakarios et al., 2018) as MAF. All experiments are performed using the Adam optimizer (Kingma & Ba, 2015). At training, the ELBO was estimated using a Monte Carlo procedure with 8 samples. All architectures were evaluated over a fixed set of 20 samples X, with 5 seeds per sample. Non-sample-amortized architectures 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.

**Termination condition** Our termination condition amounts to plateau detection across the statistics of multiple runs. In details, we ran the inference with a large, conservative number of epochs for all methods and collected the evolution of the ELBO for each run. We then took as runtime the time to reach an ELBO close to the one at convergence for the method. We averaged over 100 repetitions to avoid stochastic noise due to Monte Carlo estimation of the ELBO. The order of magnitude difference between the results is clearly above the potential noise effects due to our procedure.

### B.5.1 Plate amortization and convergence speed (5.1)

All 3 architectures (baseline, PAVI-F, PAVI-E) used:

* for the flows $F_i$, a MAF with [32, 32] hidden units;

* as encoding size, 128

For the encoder *f* in the PAVI-E scheme, we used a multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 64 inducing points.

### B.5.2 Impact of encoding size (5.2)

All architectures used:

* for the flows $F_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* as encoding size, a value varying from 2 to 16

### B.5.3 Scaling with plate cardinalities (5.3)

**ADAVI** (Rouillard & Wassermann, 2022) we used:

* for the flows $F_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* for the encoder, an encoding size of 8 with a multi-head architecture with 2 heads of 4 units each, 2 ISAB blocks with 32 inducing points.

**Cascading Flows** (Ambrogioni et al., 2021b) we used:

* a mean-field distribution over the auxiliary variables *r*

* as auxiliary size, a fixed value of 8

* as flows, *Highway Flows* as designed by the Cascading Flows authors

**PAVI-F** we used:
* for the flows $\mathcal{F}_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* an encoding size of 8

**PAVI-E** we used:

* for the flows $\mathcal{F}_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads of 8 units each, 2 ISAB blocks with 64 inducing points.

**UIVI** we used as $Card(\mathcal{P}_1) = 2 \to 20 \to 200$:

* as base distribution, a standard Gaussian with dimensionality $6 \to 42 \to 402$

* as transform $h$, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality $3 \to 6 \to 9$

* as transform weights regressor, an MLP with hidden units [32, 32] $\to$ [64, 64] $\to$ [128, 128]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-2$, $\times 0.9$ every 300 steps

### B.5.4 Gaussian mixture (B.4.1)

**ADAVI** (Rouillard & Wassermann, 2022) we used:

* for the flows $\mathcal{F}_i$, a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.

* for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads of 8 units each, 2 ISAB blocks with 8 inducing points.

**Cascading Flows** (Ambrogioni et al., 2021b) we used:

* a 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 $\mathcal{F}_i$, a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.

* an encoding size of 16

* $Card(\mathcal{P}_1) = 5$

**PAVI-E** we used:

* for the flows $\mathcal{F}_i$, a MAF with [128, 128] hidden units, after an affine block with triangular scaling matrix.
* for the flows $\mathcal{F}_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* an encoding size of 8

**PAVI-E** we used:

* for the flows $\mathcal{F}_i$, a MAF with [32, 32] hidden units, after an affine block with triangular scaling matrix.

* for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads of 8 units each, 2 ISAB blocks with 64 inducing points.

**UIVI** we used as $Card(\mathcal{P}_1) = 2 \to 20 \to 200$:

* as base distribution, a standard Gaussian with dimensionality $6 \to 42 \to 402$

* as transform $h$, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality $3 \to 6 \to 9$

* as transform weights regressor, an MLP with hidden units [32, 32] $\to$ [64, 64] $\to$ [128, 128]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-2$, $\times 0.9$ every 300 steps

### B.5.4 Gaussian mixture (B.4.1)

**ADAVI** (Rouillard & Wassermann, 2022) we used:

* for the flows $\mathcal{F}_i$, a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.

* for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads of 8 units each, 2 ISAB blocks with 8 inducing points.

**Cascading Flows** (Ambrogioni et al., 2021b) we used:

* a 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 $\mathcal{F}_i$, a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.

* an encoding size of 16

* $Card(\mathcal{P}_1) = 5$

**PAVI-E** we used:

* for the flows $\mathcal{F}_i$, a MAF with [128, 128] hidden units, after an affine block with triangular scaling matrix.
* for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 128 inducing points.

* $\text{Card}^r(\mathcal{P}_1) = 5$

UIVI we used:

* as base distribution, a standard Gaussian with dimensionality 82

* as transform $h$, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality 6

* as transform weights regressor, an MLP with hidden units [64, 64]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-2$, $\times 0.9$ every 300 steps

### B.5.5 Hierarchical Variances (B.4.2)

Cascading Flows (Ambrogioni et al., 2021b) we used:

* a 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 we used:

* for the flows $\mathcal{F}_i$, a MAF with [32, 32] hidden units, after an affine block with diagonal scaling matrix.

* an encoding size of 32

* $\text{Card}^r(\mathcal{P}_1) = 3$

* $\text{Card}^r(\mathcal{P}_0) = 3$

PAVI-E we used:

* for the flows $\mathcal{F}_i$, a MAF with [128, 128] hidden units, after an affine block with diagonal scaling matrix.

* for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 128 inducing points.

* $\text{Card}^r(\mathcal{P}_1) = 3$

* $\text{Card}^r(\mathcal{P}_0) = 3$

UIVI we used:

* as base distribution, a standard Gaussian with dimensionality 32

* as transform $h$, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality 3

* as transform weights regressor, an MLP with hidden units [64, 64]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-2$, $\times 0.9$ every 300 steps
* for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 128 inducing points.

* $\text{Card}^r(\mathcal{P}_1) = 5$

UIVI we used:

* as base distribution, a standard Gaussian with dimensionality 82

* as transform $h$, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality 6

* as transform weights regressor, an MLP with hidden units [64, 64]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-2$, $\times 0.9$ every 300 steps

### B.5.5 Hierarchical Variances (B.4.2)

Cascading Flows (Ambrogioni et al., 2021b) we used:

* a 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 we used:

* for the flows $\mathcal{F}_i$, a MAF with [32, 32] hidden units, after an affine block with diagonal scaling matrix.

* an encoding size of 32

* $\text{Card}^r(\mathcal{P}_1) = 3$

* $\text{Card}^r(\mathcal{P}_0) = 3$

PAVI-E we used:

* for the flows $\mathcal{F}_i$, a MAF with [128, 128] hidden units, after an affine block with diagonal scaling matrix.

* for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads of 32 units each, 2 ISAB blocks with 128 inducing points.

* $\text{Card}^r(\mathcal{P}_1) = 3$

* $\text{Card}^r(\mathcal{P}_0) = 3$

UIVI we used:

* as base distribution, a standard Gaussian with dimensionality 32

* as transform $h$, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality 3

* as transform weights regressor, an MLP with hidden units [64, 64]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-2$, $\times 0.9$ every 300 steps
### B.5.6 Small Neuroimaging example (B.4.3)

Cascading Flows (Ambrogioni et al., 2021b) we used:

* a 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 we used:

* for the flows $F_i$, a MAF with [32, 32] hidden units, after an affine block with diagonal scaling matrix.

* an combination of encoding sizes of 32 and 8

UIVI we used:

* as base distribution, a standard Gaussian with dimensionality 1,802

* as transform *h*, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality 16

* as transform weights regressor, an MLP with hidden units [128, 256, 512, 1024]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-3$, $\times 0.9$ every 300 steps

Optimization systematically degenerated into NaN results after around 200 optimization steps.

## B.6 Details about our Neuroimaging experiment (5.4)

### B.6.1 Data description

In this experiment, we use data from the Human Connectome Project (HCP) dataset (Van Essen et al., 2012). We randomly select a cohort of *S* = 1,000 subjects from this dataset, each subject is associated with *T* = 2 resting-state fMRI sessions (Smith et al., 2013). We minimally pre-process the signal using the *nilearn* python library (Abraham et al., 2014):

1. removing high variance confounds

2. detrending the data

3. band-filtering the data (0.01 to 0.1 Hz), with a 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* = 59,412 vertices across the whole cortex. We compare this signal with the extracted signal of *D* = 64 DiFuMo components: a dictionary of brain spatial maps allowing for an effective fMRI dimensionality reduction (Dadi et al., 2020). Specifically, we compute the one-to-one Pearson'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 ($S \times T \times N \times D$). We project this data —correlation coefficients lying in $]-1; 1[$ — in an unbounded space using an inverse sigmoid function.
### B.5.6 Small Neuroimaging example (B.4.3)

Cascading Flows (Ambrogioni et al., 2021b) we used:

* a 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 we used:

* for the flows $F_i$, a MAF with [32, 32] hidden units, after an affine block with diagonal scaling matrix.

* an combination of encoding sizes of 32 and 8

UIVI we used:

* as base distribution, a standard Gaussian with dimensionality 1,802

* as transform *h*, an affine transform with diagonal scale

* as embedding distribution, a standard Gaussian with dimensionality 16

* as transform weights regressor, an MLP with hidden units [128, 256, 512, 1024]

* to sample uncorrelated samples $\epsilon$, an HMC run with 5 burn-in steps and 5 samples

* an Adam optimizer with exponential learning rate decay, starting at $1e-3$, $\times 0.9$ every 300 steps

Optimization systematically degenerated into NaN results after around 200 optimization steps.

## B.6 Details about our Neuroimaging experiment (5.4)

### B.6.1 Data description

In this experiment, we use data from the Human Connectome Project (HCP) dataset (Van Essen et al., 2012). We randomly select a cohort of *S* = 1,000 subjects from this dataset, each subject is associated with *T* = 2 resting-state fMRI sessions (Smith et al., 2013). We minimally pre-process the signal using the *nilearn* python library (Abraham et al., 2014):

1. removing high variance confounds

2. detrending the data

3. band-filtering the data (0.01 to 0.1 Hz), with a 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* = 59,412 vertices across the whole cortex. We compare this signal with the extracted signal of *D* = 64 DiFuMo components: a dictionary of brain spatial maps allowing for an effective fMRI dimensionality reduction (Dadi et al., 2020). Specifically, we compute the one-to-one Pearson'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 ($S \times T \times N \times D$). We project this data —correlation coefficients lying in $]-1; 1[$ — in an unbounded space using an inverse sigmoid function.
## B.6.2 Model description

We use a model inspired by the work of Kong et al. (2019). We hypothesize that every vertex in the cortex belongs to either one of $L = 7$ functional networks. This number is inspired by the work of Thomas Yeo et al. (2011) and Kong et al. (2019), in which authors identify 7 general networks on resting-state fMRI data.

Each network corresponds to a pattern of connectivity with the brain, represented as the correlation of the BOLD signal with the signal from the $D = 64$ DiFuMo components. We define $L = 7$ functional networks at the population level. Each network $l$ corresponds to the connectivity fingerprint $\mu_l$.

At the population level, for a vertex $n$ we denote as $logits_n \in \mathbb{R}^L$ the logits of its probability to belong to each network $l$. Each subject $s$ is associated with the logits $logits_{s,n}$, which are considered as a perturbation of the population logits. This creates a regularization across subjects: the same vertex $n$ is encouraged to have the same label across all subjects. The variable $\gamma_l$ controls the inter-subject spatial variability across all subjects, for the network $l$.

For every subject $s$, session $t$, and vertex $n$, we denote as $X_{s,t,n}$ the observed connectivity. This connectivity is modeled via a mixture model: $X_{s,t,n}$ is a perturbation of either one of the $\mu_l$'s connectivity fingerprints. $\text{label}_{s,n}$ denotes the label in the mixture (that depends on the subject-specific logits). $\kappa_l$ controls the variability between $X_{s,t,n}$ and the mixture component $\mu_{\text{label}_{s,n}}$.

The resulting model can be described as:

$$
\begin{align}
S, T, N, D, L &= 1000, 2, 59412, 64, 7 \nonumber \\
\forall l=1..L: \quad &\mu_l \sim \mathcal{N}(0 \times \vec{1}_D, 6) \nonumber \\
\forall l=1..L: \quad &\log \kappa_l \sim \mathcal{N}(0 \times \vec{1}_L, 1) \nonumber \\
\forall n=1..N: \quad &\text{logits}_n \sim \mathcal{N}(0 \times \vec{1}_L, 6) \nonumber \\
\forall l=1..L: \quad &\log \gamma_l \sim \mathcal{N}(0 \times \vec{1}_L, 1) \tag{B.24} \\
\forall s=1..S \atop \forall n=1..N: \quad &\text{logits}_{s,n} | \text{logits}_{n, [\gamma_l]_{l=1..L}} \sim \mathcal{N}(\text{logits}_n, [\gamma_1 \dots \gamma_L]) \nonumber \\
\forall s=1..S \atop \forall n=1..N: \quad &\text{labels}_{s,n} | \text{logits}_{s,n} \sim \text{Categorical}(\text{logits}_{s,n}) \nonumber \\
\forall s=1..S \atop \forall t=1..T \atop \forall n=1..N: \quad &X_{s,t,n} | [\mu_l]_{l=1..L}, [\kappa_l]_{l=1..L}, \text{label}_{s,n} \sim \mathcal{N}(\mu_{\text{label}_{s,n}}, \kappa_{\text{label}_{s,n}}) \nonumber
\end{align}
$$

The model contains 4 plates: the *network* plate of full cardinality $L$ (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 $\mu$ and the labels $\text{label}$—represented as the parcellation in Figure 4— given the observed connectome described in Appendix B.6.1.

## B.6.3 PAVI implementation

We used in this experiment the PAVI-F scheme, using:

* for the RVs $\mu_l$:

- for the flows $F_i$, a MAF with [64, 128, 64] hidden units, following an affine block with diagonal scale

- for the encoding size: 16

* for the RVs $\kappa_l, \gamma_l$:

- a MAP estimator using a regressor with [64, 128, 64] hidden units

- for the encoding size: 8

* for the RVs $\text{logits}_n, \text{logits}_{s,n}$:

- for the flows $F_i$, an affine block with diagonal scale regressed using a [64, 128, 64] MLP
## B.6.2 Model description

We use a model inspired by the work of Kong et al. (2019). We hypothesize that every vertex in the cortex belongs to either one of $L = 7$ functional networks. This number is inspired by the work of Thomas Yeo et al. (2011) and Kong et al. (2019), in which authors identify 7 general networks on resting-state fMRI data.

Each network corresponds to a pattern of connectivity with the brain, represented as the correlation of the BOLD signal with the signal from the $D = 64$ DiFuMo components. We define $L = 7$ functional networks at the population level. Each network $l$ corresponds to the connectivity fingerprint $\mu_l$.

At the population level, for a vertex $n$ we denote as $logits_n \in \mathbb{R}^L$ the logits of its probability to belong to each network $l$. Each subject $s$ is associated with the logits $logits_{s,n}$, which are considered as a perturbation of the population logits. This creates a regularization across subjects: the same vertex $n$ is encouraged to have the same label across all subjects. The variable $\gamma_l$ controls the inter-subject spatial variability across all subjects, for the network $l$.

For every subject $s$, session $t$, and vertex $n$, we denote as $X_{s,t,n}$ the observed connectivity. This connectivity is modeled via a mixture model: $X_{s,t,n}$ is a perturbation of either one of the $\mu_l$'s connectivity fingerprints. $\text{label}_{s,n}$ denotes the label in the mixture (that depends on the subject-specific logits). $\kappa_l$ controls the variability between $X_{s,t,n}$ and the mixture component $\mu_{\text{label}_{s,n}}$.

The resulting model can be described as:

$$
\begin{align}
S, T, N, D, L &= 1000, 2, 59412, 64, 7 \nonumber \\
\forall l=1..L &: \quad \mu_l \sim \mathcal{N}(0 \times \vec{1}_D, 6) \nonumber \\
\forall l=1..L &: \quad \log \kappa_l \sim \mathcal{N}(0 \times \vec{1}_L, 1) \nonumber \\
\forall n=1..N &: \quad \text{logits}_n \sim \mathcal{N}(0 \times \vec{1}_L, 6) \nonumber \\
\forall l=1..L &: \quad \log \gamma_l \sim \mathcal{N}(0 \times \vec{1}_L, 1) \tag{B.24} \\
\forall s=1..S \atop \forall n=1..N &: \quad \text{logits}_{s,n} | \text{logits}_{n, [\gamma_l]_{l=1..L}} \sim \mathcal{N}(\text{logits}_n, [\gamma_1 \dots \gamma_L]) \nonumber \\
\forall s=1..S \atop \forall n=1..N &: \quad \text{labels}_{s,n} | \text{logits}_{s,n} \sim \text{Categorical}(\text{logits}_{s,n}) \nonumber \\
\forall s=1..S \atop \forall t=1..T \atop \forall n=1..N &: \quad X_{s,t,n} | [\mu_l]_{l=1..L}, [\kappa_l]_{l=1..L}, \text{label}_{s,n} \sim \mathcal{N}(\mu_{\text{label}_{s,n}}, \kappa_{\text{label}_{s,n}}) \nonumber
\end{align}
$$

The model contains 4 plates: the *network* plate of full cardinality $L$ (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 $\mu$ and the labels $\text{label}$—represented as the parcellation in Figure 4— given the observed connectome described in Appendix B.6.1.

## B.6.3 PAVI implementation

We used in this experiment the PAVI-F scheme, using:

* for the RVs $\mu_l$:

- for the flows $\mathcal{F}_i$, a MAF with [64, 128, 64] hidden units, following an affine block with diagonal scale

- for the encoding size: 16

* for the RVs $\kappa_l, \gamma_l$:

- a MAP estimator using a regressor with [64, 128, 64] hidden units

- for the encoding size: 8

* for the RVs $\text{logits}_n, \text{logits}_{s,n}$:

- for the flows $\mathcal{F}_i$, an affine block with diagonal scale regressed using a [64, 128, 64] MLP
- for the encoding size: 8

• for the reduced model, we used $S^r = 128$, $T^r = 1$ and $N^r = 128$.

To allow for the optimization over the discrete label$_{s,n}$ RV, we used the Gumbell-Softmax trick, using a fixed temperature of 1.0 (Jang et al., 2017; Maddison et al., 2016).

### B.6.4 Cognitive scores prediction

For the cognitive scores prediction in Figure 4, we reproduced the standard methodology from Kong et al. (2019).

We perform a 20-fold cross-validation across 1,000 subjects. We start from the logits$_{s,n}$ associated with each subject. We use PCA to project the features from the 19 training folds to their 100 first components (with 33% explained variance). We then train a linear regression to predict each of the 13 cognitive scores from the training-fold PCA features. We compute the test performance on the test fold —using the training set PCA and linear regression— averaged across the 13 cognitive measures. This process is reproduced 100 times. The reported scores are the triple average across folds, measures, and repetition, while the standard deviation is computed across the 100 repetitions only.

## C Supplemental discussion

### C.1 Plate amortization as a generalization of sample amortization

In Section 4.1, we introduced plate amortization as applying the generic concept of amortization to the granularity of plates. There is actually an even stronger connection between sample amortization and plate amortization.

An HBM $p$ models the distribution of a given observed RV $X$ —jointly with the parameters $\Theta$. Different samples $X_0, X_1, ...$ of the model $p$ are i.i.d. draws from the distribution $p(X)$. $p$ can thus be considered as the model for "one sample". Instead of $p$, consider a "macro" model for the whole population of samples one could draw from $p$. The observed RV of that macro model would be the infinite collection of samples drawn from the same distribution $p(X)$. In that light, the i.i.d. sampling of different $X$ values from $p$ could be interpreted as a plate of the macro model. Thus, we could consider sample amortization as an instance of plate amortization for the "sample plate". Or equivalently: plate amortization can be 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^{full}$ versus reduced HBM $M^{redu}$ sharing the same template $\mathcal{T}$. This formalism is 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^{redu}$- in which encodings are "plugged in" -either sliced from larger encoding arrays or as the output of an encoder applied to a data slice, see Section 4.1&4.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 as 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 $M^{redu}$ and obtain "for free" a variational posterior for the full model $M^{full}$ —see Section 4.2. In this context, our scheme is no longer a different take on hierarchical, batched SVI: the cardinality of the full model is truly independent of the cardinality of the training and is only simulated as a scaling factor in the stochastic training —see Section 4.2. We have the intuition that fruitful research directions could stem from this concept.
- for the encoding size: 8

• for the reduced model, we used $S^r = 128$, $T^r = 1$ and $N^r = 128$.

To allow for the optimization over the discrete label$_{s,n}$ RV, we used the Gumbell-Softmax trick, using a fixed temperature of 1.0 (Jang et al., 2017; Maddison et al., 2016).

### B.6.4 Cognitive scores prediction

For the cognitive scores prediction in Figure 4, we reproduced the standard methodology from Kong et al. (2019).

We perform a 20-fold cross-validation across 1,000 subjects. We start from the logits$_{s,n}$ associated with each subject. We use PCA to project the features from the 19 training folds to their 100 first components (with 33% explained variance). We then train a linear regression to predict each of the 13 cognitive scores from the training-fold PCA features. We compute the test performance on the test fold —using the training set PCA and linear regression— averaged across the 13 cognitive measures. This process is reproduced 100 times. The reported scores are the triple average across folds, measures, and repetition, while the standard deviation is computed across the 100 repetitions only.

## C Supplemental discussion

### C.1 Plate amortization as a generalization of sample amortization

In Section 4.1, we introduced plate amortization as applying the generic concept of amortization to the granularity of plates. There is actually an even stronger connection between sample amortization and plate amortization.

An HBM $p$ models the distribution of a given observed RV $X$ —jointly with the parameters $\Theta$. Different samples $X_0, X_1, ...$ of the model $p$ are i.i.d. draws from the distribution $p(X)$. $p$ can thus be considered as the model for "one sample". Instead of $p$, consider a "macro" model for the whole population of samples one could draw from $p$. The observed RV of that macro model would be the infinite collection of samples drawn from the same distribution $p(X)$. In that light, the i.i.d. sampling of different $X$ values from $p$ could be interpreted as a plate of the macro model. Thus, we could consider sample amortization as an instance of plate amortization for the "sample plate". Or equivalently: plate amortization can be 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^{\text{full}}$ versus reduced HBM $M^{\text{redu}}$ sharing the same template $\mathcal{T}$. This formalism is 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^{\text{redu}}$- in which encodings are "plugged in" -either sliced from larger encoding arrays or as the output of an encoder applied to a data slice, see Section 4.1&4.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 as 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 $M^{\text{redu}}$ and obtain "for free" a variational posterior for the full model $M^{\text{full}}$ —see Section 4.2. In this context, our scheme is no longer a different take on hierarchical, batched SVI: the cardinality of the full model is truly independent of the cardinality of the training and is only simulated as a scaling factor in the stochastic training —see Section 4.2. 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 4.1. Plate amortization is particularly useful in the context of heavily parameterized density approximators 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 applications in other highly-structured problem classes such as graphs or sequences (Wu et al., 2020; Salehinejad et al., 2018).

Central to our design, PAVI adjoins an encoding structure to an inference problem. The encodings $E_{i,n}$ embed all the "individualized" information in the problem, while the shared conditional density estimators $\mathcal{F}_i$ translate those encodings into posterior distributions. We believe there is potential in externalizing part of the inference problem into an embedding. For instance, embeddings could be learned separately from density estimators, opening up the possibility for transfer learning (Weiss et al., 2016). Or the encodings could integrate some known symmetry of the problem, such as convolution-based encodings in the case of Random Fields (Bishop, 2006).

### C.4 Connection with meta-learning

In Section 4.1, we introduced plate amortization: sharing the parameterization and learning across a model's plates. This section discusses the connection between plate amortization and meta-learning (Ravi & Beatson, 2019; Iakovleva et al., 2020; Yao et al., 2019).

Supervised learning can be seen as the mapping from a given context set $C = \{(x, y)\}$ to a predictive function $f$ (Bishop, 2006) such that $f(x) = y$. Meta-learning—or "learning to learn"—instead recovers this mapping $C \mapsto f$ in the general case. Once the meta-training is completed, a predictive function $f$ conditioned by an unseen context $C$ can be obtained in a single forward pass —without any training done on $C$. As an 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$ is then used to condition an estimator for the density $q(y|x, C)$.

This framework is similar to the PAVI-E scheme, where a combination of a deep set encoder and a normalizing-flow-based density estimator output the posterior probability of a ground RV $\theta_{i,n}$. This encoder-estimator pair is repeatedly used across a 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_{i,n}$ at a lower hierarchy plays a role similar to the context $C$ in meta-leaning. A few differences, however, exist between the two frameworks:

* meta-learning is typically concatenated to the 1-plate regime, whereas PAVI is designed for the multi-hierarchy scenario;

* meta-learning typically considers a set of i.i.d. tasks, whereas in PAVI the inference of different ground RVs $\theta_{i,n}$ are conditionally dependent through the hierarchical model $p$;

* meta-learning is trained using the forward KL loss, that is to say in the sample-amortized regime, maximizing the probability $q(y)$ of samples from the underlying generative process. In contrast, PAVI —though possible to train using the forward KL loss— is trained via the reverse KL, which requires to evaluate the density of the generative process explicitly.

We suspect there would be interesting applications for the PAVI architecture in hierarchical meta-learning scenarios.

### C.5 Conditional dependencies modeled in the variational family

Consider an inference problem with three latent parameters $\theta_1, \theta_2, \theta_3$, and a model $p(\Theta, X)$ that factorizes as $p(\theta_1, \theta_2, \theta_3) = p(\theta_1)p(\theta_2)p(\theta_3|\theta_1, \theta_2)p(X|\theta_3)$. To approximate the posterior $p(\Theta|X)$, different conditional
### C.3 Benefiting from structure in inference

Our contributions can be abstracted through the concept of plate amortization -see Section 4.1. Plate amortization is particularly useful in the context of heavily parameterized density approximators 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 applications in other highly-structured problem classes such as graphs or sequences (Wu et al., 2020; Salehinejad et al., 2018).

Central to our design, PAVI adjoins an encoding structure to an inference problem. The encodings $E_{i,n}$ embed all the "individualized" information in the problem, while the shared conditional density estimators $\mathcal{F}_i$ translate those encodings into posterior distributions. We believe there is potential in externalizing part of the inference problem into an embedding. For instance, embeddings could be learned separately from density estimators, opening up the possibility for transfer learning (Weiss et al., 2016). Or the encodings could integrate some known symmetry of the problem, such as convolution-based encodings in the case of Random Fields (Bishop, 2006).

### C.4 Connection with meta-learning

In Section 4.1, we introduced plate amortization: sharing the parameterization and learning across a model's plates. This section discusses the connection between plate amortization and meta-learning (Ravi & Beatson, 2019; Iakovleva et al., 2020; Yao et al., 2019).

Supervised learning can be seen as the mapping from a given context set $C = \{(x, y)\}$ to a predictive function $f$ (Bishop, 2006) such that $f(x) = y$. Meta-learning—or "learning to learn"—instead recovers this mapping $C \mapsto f$ in the general case. Once the meta-training is completed, a predictive function $f$ conditioned by an unseen context $C$ can be obtained in a single forward pass —without any training done on $C$. As an 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$ is then used to condition an estimator for the density $q(y|x, C)$.

This framework is similar to the PAVI-E scheme, where a combination of a deep set encoder and a normalizing-flow-based density estimator output the posterior probability of a ground RV $\theta_{i,n}$. This encoder-estimator pair is repeatedly used across a 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_{i,n}$ at a lower hierarchy plays a role similar to the context $C$ in meta-leaning. A few differences, however, exist between the two frameworks:

* meta-learning is typically concatenated to the 1-plate regime, whereas PAVI is designed for the multi-hierarchy scenario;

* meta-learning typically considers a set of i.i.d. tasks, whereas in PAVI the inference of different ground RVs $\theta_{i,n}$ are conditionally dependent through the hierarchical model $p$;

* meta-learning is trained using the forward KL loss, that is to say in the sample-amortized regime, maximizing the probability $q(y)$ of samples from the underlying generative process. In contrast, PAVI —though possible to train using the forward KL loss— is trained via the reverse KL, which requires to evaluate the density of the generative process explicitly.

We suspect there would be interesting applications for the PAVI architecture in hierarchical meta-learning scenarios.

### C.5 Conditional dependencies modeled in the variational family

Consider an inference problem with three latent parameters $\theta_1, \theta_2, \theta_3$, and a model $p(\Theta, X)$ that factorizes as $p(\theta_1, \theta_2, \theta_3) = p(\theta_1)p(\theta_2)p(\theta_3|\theta_1, \theta_2)p(X|\theta_3)$. To approximate the posterior $p(\Theta|X)$, different conditional
dependencies can be modeled in the chosen variational family $\mathcal{Q}$. The mean-field (MF) approximation fully factorizes the variational distribution as $q(\theta_1, \theta_2, \theta_3) = q(\theta_1)q(\theta_2)q(\theta_3)$. MF was originally introduced in VI to facilitate computation, allowing dedicated optimization schemes (Blei et al., 2017). Nonetheless, this approximation assumes independence between RVs in the posterior and ultimately limits the expressivity of the variational family. To remove this approximation, modeling complex conditional dependencies in the variational family is an open research subject (Ambrogioni et al., 2021b; Webb et al., 2018). In the PAVI design, we inherit our statistical dependency structure from the prior distribution $p$, as detailed in Equation (3). This amounts to factorizing $q(\theta_1, \theta_2, \theta_3) = q(\theta_1)q(\theta_2)q(\theta_3|\theta_1, \theta_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 modeling only those forward dependencies, the modeling of colliders can be an issue. More generally, one would like to model arbitrary dependencies in the variational posterior. As an example, the factorization $q(\theta_1, \theta_2, \theta_3) = q(\theta_3)q(\theta_2|\theta_3)q(\theta_1|\theta_2, \theta_3)$ would more faithfully represent the conditional dependencies that can arise in the posterior. Modeling arbitrary dependencies reduces the inference gap but can become computationally intractable as the number of RVs increases. In PAVI, we strike a particular balance between expressivity and computational tractability, as further motivated in this section.

In the case of a single plate —the 2-level case— Agrawal & Domke (2021) demonstrate that modeling only the *forward* dependencies does not result in reduced expressivity compared to the modeling of the full dependencies. Yet this result does not hold in the n-level case, as Webb et al. (2018) shows that faithful inversion features conditional dependencies in the posterior between ground RVs of the same template. We can dub those dependencies as *horizontal* dependencies across RVs in the same plate. Like Structured SVI (Hoffman & Blei, 2014) or Automatic SSVI (Ambrogioni et al., 2021a), the PAVI design thus results in reduced expressivity when stacking multiple plates in the model $p$. This issue can be partially alleviated with the usage of a *backward* encoding scheme —going in the reverse direction compared to the prior's dependencies—as 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 4.1) and the stochastic training over batches of those RVs (see Section 4.2) prevent the direct modeling *horizontal* dependencies. Put differently, the fact that we consider the inference over different ground RVs as conditionally independent inference problems is central to our design and adverse to the modeling of *horizontal* dependencies.

Injecting *horizontal* dependencies back into our variational family is therefore a non-trivial research direction, that is not at the core of this paper. This opens up promising research directions: how could arbitrary conditional dependencies be modeled in the variational posterior in the context of stochastic training?

## C.6 Towards user-friendly Variational Inference

By re-purposing the concept of amortization at the plate level, we aim to propose clear computation versus precision trade-offs in VI. Hyper-parameters such as the encoding size —as illustrated in Figure 2 (right)— allow to clearly trade inference quality in exchange for a reduced memory footprint. On the contrary, in classical VI, changing $\mathcal{Q}$'s parametric form —for instance, switching from Gaussian to Student distributions— can have a strong and complex impact on the number of weights and inference quality (Blei et al., 2017). By allowing the usage of normalizing flows in very large cardinality regimes, our contribution aims at disentangling approximation power and computational feasibility. In particular, having access to expressive density approximators for the posterior can help experimenters diversify the proposed HBMs, removing the need for properties such as conjugacy to obtain meaningful inference (Gelman et al., 2004). Combining clear hyper-parameters and scalable yet universal density approximators, we tend towards a user-friendly methodology in the context of large population studies VI.

In its current implementation, PAVI requires to set up a rather large number of hyper-parameters. This includes choosing a normalizing flow architecture, the size of the encodings at different plate levels, and the reduced model cardinalities. Those hyper-parameters are in addition to generic ones such as the number of samples from the variational family used to estimate the ELBO, or the choice of an optimizer. However,
dependencies can be modeled in the chosen variational family $\mathcal{Q}$. The mean-field (MF) approximation fully factorizes the variational distribution as $q(\theta_1, \theta_2, \theta_3) = q(\theta_1)q(\theta_2)q(\theta_3)$. MF was originally introduced in VI to facilitate computation, allowing dedicated optimization schemes (Blei et al., 2017). Nonetheless, this approximation assumes independence between RVs in the posterior and ultimately limits the expressivity of the variational family. To remove this approximation, modeling complex conditional dependencies in the variational family is an open research subject (Ambrogioni et al., 2021b; Webb et al., 2018). In the PAVI design, we inherit our statistical dependency structure from the prior distribution $p$, as detailed in Equation (3). This amounts to factorizing $q(\theta_1, \theta_2, \theta_3) = q(\theta_1)q(\theta_2)q(\theta_3|\theta_1, \theta_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 modeling only those forward dependencies, the modeling of colliders can be an issue. More generally, one would like to model arbitrary dependencies in the variational posterior. As an example, the factorization $q(\theta_1, \theta_2, \theta_3) = q(\theta_3)q(\theta_2|\theta_3)q(\theta_1|\theta_2, \theta_3)$ would more faithfully represent the conditional dependencies that can arise in the posterior. Modeling arbitrary dependencies reduces the inference gap but can become computationally intractable as the number of RVs increases. In PAVI, we strike a particular balance between expressivity and computational tractability, as further motivated in this section.

In the case of a single plate —the 2-level case— Agrawal & Domke (2021) demonstrate that modeling only the *forward* dependencies does not result in reduced expressivity compared to the modeling of the full dependencies. Yet this result does not hold in the n-level case, as Webb et al. (2018) shows that faithful inversion features conditional dependencies in the posterior between ground RVs of the same template. We can dub those dependencies as *horizontal* dependencies across RVs in the same plate. Like Structured SVI (Hoffman & Blei, 2014) or Automatic SSVI (Ambrogioni et al., 2021a), the PAVI design thus results in reduced expressivity when stacking multiple plates in the model $p$. This issue can be partially alleviated with the usage of a *backward* encoding scheme —going in the reverse direction compared to the prior's dependencies—as 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 4.1) and the stochastic training over batches of those RVs (see Section 4.2) prevent the direct modeling *horizontal* dependencies. Put differently, the fact that we consider the inference over different ground RVs as conditionally independent inference problems is central to our design and adverse to the modeling of *horizontal* dependencies.

Injecting *horizontal* dependencies back into our variational family is therefore a non-trivial research direction, that is not at the core of this paper. This opens up promising research directions: how could arbitrary conditional dependencies be modeled in the variational posterior in the context of stochastic training?

## C.6 Towards user-friendly Variational Inference

By re-purposing the concept of amortization at the plate level, we aim to propose clear computation versus precision trade-offs in VI. Hyper-parameters such as the encoding size —as illustrated in Figure 2 (right)— allow to clearly trade inference quality in exchange for a reduced memory footprint. On the contrary, in classical VI, changing $\mathcal{Q}$'s parametric form —for instance, switching from Gaussian to Student distributions— can have a strong and complex impact on the number of weights and inference quality (Blei et al., 2017). By allowing the usage of normalizing flows in very large cardinality regimes, our contribution aims at disentangling approximation power and computational feasibility. In particular, having access to expressive density approximators for the posterior can help experimenters diversify the proposed HBMs, removing the need for properties such as conjugacy to obtain meaningful inference (Gelman et al., 2004). Combining clear hyper-parameters and scalable yet universal density approximators, we tend towards a user-friendly methodology in the context of large population studies VI.

In its current implementation, PAVI requires to set up a rather large number of hyper-parameters. This includes choosing a normalizing flow architecture, the size of the encodings at different plate levels, and the reduced model cardinalities. Those hyper-parameters are in addition to generic ones such as the number of samples from the variational family used to estimate the ELBO, or the choice of an optimizer. However,
our experiments suggest that those hyper-parameters could default to conservative values. As shown in Section 5.2, the encoding size should be maximized with respect to the available memory, as should the reduced cardinalities, as detailed in Appendix B.3. As for the normalizing flow architecture, since the latter is shared across plates, the user can always default to the most expressive architecture available without incurring an exploding number of weights. As an example, in our experiments, we stuck to the MAF architecture (Papamakarios & Murray, 2018). Since all hyper-parameters can be set conservatively, a user-friendly automatic VI API can be designed, where the user only provides the system with the generative model *p*, and some observed value *X*. The variational family *q* can then automatically be derived from *p*, and directly optimized over *X*. This is the design principle we implemented in the codebase adjoined to this publication.
our experiments suggest that those hyper-parameters could default to conservative values. As shown in Section 5.2, the encoding size should be maximized with respect to the available memory, as should the reduced cardinalities, as detailed in Appendix B.3. As for the normalizing flow architecture, since the latter is shared across plates, the user can always default to the most expressive architecture available without incurring an exploding number of weights. As an example, in our experiments, we stuck to the MAF architecture (Papamakarios & Murray, 2018). Since all hyper-parameters can be set conservatively, a user-friendly automatic VI API can be designed, where the user only provides the system with the generative model *p*, and some observed value *X*. The variational family *q* can then automatically be derived from *p*, and directly optimized over *X*. This is the design principle we implemented in the codebase adjoined to this publication.