Original: PAVI:
Revised: PAVI:

Original: P LATE -A MORTIZED V ARIATIONAL I NFERENCE
Revised: P LATE -A MORTIZED V ARIATIONAL I NFERENCE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Given observed data and a probabilistic generative model, Bayesian inference aims at obtaining the distribution of a model’s latent parameters that could have yielded the data.
Revised: Given observed data and a probabilistic generative model, Bayesian inference aims at obtaining the distribution of a model’s latent parameters that could have yielded the data.

Original: This task is challenging for large population studies where thousands of measurements are performed over a cohort of hundreds of subjects, resulting in a massive latent parameter space.
Revised: This task is challenging for large population studies where thousands of measurements are performed over a cohort of hundreds of subjects, resulting in a massive latent parameter space.

Original: This large cardinality renders offthe-shelf Variational Inference (VI) computationally impractical.
Revised: This large cardinality renders offthe-shelf Variational Inference (VI) computationally impractical.

Original: In this work, we design structured VI families that can efficiently tackle large population studies.
Revised: In this work, we design structured VI families that can efficiently tackle large population studies.

Original: Our main idea is to share the parameterization and learning across the different i.i.d.
Revised: Our main idea is to share the parameterization and learning across the different i.i.d.

Original: variables in a generative model -symbolized by the model’s plates .
Revised: variables in a generative model -symbolized by the model’s plates .

Original: We name this concept plate amortization , and illustrate the powerful synergies it entitles, resulting in expressive, parsimoniously parameterized and orders of magnitude faster to train large scale hierarchical variational distributions.
Revised: We name this concept plate amortization , and illustrate the powerful synergies it entitles, resulting in expressive, parsimoniously parameterized and orders of magnitude faster to train large scale hierarchical variational distributions.

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

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Population studies analyse measurements over large cohorts of human subjects.
Revised: Population studies analyse measurements over large cohorts of human subjects.

Original: These studies are ubiquitous in health care (Fayaz et al., 2016;
Revised: These studies are ubiquitous in health care (Fayaz et al., 2016;

Original: Towsley et al., 2011), and can typically involve hundreds of subjects and measurements per subject.
Revised: Towsley et al., 2011), and can typically involve hundreds of subjects and measurements per subject.

Original: For instance in Neuroimaging (Kong et al., 2019), measurements X can correspond to signals in hundreds of locations in the brain for a thousand subjects.
Revised: For instance in Neuroimaging (Kong et al., 2019), measurements X can correspond to signals in hundreds of locations in the brain for a thousand subjects.

Original: Given this observed data X , and a generative model that can produce data given model parameters Θ , we want to recover the parameters Θ that could have yielded the observed X .
Revised: Given this observed data X , and a generative model that can produce data given model parameters Θ , we want to recover the parameters Θ that could have yielded the observed X .

Original: In our Neuroimaging example, Θ can be local labels for each brain location and subject, together with global parameters common to all subjects –such as the connectivity corresponding to each label.
Revised: In our Neuroimaging example, Θ can be local labels for each brain location and subject, together with global parameters common to all subjects –such as the connectivity corresponding to each label.

Original: We want to recover the distribution of the Θ that could have produced X .
Revised: We want to recover the distribution of the Θ that could have produced X .

Original: Following the Bayesian formalism (Gelman et al., 2004), we cast both Θ and X as Random Variables (RVs) and our goal is to recover the posterior distribution p (Θ | X ) .
Revised: Following the Bayesian formalism (Gelman et al., 2004), we cast both Θ and X as Random Variables (RVs) and our goal is to recover the posterior distribution p (Θ | X ) .

Original: Due to the nested structure of our applications we focus on the case where p corresponds to a Hierarchical Bayesian Model (HBM) (Gelman et al., 2004).
Revised: Due to the nested structure of our applications we focus on the case where p corresponds to a Hierarchical Bayesian Model (HBM) (Gelman et al., 2004).

Original: In the context of population studies, the multitude of subjects and measurements per subject implies a large dimensionality for both Θ and X .
Revised: In the context of population studies, the multitude of subjects and measurements per subject implies a large dimensionality for both Θ and X .

Original: This large dimensionality in turn creates computational hurdles that we tackle through our method.
Revised: This large dimensionality in turn creates computational hurdles that we tackle through our method.

Original: Several inference methods exist in the literature.
Revised: Several inference methods exist in the literature.

Original: Earliest works resorted to Markov Chain Monte
Revised: Earliest works resorted to Markov Chain Monte

Original: Carlo (Koller & Friedman, 2009), which tend to be slow in high dimensional settings (Blei et al., 2017).
Revised: Carlo (Koller & Friedman, 2009), which tend to be slow in high dimensional settings (Blei et al., 2017).

Original: Recent approaches, coined Variational Inference (VI), cast the inference as an optimization problem (Blei et al., 2017;
Revised: Recent approaches, coined Variational Inference (VI), cast the inference as an optimization problem (Blei et al., 2017;

Original: Zhang et al., 2019; Ruiz & Titsias, 2019).
Revised: Zhang et al., 2019; Ruiz & Titsias, 2019).

Original: Inference reduces to finding the distribution q (Θ; ϕ ) ∈ Q closest to the unknown posterior p (Θ | X ) in a variational family Q chosen by the experimenter.
Revised: Inference reduces to finding the distribution q (Θ; ϕ ) ∈ Q closest to the unknown posterior p (Θ | X ) in a variational family Q chosen by the experimenter.

Original: Historically, VI required to manually derive and optimize over Q , which remains an effective method where applicable (Dao et al., 2021).
Revised: Historically, VI required to manually derive and optimize over Q , which remains an effective method where applicable (Dao et al., 2021).

Original: In contrast, we follow the idea of automatic VI: deriving an efficient family Q directly from the HBM p (Kucukelbir et al., 2016;
Revised: In contrast, we follow the idea of automatic VI: deriving an efficient family Q directly from the HBM p (Kucukelbir et al., 2016;

Original: Ambrogioni et al., 2021a;b).
Revised: Ambrogioni et al., 2021a;b).

Original: Our method also relies on the pervasive idea of amortization in VI: deriving posteriors usable across multiple data points, which can be linked to meta-learning (Zhang et al., 2019; Ravi & Beatson, 2019; Iakovleva et al., 2020; Yao et al., 2019).
Revised: Our method also relies on the pervasive idea of amortization in VI: deriving posteriors usable across multiple data points, which can be linked to meta-learning (Zhang et al., 2019; Ravi & Beatson, 2019; Iakovleva et al., 2020; Yao et al., 2019).

Original: In particular, Neural
Revised: In particular, Neural

Original: Processes share with our method the conditioning a density estimator by the output of a permutationinvariant encoder (Garnelo et al., 2018; Dubois et al., 2020; Zaheer et al., 2018).
Revised: Processes share with our method the conditioning a density estimator by the output of a permutationinvariant encoder (Garnelo et al., 2018; Dubois et al., 2020; Zaheer et al., 2018).

Original: Meta-learning studies problems with few hierarchies, similar to many VI methods (Ravi & Beatson, 2019; Tran et al., 2017).
Revised: Meta-learning studies problems with few hierarchies, similar to many VI methods (Ravi & Beatson, 2019; Tran et al., 2017).

Original: Agrawal & Domke (2021) notably studied the case of 2-level HBMs, providing with theoretical guarantees.
Revised: Agrawal & Domke (2021) notably studied the case of 2-level HBMs, providing with theoretical guarantees.

Original: In contrast, our focus is rather computational, and we study generic HBMs with an arbitrary number of hierarchies, aiming at tackling large population studies efficiently.
Revised: In contrast, our focus is rather computational, and we study generic HBMs with an arbitrary number of hierarchies, aiming at tackling large population studies efficiently.

Original: Modern VI is effective in low-dimensional settings, but does not scale up to large population studies –that involve millions of random variables.
Revised: Modern VI is effective in low-dimensional settings, but does not scale up to large population studies –that involve millions of random variables.

Original: In this work we identify and tackle 2 challenges to enable this scale-up.
Revised: In this work we identify and tackle 2 challenges to enable this scale-up.

Original: A first challenge with scalability is a detrimental trade-off between expressivity and high-dimensionality (Rouillard & Wassermann, 2022).
Revised: A first challenge with scalability is a detrimental trade-off between expressivity and high-dimensionality (Rouillard & Wassermann, 2022).

Original: To reduce the inference gap, VI requires the variational family Q to contain distributions closely approximating p (Θ | X ) (Blei et al., 2017).
Revised: To reduce the inference gap, VI requires the variational family Q to contain distributions closely approximating p (Θ | X ) (Blei et al., 2017).

Original: Yet the form of p (Θ | X ) is usually unknown to the experimenter.
Revised: Yet the form of p (Θ | X ) is usually unknown to the experimenter.

Original: Instead of a lengthy search for a valid family, one can resort to universal density approximators: normalizing flows (Papamakarios et al., 2019).
Revised: Instead of a lengthy search for a valid family, one can resort to universal density approximators: normalizing flows (Papamakarios et al., 2019).

Original: But the cost for this generality is a heavy parameterization, and normalizing flows scale poorly with the dimensionality of Θ .
Revised: But the cost for this generality is a heavy parameterization, and normalizing flows scale poorly with the dimensionality of Θ .

Original: In large populations studies, as this dimensionality grows, the parameterization of normalizing flows becomes prohibitively large.
Revised: In large populations studies, as this dimensionality grows, the parameterization of normalizing flows becomes prohibitively large.

Original: To tackle this challenge, Rouillard & Wassermann (2022) recently proposed –in the ADAVI architecture– to partially share the parameterization of normalizing flows across the hierarchies of a generative model.
Revised: To tackle this challenge, Rouillard & Wassermann (2022) recently proposed –in the ADAVI architecture– to partially share the parameterization of normalizing flows across the hierarchies of a generative model.

Original: While ADAVI tackled the over-parameterization of VI in population studies, it still could not perform inference in very large data regimes.
Revised: While ADAVI tackled the over-parameterization of VI in population studies, it still could not perform inference in very large data regimes.

Original: This is due to a second challenge with scalability: as the size of Θ increases, the evaluation of a single gradient over the entirety of an architecture’s weights quickly requires too much memory and compute.
Revised: This is due to a second challenge with scalability: as the size of Θ increases, the evaluation of a single gradient over the entirety of an architecture’s weights quickly requires too much memory and compute.

Original: This second challenge can be overcome using stochastic VI (Hoffman et al., 2013), which subsamples the parameters Θ inferred for at each optimization step.
Revised: This second challenge can be overcome using stochastic VI (Hoffman et al., 2013), which subsamples the parameters Θ inferred for at each optimization step.

Original: However, using SVI, the weights for the posterior for a local parameter θ ∈ Θ are only updated when θ is visited by the algorithm.
Revised: However, using SVI, the weights for the posterior for a local parameter θ ∈ Θ are only updated when θ is visited by the algorithm.

Original: In the presence of hundreds of thousands of such local parameters, stochastic VI can become prohibitively slow.
Revised: In the presence of hundreds of thousands of such local parameters, stochastic VI can become prohibitively slow.

Original: In this work, we introduce the concept of plate amortization (PAVI) for fast inference in large scale HBMs.
Revised: In this work, we introduce the concept of plate amortization (PAVI) for fast inference in large scale HBMs.

Original: Instead of considering the inference over local parameters θ as separate problems, our main idea is to share both the parameterization and learning across those local parameters –or equivalently across a model’s plates .
Revised: Instead of considering the inference over local parameters θ as separate problems, our main idea is to share both the parameterization and learning across those local parameters –or equivalently across a model’s plates .

Original: We first propose an algorithm to automatically derive an expressive yet parsimoniously-parameterized variational family from a plate-enriched HBM.
Revised: We first propose an algorithm to automatically derive an expressive yet parsimoniously-parameterized variational family from a plate-enriched HBM.

Original: We then propose a hierarchical stochastic optimization scheme to train this architecture efficiently, obtaining orders of magnitude faster convergence.
Revised: We then propose a hierarchical stochastic optimization scheme to train this architecture efficiently, obtaining orders of magnitude faster convergence.

Original: We detail two variants of our method, with different trade-offs between parameterization and inference quality.
Revised: We detail two variants of our method, with different trade-offs between parameterization and inference quality.

Original: PAVI leverages the repeated structure of plate-enriched HBMs via an original combination of amortization and stochastic training.
Revised: PAVI leverages the repeated structure of plate-enriched HBMs via an original combination of amortization and stochastic training.

Original: Through this combination, our main claim is to enable inference over arbitrarily large population studies –up to a million RVs– with reduced parameterization and training time as the cardinality of the problem augments.
Revised: Through this combination, our main claim is to enable inference over arbitrarily large population studies –up to a million RVs– with reduced parameterization and training time as the cardinality of the problem augments.

Original: We illustrate this by applying PAVI to a challenging human brain cortex parcellation, featuring inference over a cohort of 1000 subjects, demonstrating a significant step towards scalable, expressive and fast VI.
Revised: We illustrate this by applying PAVI to a challenging human brain cortex parcellation, featuring inference over a cohort of 1000 subjects with tens of thousand of measurements per subject, demonstrating a significant step towards scalable, expressive and fast VI.

Original:  M ETHODS 2. H IERARCHICAL B AYESIAN M ODELS (HBM S ), TEMPLATES AND PLATES
Revised:  M ETHODS 2. H IERARCHICAL B AYESIAN M ODELS (HBM S ), TEMPLATES AND PLATES

Original: Our objective is to perform inference in large population studies.
Revised: Our objective is to perform inference in large population studies.

Original: As an example of how inference becomes impractical in this context, consider M in fig.
Revised: As an example of how inference becomes impractical in this context, consider M in fig.

Original: 1 as a model for the height distribution in a population.
Revised: 1 as a model for the height distribution in a population.

Original: θ 2 , 0 denotes the mean height across the population.
Revised: θ 2 , 0 denotes the mean height across the population.

Original: θ 1 , 0 , θ 1 , 1 , θ 1 , 2 denote the mean heights for 3 groups of subjects, distributed around the population mean.
Revised: θ 1 , 0 , θ 1 , 1 , θ 1 , 2 denote the mean heights for 3 groups of subjects, distributed around the population mean.

Original: X 0 , X 1 represent the observed heights of 2 subjects from group 0 , distributed around the group mean.
Revised: X 0 , X 1 represent the observed heights of 2 subjects from group 0 , distributed around the group mean.

Original: Given the observed subject heights X , the goal is to determine the posterior distributions of the group means p ( θ 1 ,n | X ) and population mean p ( θ 2 , 0 | X ) .
Revised: Given the observed subject heights X , the goal is to determine the posterior distributions of the group means p ( θ 1 ,n | X ) and population mean p ( θ 2 , 0 | X ) .

Original: As the number of groups and subjects per group augments, the parameterization and time required to infer the posteriors of a growing number of RVs become prohibitively large.
Revised: As the number of groups and subjects per group augments, the parameterization and time required to infer the posteriors of a growing number of RVs become prohibitively large.

Original: Our goal is to keep this inference computationally tractable.
Revised: Our goal is to keep this inference computationally tractable.

Original: In the rest of this section, we define the inference problem and the associated notations.
Revised: In the rest of this section, we define the inference problem and the associated notations.

Original: Population studies can be modelled using Hierarchical Bayesian Models (HBMs).
Revised: Population studies can be modelled using Hierarchical Bayesian Models (HBMs).

Original: Those HBMs feature samples from common conditional distributions at multiple levels.
Revised: Those HBMs feature samples from common conditional distributions at multiple levels.

Original: For instance subject heights are i.i.d.
Revised: For instance subject heights are i.i.d.

Original: given group heights, which are i.i.d.
Revised: given group heights, which are i.i.d.

Original: given population height.
Revised: given population height.

Original: Due to data being collected across hundreds of subjects (Kong et al., 2019), HBMs representing population studies feature thousands of RVs, most of which correspond to symmetrical distributions.
Revised: Due to data being collected across hundreds of subjects (Kong et al., 2019), HBMs representing population studies feature thousands of RVs, most of which correspond to symmetrical distributions.

Original: For instance, the distribution of a subject’s height given the group height is the same across all subjects.
Revised: For instance, the distribution of a subject’s height given the group height is the same across all subjects.

Original: As a result, those large HBMs M can be compactly represented via plate-enriched Directed Acyclic Graph (DAG) templates T (Gilks et al., 1994; Koller & Friedman, 2009).
Revised: As a result, those large HBMs M can be compactly represented via plate-enriched Directed Acyclic Graph (DAG) templates T (Gilks et al., 1994; Koller & Friedman, 2009).

Original: The plates in the template T translate graphically the hierarchical i.i.d.
Revised: The plates in the template T translate graphically the hierarchical i.i.d.

Original: structure in the model M .
Revised: structure in the model M .

Original: We exploit this plate-induced repeated structure to design scalable inference methods.
Revised: We exploit this plate-induced repeated structure to design scalable inference methods.

Original: Starting on the left, the graph template
Revised: Starting on the left, the graph template

Original: T is grounded into 2 separate HBMs:
Revised: T is grounded into 2 separate HBMs:

Original: M (top) and M r (down) of respective plate cardinalities (3 , 2) –large– and (2 , 1) –small.
Revised: M (top) and M r (down) of respective plate cardinalities (3 , 2) –large– and (2 , 1) –small.

Original: Based on M r , the reduced distribution q r is constructed.
Revised: Based on M r , the reduced distribution q r is constructed.

Original: We train q r over data slices of small cardinality, before performing inference over the full model of large cardinality.
Revised: We train q r over data slices of small cardinality, before performing inference over the full model of large cardinality.

Original: We denote T ’s vertices as X and Θ = { θ i } i =

Revised: We denote T ’s vertices as X and Θ = { θ i } i =

Original: I .
Revised: I .

Original: X denotes the RVs observed during inference, and Θ the parameters we infer: our goal is to approximate the posterior distribution p
Revised: X denotes the RVs observed during inference, and Θ the parameters we infer: our goal is to approximate the posterior distribution p

Original: (Θ | X ) .
Revised: (Θ | X ) .

Original: We denote T ’s plates as {P p } p =0 ..
Revised: We denote T ’s plates as {P p } p =0 ..

Original: P , and the plates θ i belongs to as Plates( θ i ) .
Revised: P , and the plates θ i belongs to as Plates( θ i ) .

Original: In the toy example from fig.
Revised: In the toy example from fig.

