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: The paper reviews the issue of partial conditioning of the amortized posterior in sequential latent variable models, typically state-space models trained with a VAE-style loss, but where the posterior used is the filtering rather than smoothing posterior. The author show that training a model with posterior with missing information can lead to a gap in estimating both the posterior and the corresponding model. They show the benefits of using the correct posteriors in simple examples.