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: This work introduces ADAVI, an approximate inference algorithm for hierarchical Bayesian models (HBMs). The approach is similar to NPE from simulation-based inference but exploits the hierarchical structure of the forward model to generate an efficient variational family automatically. Experiments demonstrate the applicability of the method on HBMs of increasing complexity, including a challenging neuroimaging model. Results indicate good performance against other amortized methods.