Original: 1, there are 2 latent RV templates: θ 1 and θ 2 , two plates P 0 , P 1 and we want to approximate p ( θ 1 , θ 2 | X ) .
Revised: 1, there are 2 latent RV templates: θ 1 and θ 2 , two plates P 0 , P 1 and we want to approximate p ( θ 1 , θ 2 | X ) .

Original: We build upon the interplay between the template T and the HBM M .
Revised: We build upon the interplay between the template T and the HBM M .

Original: To go from one representation to the other, T can be grounded into M given plate cardinalities { Card( P p ) } p =0 ..
Revised: To go from one representation to the other, T can be grounded into M given plate cardinalities { Card( P p ) } p =0 ..

Original: P (Koller & Friedman, 2009).
Revised: P (Koller & Friedman, 2009).

Original: Card( P ) represents the number of elements in the plate P , for instance the number of groups in the study.
Revised: Card( P ) represents the number of elements in the plate P , for instance the number of groups in the study.

Original: The grounding operation instantiates the repeated structures symbolized by the plates: a given RV template θ i corresponds to multiple ground RVs { θ i,n } n =0 ..
Revised: The grounding operation instantiates the repeated structures symbolized by the plates: a given RV template θ i corresponds to multiple ground RVs { θ i,n } n =0 ..

Original: N i with the same parametric form, where N i = (cid:81)
Revised: N i with the same parametric form, where N i = (cid:81)

Original: P∈ Plates(
Revised: P∈ Plates(

Original: [Equation]
Revised: [Equation]

Original: i ) Card( P ) .
Revised: i ) Card( P ) .

Original: Template grounding is illustrated in fig.
Revised: Template grounding is illustrated in fig.

Original: 1, where T is instantiated into models M .
Revised: 1, where T is instantiated into models M .

Original: In PAVI, we consider the inference over the ground RVs θ i,n as symmetrical problems, over which we share parameterization and learning.
Revised: In PAVI, we consider the inference over the ground RVs θ i,n as symmetrical problems, over which we share parameterization and learning.

Original: Our goal is to perform inference over the full model
Revised: Our goal is to perform inference over the full model

Original: M , which is associated to the density p .
Revised: M , which is associated to the density p .

Original: In p , the plate structure indicates that a RV template θ i is associated to a conditional distribution p i shared across all ground RVs θ i,n :
Revised: In p , the plate structure indicates that a RV template θ i is associated to a conditional distribution p i shared across all ground RVs θ i,n :

Original: [Equation]
Revised: [Equation]

Original: where π ( θ i,n ) are the parents of the RV θ i,n –conditioning its distribution.
Revised: where π ( θ i,n ) are the parents of the RV θ i,n –conditioning its distribution.

Original: We denote with a • X index all variables related to the observed RVs X .
Revised: We denote with a • X index all variables related to the observed RVs X .

Original: M typically features large plate cardinalities
Revised: M typically features large plate cardinalities

Original: Card( P ) , and thus many ground RVs, making it computationally intractable.
Revised: Card( P ) , and thus many ground RVs, making it computationally intractable.

Original: Instead of inferring directly over M , we will train over a smaller replica of M .
Revised: Instead of inferring directly over M , we will train over a smaller replica of M .

Original: To this end, we instantiate the template T into a second HBM M r –the reduced model– of tractable plate cardinalities Card r ( P )
Revised: To this end, we instantiate the template T into a second HBM M r –the reduced model– of tractable plate cardinalities Card r ( P )

Original: ≪ Card( P ) .
Revised: ≪ Card( P ) .

Original: M r has the same template as M , meaning the same dependency structure and the same parametric form for its distributions.
Revised: M r has the same template as M , meaning the same dependency structure and the same parametric form for its distributions.

Original: The only difference lies in M r ’s smaller cardinalities, resulting in fewer ground RVs, as visible in fig.
Revised: The only difference lies in M r ’s smaller cardinalities, resulting in fewer ground RVs, as visible in fig.

Original: We train over the tractable reduced model M r to obtain a variational distribution q usable to approximate the posterior for the intractable target model M .
Revised: We train over the tractable reduced model M r to obtain a variational distribution q usable to approximate the posterior for the intractable target model M .

Original: 2. PAVI ARCHITECTURE 2.2. F ULL VARIATIONAL FAMILY DESIGN
Revised: 2. PAVI ARCHITECTURE 2.2. F ULL VARIATIONAL FAMILY DESIGN

Original: Here we define the variational distribution q –corresponding to the full model M – introduced in the previous section.
Revised: Here we define the variational distribution q –corresponding to the full model M – introduced in the previous section.

Original: q is the distribution used for inference, but that –due to M ’s large cardinalities– we cannot directly train over.
Revised: q is the distribution used for inference, but that –due to M ’s large cardinalities– we cannot directly train over.

Original: To derive q , we push forward the prior p (Θ) using trainable normalizing flows, denoted as F (Papamakarios et al., 2019).
Revised: To derive q , we push forward the prior p (Θ) using trainable normalizing flows, denoted as F (Papamakarios et al., 2019).

Original: To every ground RV θ i,n , we associate a learnable flow F i,n to approximate its posterior distribution:
Revised: To every ground RV θ i,n , we associate a learnable flow F i,n to approximate its posterior distribution:

Original: [Equation]
Revised: [Equation]

Original: where q i,n is the push-forward of the prior p i through the flow F i,n .
Revised: where q i,n is the push-forward of the prior p i through the flow F i,n .

Original: This push-forward is illustrated in fig.
Revised: This push-forward is illustrated in fig.

Original: 2, where flows F push RVs u into θ .
Revised: 2, where flows F push RVs u into θ .

Original: This cascading scheme was first introduced by
Revised: This cascading scheme was first introduced by

Original: Ambrogioni et al.
Revised: Ambrogioni et al.

Original: (2021b), and makes q inherit the conditional dependencies of the prior p .
Revised: (2021b), and makes q inherit the conditional dependencies of the prior p .

Original: More details about the dependencies modelled in the variational distribution can be found in appendix C.5.
Revised: More details about the dependencies modelled in the variational distribution can be found in appendix C.5.

Original: 2.2. P LATE AMORTIZATION
Revised: 2.2. P LATE AMORTIZATION

Original: Here we introduce plate amortization: sharing the parameterization of density estimators across a model’s plates.
Revised: Here we introduce plate amortization: sharing the parameterization of density estimators across a model’s plates.

Original: Traditional VI searches for a distribution q (Θ; ϕ ) that best approximates the posterior of Θ given a value X 0 for X : q (Θ; ϕ 0 ) ≃ p (Θ | X = X 0 ) .
Revised: Traditional VI searches for a distribution q (Θ; ϕ ) that best approximates the posterior of Θ given a value X 0 for X : q (Θ; ϕ 0 ) ≃ p (Θ | X = X 0 ) .

Original: When presented with a new value X 1 , optimization has to be performed again to search for the weights ϕ 1 , such that q (Θ; ϕ 1 ) ≃ p (Θ | X = X 1 ) .
Revised: When presented with a new value X 1 , optimization has to be performed again to search for the weights ϕ 1 , such that q (Θ; ϕ 1 ) ≃ p (Θ | X = X 1 ) .

Original: Sample amortized inference (Zhang et al., 2019; Cremer et al., 2018) instead infers in the general case, regressing the weights ϕ using an encoder f of the observed data X : q (Θ; ϕ
Revised: Sample amortized inference (Zhang et al., 2019; Cremer et al., 2018) instead infers in the general case, regressing the weights ϕ using an encoder f of the observed data X : q (Θ; ϕ

Original: = f ( X )) ≃ p (Θ | X = X ) .
Revised: = f ( X )) ≃ p (Θ | X = X ) .

Original: The cost of learning the encoder weights is amortized since inference over any new sample X requires no additional optimization.
Revised: The cost of learning the encoder weights is amortized since inference over any new sample X requires no additional optimization.

Original: We propose to exploit the concept of amortization, but to apply it at a different granularity, leading to our notion of plate amortization .
Revised: We propose to exploit the concept of amortization, but to apply it at a different granularity, leading to our notion of plate amortization .

Original: Similar to amortizing across the different data samples X , we amortize across the different ground
Revised: Similar to amortizing across the different data samples X , we amortize across the different ground

Original: RVs { θ i,n } n =0 ..
Revised: RVs { θ i,n } n =0 ..

Original: N i corresponding to the same RV template θ i .
Revised: N i corresponding to the same RV template θ i .

Original: Instead of casting every flow F i,n –defined in eq.
Revised: Instead of casting every flow F i,n –defined in eq.

Original: (2)– as an separate, fully-parameterized flow, we will share some parameters across the F i,n .
Revised: (2)– as an separate, fully-parameterized flow, we will share some parameters across the F i,n .

Original: To the template θ i , we associate a conditional flow F i ( · ; ϕ i , • ) with weights ϕ
Revised: To the template θ i , we associate a conditional flow F i ( · ; ϕ i , • ) with weights ϕ

Original: i shared across all the { θ i,n } n =0 ..
Revised: i shared across all the { θ i,n } n =0 ..

Original: N i .
Revised: N i .

Original: The flow F i,n associated to a given ground RV θ i,n will be an instance of this conditional flow, conditioned by an encoding E
Revised: The flow F i,n associated to a given ground RV θ i,n will be an instance of this conditional flow, conditioned by an encoding E

Original: [Equation]
Revised: [Equation]

Original: The distributions q i,n thus have 2 sets of weights, ϕ i and E i,n , creating a parameterization trade-off.
Revised: The distributions q i,n thus have 2 sets of weights, ϕ i and E i,n , creating a parameterization trade-off.

Original: Concentrating all of q i,n ’s parameterization into ϕ
Revised: Concentrating all of q i,n ’s parameterization into ϕ

Original: i results in all the ground RVs θ i,n having the same posterior distribution.
Revised: i results in all the ground RVs θ i,n having the same posterior distribution.

Original: On the contrary, concentrating all of q i,n ’s parameterization into
Revised: On the contrary, concentrating all of q i,n ’s parameterization into

Original: E i,n allows the θ i,n to have completely different posterior distributions.
Revised: E i,n allows the θ i,n to have completely different posterior distributions.

Original: But in a large cardinality setting, this freedom can result in a massive number of weights, proportional to the number of ground RVs times the encoding size.
Revised: But in a large cardinality setting, this freedom can result in a massive number of weights, proportional to the number of ground RVs times the encoding size.

Original: This double parameterization is therefore efficient when the majority of the weights of q i,n is concentrated into ϕ i .
Revised: This double parameterization is therefore efficient when the majority of the weights of q i,n is concentrated into ϕ i .

Original: Using normalizing flows F i , the burden of approximating the correct parametric form for the posterior is placed onto ϕ
Revised: Using normalizing flows F i , the burden of approximating the correct parametric form for the posterior is placed onto ϕ

Original: i , while the E i,n encode lightweight summary statistics specific to each θ i,n .
Revised: i , while the E i,n encode lightweight summary statistics specific to each θ i,n .

Original: In section 2.3.2, we also see that this shared parameterization has synergies with stochastic training.
Revised: In section 2.3.2, we also see that this shared parameterization has synergies with stochastic training.

Original: 2. PAVI STOCHASTIC TRAINING 2.3. R EDUCED DISTRIBUTION AND LOSS
Revised: 2. PAVI STOCHASTIC TRAINING 2.3. R EDUCED DISTRIBUTION AND LOSS

Original: Because of the large number of RVs in M , optimizing over the full distribution q –defined in eq.
Revised: Because of the large number of RVs in M , optimizing over the full distribution q –defined in eq.

Original: (2)– is computationally intractable.
Revised: (2)– is computationally intractable.

Original: Instead, we optimize over a distribution with the smaller cardinalities of the reduced model.
Revised: Instead, we optimize over a distribution with the smaller cardinalities of the reduced model.

Original: At each optimization step t , we randomly choose inside M paths of reduced cardinality, as visible in fig.
Revised: At each optimization step t , we randomly choose inside M paths of reduced cardinality, as visible in fig.

Original: Selecting paths is equivalent to selecting from X a subset X r [ t ] of size N r X , and from Θ a subset Θ r [ t ] .
Revised: Selecting paths is equivalent to selecting from X a subset X r [ t ] of size N r X , and from Θ a subset Θ r [ t ] .

Original: For a given θ i , we denote as B i [ t ] the batch of selected ground
Revised: For a given θ i , we denote as B i [ t ] the batch of selected ground

Original: RVs, of size N r i .
Revised: RVs, of size N r i .

Original: Inferring over Θ r [ t ] , we will simulate training over the full distribution q :
Revised: Inferring over Θ r [ t ] , we will simulate training over the full distribution q :

Original: [Equation]
Revised: [Equation]

Original: where the factor N i
Revised: where the factor N i

Original: /N
Revised: /N

Original: r i simulates that we observe as many ground RVs as in M by repeating the RVs from M r (Hoffman et al., 2013).
Revised: r i simulates that we observe as many ground RVs as in M by repeating the RVs from M r (Hoffman et al., 2013).

Original: Similarly, the loss used at step t is the reduced ELBO constructed using X r [ t ] as observed RVs:
Revised: Similarly, the loss used at step t is the reduced ELBO constructed using X r [ t ] as observed RVs:

Original: [Equation]
Revised: [Equation]

Original: The reduced distribution q r featuresconditional normalizing flows F 1 and Frespectively associated to the RV templates θ 1 and θ 2 .
Revised: The reduced distribution q r featuresconditional normalizing flows F 1 and Frespectively associated to the RV templates θ 1 and θ 2 .

Original: During the stochastic training, q r is instantiated over different branchings of the full model M –highlighted in blue on the left.
Revised: During the stochastic training, q r is instantiated over different branchings of the full model M –highlighted in blue on the left.

Original: The branchings have the cardinalities of
Revised: The branchings have the cardinalities of

Original: M r and change at each stochastic training step t .
Revised: M r and change at each stochastic training step t .

Original: The branching determine the encodings E conditioning the flows F –as symbolised by the letters A, B, C– and the observed data slice –as symbolised by the letters D, E.
Revised: The branching determine the encodings E conditioning the flows F –as symbolised by the letters A, B, C– and the observed data slice –as symbolised by the letters D, E.

Original: This scheme can be viewed as the instantiation of M r over batches of M ’s ground RVs.
Revised: This scheme can be viewed as the instantiation of M r over batches of M ’s ground RVs.

Original: In fig.
Revised: In fig.

Original: 2 we see that q r has the cardinalities of M r , and replicates its conditional dependencies.
Revised: 2 we see that q r has the cardinalities of M r , and replicates its conditional dependencies.

Original: This training is analogous to stochastic VI (Hoffman et al., 2013), generalized with multiple hierarchies and minibatches of RVs.
Revised: This training is analogous to stochastic VI (Hoffman et al., 2013), generalized with multiple hierarchies and minibatches of RVs.

Original: Our novelty lies in the interaction of this stochastic scheme with plate amortization, as explained in the next section.
Revised: Our novelty lies in the interaction of this stochastic scheme with plate amortization, as explained in the next section.

Original: 2.3. S HARING LEARNING ACROSS PLATES
Revised: 2.3. S HARING LEARNING ACROSS PLATES

Original: Here we detail how our shared parameterization –detailed in section 2.2.2– combined with our stochastic training scheme, results in faster inference.
Revised: Here we detail how our shared parameterization –detailed in section 2.2.2– combined with our stochastic training scheme, results in faster inference.

Original: In traditional stochastic VI, every θ i,n corresponding to the same template θ i is associated to individual weights.
Revised: In traditional stochastic VI, every θ i,n corresponding to the same template θ i is associated to individual weights.

Original: Those weights are trained only when θ i,n is visited by the algorithm, that is to say at step t when n ∈ B i [ t ] .
Revised: Those weights are trained only when θ i,n is visited by the algorithm, that is to say at step t when n ∈ B i [ t ] .

Original: As plates become larger, this event becomes rare.
Revised: As plates become larger, this event becomes rare.

Original: If θ i,n is furthermore associated to a highly-parameterized density estimator –such as a normalizing flow– many optimization steps are required for q i,n to converge.
Revised: If θ i,n is furthermore associated to a highly-parameterized density estimator –such as a normalizing flow– many optimization steps are required for q i,n to converge.

Original: The combination of those two items leads to slow training.
Revised: The combination of those two items leads to slow training.

Original: Instead, our idea is to share the learning across the ground RVs θ i,n .
Revised: Instead, our idea is to share the learning across the ground RVs θ i,n .

Original: Due to the problem’s plate structure, we consider the inference over the θ i,n as different instances of a common density estimation task.
Revised: Due to the problem’s plate structure, we consider the inference over the θ i,n as different instances of a common density estimation task.

Original: In PAVI, a large part of the parameterization of the estimators q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) is mutualized via the plate-wide-shared weights ϕ
Revised: In PAVI, a large part of the parameterization of the estimators q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) is mutualized via the plate-wide-shared weights ϕ

Original: i .
Revised: i .

Original: This means that most of the weights of the flows
Revised: This means that most of the weights of the flows

Original: F i,n –concentrated in ϕ i – are trained at every optimization step, across all the selected batches
Revised: F i,n –concentrated in ϕ i – are trained at every optimization step, across all the selected batches

Original: B i [ t ] .
Revised: B i [ t ] .

Original: This results in drastically faster convergence compared to SVI –as seen in exp.
Revised: This results in drastically faster convergence compared to SVI –as seen in exp.

Original: 2. E NCODING SCHEMES
Revised: 2. E NCODING SCHEMES

Original: PAVI shares the parameterization and learning of density estimators across an HBM’s plates.
Revised: PAVI shares the parameterization and learning of density estimators across an HBM’s plates.

Original: In practice the distributions q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) –defined in eq.
Revised: In practice the distributions q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) –defined in eq.

Original: (3)– with different n only differ through the value of the encodings E i,n .
Revised: (3)– with different n only differ through the value of the encodings E i,n .

Original: We detail one core and one alternate scheme to derive those encodings:
Revised: We detail one core and one alternate scheme to derive those encodings:

Original: Free plate encodings (PAVI-F)
Revised: Free plate encodings (PAVI-F)

Original: In our core implementation, E i,n are free weights.
Revised: In our core implementation, E i,n are free weights.

Original: We define encoding arrays with the cardinality of the full model M , one array E i =
Revised: We define encoding arrays with the cardinality of the full model M , one array E i =

Original: [ E i,n ] n =0 ..
Revised: [ E i,n ] n =0 ..

