Learning Likelihoods with Conditional Normalizing Flows 

1 INTRODUCTION . Learning conditional distributions pY |X ( y|x ) is one of the oldest problems in machine learning . When the output y is high-dimensional this is a particularly challenging task , and the practitioner is left with many design choices . Do we factorize the conditional ? If not , do we model correlations with , say , a conditional random field ( Prince , 2012 ) ? Do we use a unimodal distribution ? How fat should the tails be ? Do we use an explicit likelihood at all , or use implicit methods ( Mohamed & Rezende , 2015 ) such as a GAN ( Goodfellow et al. , 2014 ) ? Do we quantize the output ? Ideally , the practitioner should not have to make design choices at all , and the distribution should be learned from the data . In the field of density estimation normalizing flows ( NFs ) are a relatively new family of models ( Rezende & Mohamed , 2015 ) . NFs model complicated high dimensional marginal distributions pY ( y ) by transforming a simple base distribution or prior pZ ( z ) through a learnable , invertible mapping fφ and then applying the change of variables formula . NFs are efficient in inference and sampling , are able to learn inter-dimensional correlations and multi-modality , and they are exact likelihood models , amenable to gradient-based optimization . Flow-based generative models ( Dinh et al. , 2016 ) are generally trained on the image space , and are in some cases computationally efficient in both the forward and inverse direction . These are advantageous over other likelihood based methods because i ) sampling is efficient opposed to autoregressive models ( Van Oord et al. , 2016 ) , and ii ) flows admit exact likelihood optimization in contrast with variational autoencoders ( Kingma & Welling , 2014 ) . Conditional random fields directly model correlations between pixels , and have been fused with deep learning ( Chen et al. , 2016 ) . However , they require the practitioner to choose which pixels have pairwise interactions . Another approach uses adversarial training ( Goodfellow et al. , 2014 ) . A downside is that the training procedure can be unstable and they are difficult to evaluate quantitatively . We propose to learn the likelihood of conditional distributions with few modeling choices using Conditional Normalizing Flows ( CNFs ) . CNFs can be harnessed for conditional distributions pY |X ( y|x ) by conditioning the prior and the invertible mapping on the input x . In particular , we apply conditional flows to super-resolution ( Wang et al. , 2018 ) and vessel segmentation ( Staal et al. , 2004 ) . We evaluate their performance gains on multivariate prediction tasks along side architecturally-matched factored baselines by comparing likelihood and application specific evaluation metrics . 2 BACKGROUND . In the following , we present the relevant background material on normalizing flows and structured prediction . This section covers the change of variables formula , invertible modules , variational dequantization and conventional likelihood optimization . 2.1 NORMALIZING FLOWS . A standard NF in continuous space is based on a simple change of variables formula . Given two spaces of equal dimension Z and Y ; a once-differentiable , parametric , bijective1 mapping fφ : Y → Z , where φ are the parameters of f ; and a prior distribution pZ ( z ) , we can model a complicated distribution pY ( y ) as pY ( y ) = pZ ( fφ ( y ) ) ∣∣∣∣∂fφ ( y ) ∂y ∣∣∣∣ . ( 1 ) The term |∂fφ ( y ) /∂y| is the Jacobian determinant of fφ , evaluated at y and it accounts for volume changes induced by fφ . The transformation fφ introduces correlations and multi-modality in pY . The main challenge in the field of normalizing flows is designing the transformation fφ . It has to be i ) bijective , ii ) have an efficient and tractable Jacobian determinant , iii ) be from a ‘ flexible ’ model class . In addition , iv ) for fast sampling the inverse needs to be efficiently computable . Below we briefly state which invertible modules are used in our architectures , obeying the aforementioned points . Coupling layers Affine coupling layers ( Dinh et al. , 2016 ) are invertible , nonlinear layers . They work by splitting the input z into two components z0 and z1 and nonlinearly transforming z0 as a function of z1 , before reconcatenating the result . If z = [ z0 , z1 ] and y = [ y0 , y1 ] this is y0 = s ( z1 ) · z0 + t ( z1 ) z0 = ( z0 − t ( y1 ) ) /s ( y1 ) y1 = z1 z1 = y1 Where the scale s ( · ) and translation t ( · ) functions can be any function , typically implemented with a CNN . Similar conditioning with normalizing flows has been done in previous works by Mohamed & Rezende ( 2015 ) and Kingma et al . ( 2016b ) . Invertible 1 x 1 Convolutions Proposed in Kingma & Dhariwal ( 2018 ) , invertible 1 × 1 convolutions help mix information across channel dimensions . We implement them as regular 1 × 1 convolutions and for the inverse , we convolve with the inverse of the kernel . Squeeze layers Squeeze layers ( Dinh et al. , 2016 ) are used to compress the spatial resolution of activations . These also help with increasing spatial receptive field of pixels in the deeper activations . Split Prior Split priors ( Dinh et al. , 2016 ) work by spliting a set of activations z into two components z0 and z1 . We then condition z1 on z0 using a simple base density e.g . p ( z1|z0 ) = N ( z1 ; µ ( z0 ) , σ2 ( z0 ) ) , where µ ( · ) and σ2 ( · ) are neural networks . The component z0 can be modeled by further flow layers . This prior , is useful for modeling hierarchical correlations between dimensions , and also helps reduce computation , since z0 is reduced in size . Variational dequantization When modeling discrete data , Theis et al . ( 2016 ) introduced the concept of dequantization . For this , they modeled the probability mass function over y as a latent variable model Pmodel ( y ) = ∫ V P ( y|v ) p ( v ) dv = ∫ V p ( y , v ) dv ( 2 ) 1We can also use injective mappings where the latent variables v ∈ V are continuous-valued . This is a convenient model to use , since the marginal p ( v ) , living on a continuous sample space , can be modelled with a continuous NF . The distribution P ( y|v ) is known as the quantizer and is typically an indicator function P ( y|v ) = I [ v ∈ y + [ 0 , 1 ) D ] . Other works ( Hoogeboom et al. , 2019a ; Tran et al. , 2019 ) directly model Pmodel ( y ) with a discrete-valued flow , but these are known to be difficult to optimize . As an extension of dequantization , Ho et al . ( 2019 ) introduced a variational distribution q ( v|y ) , called a dequantizer , and write a lower bound on the data log-likelihood using Jensen ’ s inequality as follows EPdata ( y ) logPmodel ( y ) = EPdata ( y ) log ∫ p ( y , v ) dv ( 3 ) = EPdata ( y ) log ∫ q ( v|y ) q ( v|y ) p ( y , v ) dv ≥ EPdata ( y ) ∫ q ( v|y ) log p ( y , v ) q ( v|y ) dv ( 4 ) Noting that the joint p ( y , v ) = I [ v ∈ y + [ 0 , 1 ) D ] p ( v ) , we see that the dequantizer distribution q ( v|y ) must be defined such that v ∈ y + [ 0 , 1 ) D , otherwise p ( y , v ) = 0 and the lower-bound is undefined . Restricting q ( v|y ) to satisfy this condition , results in the following variational dequantization bound EPdata ( y ) logPmodel ( y ) ≥ EPdata ( y ) ∫ q ( v|y ) log p ( v ) q ( v|y ) dv . ( 5 ) 2.2 STRUCTURED PREDICTION . Structured prediction tasks such as image segmentation or super-resolution , can be probabilistically framed as learning an unknown target distribution p∗ ( y|x ) , with an input x ∈ X and a target y ∈ Y . In practice with deep learning models , the unknown distribution is often learned by a factored model : p ( y|x ) = D∏ d=1 p ( yd|x ) , ( 6 ) where yd represents the dth dimension of y . Several loss-based optimization methods are a special case of this factored model . The mean squared error is equivalent to a product of normal distributions with equal and fixed standard deviation . Other examples are : cross entropy , equivalent to a product of log categorical distributions , and binary cross entropy , equivalent to a product of log Bernoulli distributions . With factorized independent likelihoods , individual dimensions of y are assumed to be conditionally independent . As a result , sampling leads to results with uncorrelated noise over the output dimensions . In the literature , a fix for this problem is to visualize the mode of the distribution and interpret that as a prediction . However , because the likelihood was optimized assuming a conditionally independent noise distribution , these modes tend to be blurry and lack crisp details . 3 METHOD . In this section we present our main innovations . i ) learning conditional likelihoods using CNFs and ii ) a variational dequantization framework for binary random variables . 3.1 CONDITIONAL NORMALIZING FLOWS . We propose to learn conditional likelihoods using conditional normalizing flows for complicated target distributions in multivariate prediction tasks . Take an input x ∈ X and a regression target y ∈ Y . We learn a complicated distribution pY |X ( y|x ) using a conditional prior pZ|X ( z|x ) and a mapping fφ : Y × X → Z , which is bijective in Y and Z . The likelihood of this model is : pY |X ( y|x ) = pZ|X ( z|x ) ∣∣∣∣ ∂z∂y ∣∣∣∣ = pZ|X ( fφ ( y , x ) |x ) ∣∣∣∣∂fφ ( y , x ) ∂y ∣∣∣∣ . ( 7 ) Notice that the difference between Equations 1 and 7 is that all distributions are conditional and the flow has a conditioning argument of x . The generative process from x to y ( shown in Figure 1 ) can be described by first sampling z ∼ pZ|X ( z|x ) from a simple base density with its parameters conditioned on x ( for us this is a diagonal Gaussian ) and then passing it through a sequence of bijective mappings f−1φ ( z ; x ) . This allows for modelling multimodal conditional distributions in y , which is typically uncommon . For the training procedure , the process runs in reverse . We begin with label y and conditioning input x . We ‘ flow ’ the label back through fφ to yield z = fφ ( y ; x ) , and then we evaluate the log-likelihood of the parameters of the prior , given this transformed label z . The flow and prior parameters can be optimized using stochastic gradient descent and training in minibatches in the usual fashion . Note that this style of training a conditional density model pY |X ( y|x ) differs fundamentally from traditional models , because we compute the log-likelihood in z-space and not y-space . As a result , we are not biasing our results with an arbitrary choice of output-space likelihood or in the case of this paper , handcrafted image loss . Instead , one could interpret this method as learning the correlational and multimodal structure of the likelihood or simply put loss-learning . Conditional modules In our work , the conditioning is introduced in the prior , the split priors , and the affine coupling modules . For the prior , we set the mean and variance as functions of x . For the split prior , we add x as a conditoning argument to the conditional p ( z1|z0 , x ) . And for the affine coupling layers , we pass x to the scale and translation networks so that Conditional Prior p ( z|x ) = N ( z ; µ ( x ) , σ2 ( x ) ) Conditional Split Prior p ( z1|z0 , x ) = N ( z1 ; µ ( z0 , x ) , σ2 ( z0 , x ) ) Conditional Coupling y0 = s ( z1 , x ) · z0 + t ( z1 , x ) ; y1 = z1 In practice these functions are implemented using deep neural networks . First the conditioning term x is transformed into a rich representation h = g ( x ) using a large network g. Subsequently , each function in the flow is applied to a concatenation [ · , · ] of h and the relevant part of z . For example , the translation of a conditional coupling is computed as t ( z1 , x ) = NN ( [ z1 , h ] ) .

Summary: This paper presented the conditional normalizing flows (CNFs) as a new kind of likelihood-based learning objective.  There are two keys in CNFs. One is the parametric mapping function f_{\phi} and the other is the conditional prior. This paper assumed the conditional prior as Gaussian distribution of x. The mapping function is invertible with x as a parameter. The prior parameter and \phi are updated by stochastic gradient descent. The latent variable z is then sampled from conditional prior. The output targe y is obtained with dependency on x and f_{\phi}.