Hybrid Memoised Wake-Sleep: Approximate Inference at the Discrete-Continuous Interface

1 INTRODUCTION . We naturally understand the world around us in terms of discrete symbols . When looking at a scene , we automatically parse what is where , and understand the relationships between objects in the scene . We understand that there is a book and a table , and that the book is on the table . Such symbolic representations are often necessary for planning , communication and abstract reasoning . They allow the specification of goals like “ book on shelf ” or preconditions like the fact that “ book in hand ” must come before “ move hand to shelf ” , both of which are necessary for high-level planning . In communication , we ’ re forced to describe such plans , among many other things we say , in discrete words . And further , abstract reasoning requires creating new symbols by composing old ones which allows us to generalize to completely new situations . A “ tower ” structure made out of books is still understood as a “ tower ” when built out of Jenga blocks . How do we represent and learn models that support such symbolic reasoning while supporting efficient inference ? We focus on a particular class of hybrid generative models pθ ( zd , zc , x ) of observed data x with discrete latent variables zd , continuous latent variables zc and learnable parameters θ with a graphical model structure shown in Figure 1 . In particular , the discrete latent variables zd represent an underlying structure present in the data , while the remaining continuous latent variables zc represent non-structural quantities . For example , in the context of compositional scene understanding , zd can represent a scene graph comprising the object identities and the relationships between them , like “ a small green pyramid is on top of a yellow cube ; a blue doughnut leans on the yellow cube ; the large pyramid is next to the yellow cube ” while zc represents the continuous poses of these objects . In this model , we assume that object identities are discrete , symbolic variables indexing into a set of learnable primitives , parameterized by a subset of the generative model parameters θ . The idea is to make these primitives learn to represent concrete objects like “ yellow cube ” or “ large green pyramid ” from data in an unsupervised fashion . Algorithms suitable for learning such models are based on variational inference or wake-sleep . However , these algorithms are either inefficient or inapplicable to general settings . First , stochastic variational inference methods that optimize the evidence lower bound ( ELBO ) using the reparameterization trick ( Kingma & Welling , 2014 ; Rezende & Mohamed , 2015 ) are not applicable to discrete kernel : WN ( 0.49 ) + SE ( 0.49 , 0.50 ) SE ( 0.50 , 0.51 ) + PerL ( 0.50 , 1.00 , 0.50 ) PerXL ( 0.49 , 1.50 , 0.50 ) + WN ( 0.49 ) ( a ) Trends in time-series : Learning Gaussian process ( GP ) kernel to fit data ( blue ) . Extrapolation ( orange ) shown with inferred kernel expression ( below ) . latent variable models . The REINFORCE gradient estimator ( Williams , 1992 ) , on the other hand , has high variance and continuous relaxations of discrete variables ( Jang et al. , 2017 ; Maddison et al. , 2016 ) don ’ t naturally apply to stochastic control flow models ( Le et al. , 2019 ) . Second , wake-sleep methods ( Hinton et al. , 1995 ; Dayan et al. , 1995 ) like reweighted wake-sleep ( RWS ) ( Bornschein & Bengio , 2015 ; Le et al. , 2019 ) require inferring discrete latent variables at every learning iteration , without saving previously performed inferences . Memoised wake-sleep ( MWS ) ( Hewitt et al. , 2020 ) addresses this issue , but is only applicable to purely discrete latent variable models . We propose hybrid memoised wake-sleep ( HMWS ) —a method for learning and amortized inference in probabilistic generative models with hybrid discrete-continuous latent variables . HMWS combines the strengths of MWS in memoising the discrete latent variables and RWS in handling continuous latent variables . The core idea in HMWS is to memoise discrete latent variables zd and learn a separate recognition model which is used for importance-sampling based approximate inference and marginalization of the continuous latent variables zc . We empirically compare HMWS with state-of-the-art ( i ) stochastic variational inference methods that use control variates to reduce the REINFORCE gradient variance , VIMCO ( Mnih & Rezende , 2016 ) , and ( ii ) a wake-sleep extension , RWS . We show that HMWS outperforms these baselines in two domains : structured time series and compositional 3D scene understanding , respectively . 2 BACKGROUND . Our goal is to learn the parameters θ of a generative model pθ ( z , x ) of latent variables z and data x , and parameters φ of a recognition model qφ ( z|x ) which acts as an approximation to the posterior pθ ( z|x ) . This can be achieved by maximizing the evidence lower bound ELBO ( x , pθ ( z , x ) , qφ ( z|x ) ) = log pθ ( x ) − KL ( qφ ( z|x ) ||pθ ( z|x ) ) ( 1 ) which maximizes the evidence log pθ ( x ) while minimizing the Kullback-Leibler ( KL ) divergence , thereby encouraging the recognition model to approximate the posterior . If the latent variables are discrete , a standard way to estimate the gradients of the ELBO with respect to the recognition model parameters involves using the REINFORCE ( or score function ) gradient estimator ( Williams , 1992 ; Schulman et al. , 2015 ) ∇φELBO ( x , pθ ( z , x ) , qφ ( z|x ) ) ≈ log pθ ( z , x ) qφ ( z|x ) · ∇φ log qφ ( z|x ) +∇φ log pθ ( z , x ) qφ ( z|x ) ( 2 ) where z ∼ qφ ( z|x ) . However , the first term is often high-variance which makes learning inefficient . This issue can be addressed by ( i ) introducing control variates ( Mnih & Gregor , 2014 ; Mnih & Rezende , 2016 ; Tucker et al. , 2017 ; Grathwohl et al. , 2018 ) which reduce gradient variance , ( ii ) continuous-relaxation of discrete latent variables to allow differentiation ( Jang et al. , 2017 ; Maddison et al. , 2016 ) , or ( iii ) introducing a separate “ wake-sleep ” objective for learning the recognition model that is trained in different steps and sidesteps the need to differentiate through discrete latent variables ( Hinton et al. , 1995 ; Dayan et al. , 1995 ; Bornschein & Bengio , 2015 ; Le et al. , 2019 ) . 2.1 MEMOISED WAKE-SLEEP . Both ELBO-based and wake-sleep based approaches to learning require re-solving the inference task by sampling from the recognition model qφ ( z|x ) at every iteration . This repeated sampling can be wasteful , especially when only a few latent configurations explain the data well . Memoised wakesleep ( MWS ) ( Hewitt et al. , 2020 ) extends wake-sleep by introducing a memory—a set of M unique discrete latent variables { zmd } Mm=1 for each data point x which induces a variational distribution qMEM ( zd|x ) = M∑ m=1 ωmδzmd ( zd ) , ( 3 ) consisting of a weighted set of delta masses δzmd centered on the memory elements ( see also Saeedi et al . ( 2017 ) ) . This variational distribution is proven to improve the evidence lower bound ELBO ( x , pθ ( zd , x ) , qMEM ( zd|x ) ) ( see Sec . 3 of ( Hewitt et al. , 2020 ) ) by a memory-update phase comprising ( i ) the proposal of N new values { z′nd } Nn=1 ∼ qφ ( zd|x ) , ( ii ) retaining the best M values from the union of the old memory elements and newly proposed values { zmd } Mm=1 ∪ { z′ n d } Nn=1 scored by pθ ( zd , x ) , and ( iii ) setting the weights to ωm = pθ ( zmd , x ) / ∑M i=1 pθ ( z i d , x ) . MWS , however , only works on models with purely discrete latent variables . If we try to use the same approach for hybrid discrete-continuous latent variable models , all proposed continuous values are will be unique and the posterior approximation will collapse onto the MAP estimate . 2.2 IMPORTANCE SAMPLING BASED APPROXIMATE INFERENCE AND MARGINALIZATION . In our proposed method , we will rely on importance sampling ( IS ) to perform approximate inference and marginalization . In general , given an unnormalized density γ ( z ) , and its corresponding normalizing constant Z = ∫ γ ( z ) dz and normalized density π ( z ) = γ ( z ) /Z , we want to estimate Z and the expectation of an arbitrary function f , Eπ ( z ) [ f ( z ) ] . To do this , we sample K values { zk } Kk=1 from a proposal distribution ρ ( z ) , and weight each sample by wk = γ ( zk ) ρ ( zk ) , leading to the estimators Z ≈ 1 K K∑ k=1 wk = : Ẑ , Eπ ( z ) [ f ( z ) ] ≈ K∑ k=1 w̄kf ( zk ) = : Î , ( 4 ) where w̄k = wk/ ( KẐ ) is the normalized weight . The estimator Ẑ is often used to estimate marginal distributions , for example p ( x ) , with γ ( z ) being the joint distribution p ( z , x ) . It is unbiased and its variance decreases linearly with K. The estimator Î is often used to estimate the posterior expectation of gradients , for example the “ wake-φ ” gradient of RWS ( Bornschein & Bengio , 2015 ; Le et al. , 2019 ) , Epθ ( z|x ) [ −∇φ log qφ ( z|x ) ] with γ ( z ) = pθ ( z , x ) , π ( z ) = pθ ( z|x ) and f ( z ) = −∇φ log qφ ( z|x ) . This estimator is asymptotically unbiased and its asymptotic variance decreases linearly with K ( Owen , 2013 , Eq . 9.8 ) which means increasing K improves the estimator . 3 HYBRID MEMOISED WAKE-SLEEP . We propose hybrid memoised wake-sleep ( HMWS ) which extends memoised wake-sleep ( MWS ) to address the issue of memoization of continuous latent variables . In HMWS , we learn a generative model pθ ( zd , zc , x ) of hybrid discrete ( zd ) and continuous ( zc ) latent variables , and a recognition model qφ ( zd , zc|x ) which factorizes into a discrete recognition model qφ ( zd|x ) and a continuous recognition model qφ ( zc|zd , x ) . Like MWS , HMWS maintains a memory ofM discrete latent variables { zmd } Mm=1 per data point x which is updated in the wake phase of every learning iteration . In the sleep : replay phase , we use the memoized discrete latents to train both the generative model and the recognition model . In the sleep : fantasy phase , we optionally train the recognition model on data generated from the generative model as well . We summarize these learning phases in Fig . 3 , the full algorithm in Alg . 1 , and describe each learning phase in detail below . For notational clarity , we present the algorithm for a single training data point x .

Summary: The paper extends the previous memoized wake-sleep method to train complex generative models with discrete/structural and continuous latent variables. The method incorporates an assumption on the conditional dependence between the discrete and continuous latent variables and essentially employs a clever importance sampling procedure to capture the dependence. The experiments show applications to spatial and temporal compositional models where the proposed method outperforms VIMCO and reweighted wake-sleep.