Original: N per template θ i .
Revised: N per template θ i .

Original: This means that an additional ground RV –for instance adding a subject in a population study– requires an additional encoding vector.
Revised: This means that an additional ground RV –for instance adding a subject in a population study– requires an additional encoding vector.

Original: The associated increment in the total number of weights is much lighter than the addition of a fully parameterized normalizing flow –as would be the case in the non-plate-amortized regime.
Revised: The associated increment in the total number of weights is much lighter than the addition of a fully parameterized normalizing flow –as would be the case in the non-plate-amortized regime.

Original: The PAVI-F scheme cannot be sample amortized: when presented with an unseen X , though ϕ i can be kept as an efficient warm start, the optimal values for the encodings E i,n have to be searched again.
Revised: The PAVI-F scheme cannot be sample amortized: when presented with an unseen X , though ϕ i can be kept as an efficient warm start, the optimal values for the encodings E i,n have to be searched again.

Original: During training , the encodings E i,n corresponding to n ∈ B i [ t ] are sliced from the arrays E i and are optimized for along with ϕ i .
Revised: During training , the encodings E i,n corresponding to n ∈ B i [ t ] are sliced from the arrays E i and are optimized for along with ϕ i .

Original: In the toy example from fig.
Revised: In the toy example from fig.

Original: 2, at t = 0 , B 1 [0] =
Revised: 2, at t = 0 , B 1 [0] =

Original: { 1 , 2 } and the trained encodings are { E 1 , 1 , E 1 , 2 } , and at t =
Revised: { 1 , 2 } and the trained encodings are { E 1 , 1 , E 1 , 2 } , and at t =

Original: B 1 [1] = { 0 , 1 } and we train { E 1 , 0 , E 1 , 1 } .
Revised: B 1 [1] = { 0 , 1 } and we train { E 1 , 0 , E 1 , 1 } .

Original: In appendix A.2, we show that the training of PAVI-F is unbiased: training over M r , we converge to the same distribution as if training directly over M .
Revised: In appendix A.2, we show that the training of PAVI-F is unbiased: training over M r , we converge to the same distribution as if training directly over M .

Original: Deep set encoder (PAVI-E)
Revised: Deep set encoder (PAVI-E)

Original: The parameterization of PAVI-F scales lightly but linearly with
Revised: The parameterization of PAVI-F scales lightly but linearly with

Original: Card( P ) .
Revised: Card( P ) .

Original: This scaling could become unaffordable in large population studies.
Revised: This scaling could become unaffordable in large population studies.

Original: We thus propose an alternate scheme –PAVI-E– with a parameterization independent of cardinalities.
Revised: We thus propose an alternate scheme –PAVI-E– with a parameterization independent of cardinalities.

Original: In this more experimental scheme, free encodings are replaced by an encoder f with weights η applied to the observed data: E = f
Revised: In this more experimental scheme, free encodings are replaced by an encoder f with weights η applied to the observed data: E = f

Original: ( X ; η ) .
Revised: ( X ; η ) .

Original: As encoder f we use a deep-set architecture exploiting the data’s plate-induced permutation invariance –detailed in appendix A.1.3 (Zaheer et al., 2018; Lee et al., 2019).
Revised: As encoder f we use a deep-set architecture exploiting the data’s plate-induced permutation invariance –detailed in appendix A.1.3 (Zaheer et al., 2018; Lee et al., 2019).

Original: This scheme allows for sample amortization across different data samples X 0 , X 1 , ... .
Revised: This scheme allows for sample amortization across different data samples X 0 , X 1 , ... .

Original: Note that an encoder will be used to generate the encodings whether inference is sample amortized or not.
Revised: Note that an encoder will be used to generate the encodings whether inference is sample amortized or not.

Original: During training, shared learning is further amplified as all the architecture’s weights – ϕ i and η – are trained at every step t .
Revised: During training, shared learning is further amplified as all the architecture’s weights – ϕ i and η – are trained at every step t .

Original: To collect the encodings to plug into q r , we build up on a property of f : set size generalization (Zaheer et al., 2018).
Revised: To collect the encodings to plug into q r , we build up on a property of f : set size generalization (Zaheer et al., 2018).

Original: Instead of encoding the full-sized data X , f is applied to the slice X r
Revised: Instead of encoding the full-sized data X , f is applied to the slice X r

Original: [ t ] .
Revised: [ t ] .

Original: This amounts to aggregating summary statistics across a subset of the observed data (Lee et al., 2019; Agrawal & Domke, 2021).
Revised: This amounts to aggregating summary statistics across a subset of the observed data (Lee et al., 2019; Agrawal & Domke, 2021).

Original: This property is intensified in the sample amortized context: we train a sample amortized family over the lightweight model M r , and use it ”for free” over the heavyweight model M .
Revised: This property is intensified in the sample amortized context: we train a sample amortized family over the lightweight model M r , and use it ”for free” over the heavyweight model M .

Original: As detailed in appendix A.2, we rely on a computationally efficient but theoretically biased training scheme for PAVI-E.
Revised: As detailed in appendix A.2, we rely on a computationally efficient but theoretically biased training scheme for PAVI-E.

Original: The negative impact of this bias on PAVI-E’s performance was seldom noticeable and always marginal throughout our experiments.
Revised: The negative impact of this bias on PAVI-E’s performance was seldom noticeable and always marginal throughout our experiments.

Original: Summary
Revised: Summary

Original: In section 2.2 we derived an architecture sharing its parameterization across a model’s plates.
Revised: In section 2.2 we derived an architecture sharing its parameterization across a model’s plates.

Original: In section 2.3 we derived a stochastic scheme to train this architecture over batches of data.
Revised: In section 2.3 we derived a stochastic scheme to train this architecture over batches of data.

Original: Our novelty lies in this original combination of amortization and stochastic training, which results in significantly faster inference, as demonstrated in the following experiments.
Revised: Our novelty lies in this original combination of amortization and stochastic training, which results in significantly faster inference, as demonstrated in the following experiments.

Original:  R ESULTS AND DISCUSSION
Revised:  R ESULTS AND DISCUSSION

Original: Throughout this section we use the ELBO as a proxy to the KL divergence between the variational posterior and the unknown true posterior.
Revised: Throughout this section we use the ELBO as a proxy to the KL divergence between the variational posterior and the unknown true posterior.

Original: ELBO is measured across 20 samples X , 5 random seeds per sample.
Revised: ELBO is measured across 20 samples X , 5 random seeds per sample.

Original: The ELBO allows to compare the relative performance of different architectures on a given inference problem.
Revised: The ELBO allows to compare the relative performance of different architectures on a given inference problem.

Original: In appendix B.1 we provide with sanity checks to assess the quality of the results.
Revised: In appendix B.1 we provide with sanity checks to assess the quality of the results.

Original: In appendix B.2 we evaluate the impact of the reduced model cardinalities on performance.
Revised: In appendix B.2 we evaluate the impact of the reduced model cardinalities on performance.

Original: In appendix B.3 we compare numerically our method against baselines over a Gaussian mixture model; a model featuring the aggregation of higher order summary statistics; and a smaller version of our Neuroimaging model used in section 3.4.
Revised: In appendix B.3 we compare numerically our method against baselines over a Gaussian mixture model; a model featuring the aggregation of higher order summary statistics; and a smaller version of our Neuroimaging model used in section 3.4.

Original: Experimentally, we found the PAVI-F scheme overall faster to train and yielding better inference quality than the PAVI-E scheme.
Revised: Experimentally, we found the PAVI-F scheme overall faster to train and yielding better inference quality than the PAVI-E scheme.

Original: When its parameterization is affordable, PAVI-F should be preferred.
Revised: When its parameterization is affordable, PAVI-F should be preferred.

Original: PAVI-E nevertheless opens up promising research directions, with the potential for parameterization-constant, time-constant sample-amortized inference as cardinality augments.
Revised: PAVI-E nevertheless opens up promising research directions, with the potential for parameterization-constant, time-constant sample-amortized inference as cardinality augments.

Original: 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 exp.
Revised: 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 exp.

Original: 3.3&B.3.
Revised: 3.3&B.3.

Original: 3. P LATE AMORTIZATION AND CONVERGENCE SPEED
Revised: 3. P LATE AMORTIZATION AND CONVERGENCE SPEED

Original: In this experiment, we illustrate how plate amortization results in faster training.
Revised: In this experiment, we illustrate how plate amortization results in faster training.

Original: We consider the following Gaussian Random Effects model (GRE):
Revised: We consider the following Gaussian Random Effects model (GRE):

Original: [Equation]
Revised: [Equation]

Original: where D represents the data X ’s feature size, with group means θ 1 and population means θ 2 as
Revised: where D represents the data X ’s feature size, with group means θ 1 and population means θ 2 as

Original: D-dimensional Gaussians.
Revised: D-dimensional Gaussians.

Original: We opt in this equation for a double indexing scheme instead of a simple indexing as in our methods.
Revised: We opt in this equation for a double indexing scheme instead of a simple indexing as in our methods.

Original: The GRE model features two nested plates: the group plate P 1 and the sample plate P 0 as in fig.
Revised: The GRE model features two nested plates: the group plate P 1 and the sample plate P 0 as in fig.

Original: Inferring over the GRE model, the objective is to retrieve the posterior distribution of the group and population means given the observed sample.
Revised: Inferring over the GRE model, the objective is to retrieve the posterior distribution of the group and population means given the observed sample.

Original: [Figure]
Revised: [Figure]

Original: Here we set D = 8 , Card( P 1 ) = 100 and Card r ( P 1 )
Revised: Here we set D = 8 , Card( P 1 ) = 100 and Card r ( P 1 )

Original: =
Revised: =

Original: We compare our PAVI architecture to a stochastic non-plate-amortized baseline with the same architecture as PAVI (Hoffman et al., 2013).
Revised: We compare our PAVI architecture to a stochastic non-plate-amortized baseline with the same architecture as PAVI (Hoffman et al., 2013).

Original: The only difference is that ground RVs θ i,n are associated in the baseline to individual fully-parameterized flows F i,n instead of sharing the same conditional flow F i –as described in section 2.2.2.
Revised: The only difference is that ground RVs θ i,n are associated in the baseline to individual fully-parameterized flows F i,n instead of sharing the same conditional flow F i –as described in section 2.2.2.

Original: Figure 3 (left) displays the evolution of the ELBO for the baseline and PAVI with free encoding (PAVI-F) and deep set encoders (PAVI-E).
Revised: Figure 3 (left) displays the evolution of the ELBO for the baseline and PAVI with free encoding (PAVI-F) and deep set encoders (PAVI-E).

Original: We see that both plate amortized methods reach asymptotic ELBO equal to the non-plate-amortized baseline’s, but with orders of magnitudes faster convergence, and more numerical stability.
Revised: We see that both plate amortized methods reach asymptotic ELBO equal to the non-plate-amortized baseline’s, but with orders of magnitudes faster convergence, and more numerical stability.

Original: This stems from the individual flows F i,n in the baseline only being trained when the corresponding θ i,n is visited by the stochastic training, while the shared flow F i is updated at every optimization step in PAVI.
Revised: This stems from the individual flows F i,n in the baseline only being trained when the corresponding θ i,n is visited by the stochastic training, while the shared flow F i is updated at every optimization step in PAVI.

Original: We also note that the PAVI-E scheme converges faster in theory than PAVI-F –in terms of number of steps– by sharing not only the training of the flows, but also of the encoder through the optimization.
Revised: We also note that the PAVI-E scheme converges faster in theory than PAVI-F –in terms of number of steps– by sharing not only the training of the flows, but also of the encoder through the optimization.

Original: In practice however, computing the encodings results in longer steps, and in slower convergence, as illustrated in section 3.3.
Revised: In practice however, computing the encodings results in longer steps, and in slower convergence, as illustrated in section 3.3.

Original: 3. I MPACT OF ENCODING SIZE
Revised: 3. I MPACT OF ENCODING SIZE

Original: Here we illustrate the role of encodings as ground RV posterior’s summary statistics –as described in section 2.2.2.
Revised: Here we illustrate the role of encodings as ground RV posterior’s summary statistics –as described in section 2.2.2.

Original: We use the GRE HBM detailed in eq.
Revised: We use the GRE HBM detailed in eq.

Original: (6), using D = 8 , Card
Revised: (6), using D = 8 , Card

Original: ( P 1 )
Revised: ( P 1 )

Original: = 20 and
Revised: = 20 and

Original: Card r ( P 1 ) =
Revised: Card r ( P 1 ) =

Original: We use a single PAVI-F architecture, varying the size of the encodings E i,n –see section 2.4.
Revised: We use a single PAVI-F architecture, varying the size of the encodings E i,n –see section 2.4.

Original: Due to plate amortization, encodings determines how much individual information each RV θ i,n is associated to.
Revised: Due to plate amortization, encodings determines how much individual information each RV θ i,n is associated to.

Original: The encoding size –varying from 2 to 16 – is to be compared with the dimensionality of the problem, D = 8 .
Revised: The encoding size –varying from 2 to 16 – is to be compared with the dimensionality of the problem, D = 8 .

Original: In GRE, D = 8 corresponds to the size of the sufficient statistics needed to reconstruct the posterior of a group mean –all other statistics such as the variance being shared between the group means.
Revised: In GRE, D = 8 corresponds to the size of the sufficient statistics needed to reconstruct the posterior of a group mean –all other statistics such as the variance being shared between the group means.

Original: Figure 3 (right) shows how the asymptotic performance steadily increases when the encoding size augments, and plateaus once reaching the sufficient summary statistic size D = 8 .
Revised: Figure 3 (right) shows how the asymptotic performance steadily increases when the encoding size augments, and plateaus once reaching the sufficient summary statistic size D = 8 .

Original: Interestingly, increasing the encoding size also leads to faster convergence: redundancy can likely be exploited in the optimization.
Revised: Interestingly, increasing the encoding size also leads to faster convergence: redundancy can likely be exploited in the optimization.

Original: Encoding size appears as a unequivocal hyperparameter allowing to trade inference quality for computational efficiency.
Revised: Encoding size appears as a unequivocal hyperparameter allowing to trade inference quality for computational efficiency.

Original: Increasing the encoding size also leads experimentally to diminishing returns in terms of performance.
Revised: Increasing the encoding size also leads experimentally to diminishing returns in terms of performance.

Original: This property can be exploited in large settings to drastically reduce the memory footprint of inference while maintaining acceptable performance –choosing the encoding size approximately equal to the expected size of the sufficient statistics.
Revised: This property can be exploited in large settings to drastically reduce the memory footprint of inference while maintaining acceptable performance –choosing the encoding size approximately equal to the expected size of the sufficient statistics.

Original: 3. S CALING WITH PLATE CARDINALITIES
Revised: 3. S CALING WITH PLATE CARDINALITIES

Original: Here we put in perspective the gains from plate amortization when scaling up an inference problem’s cardinality.
Revised: Here we put in perspective the gains from plate amortization when scaling up an inference problem’s cardinality.

Original: We consider the GRE model in eq.
Revised: We consider the GRE model in eq.

Original: (6) with D = 2 and augment the plate cardinalities
Revised: (6) with D = 2 and augment the plate cardinalities

Original: (Card( P 1 ) , Card r ( P 1 )) : (2 , 1) → (20 , 5) → (200 , 20) .
Revised: (Card( P 1 ) , Card r ( P 1 )) : (2 , 1) → (20 , 5) → (200 , 20) .

Original: In doing
Revised: In doing

Original: → 42 →
Revised: → 42 →

Original: Baselines We compare our PAVI architecture against 3 state-of-the-art baselines.
Revised: Baselines We compare our PAVI architecture against 3 state-of-the-art baselines.

Original: Cascading Flows (CF) (Ambrogioni et al., 2021b) is a non-plate-amortized structured VI architecture improving on the baseline presented in section 3.1.
Revised: Cascading Flows (CF) (Ambrogioni et al., 2021b) is a non-plate-amortized structured VI architecture improving on the baseline presented in section 3.1.

Original: CF push the prior p into the posterior q using Highway
Revised: CF push the prior p into the posterior q using Highway

Original: Flows .
Revised: Flows .

Original: CF follows a cascading dependency structure, complemented with a backward auxiliary coupling.
Revised: CF follows a cascading dependency structure, complemented with a backward auxiliary coupling.

Original: ADAVI (Rouillard & Wassermann, 2022) is a structured VI architecture similar to
Revised: ADAVI (Rouillard & Wassermann, 2022) is a structured VI architecture similar to

Original: PAVI, with constant parameterization with respect to a problem’s cardinality, but large training times and memory.
Revised: PAVI, with constant parameterization with respect to a problem’s cardinality, but large training times and memory.

Original: ADAVI has several limitations compared to PAVI: ADAVI implements a
Revised: ADAVI has several limitations compared to PAVI: ADAVI implements a

Original: Mean Field approximation (Blei et al., 2017) while PAVI implements a cascading flow; ADAVI is limited pyramidal HBMs while PAVI tackles generic plate-enriched HBMs; ADAVI is limited to a full-model sample-amortized variant.
Revised: Mean Field approximation (Blei et al., 2017) while PAVI implements a cascading flow; ADAVI is limited pyramidal HBMs while PAVI tackles generic plate-enriched HBMs; ADAVI is limited to a full-model sample-amortized variant.

Original: Unbiased Implicit VI (UIVI) is an unstructured implicit VI architecture.
Revised: Unbiased Implicit VI (UIVI) is an unstructured implicit VI architecture.

Original: UIVI infers over the full parameter space –without any SVI-amenable factorization– by reparameterizing a base distribution with a stochastic transform.
Revised: UIVI infers over the full parameter space –without any SVI-amenable factorization– by reparameterizing a base distribution with a stochastic transform.

Original: For all architectures, we indicate with the suffix (sa) sample amortization , corresponding to the classical meaning of amortization, as detailed in section 2.2.2.
Revised: For all architectures, we indicate with the suffix (sa) sample amortization , corresponding to the classical meaning of amortization, as detailed in section 2.2.2.

Original: More implementation details can be found in appendix B.4.
Revised: More implementation details can be found in appendix B.4.

Original: As the cardinality of the problem augments, fig.
Revised: As the cardinality of the problem augments, fig.

Original: 4 shows how PAVI maintains a stateof-the-art inference quality, while being more computationally attractive.
Revised: 4 shows how PAVI maintains a stateof-the-art inference quality, while being more computationally attractive.

