ADAVI: Automatic Dual Amortized Variational Inference Applied To Pyramidal Bayesian Models

1 INTRODUCTION . Inference aims at obtaining the posterior distribution p ( θ|X ) of latent model parameters θ given the observed data X . In the context of Hierarchical Bayesian Models ( HBM ) , p ( θ|X ) usually has no known analytical form , and can be of a complex shape -different from the prior ’ s ( Gelman et al. , 2004 ) . Modern normalizing-flows based techniques -universal density estimators- can overcome this difficulty ( Papamakarios et al. , 2019a ; Ambrogioni et al. , 2021 ) . Yet , in setups such as neuroimaging , featuring HBMs representing large population studies ( Kong et al. , 2018 ; Bonkhoff et al. , 2021 ) , the dimensionality of θ can go over the million . This high dimensionality hinders the usage of normalizing flows , since their parameterization usually scales quadratically with the size of the parameter space ( e.g . Dinh et al. , 2017 ; Papamakarios et al. , 2018 ; Grathwohl et al. , 2018 ) . Population studies with large dimensional features are therefore inaccessible to off-the-shelf flow-based techniques and their superior expressivity . This can in turn lead to complex , problem-specific derivations : for instance Kong et al . ( 2018 ) rely on a manually-derived Expectation Maximization ( EM ) technique . Such an analytical complexity constitutes a strong barrier to entry , and limits the wide and fruitful usage of Bayesian modelling in fields such as neuroimaging . Our main aim is to meet that experimental need : how can we derive a technique both automatic and efficient in the context of very large , hierarchically-organised data ? Approximate inference features a large corpus of methods including Monte Carlo methods ( Koller & Friedman , 2009 ) and Variational Auto Encoders ( Zhang et al. , 2019 ) . We take particular inspiration from the field of Variational Inference ( VI ) ( Blei et al. , 2017 ) , deemed to be most adapted to large parameter spaces . In VI , the experimenter posits a variational familyQ so as to approximate q ( θ ) ≈ Generative Pyramidal . p ( θ|X ) . In practice , deriving an expressive , yet computationally attractive variational family can be challenging ( Blei et al. , 2017 ) . This triggered a trend towards the derivation of automatic VI techniques ( Kucukelbir et al. , 2016 ; Ranganath et al. , 2013 ; Ambrogioni et al. , 2021 ) . We follow that logic and present a methodology that automatically derives a variational family Q . In Fig . 1 , from the HBM on the left we derive automatically a neural network architecture on the right . We aim at deriving our variational family Q in the context of amortized inference ( Rezende & Mohamed , 2016 ; Cranmer et al. , 2020 ) . Amortization is usually obtained at the cost of an amortization gap from the true posterior , that accumulates on top of a approximation gap dependent on the expressivity of the variational family Q ( Cremer et al. , 2018 ) . However , once an initial training overhead has been “ paid for ” , amortization means that our technique can be applied to a any number of data points to perform inference in a few seconds . Due to the very large parameter spaces presented above , our target applications aren ’ t amenable to the generic flow-based techniques described in Cranmer et al . ( 2020 ) or Ambrogioni et al . ( 2021 ) . We therefore differentiate ourselves in exploiting the invariance of the problem not only through the design of an adapted encoder , but down to the very architecture of our density estimator . Specifically , we focus on the inference problem for Hierarchical Bayesian Models ( HBMs ) ( Gelman et al. , 2004 ; Rodrigues et al. , 2021 ) . The idea to condition the architecture of a density estimator by an analysis of the dependency structure of an HBM has been studied in ( Wehenkel & Louppe , 2020 ; Weilbach et al . ) , in the form of the masking of a single normalizing flow . With Ambrogioni et al . ( 2021 ) , we instead share the idea to combine multiple separate flows . More generally , our static analysis of a generative model can be associated with structured VI ( Hoffman & Blei , 2014 ; Ambrogioni et al . ; 2021 ) . Yet our working principles are rather orthogonal : structured VI usually aims at exploiting model structure to augment the expressivity of a variational family , whereas we aim at reducing its parameterization . Our objective is therefore to derive an automatic methodology that takes as input a generative HBM and generates a dual variational family able to perform amortized parameter inference . This variational family exploits the exchangeability in the HBM to reduce its parameterization by orders of magnitude compared to generic methods ( Papamakarios et al. , 2019b ; Greenberg et al. , 2019 ; Am- brogioni et al. , 2021 ) . Consequently , our method can be applied in the context of large , pyramidallystructured data , a challenging setup inaccessible to existing flow-based methods and their superior expressivity . We apply our method to such a large pyramidal setup in the context of neuroimaging ( section 3.5 ) , but demonstrate the benefit of our method beyond that scope . Our general scheme is visible in Fig . 1 , a figure that we will explain throughout the course of the next section . 2 METHODS . 2.1 PYRAMIDAL BAYESIAN MODELS . We are interested in experimental setups modelled using plate-enriched Hierarchical Bayesian Models ( HBMs ) ( Kong et al. , 2018 ; Bonkhoff et al. , 2021 ) . These models feature independent sampling from a common conditional distribution at multiple levels , translating the graphical notion of plates ( Gilks et al. , 1994 ) . This nested structure , combined with large measurements -such as the ones in fMRI- can result in massive latent parameter spaces . For instance the population study in Kong et al . ( 2018 ) features multiple subjects , with multiple measures per subject , and multiple brain vertices per measure , for a latent space of around 0.4 million parameters . Our method aims at performing inference in the context of those large plate-enriched HBMs . Such HBMs can be represented with Directed Acyclic Graphs ( DAG ) templates ( Koller & Friedman , 2009 ) with vertices -corresponding to RVs- { θi } i=0 ... L and plates { Pp } p=0 ... P . We denote as Card ( P ) the -fixed- cardinality of the plate P , i.e . the number of independent draws from a common conditional distribution it corresponds to . In a template DAG , a given RV θ can belong to multiple plates Ph , . . .PP . When grounding the template DAG into a ground graph -instantiating the repeated structure symbolized by the plates P- θ would correspond to multiple RVs of similar parametric form { θih , ... , iP } , with ih = 1 . . .Card ( Ph ) , . . . , iP = 1 . . .Card ( PP ) . This equivalence visible on the left on Fig . 1 , where the template RV Γ corresponds to the ground RVs [ γ1 , γ2 ] . We wish to exploit this plate-induced exchangeability . We define the sub-class of models we specialize upon as pyramidal models , which are plate-enriched DAG templates with the 2 following differentiating properties . First , we consider a single stack of the plates P0 , . . . , PP . This means that any RV θ belonging to plate Pp also belongs to plates { Pq } q > p . We thus don ’ t treat in this work the case of colliding plates ( Koller & Friedman , 2009 ) . Second , we consider a single observed RV θ0 , with observed value X , belonging to the plate P0 ( with no other -latent- RV belonging to P0 ) . The obtained graph follows a typical pyramidal structure , with the observed RV at the basis of the pyramid , as seen in Fig . 1 . This figure features 2 plates P0 and P1 , the observed RV is X , at the basis of the pyramid , and latent RVs are Γ , λ and κ at upper levels of the pyramid . Pyramidal HBMs delineate models that typically arise as part of population studies -for instance in neuroimaging- featuring a nested group structure and data observed at the subject level only ( Kong et al. , 2018 ; Bonkhoff et al. , 2021 ) . The fact that we consider a single pyramid of plates allows us to define the hierarchy of an RV θi denoted Hier ( θi ) . An RV ’ s hierarchy is the level of the pyramid it is placed at . Due to our pyramidal structure , the observed RV will systematically be at hierarchy 0 and latent RVs at hierarchies > 0 . For instance , in the example in Fig . 1 the observed RV X is at hierarchy 0 , Γ is at hierarchy 1 and both λ and κ are at hierarchy 2 . Our methodology is designed to process generative models whose dependency structure follows a pyramidal graph , and to scale favorably when the plate cardinality in such models augments . Given the observed data X , we wish to obtain the posterior density for latent parameters θ1 , . . . , θL , exploiting the exchangeability induced by the plates P0 , . . . , PP . 2.2 AUTOMATIC DERIVATION OF A DUAL AMORTIZED VARIATIONAL FAMILY . In this section , we derive our main methodological contribution . We aim at obtaining posterior distributions for a generative model of pyramidal structure . For this purpose , we construct a family of variational distributions Q dual to the model . This architecture consists in the combination of 2 items . First , a Hierarchical Encoder ( HE ) that aggregates summary statistics from the data . Second , a set of conditional density estimators . Tensor functions We first introduce the notations for tensor functions which we define in the spirit of Magnus & Neudecker ( 1999 ) . We leverage tensor functions throughout our entire architecture to reduce its parameterization . Consider a function f : F → G , and a tensor TF ∈ FB of shape B . We denote the tensor TG ∈ GB resulting from the element-wise application of f over TF as TG =−→ f ( B ) ( TF ) ( in reference to the programming notion of vectorization in Harris et al . ( 2020 ) ) . In Fig . 1 , −→ ST ( B1 ) 0 and −−−−→ lγ ◦ Fγ ( B1 ) are examples of tensor functions . At multiple points in our architecture , we will translate the repeated structure in the HBM induced by plates into the repeated usage of functions across plates . Hierarchical Encoder For our encoder , our goal is to learn a function HE that takes as input the observed data X and successively exploits the permutation invariance across plates P0 , . . . , PP . In doing so , HE produces encodings E at different hierarchy levels . Through those encodings , our goal is to learn summary statistics from the observed data , that will condition our amortized inference . For instance in Fig . 1 , the application of HE over X produces the encodings E1 and E2 . To build HE , we need at multiple hierarchies to collect summary statistics across i.i.d samples from a common distribution . To this end we leverage SetTransformers ( Lee et al. , 2019 ) : an attention-based , permutation-invariant architecture . We use SetTransformers to derive encodings across a given plate , repeating their usage for all larger-rank plates . We cast the observed data X as the encoding E0 . Then , recursively for every hierarchy h = 1 . . . P + 1 , we define the encoding Eh as the application to the encoding Eh−1 of the tensor function corresponding to the set transformer STh−1 . HE ( X ) then corresponds to the set of encodings { E1 , . . . , EP+1 } obtained from the successive application of { STh } h=0 , ... , P . If we denote the batch shape Bh = Card ( Ph ) × . . .× Card ( PP ) : Eh = −→ ST ( Bh ) h−1 ( Eh−1 ) HE ( X ) = { E1 , . . . , EP+1 } ( 1 ) In collecting summary statistics across the i.i.d . samples in plate Ph−1 , we decrease the order of the encoding tensor Eh−1 . We repeat this operation in parallel on every plate of larger rank than the rank of the contracted plate . We consequently produce an encoding tensor Eh with the batch shape Bh , which is the batch shape of every RV of hierarchy h. In that line , successively summarizing plates P0 , , . . . , PP , of increasing rank results in encoding tensors E1 , . . . , EP+1 of decreasing order . In Fig . 1 , there are 2 plates P0 and P1 , hence 2 encodings E1 = −→ ST ( B1 ) 0 ( X ) and E2 = ST1 ( E1 ) . E1 is an order 2 tensor : it has a batch shape of B1 = Card ( P1 ) -similar to Γ- whereas E2 is an order 1 tensor . We can decompose E1 = [ e11 , e21 ] = [ ST0 ( [ X1,1 , X1,2 ] ) , ST0 ( [ X2,1 , X2,2 ] ) ] . Conditional density estimators We now will use the encodings E , gathering hierarchical summary statistics on the data X , to condition the inference on the parameters θ . The encodings { Eh } h=1 ... P+1 will respectively condition the density estimators for the posterior distribution of parameters sharing their hierarchy { { θi : Hier ( θi ) = h } } h=1 ... P+1 . Consider a latent RV θi of hierarchy hi = Hier ( θi ) . Due to the plate structure of the graph , θi can be decomposed in a batch of shape Bhi = Card ( Phi ) × . . . × Card ( PP ) of multiple similar , conditionally independent RVs of individual size Sθi . This decomposition is akin to the grounding of the considered graph template ( Koller & Friedman , 2009 ) . A conditional density estimator is a 2-step diffeomorphism from a latent space onto the event space in which the RV θi lives . We initially parameterize every variational density as a standard normal distribution in the latent space RSθi . First , this latent distribution is reparameterized by a conditional normalizing flowFi ( Rezende & Mohamed , 2016 ; Papamakarios et al. , 2019a ) into a distribution of more complex density in the space RSθi . The flow Fi is a diffeomorphism in the space RSθi conditioned by the encoding Ehi . Second , the obtained latent distribution is projected onto the event space in which θi lives by the application of a link function diffeomorphism li . For instance , if θi is a variance parameter , the link function would map R onto R+∗ ( li = Exp as an example ) . The usage of Fi and the link function li is repeated on plates of larger rank than the hierarchy hi of θi . The resulting conditional density estimator qi for the posterior distribution p ( θi|X ) is given by : ui ∼ N ( −→ 0 Bhi×Sθi , IBhi×Sθi ) θ̃i = −−−→ li ◦ Fi ( Bhi ) ( ui ; Ehi ) ∼ qi ( θi ; Ehi ) ( 2 ) In Fig . 1 Γ = [ γ1 , γ2 ] is associated to the diffeomorphism −−−−→ lγ ◦ Fγ ( B1 ) . This diffeomorphism is conditioned by the encoding E1 . Both Γ and E1 share the batch shape B1 = Card ( P1 ) . Decomposing the encoding E1 = [ e11 , e21 ] , e11 is used to condition the inference on γ1 , and e21 for γ2 . λ is associated to the diffeomorphism lλ ◦ Fλ , and κ to lκ ◦ Fκ , both conditioned by E2 . Parsimonious parameterization Our approach produces a parameterization effectively independent from plate cardinalities . Consider the latent RVs θ1 , . . . , θL . Normalizing flow-based density estimators have a parameterization quadratic with respect to the size of the space they are applied to ( e.g . Papamakarios et al. , 2018 ) . Applying a single normalizing flow to the total event space of θ1 , . . . , θL would thus result in O ( [ ∑L i=1 Sθi ∏P p=hi Card ( Pp ) ] 2 ) weights . But since we instead apply multiple flows on the spaces of size Sθi and repeat their usage across all plates Phi , . . . , PP , we effectively reduce this parameterization to : # weightsADAVI = O ( L∑ i=1 S2θi ) ( 3 ) As a consequence , our method can be applied to HBMs featuring large plate cardinalities without scaling up its parameterization to impractical ranges , preventing a computer memory blow-up .

Summary: The paper tackles approximate inference for hierarchical Bayesian models with fully nested structure. The specific approach taken is variational inference with a q-distribution that iteratively applies conditional normalizing flows to derive a hierarchical representation, and factorizes in a manner parallel to the generative model by reusing flow parameters, thus having a number of parameters that does not grow with cardinality of each plate. The benefits of the model are illustrated in a few synthetic experiments as well as in application to a human neuroimaging dataset.