Mind the Gap when Conditioning Amortised Inference in Sequential Latent-Variable Models

1 INTRODUCTION . Variational inference has paved the way towards learning deep latent-variable models ( LVMs ) : maximising the evidence lower bound ( ELBO ) approximates maximum likelihood learning ( Jordan et al. , 1999 ; Beal , 2003 ; MacKay , 2003 ) . An efficient variant is amortised variational inference where the approximate posteriors are represented by a deep neural network , the inference network ( Hinton et al. , 1995 ) . It produces the parameters of the variational distribution for each observation by a single forward pass—in contrast to classical variational inference with a full optimisation process per sample . The framework of variational auto-encoders ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) adds the reparameterisation trick for low-variance gradient estimates of the inference network . Learning of deep generative models is hence both tractable and flexible , as all its parts can be represented by neural networks and fit by stochastic gradient descent . The quality of the solution is largely determined by how well the true posterior is approximated : the gap between the ELBO and the log marginal likelihood is the KL divergence from the approximate to the true posterior . Recent works have proposed ways of closing this gap , suggesting tighter alternatives to the ELBO ( Burda et al. , 2016 ; Mnih & Rezende , 2016 ) or more expressive variational posteriors ( Rezende & Mohamed , 2015 ; Mescheder et al. , 2017 ) . Cremer et al . ( 2018 ) provide an analysis of this gap , splitting it in two . The approximation gap is caused by restricting the approximate posterior to a specific parametric form , the variational family . The amortisation gap comes from the inference network failing to produce the optimal parameters within the family . In this work , we address a previously unstudied aspect of amortised posterior approximations : successful approximate posterior design is not merely about picking a parametric form of a probability density . It is also about carefully deciding which conditions to feed into the inference network . This is a particular problem in sequential LVMs , where it is common practice to feed only a restricted set of inputs into the inference network : not conditioning on latent variables from other time steps can vastly improve efficiency ( Bayer & Osendorfer , 2014 ; Lai et al. , 2019 ) . Further , leaving out ∗Equal contribution . p ( z | C , C̄ = 0 ) p ( z | C , C̄ = 1 ) p ( z | C ) ν ( z | C ) ≈ p ( z | C ) ω ( z | C ) Figure 1 : Illustration of the effect of partial conditioning . Consider a latent-variable model p ( C , C ∣∣ z ) × p ( z ) , with binary C and arbitrary C. We omit C from the amortised approximate posterior q ( z | C ) . Shown are two cases of separated ( left ) and overlapping ( right ) Gaussian true posteriors . Top row : the full posteriors p ( z ∣∣ C , C = 0 ) and p ( z ∣∣ C , C = 1 ) as well as their average , the marginal posterior p ( z | C ) . Middle row : Variational Gaussian approximation ν ( z | C ) to the marginal posterior , which was obtained by stochastic gradient descent on the reverse , modeseeking KL-divergence ( Hoffman et al. , 2013 ) . Bottom row : The optimal ω ( z | C ) obtained by optimising the ELBO with a partially-conditioned amortised approximate posterior w. r. t. q . It is located far away from the modes , sharply peaked and shares little mass with the true full posteriors , the marginal posterior as well as the approximate marginal posterior . observations from the future structurally mimics the Bayesian filter , useful for state estimation in the control of partially observable systems ( Karl et al. , 2017b ; Hafner et al. , 2019 ; Lee et al. , 2020 ) . We analyse the emerging conditioning gap and the resulting suboptimality for sequential amortised inference models . If the variational posterior is only partially conditioned , all those posteriors equal w. r. t. included conditions but different on the excluded ones will need to share an approximate posterior . The result is a suboptimal compromise over many different true posteriors which can not be mitigated by improving variational family or inference network capacity . We empirically show its effects in an extensive study on three real-world data sets on the common use case of variational state-space models . 2 AMORTISED VARIATIONAL INFERENCE . 2.1 VARIATIONAL AUTO-ENCODERS . Consider the task of fitting the parameters of a latent-variable model pθ ( x , z ) = pθ ( x | z ) × pθ ( z ) to a target distribution p̂ ( x ) by maximum likelihood : argmax θ Ex∼p̂ ( x ) [ log pθ ( x ) ] . ( 1 ) In the case of variational auto-encoders ( VAEs ) , the prior distribution pθ ( z ) is often simple , such as a standard Gaussian . The likelihood pθ ( x | z ) on the other hand is mostly represented by a deep neural network producing the likelihood parameters as a function of z . The log marginal likelihood log pθ ( x ) = log ∫ pθ ( x , z ) dz contains a challenging integral within a logarithm . The maximisation of eq . ( 1 ) is thus generally intractable . Yet , for a single sample x the log marginal likelihood can be bounded from below by the evidence lower bound ( ELBO ) : log pθ ( x ) ≥ log pθ ( x ) −KL ( q ( z ) || pθ ( z | x ) ) ( 2 ) = Ez∼q [ log pθ ( x | z ) ] −KL ( q ( z ) || pθ ( z ) ) = : − ` ( θ , q , x ) , where a surrogate distribution q , the variational posterior , is introduced . The bound is tighter the better q ( z ) approximates pθ ( z | x ) . The gap vanishes with the posterior KL divergence in eq . ( 2 ) . Typically , the target distribution is given by a finite data set D = { x ( n ) } Nn=1 . In classical variational inference , equally many variational posteriors { q ( n ) ( z ) } Nn=1 are determined by direct optimisation . For instance , each Gaussian q ( n ) ( z ) is represented by its mean µ ( n ) and standard deviation σ ( n ) . The set of all possible distributions of that parametric form is called the variational family Q 3 q ( z ) . Learning the generative model parameters θ can then be performed by an empirical risk minimisation approach , where the average negative ELBO over a data set is minimised : arg min θ , { q ( n ) } Nn=1 N∑ n=1 ` ( θ , q ( n ) ( z ) , x ( n ) ) . The situation is different for VAEs . The parameters of the approximate posterior for a specific sample x are produced by another deep neural network , the inference network qφ ( z | x ) . Instead of directly optimising for sample-wise parameters , a global set of weights φ for the inference network is found by stochastic gradient-descent of the expected negative ELBO : argmin θ , φ Ex∼p̂ ( x ) [ ` ( θ , qφ ( z | x ) , x ) ] . This approach tends to scale more gracefully to large data sets . The Sequential Case The posterior of a sequential latent-variable model p ( x1 : T , z1 : T ) = T∏ t=1 p ( xt | x1 : t−1 , z1 : t ) p ( zt | z1 : t−1 , x1 : t−1 ) . of observations ( x1 , . . . , xT ) = x1 : T and latents ( z1 , . . . , zT ) = z1 : T , with z1:0 = ∅ , factorises as p ( z1 : T | x1 : T ) = T∏ t=1 p ( zt | z1 : t−1 , x1 : T ) . ( 3 ) The posterior KL divergence is a sum of step-wise divergences ( Cover & Thomas , 2006 ) : KL ( q ( z1 : T | · ) || p ( z1 : T | x1 : T ) ) = T∑ t=1 Ez1 : t−1∼q [ KL ( q ( zt | · ) || p ( zt | z1 : t−1 , x1 : T ) ) ] ( 4 ) where we have left out the conditions of q ( z1 : T | · ) = ∏ t q ( zt | · ) for now . This decomposition allows us to focus on the non-sequential case in our theoretical analysis , even though the phenomenon most commonly occurs with sequential LVMs . 2.2 THE APPROXIMATION AND AMORTISATION GAPS . The ELBO and the log marginal likelihood are separated by the posterior Kullback-Leibler divergence KL ( q ( z ) || pθ ( z | x ) ) , cf . eq . ( 2 ) . If the optimal member of the variational family q ? ∈ Q is not equal to the true posterior , or p ( z | x ) /∈ Q , then KL ( q ? ( z ) || pθ ( z | x ) ) > 0 , so that the ELBO is strictly lower than the log marginal likelihood—the approximation gap . Using an inference network we must expect that even the optimal variational posterior within the variational family is not found in general . The gap widens : KL ( qφ ( z | x ) || pθ ( z | x ) ) ≥ KL ( q ? ( z ) || pθ ( z | x ) ) . This phenomenon was first discussed by Cremer et al . ( 2018 ) . The additional gap KL ( qφ ( z | x ) || pθ ( z | x ) ) −KL ( q ? ( z ) || pθ ( z | x ) ) ≥ 0. is called the amortisation gap . The two gaps represent two distinct suboptimality cases when training VAEs . The former is caused by the choice of variational family , the latter by the success of the amortised search for a good candidate q ∈ Q . 3 PARTIAL CONDITIONING OF INFERENCE NETWORKS . VAEs add an additional consideration to approximate posterior design compared to classical variational inference : the choice of inputs to the inference networks . Instead of feeding the entire observation x , one can choose to feed a strict subset C ⊂ x . This is common practice in many sequential variants of the VAE , cf . table 1 , typically motivated by efficiency or downstream applications . The discrepancy between inference models inputs and the true posterior conditions has been acknowledged before , e. g. by Krishnan et al . ( 2017 ) , where they investigate different variants . Their analysis does not go beyond quantitative empirical comparisons . In the following , we show that such design choices lead to a distinct third cause of inference suboptimality , the conditioning gap .

Summary: This paper investigates the effect of partial conditioning on amortized inference in variational auto-encoders, focusing specifically on sequential data sources where it is common practice to have a posterior that is factorized in such a way that conditioning is partial (usually only conditioning on past signals in the sequence). Given a true posterior that is conditioned on the entire observed datapoint, the authors discuss the effect of having an approximate posterior that is only conditioned on part of the input. As the approximate posterior cannot adapt to the part of the input that is left out of the conditioning, the evidence lower bound becomes less tight, due to the larger KL divergence between the approximate posterior and the true posterior. The authors compare this to the work by Cramer et al. [1], where the distinction was made between having a restricted family of possible distributions for the approximate posterior (approximation gap) and the gap between an amortized approximate posterior with an inference network shared for all datapoints and a non-amortized approximate posterior that is optimized for each datapoint separately (amortisation gap). They argue that partial conditioning leads to a third type of gap which is distinct of the aforementioned inference gaps. Through an example with discrete observations the authors derive that when the true posterior is conditioned on the full data, and the approximate posterior is only partially conditioned, the optimal approximate posterior is something akin to a product of true posteriors over the unconditioned information, and not a mixture where the left out information is marginalized out. Through a 1D example they show that this could lead to overly sharp posteriors that have high densities in regions where the true posterior has very low density.