Original: Specifically, in terms of parameterization, both ADAVI and PAVI-E provide with a heavyweight but constant parameterization as the cardinality Card( P 1 ) of the problem augments.
Revised: Specifically, in terms of parameterization, both ADAVI and PAVI-E provide with a heavyweight but constant parameterization as the cardinality Card( P 1 ) of the problem augments.

Original: Comparatively, both CF and PAVI-F’s parameterization scale linearly with Card( P 1 ) , but with a drastically lighter augmentation for PAVI-F. For an additional ground RV, CF requires an additional fully parameterized normalizing flow, whereas PAVI-F only requires an additional lightweight encoding vector.
Revised: Comparatively, both CF and PAVI-F’s parameterization scale linearly with Card( P 1 ) , but with a drastically lighter augmentation for PAVI-F. For an additional ground RV, CF requires an additional fully parameterized normalizing flow, whereas PAVI-F only requires an additional lightweight encoding vector.

Original: UIVI’s parameterization scales quadratically with the size of the parameter space Θ , due to a neural network regressing a transform applied to a base distribution with the size of Θ .
Revised: UIVI’s parameterization scales quadratically with the size of the parameter space Θ , due to a neural network regressing a transform applied to a base distribution with the size of Θ .

Original: In detail,
Revised: In detail,

Original: PAVI-F’s parameterization due to the plate-wide-shared ϕ 1 represents a constant ≈ 2 k weights, while the part due to the encodings E 1 ,n grows linearly from 16 to 160 to 1 .
Revised: PAVI-F’s parameterization due to the plate-wide-shared ϕ 1 represents a constant ≈ 2 k weights, while the part due to the encodings E 1 ,n grows linearly from 16 to 160 to 1 .

Original: 6 k weights.
Revised: 6 k weights.

Original: Note that
Revised: Note that

Original: PAVI’s stochastic training also allows for a controlled GPU memory during optimization, removing the need for a larger memory as the cardinality of the problem augments –a hardware constraint that can become unaffordable at very large cardinalities.
Revised: PAVI’s stochastic training also allows for a controlled GPU memory during optimization, removing the need for a larger memory as the cardinality of the problem augments –a hardware constraint that can become unaffordable at very large cardinalities.

Original: To remove this memory constraint, CF could be trained stochastically, but –without plate amortization– would suffer from slower inference, as illustrated in section 3.1.
Revised: To remove this memory constraint, CF could be trained stochastically, but –without plate amortization– would suffer from slower inference, as illustrated in section 3.1.

Original: In contrast, UIVI could not be trained stochastically, as it infers over the full Θ at once instead of factorizing it.
Revised: In contrast, UIVI could not be trained stochastically, as it infers over the full Θ at once instead of factorizing it.

Original: As a result UIVI would be ultimately limited by memory to infer over larger problems.
Revised: As a result UIVI would be ultimately limited by memory to infer over larger problems.

Original: In terms of convergence speed , PAVI benefits from plate amortization to have orders of magnitude faster convergence compared to structured VI baselines CF and ADAVI.
Revised: In terms of convergence speed , PAVI benefits from plate amortization to have orders of magnitude faster convergence compared to structured VI baselines CF and ADAVI.

Original: For UIVI, as the cardinality augments, training amounts to evaluating a neural network with increasingly larger layers.
Revised: For UIVI, as the cardinality augments, training amounts to evaluating a neural network with increasingly larger layers.

Original: GPU training time is therefore constant, but this property would not translate to larger problems as the GPU memory would become insufficient.
Revised: GPU training time is therefore constant, but this property would not translate to larger problems as the GPU memory would become insufficient.

Original: Plate amortization is particularly significant for the PAVI-E(sa) scheme, in which a sample-amortized variational family is trained over a dataset of reduced cardinality, yet performs ”for free” inference over a HBM of large cardinality.
Revised: Plate amortization is particularly significant for the PAVI-E(sa) scheme, in which a sample-amortized variational family is trained over a dataset of reduced cardinality, yet performs ”for free” inference over a HBM of large cardinality.

Original: Maintaining Card r ( P 1 ) constant while Card( P 1 ) augments allows for a constant parameterization and training time as the cardinality of the problem augments –without any maximum cardinality, contrary to UIVI.
Revised: Maintaining Card r ( P 1 ) constant while Card( P 1 ) augments allows for a constant parameterization and training time as the cardinality of the problem augments –without any maximum cardinality, contrary to UIVI.

Original: The effect of plate amortization is particularly noticeable at
Revised: The effect of plate amortization is particularly noticeable at

Original: Card( P 1 ) = 200 between the PAVI(sa) and CF(sa) architectures, where PAVI performs amortized inference with 10 × fewer weights and 100 × lower training time.
Revised: Card( P 1 ) = 200 between the PAVI(sa) and CF(sa) architectures, where PAVI performs amortized inference with 10 × fewer weights and 100 × lower training time.

Original: Scaling even higher the cardinality of the problem – Card( P 1 ) = 2000 for instance– renders ADAVI, CF and UIVI computationally intractable, while PAVI maintains a light memory footprint, and a short training time, as exemplified in the next experiment.
Revised: Scaling even higher the cardinality of the problem – Card( P 1 ) = 2000 for instance– renders ADAVI, CF and UIVI computationally intractable, while PAVI maintains a light memory footprint, and a short training time, as exemplified in the next experiment.

Original: 3. N EUROIMAGING APPLICATION : B ROCA ’ S AREA PARCELLATION
Revised: 3. N EUROIMAGING APPLICATION : B ROCA ’ S AREA PARCELLATION

Original: To illustrate its usefulness, we apply PAVI to a challenging population study for Broca’s area’s functional parcellation .
Revised: To illustrate its usefulness, we apply PAVI to a challenging population study for Broca’s area’s functional parcellation .

Original: A parcellation clusters brain vertices into different connectivity networks : labels describing co-activation with the rest of the brain –as measured using functional Magnetic
Revised: A parcellation clusters brain vertices into different connectivity networks : labels describing co-activation with the rest of the brain –as measured using functional Magnetic

Original: Resonance Imaging (fMRI).
Revised: Resonance Imaging (fMRI).

Original: Different subjects exhibit a strong variability, as in fig.
Revised: Different subjects exhibit a strong variability, as in fig.

Original: However, fMRI is costly to acquire: few noisy data is usually gathered for a given subject.
Revised: However, fMRI is costly to acquire: few noisy data is usually gathered for a given subject.

Original: It is thus essential to combine information across subjects and to display uncertainty in the results.
Revised: It is thus essential to combine information across subjects and to display uncertainty in the results.

Original: Those 2 points motivate Hierarchical Bayesian Modelling and VI in Neuroimaging (Kong et al., 2019): we search the posteriors of connectivity networks and vertex labels, measuring fMRI over a large cohort of subjects.
Revised: Those 2 points motivate Hierarchical Bayesian Modelling and VI in Neuroimaging (Kong et al., 2019): we search the posteriors of connectivity networks and vertex labels, measuring fMRI over a large cohort of subjects.

Original: We use the HCP dataset (Van Essen et al., 2012): 2 acquisitions from a 1000 subjects, with thousands of measures per acquisition, for over a million parameters Θ .
Revised: We use the HCP dataset (Van Essen et al., 2012): 2 acquisitions from a 1000 subjects, with thousands of measures per acquisition, for over a million parameters Θ .

Original: We use a model with 3 plates: subjects, sessions and brain vertices.
Revised: We use a model with 3 plates: subjects, sessions and brain vertices.

Original: None of the baselines presented in section 3.
Revised: None of the baselines presented in section 3.

Original: –CF
Revised: –CF

Original: , ADAVI, UIVI– can computationally tackle this high cardinality problem.
Revised: , ADAVI, UIVI– can computationally tackle this high cardinality problem.

Original: We nevertheless show superior performance over those baselines over a tractable problem size  in appendix B.3.3.
Revised: We nevertheless show superior performance over those baselines over a tractable problem size with 2 thousand parameters in appendix B.3.3.

Original: Despite the massive dimensionality of the problem, thanks to plate amortization PAVI converges in a dozen epochs, under an hour of GPU time.
Revised: Despite the massive dimensionality of the problem, thanks to plate amortization PAVI converges in a dozen epochs, under an hour of GPU time.

Original: Results are visible in fig.
Revised: Results are visible in fig.

Original: 5, supporting the hypothesis of a functional bi-partition of Broca’s area into a posterior part involved in phonology and an anterior part in lexical/semantic processing –following the anatomy of pars opercularis and triangularis (Heim et al., 2009; Zhang et al., 2020).
Revised: 5, supporting the hypothesis of a functional bi-partition of Broca’s area into a posterior part involved in phonology and an anterior part in lexical/semantic processing –following the anatomy of pars opercularis and triangularis (Heim et al., 2009; Zhang et al., 2020).

Original: Conclusion In this work we present the novel PAVI architecture, combining a structured variational family and a stochastic training scheme.
Revised: Conclusion In this work we present the novel PAVI architecture, combining a structured variational family and a stochastic training scheme.

Original: PAVI is based the concept of plate amortization, allowing to share parameterization and learning across a model’s plates.
Revised: PAVI is based the concept of plate amortization, allowing to share parameterization and learning across a model’s plates.

Original: We demonstrated the positive impact of plate amortization on training speed and scaling to large plate cardinality regimes, making a significant step towards scalable, expressive Variational Inference.
Revised: We demonstrated the positive impact of plate amortization on training speed and scaling to large plate cardinality regimes, making a significant step towards scalable, expressive Variational Inference.

Original: R EPRODUCIBILITY STATEMENT
Revised: R EPRODUCIBILITY STATEMENT

Original: All experiments were performed in Python using the Tensorflow Probability library (Dillon et al., 2017).
Revised: All experiments were performed in Python using the Tensorflow Probability library (Dillon et al., 2017).

Original: All experiments were conducted on computational cluster nodes equipped with a Tesla V10016Gb GPU and 4 AMD EPYC 7742 64-Core processors.
Revised: All experiments were conducted on computational cluster nodes equipped with a Tesla V10016Gb GPU and 4 AMD EPYC 7742 64-Core processors.

Original: VRAM intensive experiments in fig.were performed on an Ampere 100 PCIE-40Gb GPU.
Revised: VRAM intensive experiments in fig.were performed on an Ampere 100 PCIE-40Gb GPU.

Original: Appendix B.4 lists implementation details for all our synthetic experiments.
Revised: Appendix B.4 lists implementation details for all our synthetic experiments.

Original: Appendix B.5 is related to our Neuroimaging experiment 3.4, and details both our data pre-processing steps and our implementation.
Revised: Appendix B.5 is related to our Neuroimaging experiment 3.4, and details both our data pre-processing steps and our implementation.

Original: As part of our submission we furthermore packaged and release the code associated to our experiments.
Revised: As part of our submission we furthermore packaged and release the code associated to our experiments.

Original: E THICS STATEMENT
Revised: E THICS STATEMENT

Original: Our Neuroimaging data come from the Human Connectome Project dataset (Van Essen et al., 2012).
Revised: Our Neuroimaging data come from the Human Connectome Project dataset (Van Essen et al., 2012).

Original: All data in the HCP is strongly anonymized, as per the HCP protocols.
Revised: All data in the HCP is strongly anonymized, as per the HCP protocols.

Original: We used in this paper only
Revised: We used in this paper only

Original: Open Access imaging data data, following the HCP data use terms.
Revised: Open Access imaging data data, following the HCP data use terms.

Original: S UPPLEMENTAL M ATERIAL A S UPPLEMENTAL METHODS A. PAVI IMPLEMENTATION DETAILS A.1. P LATE BRANCHINGS AND STOCHASTIC TRAINING
Revised: S UPPLEMENTAL M ATERIAL A S UPPLEMENTAL METHODS A. PAVI IMPLEMENTATION DETAILS A.1. P LATE BRANCHINGS AND STOCHASTIC TRAINING

Original: As exposed in section 2.3.1, at each optimization step t we randomly select branchings inside the full model M , branchings over which we instantiate the reduced model M r .
Revised: As exposed in section 2.3.1, at each optimization step t we randomly select branchings inside the full model M , branchings over which we instantiate the reduced model M r .

Original: In doing so, we define batches B i [ t ] for the RV templates θ i .
Revised: In doing so, we define batches B i [ t ] for the RV templates θ i .

Original: Those batches have to be coherent with one another: they have to respect the conditional dependencies of the original model M .
Revised: Those batches have to be coherent with one another: they have to respect the conditional dependencies of the original model M .

Original: As an example, if a ground
Revised: As an example, if a ground

Original: RV is selected as part of M r , then its parent RV needs to be selected as well.
Revised: RV is selected as part of M r , then its parent RV needs to be selected as well.

Original: To ensure this, during the stochastic training we do not sample RVs directly but plates:
Revised: To ensure this, during the stochastic training we do not sample RVs directly but plates:

Original: For every plate P p , we sample without replacement Card r ( P p ) indices amongst the Card( P p ) possible indices.
Revised: For every plate P p , we sample without replacement Card r ( P p ) indices amongst the Card( P p ) possible indices.

Original: Then, for every RV template θ i , we select the ground RVs θ i,n corresponding to the sampled indices for the plates
Revised: Then, for every RV template θ i , we select the ground RVs θ i,n corresponding to the sampled indices for the plates

Original: Plates( θ i ) .
Revised: Plates( θ i ) .

Original: The selected ground RVs θ i,n constitute the set Θ r [ t ] of parameters appearing in eq.
Revised: The selected ground RVs θ i,n constitute the set Θ r [ t ] of parameters appearing in eq.

Original: The same procedure yields the observed RV subset X r [ t ] and the data slice X
Revised: The same procedure yields the observed RV subset X r [ t ] and the data slice X

Original: r
Revised: r

Original: [ t ] .
Revised: [ t ] .

Original: For instance, in the toy example from fig. 1, X 2 will be chosen iff the index 1 is selected as part of sub-sampling P 1 and the index 0 is selected as part of sub-sampling P 0 .
Revised: For instance, in the toy example from fig. 1, X 2 will be chosen iff the index 1 is selected as part of sub-sampling P 1 and the index 0 is selected as part of sub-sampling P 0 .

Original: Less formally, this is equivalent to going middle , then left in the full graph representing M .
Revised: Less formally, this is equivalent to going middle , then left in the full graph representing M .

Original: This stochastic choice is illustrated in fig.
Revised: This stochastic choice is illustrated in fig.

Original: 2 at t = 1 where X 2 corresponds to the node E .
Revised: 2 at t = 1 where X 2 corresponds to the node E .

Original: This stochastic strategy also applies to the selected encoding scheme –described in section 2.4&2.3.2– as detailed in the next sections.
Revised: This stochastic strategy also applies to the selected encoding scheme –described in section 2.4&2.3.2– as detailed in the next sections.

Original: A.1. PAVI-F DETAILS
Revised: A.1. PAVI-F DETAILS

Original: In section 2.4 we refer to encodings E i =
Revised: In section 2.4 we refer to encodings E i =

Original: [ E i,n ] n =0 ..
Revised: [ E i,n ] n =0 ..

Original: N corresponding to RV templates θ i .
Revised: N corresponding to RV templates θ i .

Original: In practice, we have some amount of sharing for those encodings: instead of defining separate encodings for every RV template, we define encodings for every plate level .
Revised: In practice, we have some amount of sharing for those encodings: instead of defining separate encodings for every RV template, we define encodings for every plate level .

Original: A plate level is a combination of plates with at least one parameter RV template θ i belonging to it:
Revised: A plate level is a combination of plates with at least one parameter RV template θ i belonging to it:

Original: [Equation]
Revised: [Equation]

Original: For every plate level, we construct a large encoding array with the cardinalities of the full model M :
Revised: For every plate level, we construct a large encoding array with the cardinalities of the full model M :

Original: [Equation]
Revised: [Equation]

Original: Where D is an encoding size that we kept constant to de-clutter the notation but can vary between plate levels.
Revised: Where D is an encoding size that we kept constant to de-clutter the notation but can vary between plate levels.

Original: The encodings for a given ground RV θ i,n then correspond to an element from the encoding array E i .
Revised: The encodings for a given ground RV θ i,n then correspond to an element from the encoding array E i .

Original: A.1. PAVI-E DETAILS
Revised: A.1. PAVI-E DETAILS

Original: In the PAVI-E scheme, encodings are not free weights but the output of en encoder f ( · , η ) applied to the observed data X .
Revised: In the PAVI-E scheme, encodings are not free weights but the output of en encoder f ( · , η ) applied to the observed data X .

Original: In this section we detail the design of this encoder.
Revised: In this section we detail the design of this encoder.

Original: As in the previous section, the role of the encoder will be to produce one encoding per plate level.
Revised: As in the previous section, the role of the encoder will be to produce one encoding per plate level.

Original: We start from a dependency structure for the plate levels:
Revised: We start from a dependency structure for the plate levels:

Original: [Equation]
Revised: [Equation]

Original: note that this dependency structure is in the backward direction: a plate level will be the parent of another plate level, if the former contains a RV who has a child in the latter.
Revised: note that this dependency structure is in the backward direction: a plate level will be the parent of another plate level, if the former contains a RV who has a child in the latter.

Original: We therefore obtain a plate level dependency structure that reverts the conditional dependency structure of the graph template T .
Revised: We therefore obtain a plate level dependency structure that reverts the conditional dependency structure of the graph template T .

Original: To avoid redundant paths in this dependency structure, we take the maximum branching of the obtained graph.
Revised: To avoid redundant paths in this dependency structure, we take the maximum branching of the obtained graph.

Original: Given the plate level dependency structure, we will recursively construct the encodings, starting from the observed data:
Revised: Given the plate level dependency structure, we will recursively construct the encodings, starting from the observed data:

Original: [Equation]
Revised: [Equation]

Original: where x is the observed data for the RV x , and ρ is a simple encoder that processes every observed ground
Revised: where x is the observed data for the RV x , and ρ is a simple encoder that processes every observed ground

Original: RV’s value independently through an identical multi-layer perceptron.
Revised: RV’s value independently through an identical multi-layer perceptron.

Original: Then, until we have exhausted all plate levels, we process existing encodings to produce new encodings:
Revised: Then, until we have exhausted all plate levels, we process existing encodings to produce new encodings:

Original: [Equation]
Revised: [Equation]

