VISCOS Flows: Variational Schur Conditional Sampling with Normalizing Flows

1 INTRODUCTION . Conditional data generation is a ubiquitous and challenging problem , even more so when the data is high dimensional . If the partitioning of the conditioned data and conditioning itself can be established in advance , a wide variety of inference tools can be used ranging from Bayesian inference ( Gelman , 2013 ) and its approximations ( e.g . variational inference ( Klys et al. , 2018 ) ) to Gaussian ( Williams & Rasmussen , 2006 ) or Neural Processes ( Garnelo et al. , 2018 ) . The task is considerably more challenging when the partitioning can not be anticipated . This is met in a wide range of applications , from few shots learning ( Wang et al. , 2020 ) to super resolution image generation ( Bashir et al. , 2021 ) , high-dimensional Bayesian optimisation ( Moriconi et al. , 2020 ; Gomez-Bombarelli et al. , 2018 ) , learning with incomplete datasets ( Yoon et al. , 2018 ; Li et al. , 2019 ; Richardson et al. , 2020 ) , image inpainting ( Elharrouss et al. , 2020 ) and many more . If we can assume that the data is normally distributed , the conditional distribution can be computed in closed form using Gaussian elimination and the Schur complement . This observation paved the way for Gaussian Processes ( GPs ) solutions ( Williams & Rasmussen , 2006 ) and alike . Although they constitute a solid component of the modern machine learning toolbox , the Gaussian assumption is quite restrictive in nature . Indeed : GPs in their vanilla form are used when the observations are uni-dimensional and when inference has to be performed on a subset of these variables . As a consequence , alternative tools may need to be used when dealing with even moderately dimensional data , or when part of the input data can be missing . It is well-established that normalizing flows ( NFs ) ( Rezende & Mohamed , 2015 ) can model arbitrarily complex joint distributions . Furthermore , they naturally embed together the joint probability of any partition of data that could occur during training or at test time . Hence , it is tempting to make use of NFs to perform conditional sampling , but only a few works have provided solutions each with some limitations . ( Cannella et al. , 2020 ) proposed Projected Latent Monte Carlo Markov Chains ( PL-MCMC ) to generate samples whose probability is guaranteed to converge to the desired conditional distribution as the chain gets longer . While such convergence is certainly an advantage , one can identify several important drawbacks : first , as for every other Monte Carlo sampling method , the mixing time can not be known in advance , and hence one may wonder whether the samples that are gathered truly belong to the conditional distribution . Second , although the observed part of the imputed data converges towards the true values , a difference between the two may persist . Third , training a model comprising of missing data with PL-MCMC is achieved through a Monte Carlo Expectation Maximisation ( MCEM ) scheme ( Dempster et al. , 1977 ; Neath , 2013 ) . Under appropriate assumptions , MCEM is guaranteed to converge to a local maximiser of the log-likelihood function , but this is conditioned on the quality of the data generated by the Monte Carlo algorithm . Consequently , as the optimisation progresses towards the optimum , longer chains may be required to ensure convergence or even to obtain convincing results ( Neath , 2013 ) . MCFlow ( Richardson et al. , 2020 ) relies on an auxiliary feedforward neural network whose role is to produce latent embeddings with maximum a posteriori likelihood values . Those values are constrained to lie on the manifold of latent vectors whose mapping to the observed space match the observations . MCFlow produces state-of-the-art results in terms of image quality or classification accuracy when compared to GAN-based methods such as ( Yoon et al. , 2018 ; Li et al. , 2019 ) . However , this method requires a set of adjustments to the model and needs retraining for any additional incomplete data . On the other hand , ACFlow Li et al . ( 2020 ) learns all conditional distributions for all possible masking operations , which can be quite computationally expensive . As such , both MCFlow and ACFlow can not be applied to post-training data completion as PL-MCMC . Several conditioning techniques have been used with NFs in contexts where the joint probability distribution of the conditioning and conditioned random variable is of no interest . For instance , ( Rezende & Mohamed , 2015 ; van den Berg et al. , 2018 ) extended the amortisation technique used in Variational Auto-Encoders ( VAEs ) ( Kingma & Welling , 2013 ) to the parameters of the flow . ( Winkler et al. , 2019 ) extended this use to the scenario of conditional likelihood estimation , while ( Trippe & Turner , 2018 ) provided a Variational Bayes view on this problem by using a Bayesian Neural Network for conditioning feature embedding . On a different line of work , ( Kingma & Dhariwal , 2018 ) use a heuristic form of a posteriori conditioning . The generative flow is trained with a classifier over the latent space , which forces the latent representation location to be indicative of the class they belong to . At test time , some parametric distribution is fitted on the latent representations to account for the attributes of the images that are to be generated . ( Nguyen et al. , 2019 ) used a hybrid approach , whereby the conditioning and conditioned variables joint distribution is modelled by a normalizing flow in such a way that conditional sampling of one given the other is straightforward to apply . Still , either with this approach or the others above , the knowledge of what part of the data constitutes the conditioning random variable is required beforehand . Our Contribution : We propose a conditional NF sampling method in the following setting : a ) conditional data generation is performed after a model has been trained without taking conditional sampling into account ; b ) training data may be itself incomplete , in the sense that some training features might be missing from examples in the training dataset ; c ) the subset of missing data may not be known in advance , and it could also be randomly distributed across input data . Importantly , we are interested in deriving a method whereby the distribution of the generated data faithfully reflects the Bayesian perspective . Our derivations heavily rely on the Schur complement properties in the spirit of the conditional Gaussian distributions . To highlight this feature we call our approach VISCOS Flows : VarIational Schur COnditional Sampling with normalizing Flows . The use of a variational posterior brings some advantages : for example , with a single fitted posterior , multiple samples can quickly be recovered . We also show how to amortise the cost of inference across multiple partially observed items by using inference networks ( Kingma & Welling , 2013 ) . 2 VISCOS FLOWS : VARIATIONAL SCHUR CONDITIONAL SAMPLING WITH NORMALIZING FLOWS . 2.1 PRELIMINARIES AND PROBLEM FORMULATION . Preliminaries : Consider a C1-diffeomorphism f ( X ) : X → Y where X ⊆ Rd , Y ⊆ Rd are open sets and define the inverse of f to be g ( y ) ≡ f−1 : Y → X . Consider the case where the random variable Y = f ( X ) while X is distributed according to the base distribution P0 ( X ) with density p0 ( x ) . According to the change of variable rule Y has a log-density given by the following formula : log p ( y ) = log p0 ( g ( y ) ) + log |det ∇yg ( y ) | = log p0 ( x ) − log |det∇xf ( x ) | . ( 1 ) We will use two set of complementary indexes covering [ d ] : observable O ⊂ [ d ] with cardinality # O = dO , and hidden H ⊂ [ d ] with # H = dH and O ∪H = [ d ] . In what follows , without loss of generality we assume O = { 0 , . . . , dO − 1 } ) and H = { dO , . . . , d } , and use the following notation : f ( x ) = ( fO ( xO ; xH ) fH ( xH ; xO ) ) , ∇xf ( x ) = J ( x ) = ( JOO ( x ) JOH ( x ) JH O ( x ) JHH ( x ) ) , and we will use a similar partition for g ( y ) and its Jacobian G ( y ) . We will use the notation mOO ( J ) = JOO to denote sub-matrix masking . The central to our approach is the Schur complements and its properties and , in particular , the following well-known identities : GHH ( f ( x ) ) = ( JHH ( x ) − JH O ( x ) ( JOO ( x ) ) −1 JOH ( x ) ) −1 , ( 2 ) det ( J ) = det ( JOO ) det ( JHH − JH O ( JOO ) −1 JOH ) = det ( JOO ) det ( GHH ) . ( 3 ) We denote by A the “ detached ” version of a matrix-valued function A ( x ) , i.e . a matrix whose values match those ofA but have a null gradient . Problem formulation : Consider a set Y containing some samples of interest , for example , images . Assume also a pre-trained normalizing flow , i.e. , a map f : X → Y with X , Y ⊆ Rd . In what follows we relax the assumption on the pre-trained flow , but for streamlining the presentation we assume that the normalizing flow is given at this point . Let us now assume that the observation sample y is partially missing or masked ( for instance , photographer ’ s finger is covering part of the image ) . In particular , the observation is split into observed YO = { yO | y ∈ Y } and hidden YH = { yH | y ∈ Y } regions , whereO , H are observed and hidden indexes , respectively . Our goal is to sample yH from the conditional distribution over a partial observation Y O ∈ YO . While this distribution has a density given by p ( yH | yO ) = p ( y ) p ( yO ) , computing p ( yO ) = ∫ p ( yO , yH ) dyH is in general intractable . Since it is often easier to sample from and optimise in the latent space we will aim to express all distributions in the space X . First , we introduce a partition into observed and hidden indexes in the latent space X . Assumption A0 . Both spaces X and Y are partitioned as O = { 0 , . . . , dO − 1 } and H = { dO , . . . , d } . We stress that Assumption A0 is made only to simplify the exposition . In general , we can have arbitrary partitions of the observed space , which will not affect our algorithm . While this partition seems somewhat artificial , it enables the use of the Schur complement in our derivations . Choosing the same partition in X and Y is natural for the invertible residual networks ( iResNet ) ( Behrmann et al. , 2019 ; Chen et al. , 2019 ) due the structure of the flow . This observation also holds for any flows sharing a residual structure x+h ( x ) with an activation h ( e.g . ReLU ) , such as Planar and Radial flows ( Rezende & Mohamed , 2015 ) , Sylvester flows ( van den Berg et al. , 2018 ) , continuous flows ( Chen et al. , 2018 ; Onken et al. , 2021 ; Mathieu & Nickel , 2020 ) or – to some extent – Implicit NFs ( Lu et al. , 2021 ) , where at each layer t of the flow , the derivative of the output xit+1 with respect to the input layer xt will be dominated by the ith component of this input . We note that the partition in latent space X can be made different from the partition in the observed Y for numerical purposes . For example , picking a latent space partitioning with a maximal value of log |det ( JOO ) | will ensure that the matrix JOO is not ill-conditioned . Indeed , for any matrix A we have 2 log |det ( A ) | = log ( det ( AAT ) ) = ∑ i log ( σi ( A ) ) , where σi are the singular values of A . Thus promoting larger singular values will improve conditioning . We empirically found that across several architectures ( iResNet , Planar , Radial , Sylvester and Implicit flows ) , this approach leads to a well-behaving algorithm . Maximising log |det ( JOO ) | to choose partitioning can also potentially be used for coupling flow architectures ( e.g . Glow ( Kingma & Dhariwal , 2018 ) or RealNVPs ( Dinh et al. , 2016 ) ) , however , this requires additional research . While Assumption A0 is made for convenience of notation , the following assumption on f is necessary . Assumption A1 . The map f : X → Y and its restriction fO ( XO ; XH ) : XO → YO for all xH ∈ XH are C1-diffeomorphisms . Let us now formally define our reparametrisation from Y space to X as follows : xH ( yH , yO ) : = gH ( yH ; yO ) , ( 4 ) xO ( xH , yO ) such that fO ( xO ; xH ) − yO = 0 . ( 5 ) Our reparametrisation is well-posed as stated in the following lemma with the proof in Appendix A. Lemma 1 If Assumption A1 holds then the reparametrisation in Equations 4 and 5 is well-posed , i.e , gH ( yH , yO ) is a diffeomorphism , while the map xO ( xH , yO ) exists and it is differentiable . Now we can rewrite Equation 1 using Equation 3 in simpler terms : log p ( y ) = log p0 ( x H ) + log ∣∣detGHH ∣∣+ log p0 ( xO ) − log ∣∣det JOO∣∣ . ( 6 ) Here we used the fact that the base distribution p0 ( x ) can be factorised as a product of independent univariate distributions , and hence p0 ( x ) = p0 ( xO ) p0 ( xH ) . Now we can define the variational distribution in terms of the latent samples as follows : q ( yH ; yO ) = q ( xH ) ∣∣detGHH ∣∣ . Combining Equation 6 with our variational distributions we obtain the evidence lower bound : LELBO = Eq ( yH ; yO ) [ log p ( yH , yO ) ] +H [ q ( yH ; yO ) ] = E q ( xH ) [ log p0 ( xH ) q ( xH ) + log p0 ( xO ) − log ∣∣∣det JOO ( xO , xH ) ∣∣∣ ] , ( 7 ) where H [ · ] is the entropy and we have taken advantage of the fact that the log-absolute determinant ( LAD ) of the Jacobian belonging to the approximate posterior q ( yH ; yO ) cancels with the one from the joint log-probability given by the model ( i.e. , the terms withGHH cancel each other out ) . Furthermore , we have eliminated the dependence on the variable yH completely instead we treat xH as a variational variable . We note here that our , at the first sight , artificial partition of the latent space into observable xO and hidden xH parts was the key allowing for the reparametrisation from Y space to X space . In particular , we used the invertibility of sub-jacobiansGHH and JOO allowing to map between observable and hidden samples in the X and Y spaces . Our approach is in the spirit of Gaussian conditional sampling and Schur complement , which justifies the method ’ s name . Note that this ELBO loss can be added to other losses , e.g. , to a classifier loss as we do in what follows . Note that this ELBO can be used in combination with other losses . Thus the normalizing flow can be trained simultaneously while performing conditional sampling of missing data . This way we can treat masked data during training as well as after training . There are still two technical issues that we cover in the following subsections before presenting our algorithm . First , we need to get xO by solving Equation 5 . Second , we need to discuss the computation of the ELBO gradient with respect to xH . Specifically , we need to reparametrise xO and compute the LAD gradient . As the following two sections are quite technical the reader interested in the algorithm itself can skip to Subsection 2.4 .

Summary: This paper proposes a technique for performing data imputation using pre-trained normalizing flows. They do so by fitting a variational distribution over a subset of the 'base' variables of the normalizing flow. Along with the observations, a sample from this variational distribution is sufficient to fully specify a sample from the normalizing flow model. The authors present techniques for constructing samples in this way, and computing the derivatives required to optimise the variational distribution.