Multiscale Invertible Generative Networks for High-Dimensional Bayesian Inference

1 INTRODUCTION . Bayesian inference provides a powerful framework to blend prior knowledge , data generation process and ( possibly small ) data for statistical inference . With some prior knowledge ⇢ ( distribution ) for the quantity of interest x 2 Rd , and some ( noisy ) measurement y 2 Rdy , it casts on x a posterior q ( x|y ) / ⇢ ( x ) L ( y|x ) , where L ( y|x ) = N ( y F ( x ) ; 0 , `` ) . ( 1 ) where L ( y|x ) is the likelihood that compares the data y with system prediction F ( x ) from the candidate x , here F denotes the forward process . We can use different distributions to model the mismatch `` = y F ( x ) , and for illustration simplicity , we assume Gaussian in Equation 1 . For example , Bayesian deep learning generates model predicted logits F ( x ) from model parameters x , and compares it with discrete labels y through binomial or multinomial distribution . Sampling or inferring from q is a long-standing challenge , especially for high-dimensional ( high-d ) cases . An arbitrary high-d posterior can have its importance regions ( also called “ modes ” ) anywhere in the high-d space , and finding these modes requires computational cost that grows exponentially with the dimension d. This intrinsic difficulty is the consequence of “ the curse of dimensionality ” , which all existing Bayesian inference methods suffer from , e.g. , MCMC-based methods ( Neal et al. , 2011 ; Welling & Teh , 2011 ; Cui et al. , 2016 ) , SVGD-type methods ( Liu & Wang , 2016 ; Chen et al. , 2018 ; 2019a ) , and generative modeling ( Morzfeld et al. , 2012 ; Parno et al. , 2016 ; Hou et al. , 2019 ) . In this paper , we focus on Bayesian inference problems with multiscale structure and exploit this structure to sample from a high-d posterior . While the original problem has a high spatial resolution ( fine-scale ) , its low resolution ( coarse-scale ) analogy is computationally attractive because it lies in a low-dimension ( low-d ) space . A problem has the multiscale structure if such coarse-scale low-d surrogate exists and gives good approximation to the fine-scale high-d problem , see Section 2.1 . Such multiscale property is very common in high-d Bayesian inference problems . For example , inferring 3-D permeability field of subsurface at the scale of meters is a reasonable approximation of itself at the scale of centimeters , while the problem dimension is 106-times fewer . We propose a Multiscale Invertible Generative Network ( MsIGN ) to sample from high-d Bayesian inference problems with multiscale structure . MsIGN is a flow-based generative network that can both generate samples and give density evaluation . It consists of multiple scales that recursively lifts up samples to a finer-scale ( higher-resolution ) , except that the coarsest scale directly samples from a low-d ( low resolution ) distribution . At each scale , a fixed prior conditioning layer combines coarse-scale samples with some random noise according to the prior to enhance the resolution , and then an invertible flow modifies the samples for better accuracy , see Figure 1 . The architecture of MsIGN makes it fully invertible between the final sample and random noise at all scales . MsIGN undergoes a multi-stage training that learns a hierarchy of distributions with dimensions growing from the lowest to the highest ( the target posterior ) . Each stage gives a good initialization to the next stage thanks to the multiscale property . To capture multiple modes , we choose Jeffreys divergence DJ ( pkq ) as the training objective at each stage , which is defined as DJ ( pkq ) = DKL ( pkq ) +DKL ( qkp ) = Ex⇠p [ log ( p ( x ) /q ( x ) ) ] + Ex⇠q [ log ( q ( x ) /p ( x ) ) ] . ( 2 ) Jeffreys divergence removes bad local minima of single-sided Kullback-Leibler ( KL ) divergence to avoid mode missing . We build an unbiased estimation of it by leveraging prior conditioning layer in importance sampling . Proper loss function and good initialization from multi-stage training solve the non-convex optimization stably and capture multi-modes of the high-d distribution . In summary , we claim four contributions in this work . First , we propose a Multiscale Invertible deep Generative Network ( MsIGN ) with a novel prior conditioning layer , which can be trained in a coarse-to-fine scale manner . Second , Jeffreys divergence is used as the objective function to avoid mode collapse , and is estimated by importance sampling based on the prior conditioning layer . Third , when applied to two Bayesian inverse problems , MsIGN clearly captures multiple modes in the high-d posterior and approximates the posterior accurately , demonstrating its superior performance compared with previous methods via the generative modeling approach . Fourth , we also apply MsIGN to image synthesis tasks , where it achieves superior performance in bits-perdimension among our baseline models , like Glow ( Kingma & Dhariwal , 2018 ) , FFJORD ( Grathwohl et al. , 2018 ) , Flow++ ( Ho et al. , 2019 ) , i-ResNet ( Behrmann et al. , 2019 ) , and Residual Flow ( Chen et al. , 2019b ) . MsIGN also yields great interpret-ability of its neurons in intermediate layers . 2 METHOLOGY . We will abbreviate q ( x|y ) in Equation 1 as q ( x ) for simplicity in the following context , because y only plays the role of defining the target distribution q ( x ) in MsIGN . In Section 2.1 , we discuss the multiscale structure in detail of the posterior q ( x ) and derive a scale decoupling that can be utilized to divide and conquer the high-d challenge of Bayesian inference . As a flow-based generative model like in Dinh et al . ( 2016 ) , MsIGN models a bijective that maps Gaussian noise z to a sample x whose distribution is denoted as p✓ ( x ) , where ✓ is the network parameters . MsIGN allows fast generation of samples x and density evaluation p✓ ( x ) , so we train our working distribution p✓ ( x ) to approximate the target distribution q ( x ) . We present the architecture of MsIGN in Section 2.2 and the training algorithm in Section 2.3 . 2.1 MULTISCALE STRUCTURE AND SCALE DECOUPLING . We say a Bayesian inference problem has multiscale structure if the associated coarse-scale likelihood Lc approximates the original likelihood L well : L ( y|x ) ⇡ Lc ( y|xc ) , where Lc ( y|xc ) : = N ( y Fc ( xc ) ; 0 , `` ) . ( 3 ) Here xc 2 Rdc is a coarse-scale version of the fine-scale quantity x 2 Rd ( dc < d ) , given by a deterministic pooling operator A : xc = A ( x ) . The map Fc : Rdc ! Rdy is a forward process that gives system prediction based on the coarse-scale information xc . A popular case of the multiscale structure is when A is the average pooling operator , and F ( x ) ⇡ Fc ( xc ) , meaning that the system prediction mainly depends on the lower-resolution information xc . Equation 3 motivates us to define a surrogate distribution q̃ ( x ) / ⇢ ( x ) Lc ( y|A ( x ) ) that approximates the target posterior q ( x ) well1 : q̃ ( x ) = ⇢ ( x ) Lc ( y|A ( x ) ) = ⇢ ( x ) Lc ( y|xc ) ⇡ ⇢ ( x ) L ( y|x ) = q ( x ) . ( 4 ) We also notice that the prior ⇢ allows an exact scale decoupling . To generate a sample x from ⇢ , one can first sample its coarse-scale version xc = A ( x ) , and then replenish missing fine-scale details without changing the coarse-scale structure by sampling from the conditional distribution ⇢ ( x|xc ) = ⇢ ( x|A ( x ) = xc ) . Using ⇢c to denote the distribution of xc = A ( x ) , the conditional probability calculation summarizes this scale decoupling process as ⇢ ( x ) = ⇢ ( x|xc ) ⇢c ( xc ) . Combining the scale effect in the likelihood and the scale decoupling in the prior , we decouple the surrogate q̃ ( x ) = ⇢ ( x ) Lc ( y|A ( x ) ) into the prior conditional distribution ⇢ ( x|xc ) and a coarse-scale posterior , defined as qc ( xc ) : = ⇢c ( xc ) L ( y|xc ) . The decoupling goes as q̃ ( x ) = ⇢ ( x ) Lc ( y|xc ) = ⇢ ( x|xc ) ⇢c ( xc ) Lc ( y|xc ) = ⇢ ( x|xc ) qc ( xc ) , ( 5 ) The prior conditional distribution ⇢ ( x|xc ) bridges the coarse-scale posterior qc ( xc ) and the surrogate q̃ ( x ) , which in turn approximates the original fine-scale posterior q ( x ) . Parno et al . ( 2016 ) proposed a similar scale decoupling relation , and we leave the discussion and comparison to Appendix A . Figure 1 shows the integrated sampling strategy . To sample an x from q , we start with an xc from qc . The prior conditioning layer then performs random upsampling from the prior conditional distribution ⇢ ( ·|xc ) , and the output will be a sample x̃ of the surrogate q̃ . Due to the approximation q̃ ⇡ q from Equation 4 , we stack multiple invertible blocks for the invertible flow F to modify the sample x̃ ⇠ q̃ to a sample x ⇠ q : x = F ( x̃ ) . F is initialized as an identity map in training . Finally , to obtain the xc from qc , we apply the above procedure recursively until the dimension of the coarsest scale is small enough so that qc can be easily sampled by a standard method . 2.2 MULTISCALE INVERTIBLE GENERATIVE NETWORK : ARCHITECTURE . Our proposed MsIGN has multiple levels to recursively apply the above strategy . We denote L the number of levels , xl 2 Rdl the sample at level l , and Al : Rdl ! Rdl 1 the pooling operator from level l to l 1 : xl 1 = Al ( xl ) . Following the idea in Section 2.1 , we can define the l-th level target ql ( xl ) and surrogate q̃l ( x̃l ) , and the last-level target qL is our original target q in Equation 1 . The l-th level of MsIGN uses a prior conditioning layer PCl and an inverse transform Fl to capture ql . Prior conditioning layer . The prior conditioning layer PCl at level l lifts a coarse-scale sample xl 1 2 Rdl 1 up to a random fine-scale one xl 2 Rdl following the conditional distribution ⇢ ( xl|xl 1 ) . The difference in dimension is compensated by a Gaussian noise zl 2 Rdl dl 1 , which is the source of randomness : xl = PCl ( xl 1 , zl ) . PCl depends only on the prior conditional distribution ⇢ ( xl|xl 1 ) , and thus can be pre-computed independently for different levels regardless of the likelihood L. When the prior is Gaussian and the pooling operators are linear ( e.g. , average pooling ) , the prior conditional distribution is still Gaussian with moments specified as follows . Lemma 2.1 Suppose that ⇢ ( xl ) = N ( xl ; 0 , ⌃l ) , and Al ( xl ) = Alxl for some Al 2 Rdl 1⇥dl , then with Ul 1 : = ⌃lATl ( Al⌃lA T l ) 1 and ⌃l|l 1 : = ⌃l ⌃lATl ( Al⌃lATl ) 1Al⌃l , we have ⇢ ( xl|xl 1 = Alxl ) = N ( xl ; Ul 1xl 1 , ⌃l|l 1 ) . 1We omit normalizing constants . Equivalence and approximation are up to normalization in the following . With the Cholesky decomposition ( or eigen-decomposition ) ⌃l|l 1 = BlBTl , we design the prior conditioning layer PCl as below , which is invertible between xl and ( xl 1 , zl ) : xl = PCl ( xl 1 , zl ) : = Ul 1xl 1 +Blzl , zl ⇠ N ( 0 , Idl dl 1 ) . ( 6 ) We refer readers to Appendix B for proof of Lemma 2.1 and the invertibility in Equation 6 . When the prior is non-Gaussian or the pooling operators are nonlinear , there exists a nonlinear invertible prior conditioning operator xl = PCl ( xl 1 , zl ) such that xl follows the prior conditional distribution ⇢ ( xl|xl 1 ) given xl 1 and zl ⇠ N ( 0 , Idl dl 1 ) . We can pre-train an invertible network to approximate this sampling process , and fix it as the prior conditioning layer . Invertible flow . The invertible flow Fl at level l modifies the surrogate q̃l towards the target ql . The more accurate the multiscale structure in Equation 3 is , the better q̃l approximates ql , and the closer Fl is to the identity map . Therefore , we parameterize Fl by some flow-based generative model and initialize it as an identity map . In practice , we utilize the invertible block of Glow ( Kingma & Dhariwal , 2018 ) , which consists of actnorm , invertible 1⇥ 1 convolution , and affine coupling layer , and stack several blocks as the inverse flow Fl in MsIGN . Overall model . MsIGN is a bijective map between random noise inputs at different scales { zl } Ll=1 and the finest-scale sample xL . The forward direction of MsIGN maps { zl } Ll=1 to xL as below : x1 = F1 ( z1 ) , x̃l = PCl ( xl 1 , zl ) , xl = Fl ( x̃l ) , 2 l L . ( 7 ) As a flow-based generative model , sample generation as in Equation 7 and density evaluation p✓ ( x ) by the change-of-variable rule is accessible and fast for MsIGN . In scenarios when certain bound needs enforcing to the output , we can append element-wise output activations at the end of MsIGN . For example , image synthesis can use the sigmoid function so that pixel values lie in [ 0 , 1 ] . Such activations should be bijective to keep the invertible relation between random noise to the sample .

Summary: This paper presents a model and a corresponding training approach for multi-scale invertible models. The presented model is defined on multiple scales with information on finer scales being conditioned on coarser scales. Data generation is hence done sequentially from a coarser to finer scale. The authors argue that this multi-scale sampling helps in addressing the curse of dimensionality problem by allowing to sample from high density regions more efficiently.