Original: where g is the composition of attention-based deep-set networks called Set Transformers (Lee et al., 2019; Zaheer et al., 2018).
Revised: where g is the composition of attention-based deep-set networks called Set Transformers (Lee et al., 2019; Zaheer et al., 2018).

Original: For every plate P p present in the parent plate level but absent in the child plate level, g will compute summary statistics across that plate, effectively contracting the corresponding batch dimensionality in the parent encoding (Rouillard & Wassermann, 2022).
Revised: For every plate P p present in the parent plate level but absent in the child plate level, g will compute summary statistics across that plate, effectively contracting the corresponding batch dimensionality in the parent encoding (Rouillard & Wassermann, 2022).

Original: In the case of multiple observed RVs, we run this ”backward pass” independently for each observed data –with one encoder per observed RV.
Revised: In the case of multiple observed RVs, we run this ”backward pass” independently for each observed data –with one encoder per observed RV.

Original: We then concatenate the resulting encodings corresponding to the same plate level.
Revised: We then concatenate the resulting encodings corresponding to the same plate level.

Original: For more precise implementation details, we invite the reader to consult the codebase released with this supplemental material.
Revised: For more precise implementation details, we invite the reader to consult the codebase released with this supplemental material.

Original: A. A NALYSIS OF BIAS IN THE STOCHASTIC TRAINING
Revised: A. A NALYSIS OF BIAS IN THE STOCHASTIC TRAINING

Original: A key concern in our stochastic training scheme is its unbiasedness: we want our stochastic optimization to converge to the same variational posterior as if we trained over the full model directly –without any stochasticity.
Revised: A key concern in our stochastic training scheme is its unbiasedness: we want our stochastic optimization to converge to the same variational posterior as if we trained over the full model directly –without any stochasticity.

Original: In this section we first show that the PAVI-F is unbiased.
Revised: In this section we first show that the PAVI-F is unbiased.

Original: Second, we identify strategies to obtain an unbiased PAVI-E scheme, yet show how the approximations we do in practice can theoretically result in biased training.
Revised: Second, we identify strategies to obtain an unbiased PAVI-E scheme, yet show how the approximations we do in practice can theoretically result in biased training.

Original: As an important note, the negative impact of this bias on the performance PAVI-E remained limited throughout our experiments –as seen in experiments 3.3, B.2, B.1 and B.3.
Revised: As an important note, the negative impact of this bias on the performance PAVI-E remained limited throughout our experiments –as seen in experiments 3.3, B.2, B.1 and B.3.

Original: A.2. G ENERAL DERIVATION ( APPLICABLE TO BOTH THE PAVI-F AND PAVI-E SCHEMES )
Revised: A.2. G ENERAL DERIVATION ( APPLICABLE TO BOTH THE PAVI-F AND PAVI-E SCHEMES )

Original: We first formalize the plate sampling strategy described in appendix A.1.1.
Revised: We first formalize the plate sampling strategy described in appendix A.1.1.

Original: To every plate P we associate the RV I P corresponding to the Card r ( P ) -sized set of indices sampled without replacement from the Card( P ) possible index values.
Revised: To every plate P we associate the RV I P corresponding to the Card r ( P ) -sized set of indices sampled without replacement from the Card( P ) possible index values.

Original: As an example, with a plate P 0 with Card( P 0 ) =
Revised: As an example, with a plate P 0 with Card( P 0 ) =

Original: and Card r ( P 0 )
Revised: and Card r ( P 0 )

Original: = 2 , { 0 , 2 } or { 2 , 3 } can be 2 different samples from I P 0 .
Revised: = 2 , { 0 , 2 } or { 2 , 3 } can be 2 different samples from I P 0 .

Original: At a given optimization step t , we sample independently from the RVs { I P p } p =0 ..P .
Revised: At a given optimization step t , we sample independently from the RVs { I P p } p =0 ..P .

Original: This defines the batches B i [ t ] and the distribution q r in eq.
Revised: This defines the batches B i [ t ] and the distribution q r in eq.

Original: To check the unbiasedness of our stochastic training, we need to show that:
Revised: To check the unbiasedness of our stochastic training, we need to show that:

Original: [Equation]
Revised: [Equation]

Original: Where:
Revised: Where:

Original: [Equation]
Revised: [Equation]

Original: And ELBO r [ t ] is defined in eq.
Revised: And ELBO r [ t ] is defined in eq.

Original: In that expression, q and p have symmetrical roles.
Revised: In that expression, q and p have symmetrical roles.

Original: As the
Revised: As the

Original: ELBO amounts to the difference between the logarithms of distributions p and q , we can prove the equality in eq.
Revised: ELBO amounts to the difference between the logarithms of distributions p and q , we can prove the equality in eq.

Original: (A.6) if we prove that the expectation of each reduced distribution is equal to the corresponding full distribution.
Revised: (A.6) if we prove that the expectation of each reduced distribution is equal to the corresponding full distribution.

Original: To prove the equality in eq.
Revised: To prove the equality in eq.

Original: (A.6), a sufficient condition is therefore to prove that:
Revised: (A.6), a sufficient condition is therefore to prove that:

Original: [Equation]
Revised: [Equation]

Original: where to de-clutter the notations we denote the expectation over the collection of RVs
Revised: where to de-clutter the notations we denote the expectation over the collection of RVs

Original: Consider a given ground RV θ i,n corresponding to the RV template θ i and to the plates
Revised: Consider a given ground RV θ i,n corresponding to the RV template θ i and to the plates

Original: Plates( θ i ) .
Revised: Plates( θ i ) .

Original: At a given stochastic step t , θ i,n will be chosen if and only if its corresponding branching is chosen.
Revised: At a given stochastic step t , θ i,n will be chosen if and only if its corresponding branching is chosen.

Original: Recall that when sampling equiprobably without replacement a set of k elements from a population of n elements, a given element will be present in the set with probability k/n .
Revised: Recall that when sampling equiprobably without replacement a set of k elements from a population of n elements, a given element will be present in the set with probability k/n .

Original: We can apply this reasoning to the choice of branching corresponding to a given ground RV.
Revised: We can apply this reasoning to the choice of branching corresponding to a given ground RV.

Original: For instance, in fig.
Revised: For instance, in fig.

Original: 1, Xwill be chosen iff the index 1 is selected as part of sub-sampling P 1 and the index 0 is selected as part of sub-sampling P 0 .
Revised: 1, Xwill be chosen iff the index 1 is selected as part of sub-sampling P 1 and the index 0 is selected as part of sub-sampling P 0 .

Original: As Card r ( P 1 ) = 2 indices are chosen inside the plate P 1 of full cardinality 3, and Card r ( P 0 )
Revised: As Card r ( P 1 ) = 2 indices are chosen inside the plate P 1 of full cardinality 3, and Card r ( P 0 )

Original: = 1 indices are chosen inside the plate P 0 of full cardinality 2, X 2 is therefore chosen with probability 2 / 3 × 1 / 2 .
Revised: = 1 indices are chosen inside the plate P 0 of full cardinality 2, X 2 is therefore chosen with probability 2 / 3 × 1 / 2 .

Original: More formally, for a ground RV θ i,n we have:
Revised: More formally, for a ground RV θ i,n we have:

Original: [Equation]
Revised: [Equation]

Original: Plates( θ i )
Revised: Plates( θ i )

Original: [Equation]
Revised: [Equation]

Original: Applying this reasoning to every RV template θ i , we have that:
Revised: Applying this reasoning to every RV template θ i , we have that:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: where we exploited the fact that the expectation of the sum of RVs is the sum of the expectations, even in the case of dependent RVs.
Revised: where we exploited the fact that the expectation of the sum of RVs is the sum of the expectations, even in the case of dependent RVs.

Original: The term 1 n ∈B i [ t ] × log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) is the product of 2 RVs –related to the stochastic choice of plate indices:
Revised: The term 1 n ∈B i [ t ] × log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) is the product of 2 RVs –related to the stochastic choice of plate indices:

Original: • the RV 1 n ∈B i [ t ] is an indicator that θ i,n ’s branching has been chosen via the stochastic sampling of plate indices.
Revised: • the RV 1 n ∈B i [ t ] is an indicator that θ i,n ’s branching has been chosen via the stochastic sampling of plate indices.

Original: By construction, this RV depends only on the indices of the plates
Revised: By construction, this RV depends only on the indices of the plates

Original: P ∈ Plates( θ i ) .
Revised: P ∈ Plates( θ i ) .

Original: • the RV log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) depends on E i,n , whose construction depends on the encoding scheme: – In the PAVI-F scheme, E i,n is a constant.
Revised: • the RV log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) depends on E i,n , whose construction depends on the encoding scheme: – In the PAVI-F scheme, E i,n is a constant.

Original: – In the PAVI-E scheme, E i,n results of the application of an encoder to the observed data of a subset of θ i,n ’s descendants.
Revised: – In the PAVI-E scheme, E i,n results of the application of an encoder to the observed data of a subset of θ i,n ’s descendants.

Original: By construction, this subset will only depend on the indices of plates containing θ
Revised: By construction, this subset will only depend on the indices of plates containing θ

Original: i ’s descendants, but not containing θ i .
Revised: i ’s descendants, but not containing θ i .

Original: The value of E i,n therefore only depends on the indices of plates P / ∈ Plates( θ i )
Revised: The value of E i,n therefore only depends on the indices of plates P / ∈ Plates( θ i )

Original: As an example of this reasoning, consider the model M illustrated in fig.
Revised: As an example of this reasoning, consider the model M illustrated in fig.

Original: We can evaluate both terms for the ground RV θ 1 , 2 in the PAVI-E scheme:
Revised: We can evaluate both terms for the ground RV θ 1 , 2 in the PAVI-E scheme:

Original: • 1 2 ∈B 1 [ t ] depends on whether the index 2 is chosen as part of sub-sampling the plate P 1 , and therefore only depends on the RV I P 1 .
Revised: • 1 2 ∈B 1 [ t ] depends on whether the index 2 is chosen as part of sub-sampling the plate P 1 , and therefore only depends on the RV I P 1 .

Original: In this case the associated probability is 2 / 3 ; • to evaluate log q 1 , 2 ( θ 1 , 2 | θ 2 , 0 ; ϕ 1 , E 1 , 2 ) , the value of E 1 , 2 will result from the application of the encoder f over the value of either X 4 or X 5 .
Revised: In this case the associated probability is 2 / 3 ; • to evaluate log q 1 , 2 ( θ 1 , 2 | θ 2 , 0 ; ϕ 1 , E 1 , 2 ) , the value of E 1 , 2 will result from the application of the encoder f over the value of either X 4 or X 5 .

Original: This choice
Revised: This choice

Original: depends on whether the index 0 or 1 is chosen as part of sub-sampling the plate P 0 .
Revised: depends on whether the index 0 or 1 is chosen as part of sub-sampling the plate P 0 .

Original: Therefore, the value of the term log q 1 , 2 only depends on the RV I P 0 .
Revised: Therefore, the value of the term log q 1 , 2 only depends on the RV I P 0 .

Original: In both PAVI-F and PAVI-E– the terms 1 n ∈B i [ t ] and log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) are therefore independent, meaning that the expectation of their product can be rewritten as the product of their expectations:
Revised: In both PAVI-F and PAVI-E– the terms 1 n ∈B i [ t ] and log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n ) are therefore independent, meaning that the expectation of their product can be rewritten as the product of their expectations:

Original: [Equation]
Revised: [Equation]

Original: This equality can be further simplified in the PAVI-F case –proving its unbiasedness– but not in the
Revised: This equality can be further simplified in the PAVI-F case –proving its unbiasedness– but not in the

Original: PAVI-E case, as detailed in the sections below.
Revised: PAVI-E case, as detailed in the sections below.

Original: A.2. U NBIASEDNESS OF THE PAVI-F SCHEME
Revised: A.2. U NBIASEDNESS OF THE PAVI-F SCHEME

Original: In the PAVI-F scheme, detailed in section 2.4, the encodings E i,n are constants with respect to the branching choice, therefore we have:
Revised: In the PAVI-F scheme, detailed in section 2.4, the encodings E i,n are constants with respect to the branching choice, therefore we have:

Original: [Equation]
Revised: [Equation]

Original: which proves eq.
Revised: which proves eq.

Original: (A.8) and eq.
Revised: (A.8) and eq.

Original: (A.6).
Revised: (A.6).

Original: In the above example of θ 1 , 2 in M , in the PAVI-F scheme the expression E
Revised: In the above example of θ 1 , 2 in M , in the PAVI-F scheme the expression E

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n )] can be evaluated into log q 1 , 2 ( θ 1 , 2 | θ 2 , 0 ; ϕ 1 , E 1 , 2 ) .
Revised: i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n )] can be evaluated into log q 1 , 2 ( θ 1 , 2 | θ 2 , 0 ; ϕ 1 , E 1 , 2 ) .

Original: This demonstrates that the PAVI-F scheme is unbiased: training over stochastically chosen subgraphs for q r is in expectation equal to training over the full graph of q .
Revised: This demonstrates that the PAVI-F scheme is unbiased: training over stochastically chosen subgraphs for q r is in expectation equal to training over the full graph of q .

Original: A.2. A PPROXIMATIONS IN THE PAVI-E SCHEME
Revised: A.2. A PPROXIMATIONS IN THE PAVI-E SCHEME

Original: In the PAVI-E scheme, detailed in section 2.3.2, the encodings E i,n are computed from the observed data X .
Revised: In the PAVI-E scheme, detailed in section 2.3.2, the encodings E i,n are computed from the observed data X .

Original: Specifically, considering the ground RV θ i,n , we have E i,n = f
Revised: Specifically, considering the ground RV θ i,n , we have E i,n = f

Original: ( X r i,n [ t ]) where
Revised: ( X r i,n [ t ]) where

Original: X r i,n [ t ] corresponds to the observed data of a subset of θ i,n ’s descendants.
Revised: X r i,n [ t ] corresponds to the observed data of a subset of θ i,n ’s descendants.

Original: Depending on the chosen branching downstream of θ i,n , the value of E i,n can therefore vary.
Revised: Depending on the chosen branching downstream of θ i,n , the value of E i,n can therefore vary.

Original: This means we cannot further simplify eq.
Revised: This means we cannot further simplify eq.

Original: (A.11): the terms log q i,n ( θ i,n | π ( θ i,n ); E i,n ) are not constants with respect to the RVs I P .
Revised: (A.11): the terms log q i,n ( θ i,n | π ( θ i,n ); E i,n ) are not constants with respect to the RVs I P .

Original: In the above example of θ 1 , 2 in M , in the PAVI-E scheme the expression E
Revised: In the above example of θ 1 , 2 in M , in the PAVI-E scheme the expression E

Original: I P [log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n )] can be evaluated into:
Revised: I P [log q i,n ( θ i,n | π ( θ i,n ); ϕ i , E i,n )] can be evaluated into:

Original: [Equation]
Revised: [Equation]

Original: How could the PAVI-E scheme be made unbiased?
Revised: How could the PAVI-E scheme be made unbiased?

Original: Specifically, by making the value of E i,n independent of the choice of downstream branching.
Revised: Specifically, by making the value of E i,n independent of the choice of downstream branching.

Original: A possibility would be to parameterize E i,n as an average –an expectation– over all the possible sub-branchings downstream of θ i,n .
Revised: A possibility would be to parameterize E i,n as an average –an expectation– over all the possible sub-branchings downstream of θ i,n .

Original: Yet, in practical cases, the cardinalities of the reduced model are much inferior to the ones of the full model:
Revised: Yet, in practical cases, the cardinalities of the reduced model are much inferior to the ones of the full model:

Original: Card r ( P ) ≪ Card( P ) .
Revised: Card r ( P ) ≪ Card( P ) .

Original: This means that numerous Card r ( P ) -sized subsets can be chosen inside the Card( P ) possible descendants.
Revised: This means that numerous Card r ( P ) -sized subsets can be chosen inside the Card( P ) possible descendants.

Original: In order to average over all those subset choices to compute E i,n , numerous encoding calculations would be required at each stochastic training step.
Revised: In order to average over all those subset choices to compute E i,n , numerous encoding calculations would be required at each stochastic training step.

Original: For large-scale cases, we deemed this possibility impractical.
Revised: For large-scale cases, we deemed this possibility impractical.

Original: Other possibilities could exist, all revolving around the problematic of aggregating collections of stochastic estimators into one general estimator –in an unbiased and efficient manner.
Revised: Other possibilities could exist, all revolving around the problematic of aggregating collections of stochastic estimators into one general estimator –in an unbiased and efficient manner.

Original: To our knowledge, this is a complex and still open research question, whose advancement could much benefit our applications.
Revised: To our knowledge, this is a complex and still open research question, whose advancement could much benefit our applications.

Original: Practical approximation for the PAVI-E scheme In practice, we compute the encoding E i,n based on the single downstream branching corresponding to the sampling of the RVs I P .
Revised: Practical approximation for the PAVI-E scheme In practice, we compute the encoding E i,n based on the single downstream branching corresponding to the sampling of the RVs I P .

Original: Compared to the previous paragraph, this amounts to estimating the expectation of E i,n –over all downstream branchings– using a single one of those branchings.
Revised: Compared to the previous paragraph, this amounts to estimating the expectation of E i,n –over all downstream branchings– using a single one of those branchings.

Original: Note that, even if this encoding estimate was unbiased, log q i,n would remain an highly non-linear function of E i,n .
Revised: Note that, even if this encoding estimate was unbiased, log q i,n would remain an highly non-linear function of E i,n .

Original: As a consequence, we need to rely on the approximation:
Revised: As a consequence, we need to rely on the approximation:

Original: [Equation]
Revised: [Equation]

Original: which can theoretically introduce some bias in our gradients.
Revised: which can theoretically introduce some bias in our gradients.

Original: The approximation eq.
Revised: The approximation eq.

Original: (A.13) can be interpreted as follow: ”the expectation of the density of θ i,n when collecting summary statistics over a stochastic subset of θ i,n ’s descendants is approximately equal to the density of θ i,n when collecting summary statistics over the entirety of θ i,n ’s descendants”.
Revised: (A.13) can be interpreted as follow: ”the expectation of the density of θ i,n when collecting summary statistics over a stochastic subset of θ i,n ’s descendants is approximately equal to the density of θ i,n when collecting summary statistics over the entirety of θ i,n ’s descendants”.

Original: Another interpretation is that the distribution associated to the summary of the full data can be approximated by an annealing of the distributions associated to summaries of subsets of this data.
Revised: Another interpretation is that the distribution associated to the summary of the full data can be approximated by an annealing of the distributions associated to summaries of subsets of this data.

Original: In practice, this approximation did not yield significantly worse performance for the PAVI-E scheme over the generative models we tested.
Revised: In practice, this approximation did not yield significantly worse performance for the PAVI-E scheme over the generative models we tested.

Original: At the same time, computing the encodings over a single branching allows to compute all the E i,n encodings in a single lightweight pass over the data X r [ t ] .
Revised: At the same time, computing the encodings over a single branching allows to compute all the E i,n encodings in a single lightweight pass over the data X r [ t ] .

Original: This simple solution therefore provided with a substantial increase in training speed with seldom noticeable bias.
Revised: This simple solution therefore provided with a substantial increase in training speed with seldom noticeable bias.

Original: Yet, we do not bar the existence of pathological generative HBMs where this approximation would become coarse.
Revised: Yet, we do not bar the existence of pathological generative HBMs where this approximation would become coarse.

Original: Experimenters should bear in mind this possibility when using the PAVI-E scheme.
Revised: Experimenters should bear in mind this possibility when using the PAVI-E scheme.

Original: 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 onto the real problem instantiating the same generative model with a larger dimensionality.
Revised: 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 onto the real problem instantiating the same generative model with a larger dimensionality.

Original: A. PAVI ALGORITHMS
Revised: A. PAVI ALGORITHMS

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

Original: A.3. A RCHITECTURE BUILD
Revised: A.3. A RCHITECTURE BUILD

Original: Output: for i = 1 ..
Revised: Output: for i = 1 ..

Original: I do
Revised: I do

Original: Construct conditional flow F i ;
Revised: Construct conditional flow F i ;

Original: Define conditional posterior distributions q i,n as the push-forward of the prior via F i , following eq.
Revised: Define conditional posterior distributions q i,n as the push-forward of the prior via F i , following eq.

Original: Combine the q i,n distributions following the cascading flows scheme, as in section 2.2.
Revised: Combine the q i,n distributions following the cascading flows scheme, as in section 2.2.

Original: (Ambrogioni et al., 2021b) ; if PAVI-F encoding scheme then
Revised: (Ambrogioni et al., 2021b) ; if PAVI-F encoding scheme then

Original: Construct encoding arrays { E i =
Revised: Construct encoding arrays { E i =

Original: [ E i,n ] n =0 ..
Revised: [ E i,n ] n =0 ..

Original: N i } i =1 ..
Revised: N i } i =1 ..

Original: I as in appendix A.1.2 ; else if PAVI-E encoding scheme then
Revised: I as in appendix A.1.2 ; else if PAVI-E encoding scheme then

Original: Construct encoder f as in appendix A.1.3 ;
Revised: Construct encoder f as in appendix A.1.3 ;

Original: A.3. S TOCHASTIC TRAINING
Revised: A.3. S TOCHASTIC TRAINING

Original: if then batches B else if PAVI-E encoding scheme then encodings as E f
Revised: if then batches B else if PAVI-E encoding scheme then encodings as E f

Original: X r [ t
Revised: X r [ t

Original: Feed obtained encodings
Revised: Feed obtained encodings

Original: Update conditional if then
Revised: Update conditional if then

Original: Update encodings
Revised: Update encodings

Original: ; else
Revised: ; else

Original: Update
Revised: Update

Original: . I NFERENCE
Revised: . I NFERENCE

Original: Input: trained architecture q , observed data X , encoding scheme
Revised: Input: trained architecture q , observed data X , encoding scheme

Original: Output: approximate posterior distribution if PAVI-F encoding scheme then
Revised: Output: approximate posterior distribution if PAVI-F encoding scheme then

Original: Collect full encoding arrays
Revised: Collect full encoding arrays

Original: E i ; else if PAVI-E encoding scheme then
Revised: E i ; else if PAVI-E encoding scheme then

Original: Compute encodings set size generalization ; encodings
Revised: Compute encodings set size generalization ; encodings

Original: A. I NFERENCE GAPS
Revised: A. I NFERENCE GAPS

Original: In terms of inference quality, the impact of our architecture can be formalized following the gaps terminology (Cremer et al., 2018).
Revised: In terms of inference quality, the impact of our architecture can be formalized following the gaps terminology (Cremer et al., 2018).

Original: Consider a joint distribution p (Θ , X ) , and a value X for the RV template X .
Revised: Consider a joint distribution p (Θ , X ) , and a value X for the RV template X .

Original: We pick a variational family Q , and in this family look for the parametric distribution q (Θ; ϕ ) that best approximates p (Θ | X = X ) .
Revised: We pick a variational family Q , and in this family look for the parametric distribution q (Θ; ϕ ) that best approximates p (Θ | X = X ) .

Original: 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.
Revised: 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.

Original: (2018) refer to as the gap G :
Revised: (2018) refer to as the gap G :

Original: [Equation]
Revised: [Equation]

Original: We denote q ∗ (Θ; ϕ ∗ ) the optimal distribution inside Q that minimizes the KL divergence with the true posterior:
Revised: We denote q ∗ (Θ; ϕ ∗ ) the optimal distribution inside Q that minimizes the KL divergence with the true posterior:

Original: [Equation]
Revised: [Equation]

Original: The approximation gap G approx depends on the expressivity of the variational family Q , specifically whether Q contains distributions arbitrarily close to the posterior –in the KL sense.
Revised: The approximation gap G approx depends on the expressivity of the variational family Q , specifically whether Q contains distributions arbitrarily close to the posterior –in the KL sense.

Original: Note: G approx is a property of the variational family Q .
Revised: Note: G approx is a property of the variational family Q .

Original: G approx is an asymptotic bound for the KL divergence between any distribution q ∈ Q and the true posterior.
Revised: G approx is an asymptotic bound for the KL divergence between any distribution q ∈ Q and the true posterior.

Original: This gap is therefore a form of bias, but is not to be mistaken with the stochasticity-induced bias studied in appendix A.2.
Revised: This gap is therefore a form of bias, but is not to be mistaken with the stochasticity-induced bias studied in appendix A.2.

Original: The bias in appendix A.2 relates to whether q ∗ can be found by training stochastically over Q , whereas G approx relates to the the bias between q ∗ and the true posterior.
Revised: The bias in appendix A.2 relates to whether q ∗ can be found by training stochastically over Q , whereas G approx relates to the the bias between q ∗ and the true posterior.

Original: Cremer et al.
Revised: Cremer et al.

Original: (2018) demonstrate that, in the case of sample amortized inference, when the weights ϕ no longer are free but the output of an encoder f ∈ F , inference cannot be better than in the non-sample-amortized case, and a positive amortization gap is introduced:
Revised: (2018) demonstrate that, in the case of sample amortized inference, when the weights ϕ no longer are free but the output of an encoder f ∈ F , inference cannot be better than in the non-sample-amortized case, and a positive amortization gap is introduced:

Original: [Equation]
Revised: [Equation]

Original: Where we denote as η ∗ the optimal weights for the encoder f inside the function family F .
Revised: Where we denote as η ∗ the optimal weights for the encoder f inside the function family F .

Original: 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 X .”
Revised: 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 X .”

Original: Using the same gap terminology, we can define gaps implied by our PAVI architecture.
Revised: Using the same gap terminology, we can define gaps implied by our PAVI architecture.

Original: Instead of picking the distribution q inside the family Q , consider picking q from the plate-amortized family
Revised: Instead of picking the distribution q inside the family Q , consider picking q from the plate-amortized family

Original: Q pa corresponding to Q .
Revised: Q pa corresponding to Q .

Original: Distributions in Q pa are distributions from Q with the additional constraints that some weights have to be equal.
Revised: Distributions in Q pa are distributions from Q with the additional constraints that some weights have to be equal.

Original: Consequently, Q pa is a subset of Q :
Revised: Consequently, Q pa is a subset of Q :

Original: [Equation]
Revised: [Equation]

Original: As such, looking for the optimal distribution inside Q pa instead of inside Q cannot result in better performance, leading to a plate amortization gap :
Revised: As such, looking for the optimal distribution inside Q pa instead of inside Q cannot result in better performance, leading to a plate amortization gap :

Original: [Equation]
Revised: [Equation]

Original: Where we denote as ψ ∗ the optimal weights for a variational distribution q inside Q pa –in the KL sense.
Revised: Where we denote as ψ ∗ the optimal weights for a variational distribution q inside Q pa –in the KL sense.

Original: The equation A.18 is valid for the PAVI-F scheme –see section 2.4.
Revised: The equation A.18 is valid for the PAVI-F scheme –see section 2.4.

Original: We can interpret it as follow: ”theoretically, plate amortization cannot be beneficial in terms of KL divergence for the inference over a given sample X ”.
Revised: We can interpret it as follow: ”theoretically, plate amortization cannot be beneficial in terms of KL divergence for the inference over a given sample X ”.

Original: Now consider that encodings are no longer free parameters but the output of an encoder f .
Revised: Now consider that encodings are no longer free parameters but the output of an encoder f .

Original: Similar to the case presented in eq.
Revised: Similar to the case presented in eq.

Original: (A.16), using an encoder cannot result in better performance, leading to an encoder gap :
Revised: (A.16), using an encoder cannot result in better performance, leading to an encoder gap :

Original: [Equation]
Revised: [Equation]

Original: The equation eq.
Revised: The equation eq.

Original: (A.19) is valid for the PAVI-E scheme –see section 2.4.
Revised: (A.19) is valid for the PAVI-E scheme –see section 2.4.

Original: The most complex case is the PAVI-E(sa) scheme, where we combine both plate and sample amortization.
Revised: The most complex case is the PAVI-E(sa) scheme, where we combine both plate and sample amortization.

Original: Our argument cannot account for the resulting G PAVI-E(sa) gap: both the PAVI-E and PAVIE(sa) schemes rely upon the same encoder f .
Revised: Our argument cannot account for the resulting G PAVI-E(sa) gap: both the PAVI-E and PAVIE(sa) schemes rely upon the same encoder f .

Original: In the PAVI-E scheme, f is overfit over a dataset composed of the slices of a given data sample X .
Revised: In the PAVI-E scheme, f is overfit over a dataset composed of the slices of a given data sample X .

Original: In the PAVI-E(sa) scheme, the encoder is trained over the whole distribution of the samples of the reduced model M r .
Revised: In the PAVI-E(sa) scheme, the encoder is trained over the whole distribution of the samples of the reduced model M r .

Original: 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.
Revised: 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.

Original: Comparing previous equations, we therefore have:
Revised: Comparing previous equations, we therefore have:

Original: [Equation]
Revised: [Equation]

Original: Note that those are theoretical results, that do not necessarily pertain to optimization in practice.
Revised: Note that those are theoretical results, that do not necessarily pertain to optimization in practice.

Original: In particular, in section 3.1&3.3, this theoretical performance loss is not observed empirically over the studied examples.
Revised: In particular, in section 3.1&3.3, this theoretical performance loss is not observed empirically over the studied examples.

Original: 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 fig.
Revised: 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 fig.

Original: 4 or experiments B.3.
Revised: 4 or experiments B.3.

Original: 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.
Revised: 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.

Original: 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 .
Revised: 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 .

Original: A less expressive architecture –using plate amortization– may in practice yield better performance.
Revised: A less expressive architecture –using plate amortization– may in practice yield better performance.

Original: 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.
Revised: 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.

Original: B S UPPLEMENTAL RESULTS B. GRE RESULTS SANITY CHECK
Revised: B S UPPLEMENTAL RESULTS B. GRE RESULTS SANITY CHECK

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

Original: [Equation]
Revised: [Equation]

Original: Given that the likelihood term log p ( X ) does not depend on the variational family q , differences in ELBOs directly transfer in differences in KL divergence, and provide with a straightforward metric to compare different variational posteriors.
Revised: Given that the likelihood term log p ( X ) does not depend on the variational family q , differences in ELBOs directly transfer in differences in KL divergence, and provide with a straightforward metric to compare different variational posteriors.

Original: Nonetheless, the ELBO doesn’t provide with an absolute metric of quality.
Revised: Nonetheless, the ELBO doesn’t provide with an absolute metric of quality.

Original: As a sanity check, we want to assert the quality of the results presented in section 3.3 –that are transferable to section 3.1&3.2, based on the same model.
Revised: As a sanity check, we want to assert the quality of the results presented in section 3.3 –that are transferable to section 3.1&3.2, based on the same model.

Original: In fig.
Revised: In fig.

Original: B.1 we plot the posterior samples of various methods against analytical ground truths, using the Card( P 1 ) =case.
Revised: B.1 we plot the posterior samples of various methods against analytical ground truths, using the Card( P 1 ) =case.

Original: All the method’s results are aligned with the analytical ground truth, with differences in ELBO translating meaningful qualitative differences in terms of inference quality.
Revised: All the method’s results are aligned with the analytical ground truth, with differences in ELBO translating meaningful qualitative differences in terms of inference quality.

Original: B. E FFECT OF THE REDUCED MODEL CARDINALITIES ON THE TRAINING EFFICIENCY
Revised: B. E FFECT OF THE REDUCED MODEL CARDINALITIES ON THE TRAINING EFFICIENCY

Original: In fig.
Revised: In fig.

Original: B.2 we show the impact of the augmentation of Card r ( P 1 ) on the efficiency of the variational posterior’s optimization.
Revised: B.2 we show the impact of the augmentation of Card r ( P 1 ) on the efficiency of the variational posterior’s optimization.

Original: 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.
Revised: 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.

Original: Indeed, training over larger cardinalities does make each optimization step slightly slower but also make the ELBO gradient estimates less noisy and allows to train more encoding vectors
Revised: Indeed, training over larger cardinalities does make each optimization step slightly slower but also make the ELBO gradient estimates less noisy and allows to train more encoding vectors

Original: E i,n at a given optimization step.
Revised: E i,n at a given optimization step.

Original: In the case of the PAVI-E scheme, the training speed is constant with respect to Card r ( P 1 ) .
Revised: In the case of the PAVI-E scheme, the training speed is constant with respect to Card r ( P 1 ) .

Original: This is due to the compute of the encodings being vectorized across plates in our deep-set encoder – see appendix A.1.3 (Lee et al., 2019).
Revised: This is due to the compute of the encodings being vectorized across plates in our deep-set encoder – see appendix A.1.3 (Lee et al., 2019).

Original: We observe a slight if barely noticeable reduction of the inference bias when augmenting Card r ( P 1 ) .
Revised: We observe a slight if barely noticeable reduction of the inference bias when augmenting Card r ( P 1 ) .

Original: Similar to the PAVI-F scheme, Card r ( P 1 ) should be maximized with respect to the GPU memory constraint.
Revised: Similar to the PAVI-F scheme, Card r ( P 1 ) should be maximized with respect to the GPU memory constraint.

Original: The intuition behind this choice is that the generalization of the learning of summary statistics across sets of data points is easier the closer the reduced set size is to the full set size.
Revised: The intuition behind this choice is that the generalization of the learning of summary statistics across sets of data points is easier the closer the reduced set size is to the full set size.

Original: 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.
Revised: 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.

Original: This example also displays a pathological choice for Card r
Revised: This example also displays a pathological choice for Card r

Original: ( P 1 ) : the encoder fails to learn to compute the correct summary statistics over sets of size 1.
Revised: ( P 1 ) : the encoder fails to learn to compute the correct summary statistics over sets of size 1.

Original: 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.
Revised: 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.

Original: When generalizing to moments of higher order –such as the variance– there is therefore a theoretical lower bound to consider when fixating the value of Card r ( P 1 ) .
Revised: When generalizing to moments of higher order –such as the variance– there is therefore a theoretical lower bound to consider when fixating the value of Card r ( P 1 ) .

Original: Our intuition is that the more complex to estimate the statistic is, the larger Card r ( P 1 ) should be.
Revised: Our intuition is that the more complex to estimate the statistic is, the larger Card r ( P 1 ) should be.

Original: Note that in practice non-stochastic VI is intractable for large-scale models, because of its memory requirement.
Revised: Note that in practice non-stochastic VI is intractable for large-scale models, because of its memory requirement.

Original: In large scale experiments such as the one presented in fig.
Revised: In large scale experiments such as the one presented in fig.

Original: 5, stochastic training is necessary.
Revised: 5, stochastic training is necessary.

Original: Our claim is to be faster than non-plate-amortized stochastic VI in those large-scale contexts –but not necessarily to be faster then non-stochastic VI in the small-scale regime of this experiment.
Revised: Our claim is to be faster than non-plate-amortized stochastic VI in those large-scale contexts –but not necessarily to be faster then non-stochastic VI in the small-scale regime of this experiment.

Original: For PAVI-F we nonetheless obtain similar training speed compared to non-stochastic
Revised: For PAVI-F we nonetheless obtain similar training speed compared to non-stochastic

Original: VI with Card r ( P 1 ) as small as 8 .
Revised: VI with Card r ( P 1 ) as small as 8 .

Original: For PAVI-E, the stochastic training is as fast as the non-stochastic one.
Revised: For PAVI-E, the stochastic training is as fast as the non-stochastic one.

Original: [Figure]
Revised: [Figure]

Original: 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).
Revised: 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).

Original: We furthermore observe a slight reduction of the inference bias as Card r ( P 1 ) augments.
Revised: We furthermore observe a slight reduction of the inference bias as Card r ( P 1 ) augments.

Original: 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 )
Revised: 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 )

Original: = 1 .
Revised: = 1 .

Original: Interestingly, the training speed is approximately constant across all values of Card r ( P 1 ) .
Revised: Interestingly, the training speed is approximately constant across all values of Card r ( P 1 ) .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: . A DDITIONAL COMPARATIVE EXPERIMENTS B.3. G AUSSIAN M IXTURE MODEL
Revised: . A DDITIONAL COMPARATIVE EXPERIMENTS B.3. G AUSSIAN M IXTURE MODEL

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

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where Mixture([ D 1 , . . .
Revised: where Mixture([ D 1 , . . .

Original: , D L ] , π ) denotes a mixture between the distributions [ D 1 , . . .
Revised: , D L ] , π ) denotes a mixture between the distributions [ D 1 , . . .

Original: , D L ] with mixture weights π .
Revised: , D L ] with mixture weights π .

Original: Results are visible in table
Revised: Results are visible in table

Original: [Equation]
Revised: [Equation]

Original: where PAVI displays the best asymptotic ELBO, as well as the shortest optimization time.
Revised: where PAVI displays the best asymptotic ELBO, as well as the shortest optimization time.

Original: 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.
Revised: 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.

Original: B.3. H IERARCHICAL V ARIANCE MODEL
Revised: B.3. H IERARCHICAL V ARIANCE MODEL

Original: 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.
Revised: 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.

Original: Our goal is in particular to evaluate a potential empirical bias for the PAVI-E scheme –as studied in appendix A.2.
Revised: Our goal is in particular to evaluate a potential empirical bias for the PAVI-E scheme –as studied in appendix A.2.

Original: We will refer the following HBM as the Hierarchical Variance model:
Revised: We will refer the following HBM as the Hierarchical Variance model:

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: This is the context in which we would expect to observe the most bias due to the approximation in eq.
Revised: This is the context in which we would expect to observe the most bias due to the approximation in eq.

Original: (A.13).
Revised: (A.13).

Original: The performance of
Revised: The performance of

Original: PAVI-E remains nonetheless competitive.
Revised: PAVI-E remains nonetheless competitive.

Original: This illustrates how PAVI-E’s theoretical bias –introduced in the stochastic training– does not result in significantly worse inference compared to state-of-theart architectures.
Revised: This illustrates how PAVI-E’s theoretical bias –introduced in the stochastic training– does not result in significantly worse inference compared to state-of-theart architectures.

Original: PAVI also displays the best ELBO, in conjunction with UIVI, but in a 5 times shorter optimization time.
Revised: PAVI also displays the best ELBO, in conjunction with UIVI, but in a 5 times shorter optimization time.

Original: B.3. S MALL DIMENSION VERSION OF OUR N EUROIMAGING MODEL
Revised: B.3. S MALL DIMENSION VERSION OF OUR N EUROIMAGING MODEL

Original: In this experiment we validate the performance of our architecture on synthetic data generated using a small dimension version of the model presented in eq.
Revised: In this experiment we validate the performance of our architecture on synthetic data generated using a small dimension version of the model presented in eq.

Original: (B.24).
Revised: (B.24).

Original: To allow the use of the comparative
Revised: To allow the use of the comparative

Original: [Table] [Table]
Revised: [Table] [Table]

Original: baselines, we fixate S, S r = 5 , 3 , T, T r = 6 , 3 , N, N r = 12 , 3 , D = 4 , L = 7 .
Revised: baselines, we fixate S, S r = 5 , 3 , T, T r = 6 , 3 , N, N r = 12 , 3 , D = 4 , L = 7 .

Original: are visible in table 3: PAVI can deal with HBMs unavailable to the ADAVI architecture (Rouillard & Wassermann, 2022), and so so with a performance superior to the CF baseline (Ambrogioni et al., 2021b).
Revised: are visible in table 3: PAVI can deal with HBMs unavailable to the ADAVI architecture (Rouillard & Wassermann, 2022), and so so with a performance superior to the CF baseline (Ambrogioni et al., 2021b).

Original: Despite intensive efforts we did not manage to obtain a numerically stable UIVI optimization.
Revised: Despite intensive efforts we did not manage to obtain a numerically stable UIVI optimization.

Original: 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.
Revised: 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.

Original: In this context, automatic structured VI baselines such as CF or PAVI exploit the parametric form of the prior p to help with inference.
Revised: In this context, automatic structured VI baselines such as CF or PAVI exploit the parametric form of the prior p to help with inference.

Original: PAVI’s plate amortization likely facilitates the inference, resulting in a reduced optimization error (Bottou & Bousquet, 2007)
Revised: PAVI’s plate amortization likely facilitates the inference, resulting in a reduced optimization error (Bottou & Bousquet, 2007)

Original: . B. E XPERIMENTAL DETAILS - ANALYTICAL EXAMPLES
Revised: . B. E XPERIMENTAL DETAILS - ANALYTICAL EXAMPLES

Original: All experiments were performed in Python, using the Tensorflow Probability library (Dillon et al., 2017).
Revised: All experiments were performed in Python, using the Tensorflow Probability library (Dillon et al., 2017).

Original: Throughout this section we refer to Masked Autoregressive Flows (Papamakarios et al., 2018) as MAF .
Revised: Throughout this section we refer to Masked Autoregressive Flows (Papamakarios et al., 2018) as MAF .

Original: All experiments are performed using the Adam optimizer (Kingma & Ba, 2015).
Revised: All experiments are performed using the Adam optimizer (Kingma & Ba, 2015).

Original: At training, the ELBO was estimated using a Monte Carlo procedure with 8 samples.
Revised: At training, the ELBO was estimated using a Monte Carlo procedure with 8 samples.

Original: All architectures were evaluated over a fixed set of 20 samples X , with 5 seeds per sample.
Revised: All architectures were evaluated over a fixed set of 20 samples X , with 5 seeds per sample.

Original: Non-sample-amortized architectures were trained and evaluated on each of those points.
Revised: Non-sample-amortized architectures were trained and evaluated on each of those points.

Original: 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.
Revised: 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.

Original: B.4. P LATE AMORTIZATION AND CONVERGENCE SPEED (3.1) All
Revised: B.4. P LATE AMORTIZATION AND CONVERGENCE SPEED (3.1) All

Original: architectures
Revised: architectures

Original: (baseline, PAVI-F, PAVI-E) used:
Revised: (baseline, PAVI-F, PAVI-E) used:

Original: • for the flows F i , a MAF with [32 , 32] hidden units; • as encoding size,
Revised: • for the flows F i , a MAF with [32 , 32] hidden units; • as encoding size,

Original: 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.
Revised: 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.

Original: B.4. I MPACT OF ENCODING SIZE (3.2)
Revised: B.4. I MPACT OF ENCODING SIZE (3.2)

Original: All architectures used:
Revised: All architectures used:

Original: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.
Revised: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.

Original: • as encoding size, a value varying from 2 to
Revised: • as encoding size, a value varying from 2 to

Original: B.4. S CALING WITH PLATE CARDINALITIES (3.3)
Revised: B.4. S CALING WITH PLATE CARDINALITIES (3.3)

Original: ADAVI (Rouillard & Wassermann, 2022) we used:
Revised: ADAVI (Rouillard & Wassermann, 2022) we used:

Original: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.
Revised: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.

Original: • 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.
Revised: • 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.

Original: Cascading Flows (Ambrogioni et al., 2021b)
Revised: Cascading Flows (Ambrogioni et al., 2021b)

Original: we used:
Revised: we used:

Original: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the
Revised: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the

Original: Cascading Flows
Revised: Cascading Flows

Original: authors
Revised: authors

Original: PAVI-F we used:
Revised: PAVI-F we used:

Original: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.
Revised: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.

Original: • an encoding size of
Revised: • an encoding size of

Original: PAVI-E we used:
Revised: PAVI-E we used:

Original: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.
Revised: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with triangular scaling matrix.

Original: • for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads ofunits each, 2 ISAB blocks with 64 inducing points.
Revised: • for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads ofunits each, 2 ISAB blocks with 64 inducing points.

Original: UIVI we used as Card( P 1 ) = 2 → 20 → 200 :
Revised: UIVI we used as Card( P 1 ) = 2 → 20 → 200 :

Original: • as base distribution, a standard Gaussian with dimensionality 6 → 42 →• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality 3 → 6 →• as transform weights regressor, a MLP with hidden units [32 , 32] → [64 , 64] → [128 , 128] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e −
Revised: • as base distribution, a standard Gaussian with dimensionality 6 → 42 →• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality 3 → 6 →• as transform weights regressor, a MLP with hidden units [32 , 32] → [64 , 64] → [128 , 128] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e −

Original: 2 , × 0 .
Revised: 2 , × 0 .

Original: 9 everysteps
Revised: 9 everysteps

Original: B.4. G AUSSIAN MIXTURE (B.3.1)
Revised: B.4. G AUSSIAN MIXTURE (B.3.1)

Original: ADAVI (Rouillard & Wassermann, 2022) we used:
Revised: ADAVI (Rouillard & Wassermann, 2022) we used:

Original: • for the flows F i , a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.
Revised: • for the flows F i , a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.

Original: • for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads ofunits each, 2 ISAB blocks with 8 inducing points.
Revised: • for the encoder, an encoding size of 16 with a multi-head architecture with 2 heads ofunits each, 2 ISAB blocks with 8 inducing points.

Original: Cascading Flows (Ambrogioni et al., 2021b)
Revised: Cascading Flows (Ambrogioni et al., 2021b)

Original: we used:
Revised: we used:

Original: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the Cascading Flows authors
Revised: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the Cascading Flows authors

Original: PAVI-F we used:
Revised: PAVI-F we used:

Original: • for the flows F i , a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.
Revised: • for the flows F i , a MAF with [32] hidden units, after an affine block with diagonal scaling matrix.

Original: • an encoding size of• Card r
Revised: • an encoding size of• Card r

Original: ( P 1 ) =
Revised: ( P 1 ) =

Original: PAVI-E
Revised: PAVI-E

Original: we
Revised: we

Original: used:
Revised: used:

Original: • for the flows F i , a MAF with [128 , 128] hidden units, after an affine block with traingular scaling matrix.
Revised: • for the flows F i , a MAF with [128 , 128] hidden units, after an affine block with traingular scaling matrix.

Original: • for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads ofunits each, 2 ISAB blocks with 128 inducing points.
Revised: • for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads ofunits each, 2 ISAB blocks with 128 inducing points.

Original: • Card r ( P 1 ) =
Revised: • Card r ( P 1 ) =

Original: UIVI we used:
Revised: UIVI we used:

Original: • as base distribution, a standard Gaussian with dimensionality• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality• as transform weights regressor, a MLP with hidden units [64 , 64] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e − 2 , × 0 .
Revised: • as base distribution, a standard Gaussian with dimensionality• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality• as transform weights regressor, a MLP with hidden units [64 , 64] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e − 2 , × 0 .

Original: 9 everysteps
Revised: 9 everysteps

Original: B.4. H IERARCHICAL V ARIANCES (B.3.2)
Revised: B.4. H IERARCHICAL V ARIANCES (B.3.2)

Original: Cascading Flows (Ambrogioni et al., 2021b)
Revised: Cascading Flows (Ambrogioni et al., 2021b)

Original: we used:
Revised: we used:

Original: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the Cascading Flows authors
Revised: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the Cascading Flows authors

Original: PAVI-F we used:
Revised: PAVI-F we used:

Original: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with diagonal scaling matrix.
Revised: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with diagonal scaling matrix.

Original: • an encoding size of• Card r
Revised: • an encoding size of• Card r

Original: ( P 1 ) =
Revised: ( P 1 ) =

Original: • Card r
Revised: • Card r

Original: ( P 0 ) =
Revised: ( P 0 ) =

Original: PAVI
Revised: PAVI

Original: -E we used
Revised: -E we used

Original: • for the flows F i , a MAF with [128 , 128] hidden units, after an affine block with diagonal scaling matrix.
Revised: • for the flows F i , a MAF with [128 , 128] hidden units, after an affine block with diagonal scaling matrix.

Original: • for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads ofunits each, 2 ISAB blocks with 128 inducing points.
Revised: • for the encoder, an encoding size of 128 with a multi-head architecture with 4 heads ofunits each, 2 ISAB blocks with 128 inducing points.

Original: • Card r ( P 1 ) =
Revised: • Card r ( P 1 ) =

Original: • Card r
Revised: • Card r

Original: ( P 0 ) =
Revised: ( P 0 ) =

Original: UIVI we used:
Revised: UIVI we used:

Original: • as base distribution, a standard Gaussian with dimensionality• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality• as transform weights regressor, a MLP with hidden units [64 , 64] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e − 2 , × 0 .
Revised: • as base distribution, a standard Gaussian with dimensionality• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality• as transform weights regressor, a MLP with hidden units [64 , 64] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e − 2 , × 0 .

Original: 9 everysteps
Revised: 9 everysteps

Original: B.4. S MALL N EUROIMAGING EXAMPLE (B
Revised: B.4. S MALL N EUROIMAGING EXAMPLE (B

Original: Cascading Flows (Ambrogioni et al., 2021b)
Revised: Cascading Flows (Ambrogioni et al., 2021b)

Original: we used:
Revised: we used:

Original: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the Cascading Flows authors
Revised: • a mean-field distribution over the auxiliary variables r • as auxiliary size, a fixed value of• as flows, Highway Flows as designed by the Cascading Flows authors

Original: PAVI-F we used:
Revised: PAVI-F we used:

Original: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with diagonal scaling matrix.
Revised: • for the flows F i , a MAF with [32 , 32] hidden units, after an affine block with diagonal scaling matrix.

Original: • an combination of encoding sizes of 32 and
Revised: • an combination of encoding sizes of 32 and

Original: UIVI we used:
Revised: UIVI we used:

Original: • as base distribution, a standard Gaussian with dimensionality 1 ,• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality• as transform weights regressor, a MLP with hidden units [128 , 256 , 512 , 1024] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e − 3 , × 0 .
Revised: • as base distribution, a standard Gaussian with dimensionality 1 ,• as transform h , an affine transform with diagonal scale • as embedding distribution, a standard Gaussian with dimensionality• as transform weights regressor, a MLP with hidden units [128 , 256 , 512 , 1024] • to sample uncorrelated samples ϵ , a HMC run with 5 burn-in steps and 5 samples • an Adam optimizer with exponential learning rate decay, starting at 1 e − 3 , × 0 .

Original: 9 everysteps
Revised: 9 everysteps

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

Original: B. D ETAILS ABOUT OUR N EUROIMAGING EXPERIMENT (3.4) B.5. D ATA DESCRIPTION
Revised: B. D ETAILS ABOUT OUR N EUROIMAGING EXPERIMENT (3.4) B.5. D ATA DESCRIPTION

Original: In this experiment we use data from the Human Connectome Project (HCP) dataset (Van Essen et al., 2012).
Revised: In this experiment we use data from the Human Connectome Project (HCP) dataset (Van Essen et al., 2012).

Original: We randomly select a cohort of S = 1 , 000 subjects from this dataset, each subject being associated with T = 2 resting state fMRI sessions (Smith et al., 2013).
Revised: We randomly select a cohort of S = 1 , 000 subjects from this dataset, each subject being associated with T = 2 resting state fMRI sessions (Smith et al., 2013).

Original: We minimally pre-process the signal using the nilearn python library (Abraham et al., 2014):
Revised: We minimally pre-process the signal using the nilearn python library (Abraham et al., 2014):

Original: 1. removing high variance confounds 2.
Revised: 1. removing high variance confounds 2.

Original: detrending the data 3.
Revised: detrending the data 3.

Original: band-filtering the data (0.01 to 0.1 Hz), with a repetition time of 0.74 seconds 4.
Revised: band-filtering the data (0.01 to 0.1 Hz), with a repetition time of 0.74 seconds 4.

Original: spatially smoothing the data with a 4mm Full-Width at Half Maximum
Revised: spatially smoothing the data with a 4mm Full-Width at Half Maximum

Original: For every subject, we extract the surface Blood Oxygenation Level Dependent (BOLD) signal of
Revised: For every subject, we extract the surface Blood Oxygenation Level Dependent (BOLD) signal of

Original: N = 314 vertices corresponding to an average Broca’s area (Heim et al., 2009).
Revised: N = 314 vertices corresponding to an average Broca’s area (Heim et al., 2009).

Original: 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).
Revised: 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).

Original: Specifically, we compute the one-to-one Pearson’s correlation coefficient of every vertex with every DiFuMo component.
Revised: Specifically, we compute the one-to-one Pearson’s correlation coefficient of every vertex with every DiFuMo component.

Original: The resulting connectome, with S subjects, T sessions, N vertices and a connectivity signal with D dimensions, is of shape ( S × T × N × D ) .
Revised: The resulting connectome, with S subjects, T sessions, N vertices and a connectivity signal with D dimensions, is of shape ( S × T × N × D ) .

Original: We project this data –correlation coefficients lying in ] − 1; 1[ – in an unbounded space using an inverse sigmoid function.
Revised: We project this data –correlation coefficients lying in ] − 1; 1[ – in an unbounded space using an inverse sigmoid function.

Original: B.5. M ODEL DESCRIPTION
Revised: B.5. M ODEL DESCRIPTION

Original: We use a model inspired from the work of Kong et al.
Revised: We use a model inspired from the work of Kong et al.

Original: We hypothesize that every vertex in
Revised: We hypothesize that every vertex in

Original: Broca’s area belongs to either one of L = 2 functional networks.
Revised: Broca’s area belongs to either one of L = 2 functional networks.

Original: This functional bi-partition would reflect the anatomical partition between pars opercularis and pars triangularis (Heim et al., 2009;
Revised: This functional bi-partition would reflect the anatomical partition between pars opercularis and pars triangularis (Heim et al., 2009;

Original: Zhang et al., 2020).
Revised: Zhang et al., 2020).

Original: Each network is a pattern of connectivity with the brain cortex, represented as a the correlation of the BOLD signal with the signal from the D = 64 DiFuMo components.
Revised: Each network is a pattern of connectivity with the brain cortex, represented as a the correlation of the BOLD signal with the signal from the D = 64 DiFuMo components.

Original: We define L = 2 such functional networks at the population level, that correspond to some ”average” across the cohort of subjects.
Revised: We define L = 2 such functional networks at the population level, that correspond to some ”average” across the cohort of subjects.

Original: Every subject has an individual connectivity, and therefore individual L = 2 networks, that are considered as a Gaussian perturbation of the population networks, with variance ϵ .
Revised: Every subject has an individual connectivity, and therefore individual L = 2 networks, that are considered as a Gaussian perturbation of the population networks, with variance ϵ .

Original: The connectivity of a given subject also evolves through time, giving rise to session-specific networks, that are a Gaussian perturbation of the subject networks with variance σ .
Revised: The connectivity of a given subject also evolves through time, giving rise to session-specific networks, that are a Gaussian perturbation of the subject networks with variance σ .

Original: Finally, every vertex in
Revised: Finally, every vertex in

Original: Broca’s area has its individual connectivity, and is a perturbation of one network’s connectivity or the other’s.
Revised: Broca’s area has its individual connectivity, and is a perturbation of one network’s connectivity or the other’s.

Original: We model this last step as a Gaussian mixture distribution with variance κ .
Revised: We model this last step as a Gaussian mixture distribution with variance κ .

Original: We explicitly model the label label of a given vertex, and we consider this label constant across sessions.
Revised: We explicitly model the label label of a given vertex, and we consider this label constant across sessions.

Original: The resulting model can be described as:
Revised: The resulting model can be described as:

Original: [Equation]
Revised: [Equation]

Original: ∀ l =1 ..
Revised: ∀ l =1 ..

Original: L ∀ s =1 ..S : µ l,s | µ l , ϵ l ∼ N ( µ l , ϵ l ) ∀ l =1 ..
Revised: L ∀ s =1 ..S : µ l,s | µ l , ϵ l ∼ N ( µ l , ϵ l ) ∀ l =1 ..

Original: L : log σ l ∼ Uniform( s − ×
Revised: L : log σ l ∼ Uniform( s − ×

Original: ⃗ 1 D , s + × ⃗ 1 D ) ∀ l =1 ..
Revised: ⃗ 1 D , s + × ⃗ 1 D ) ∀ l =1 ..

Original: L ∀ s =1 ..S ∀ t =1 ..T : µ l,s,t | µ l,s , σ l ∼ N ( µ l,s , σ l ) ∀ l =1 ..
Revised: L ∀ s =1 ..S ∀ t =1 ..T : µ l,s,t | µ l,s , σ l ∼ N ( µ l,s , σ l ) ∀ l =1 ..

Original: L : log κ l ∼ Uniform( s − × ⃗ 1 D , s + × ⃗ 1 D ) ∀ s =1 ..S ∀ n =1 ..N : probs s,n ∼ Dirichlet(1 × ⃗ 1 L ) ∀ s =1 ..S ∀ n =1 ..N : label s,n | probs s,n ∼ Categorical(probs s,n ) ∀ s =1 ..
Revised: L : log κ l ∼ Uniform( s − × ⃗ 1 D , s + × ⃗ 1 D ) ∀ s =1 ..S ∀ n =1 ..N : probs s,n ∼ Dirichlet(1 × ⃗ 1 L ) ∀ s =1 ..S ∀ n =1 ..N : label s,n | probs s,n ∼ Categorical(probs s,n ) ∀ s =1 ..

Original: S ∀ t =1 ..T ∀ n =1 ..N :
Revised: S ∀ t =1 ..T ∀ n =1 ..N :

Original: X s,t,n | [ µ l,s,t ] l =1 ..L , [ κ l ] l =1 ..
Revised: X s,t,n | [ µ l,s,t ] l =1 ..L , [ κ l ] l =1 ..

Original: L , label s,n ∼ N ( µ label s,n ,s,t , κ label s,n ) (B.24)
Revised: L , label s,n ∼ N ( µ label s,n ,s,t , κ label s,n ) (B.24)

Original: 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 .
Revised: 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 .

Original: Our goal is to recover the posterior distribution of the networks µ –represented as networks in fig.
Revised: Our goal is to recover the posterior distribution of the networks µ –represented as networks in fig.

Original: 5– and the labels label –represented as the parcellation in fig.
Revised: 5– and the labels label –represented as the parcellation in fig.

Original: 5– given the observed connectome described in appendix B.5.1.
Revised: 5– given the observed connectome described in appendix B.5.1.

Original: B.5. PAVI IMPLEMENTATION
Revised: B.5. PAVI IMPLEMENTATION

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

Original: • for the RVs µ l , µ l,s , µ l,s,t : – for the flows F i , a MAF with [128 , 128] hidden units, following an affine block with diagonal scale – for the encoding size:• for the RVs ϵ l , σ l , κ l , probs s
Revised: • for the RVs µ l , µ l,s , µ l,s,t : – for the flows F i , a MAF with [128 , 128] hidden units, following an affine block with diagonal scale – for the encoding size:• for the RVs ϵ l , σ l , κ l , probs s

Original: ,n , labels s,n :
Revised: ,n , labels s,n :

Original: – for the flows F i , a MAF with [8 , 8] hidden units, following an affine block with diagonal scale – for the encoding size:• for the reduced model, we used S r = 30 , T r = 1 and N r = 32 .
Revised: – for the flows F i , a MAF with [8 , 8] hidden units, following an affine block with diagonal scale – for the encoding size:• for the reduced model, we used S r = 30 , T r = 1 and N r = 32 .

Original: To allow for the optimization over the discrete label s,n RV, we used the Gumbell-Softmax trick, using a fixed temperature of 1 . 

Revised: To allow for the optimization over the discrete label s,n RV, we used the Gumbell-Softmax trick, using a fixed temperature of 1 .

Original: (Jang et al., 2017; Maddison et al., 2016).
Revised: (Jang et al., 2017; Maddison et al., 2016).

Original: C S UPPLEMENTAL DISCUSSION C. P LATE AMORTIZATION AS A GENERALIZATION OF SAMPLE AMORTIZATION
Revised: C S UPPLEMENTAL DISCUSSION C. P LATE AMORTIZATION AS A GENERALIZATION OF SAMPLE AMORTIZATION

Original: In section 2.2.2 we introduced plate amortization as the application of the generic concept of amortization to the granularity of plates.
Revised: In section 2.2.2 we introduced plate amortization as the application of the generic concept of amortization to the granularity of plates.

Original: There is actually an even stronger connection between sample amortization and plate amortization.
Revised: There is actually an even stronger connection between sample amortization and plate amortization.

Original: A HBM p models the distribution of a given observed RV X –jointly with the parameters Θ .
Revised: A HBM p models the distribution of a given observed RV X –jointly with the parameters Θ .

Original: Different samples X 0 , X 1 , ... of the model p are i.i.d.
Revised: Different samples X 0 , X 1 , ... of the model p are i.i.d.

Original: draws from the distribution p ( X ) .
Revised: draws from the distribution p ( X ) .

Original: p can thus be considered as the model for ”one sample”.
Revised: p can thus be considered as the model for ”one sample”.

Original: Consider, instead of p , a ”macro” model for the whole population of samples one could draw from p .
Revised: Consider, instead of p , a ”macro” model for the whole population of samples one could draw from p .

Original: The observed RV of that macro model would be the infinite collection of samples drawn from the same distribution p
Revised: The observed RV of that macro model would be the infinite collection of samples drawn from the same distribution p

Original: ( X ) .
Revised: ( X ) .

Original: In that light, the i.i.d.
Revised: In that light, the i.i.d.

Original: sampling of different X values from p could be interpreted as a plate of the macro model.
Revised: sampling of different X values from p could be interpreted as a plate of the macro model.

Original: Thus, we could consider sample amortization as a instance of plate amortization for the ”sample plate”.
Revised: Thus, we could consider sample amortization as a instance of plate amortization for the ”sample plate”.

Original: Or equivalently: plate amortization can be seen as the generalization of amortization beyond the particular case of sample amortization.
Revised: Or equivalently: plate amortization can be seen as the generalization of amortization beyond the particular case of sample amortization.

Original: C. A LTERNATE FORMALISM FOR SVI – PAVI-E( SA ) SCHEME
Revised: C. A LTERNATE FORMALISM FOR SVI – PAVI-E( SA ) SCHEME

Original: In this work, we propose a different formalism for SVI, based around the concept of full HBM
Revised: In this work, we propose a different formalism for SVI, based around the concept of full HBM

Original: M full versus reduced HBM M redu sharing the same template T .
Revised: M full versus reduced HBM M redu sharing the same template T .

Original: This formalism is helpful to set up
Revised: This formalism is helpful to set up

Original: 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 2.2.1&2.3.2.
Revised: 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 2.2.1&2.3.2.

Original: 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).
Revised: 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).

Original: 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 2.3.2.
Revised: 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 2.3.2.

Original: In this context, our scheme is no longer a different take on hierarchical, batched
Revised: In this context, our scheme is no longer a different take on hierarchical, batched

Original: SVI: the cardinality of the full model is truly independent from the cardinality of the training, and is only simulated as a scaling factor in the stochastic training –see section 2.3.1.
Revised: SVI: the cardinality of the full model is truly independent from the cardinality of the training, and is only simulated as a scaling factor in the stochastic training –see section 2.3.1.

Original: We have the intuition that fruitful research directions could stem from this concept.
Revised: We have the intuition that fruitful research directions could stem from this concept.

Original: C. B ENEFITING FROM STRUCTURE IN INFERENCE
Revised: C. B ENEFITING FROM STRUCTURE IN INFERENCE

Original: Our contributions can be abstracted through the concept of plate amortization -see section 2.2.2.
Revised: Our contributions can be abstracted through the concept of plate amortization -see section 2.2.2.

Original: 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),
Revised: 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),

Original: ASVI (Ambrogioni et al., 2021a), or implicit (Yin & Zhou, 2018; Titsias & Ruiz, 2019) schemes are also possible to use.
Revised: ASVI (Ambrogioni et al., 2021a), or implicit (Yin & Zhou, 2018; Titsias & Ruiz, 2019) schemes are also possible to use.

Original: Plate amortization can be viewed as the amortization of common density approximators across different sub-structures of a problem.
Revised: Plate amortization can be viewed as the amortization of common density approximators across different sub-structures of a problem.

Original: This general concept could have applications in other highly-structured problem classes such as graphs or sequences (Wu et al., 2020;
Revised: This general concept could have applications in other highly-structured problem classes such as graphs or sequences (Wu et al., 2020;

Original: Salehinejad et al., 2018).
Revised: Salehinejad et al., 2018).

Original: C. C ONNECTION WITH META - LEARNING
Revised: C. C ONNECTION WITH META - LEARNING

Original: In section 2.2.2 we introduced plate amortization: sharing the parameterization and learning across a model’s plates.
Revised: In section 2.2.2 we introduced plate amortization: sharing the parameterization and learning across a model’s plates.

Original: In this section we discuss the connection between plate amortization and metalearning(Ravi & Beatson
Revised: In this section we discuss the connection between plate amortization and metalearning(Ravi & Beatson

Original: Supervised learning can be seen as the mapping from a given context set C
Revised: Supervised learning can be seen as the mapping from a given context set C

Original: = { ( x, y ) } to a predictive function f (Bishop, 2006) such that f ( x ) =
Revised: = { ( x, y ) } to a predictive function f (Bishop, 2006) such that f ( x ) =

Original: y .
Revised: y .

Original: Meta-learning –or ”learning to learn”– instead recovers this mapping C
Revised: Meta-learning –or ”learning to learn”– instead recovers this mapping C

Original: (cid:55)→ f in the general case.
Revised: (cid:55)→ f in the general case.

Original: 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 .
Revised: 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 .

Original: 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).
Revised: 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).

Original: The encoded context, along with the data point x are then used to condition an estimator for the density q ( y | x, C ) .
Revised: The encoded context, along with the data point x are then used to condition an estimator for the density q ( y | x, C ) .

Original: 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 θ i,n .
Revised: 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 θ i,n .

Original: This encoder-estimator pair is repeatedly used across a plate.
Revised: This encoder-estimator pair is repeatedly used across a plate.

Original: 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.
Revised: 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.

Original: In PAVI-E, the encoding E
Revised: In PAVI-E, the encoding E

Original: i,n at a lower hierarchy play a role similar to the context C in meta-leaning.
Revised: i,n at a lower hierarchy play a role similar to the context C in meta-leaning.

Original: A few differences however exist between the two frameworks:
Revised: A few differences however exist between the two frameworks:

Original: • 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
Revised: • 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

Original: d. tasks, whereas in PAVI the inference of different ground RVs θ 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.
Revised: d. tasks, whereas in PAVI the inference of different ground RVs θ 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.

Original: In contrast PAVI –though possible to train using the forward KL loss– is trained via the reverse KL, needing to explicitly evaluate the density of the generative process.
Revised: In contrast PAVI –though possible to train using the forward KL loss– is trained via the reverse KL, needing to explicitly evaluate the density of the generative process.

Original: We suspect there would be interesting applications for the PAVI architecture in hierarchical metalearning scenarios.
Revised: We suspect there would be interesting applications for the PAVI architecture in hierarchical metalearning scenarios.

Original: C. C ONDITIONAL DEPENDENCIES MODELLED IN THE VARIATIONAL FAMILY
Revised: C. C ONDITIONAL DEPENDENCIES MODELLED IN THE VARIATIONAL FAMILY

Original: The Mean-Field approximation was originally introduced in VI to facilitate computation, allowing dedicated optimization schemes (Blei et al., 2017).
Revised: The Mean-Field approximation was originally introduced in VI to facilitate computation, allowing dedicated optimization schemes (Blei et al., 2017).

Original: Nonetheless, this approximation assumes independence between RVs in the posterior, and ultimately limits the expressivity of the variational family.
Revised: Nonetheless, this approximation assumes independence between RVs in the posterior, and ultimately limits the expressivity of the variational family.

Original: To remove this approximation, modelling complex conditional dependencies in the variational family is an open research subject (Ambrogioni et al., 2021b;
Revised: To remove this approximation, modelling complex conditional dependencies in the variational family is an open research subject (Ambrogioni et al., 2021b;

Original: Webb et al., 2018).
Revised: Webb et al., 2018).

Original: In the PAVI design, we inherit our statistical dependency structure from the prior distribution p , as detailed in eq.
Revised: In the PAVI design, we inherit our statistical dependency structure from the prior distribution p , as detailed in eq.

Original: This choice of dependencies follows the line of research from Hoffman & Blei (2014) and
Revised: This choice of dependencies follows the line of research from Hoffman & Blei (2014) and

Original: Ambrogioni et al.
Revised: Ambrogioni et al.

Original: (2021a).
Revised: (2021a).

Original: As pointed out by Ambrogioni et al.
Revised: As pointed out by Ambrogioni et al.

Original: (2021b), when modelling only those forward dependencies, the modelling of colliders can be an issue.
Revised: (2021b), when modelling only those forward dependencies, the modelling of colliders can be an issue.

Original: In the case of a single plate –the 2-level case– Agrawal & Domke (2021) demonstrate that modelling only the forward dependencies does not result in reduced expressivity compared to the modelling of the full dependencies.
Revised: In the case of a single plate –the 2-level case– Agrawal & Domke (2021) demonstrate that modelling only the forward dependencies does not result in reduced expressivity compared to the modelling of the full dependencies.

Original: Yet this result does not hold in the n-level case, as Webb et al.
Revised: Yet this result does not hold in the n-level case, as Webb et al.

Original: (2018) show that faithful inversion features conditional dependencies in the posterior between ground RVs of the same template.
Revised: (2018) show that faithful inversion features conditional dependencies in the posterior between ground RVs of the same template.

Original: We can dub those dependencies as horizontal dependencies, across RVs in the same plate.
Revised: We can dub those dependencies as horizontal dependencies, across RVs in the same plate.

Original: Similar to Structured SVI (Hoffman & Blei, 2014) or Automatic SSVI (Ambrogioni et al., 2021a), the PAVI design therefore results in reduced expressivity when stacking multiple plates in the model p .
Revised: Similar to Structured SVI (Hoffman & Blei, 2014) or Automatic SSVI (Ambrogioni et al., 2021a), the PAVI design therefore results in reduced expressivity when stacking multiple plates in the model p .

Original: This issue can be partially alleviated with the usage of a backward encoding scheme –going in the reverse direction compared to the prior’s dependencies– as in the PAVI-E design (see section 2.3.2) or in Cascading Flows (Ambrogioni et al., 2021b).
Revised: This issue can be partially alleviated with the usage of a backward encoding scheme –going in the reverse direction compared to the prior’s dependencies– as in the PAVI-E design (see section 2.3.2) or in Cascading Flows (Ambrogioni et al., 2021b).

Original: In practice, though limiting our expressivity, horizontal dependencies are difficult to inject back into our architecture.
Revised: In practice, though limiting our expressivity, horizontal dependencies are difficult to inject back into our architecture.

Original: Critically, in the PAVI design, the use of a common density estimator across the ground RVs of the same template (see section 2.2), and the stochastic training over batches of those
Revised: Critically, in the PAVI design, the use of a common density estimator across the ground RVs of the same template (see section 2.2), and the stochastic training over batches of those

Original: RVs (see section 2.3.1) prevent the direct modeling horizontal dependencies.
Revised: RVs (see section 2.3.1) prevent the direct modeling horizontal dependencies.

Original: 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.
Revised: 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.

Original: Injecting horizontal dependencies back into our variational family is therefore a non-trivial research direction, that is not at the core of this paper.
Revised: Injecting horizontal dependencies back into our variational family is therefore a non-trivial research direction, that is not at the core of this paper.

Original: This opens up promising research directions: how could arbitrary conditional dependencies be modelled in the variational posterior in the context of stochastic training?
Revised: This opens up promising research directions: how could arbitrary conditional dependencies be modelled in the variational posterior in the context of stochastic training?

Original: C. T OWARDS USER - FRIENDLY V ARIATIONAL I NFERENCE
Revised: C. T OWARDS USER - FRIENDLY V ARIATIONAL I NFERENCE

Original: By re-purposing the concept of amortization at the plate level, our goal is to propose clear computation versus precision trade-offs in VI.
Revised: By re-purposing the concept of amortization at the plate level, our goal is to propose clear computation versus precision trade-offs in VI.

Original: Hyper-parameters such as the encoding size –as illustrated in fig.
Revised: Hyper-parameters such as the encoding size –as illustrated in fig.

Original: 3 (right)– allow to clearly trade inference quality in exchanged for a reduced memory footprint.
Revised: 3 (right)– allow to clearly trade inference quality in exchanged for a reduced memory footprint.

Original: On the contrary, in classical VI, changing Q ’s parametric form –for instance switching from
Revised: On the contrary, in classical VI, changing Q ’s parametric form –for instance switching from

Original: Gaussian to Student distributions– can have a strong and complex impact both on number of weights and inference quality (Blei et al., 2017).
Revised: Gaussian to Student distributions– can have a strong and complex impact both on number of weights and inference quality (Blei et al., 2017).

Original: By allowing the usage of normalizing flows in very large cardinality regimes, our contribution aims at disentangling approximation power and computational feasibility.
Revised: By allowing the usage of normalizing flows in very large cardinality regimes, our contribution aims at disentangling approximation power and computational feasibility.

Original: In particular, having access to expressive density approximators for the posterior can help experimenters diversify the proposed HBMs, removing the need of properties such as conjugacy to obtain meaningful inference (Gelman et al., 2004).
Revised: In particular, having access to expressive density approximators for the posterior can help experimenters diversify the proposed HBMs, removing the need of properties such as conjugacy to obtain meaningful inference (Gelman et al., 2004).

Original: 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.
Revised: 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.