HYBRID MEMOISED WAKE-SLEEP: APPROXIMATE IN-
Year: 2021

Abstract:
Modeling complex phenomena typically involves the use of both discrete and continuous variables. Such a setting applies across a wide range of problems, from identifying trends in time-series data to performing effective compositional scene understanding in images. Here, we propose Hybrid Memoised Wake-Sleep (HMWS), an algorithm for effective inference in such hybrid discrete-continuous models. Prior approaches to learning suffer as they need to perform repeated expensive inner-loop discrete inference. We build on a recent approach, Memoised Wake-Sleep (MWS), which alleviates part of the problem by memoising discrete variables, and extend it to allow for a principled and effective way to handle continuous variables by learning a separate recognition model used for importancesampling based approximate inference and marginalization. We evaluate HMWS in the GP-kernel learning and 3D scene understanding domains, and show that it outperforms current state-of-the-art inference methods.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND
text: 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).

heading: 2.1 MEMOISED WAKE-SLEEP
text: 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.

heading: 2.2 IMPORTANCE SAMPLING BASED APPROXIMATE INFERENCE AND MARGINALIZATION
text: 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.

heading: 3 HYBRID MEMOISED WAKE-SLEEP
text: 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.

heading: 3.1 WAKE
text: Given a set of M memory elements {zmd }Mm=1, we define the memory-induced variational posterior qMEM(zd|x) = ∑M m=1 ωmδzmd (zd) similarly as in MWS (eq. (3)). If we knew how to evaluate pθ(z m d , x), the wake-phase for updating the memory of HMWS would be identical to MWS. Here, we use an IS-based estimator of this quantity based on K samples from the continuous recognition model
p̂θ(z m d , x) =
1
K K∑ k=1 wmk, wmk = pθ(z m d , z mk c , x) qφ(zmkc |zmd , x) , zmkc ∼ qφ(zc|zmd , x). (5)
The IS weights {{wmk}Kk=1}Mm=1 are used for marginalizing zc here but they will later be reused for estimating expected gradients by approximating the continuous posterior pθ(zc|zmd , x). Given the estimation of pθ(zmd , x), the weights {ωm}Mm=1 of the memory-induced posterior qMEM(zd|x) are
ωm = p̂θ(z m d , x)∑M
i=1 p̂θ(z i d, x)
= 1 K ∑K k=1 wmk∑M
i=1 1 K ∑K j=1 wij =
∑K k=1 wmk∑M
i=1 ∑K j=1 wij . (6)
The resulting wake-phase memory update is identical to that of MWS (Sec. 2.1), except that the evaluation of pθ(zmd , x) and ωm is done by (5) and (6) instead. For large K, these estimators are more accurate, and if exact, guaranteed to improve ELBO(x, pθ(zd, x), qMEM(zd|x)) like in MWS.

heading: 3.2 SLEEP: REPLAY
text: In this phase, we use the memoized discrete latent variables {zmd }Mm=1 (“replay from the memory”) to learn the generative model pθ(zd, zc, x), and the discrete and continuous recognition models qφ(zd|x) and qφ(zc|zd, x). We now derive the gradient estimators used for learning each of these distributions in turn and summarize how they are used within HMWS in the “Sleep: Replay” part of Alg. 1.
Algorithm 1 Hybrid Memoised Wake-Sleep (a single learning iteration) Input: Generative model pθ(zd, zc, x), recognition model qφ(zd|x)qφ(zc|zd, x), data point x, memory elements {zmd }Mm=1, replay factor λ ∈ [0, 1], # of importance samples K, # of proposals N . Output: Gradient estimators for updating θ and φ: gθ and gφ; updated memory elements {zmd }Mm=1. 1: Propose {z′nd}Nn=1 ∼ qφ(zd|x) . Wake 2: Select {zid}Li=1 ← unique({zmd }Mm=1 ∪ {z′ n d}Nn=1) where L ≤M +N
3: Sample {zikc }Kk=1 ∼ qφ(zc|zid, x) for i = 1, . . . , L 4: Compute {{wik}Kk=1}Li=1, {p̂θ(zid, x)}Li=1 (eq. (5)), and {ωi}Li=1 (eq. (6)) 5: Select {zmd , {zmkc }Kk=1, {wmk}Kk=1, ωm}Mm=1 ← best M values from {zid}Li=1, {{zikc }Kk=1}Li=1, {{wik}Kk=1}Li=1 and {ωi}Li=1 according to {p̂θ(zid, x)}Li=1 6: Compute generative model gradient gθ ← gθ({wmk, zmkc }M,Km=1,k=1) (eq. (7)) . Sleep: Replay 7: Compute gφd,REPLAY ← g φ d,REPLAY({z m d , ωm}Mm=1) (eq. (9)) 8: Compute gφc,REPLAY ← gφc,REPLAY({wmk, zmkc }M,Km=1,k=1) (eq. (10)) 9: Compute gφFANTASY ← − 1K ∑K k=1∇φ log qφ(zk|xk) from zk, xk ∼ pθ(z, x) . Sleep: Fantasy
10: Aggregate recognition model gradient gφ ← λ(gφd,REPLAY + g φ c,REPLAY) + (1− λ)gφFANTASY 11: Return: gθ , gφ, {zmd }Mm=1.
Learning the generative model. We want to minimize the negative evidence − log pθ(x). We express its gradient using Fisher’s identity ∇θ log pθ(x) = Epθ(z|x) [∇θ log pθ(z, x)] which we estimate using a combination of the memory-based approximation of pθ(zd|x) and the IS-based approximation of pθ(zc|x). The resulting estimator reuses the samples zmkc and weights wmk from the wake phase (see Appendix A for derivation):
−∇θ log pθ(x) ≈ − M∑ m=1 K∑ k=1 vmk∇θ log pθ(zmd , zmkc , x) =: gθ({wmk, zmkc } M,K m=1,k=1), (7)
where vmk = wmk∑M
i=1 ∑K j=1 wij . (8)
Learning the discrete recognition model. We want to minimize KL(pθ(zd|x)||qφ(zd|x)) whose gradient can be estimated using a memory-based approximation of pθ(zd|x) (eq. (3))
∇φKL(pθ(zd|x)||qφ(zd|x)) = Epθ(zd|x) [−∇φ log qφ(zd|x)]
≈ − M∑ m=1 ωm∇φ log qφ(zmd |x) =: g φ d,REPLAY({z m d , ωm}Mm=1), (9)
where memory elements {zmd }Mm=1 and weights {ωm}Mm=1 are reused from the wake phase (6).
Learning the continuous recognition model. We want to minimize the average of KL(pθ(zc|zmd , x)||qφ(zc|zmd , x)) over the memory elements zmd . The gradient of this KL is estimated by an IS-based estimation of pθ(zc|zmd , x)
1
M M∑ m=1 ∇φKL(pθ(zc|zmd , x)||qφ(zc|zmd , x)) = 1 M M∑ m=1 Epθ(zc|zmd ,x) [−∇φ log qφ(zc|z m d , x)]
≈ − 1 M M∑ m=1 K∑ k=1 w̄mk∇φ log qφ(zmkc |zmd , x) =: g φ c,REPLAY({wmk, zmkc } M,K m=1,k=1), (10)
where w̄mk = wmk/ ∑K i=1 wmi. Both weights and samples z mk c are reused from the wake phase (5).

heading: 3.3 SLEEP: FANTASY
text: This phase is identical to “sleep” in the original wake-sleep algorithm (Hinton et al., 1995; Dayan et al., 1995) which minimizes the Monte Carlo estimation of Epθ(zd,zc,x)[−∇φ log qφ(zd, zc|x)], equivalent to training the recognition model on samples from the generative model.
Together, the gradient estimators in the Wake and the Sleep: Replay and the Sleep: Fantasy phases are used to learn the generative model, the recognition model and the memory. We additionally introduce a model-dependent replay factor λ ∈ [0, 1] which modulates the relative importance of the Sleep: Replay versus the Sleep: Fantasy gradients for the recognition model. This is similar to modulating between the “wake-” and “sleep-” updates of the recognition model in RWS (Bornschein & Bengio, 2015). We provide a high level overview of HMWS in Fig. 3 and its full description in Alg. 1.
Factorization of the recognition model. The memory-induced variational posterior qMEM(zd|x) is a weighted sum of delta masses on {zmd }Mm=1 which cannot model the conditional posterior pθ(zd|x, zc) as it would require a continuous index space zc. We therefore use qMEM(zd|x) to approximate pθ(zd|x) and qφ(zc|zd, x) to approximate pθ(zc|zd, x), which allows capturing all dependencies in the posterior pθ(zc, zd|x). Relaxing the restriction on the factorization of qφ(zc, zd|x) is possible, but, we won’t be able to model any dependencies of discrete latent variables on preceding continuous latent variables.

heading: 4 EXPERIMENTS
text: In our experiments, we compare HMWS on hybrid generative models against state-of-the-art wakesleep-based RWS and variational inference based VIMCO. To ensure a fair comparison, we match the number of likelihood evaluations – typically the most time-consuming part of training – across
all models: O(K(N + M)) for HMWS, and O(S) for RWS and VIMCO where S is the number of particles for RWS and VIMCO. Note that in practice, HMWS is faster and more memory-efficient than the other algorithms as the total number of likelihood evaluations is typically smaller than the upper bound of K(N + M) (see App. C). For training, we use the Adam optimizer with default hyperparameters. We judge the generative-model quality using an Stest = 100 sample importance weighted autoencoder (IWAE) estimation of the log marginal likelihood log pθ(x). We also tried the variational-inference based REINFORCE but found that it underperforms VIMCO.
We focus on generative models in which (i) there are interpretable, learnable symbols, (ii) the discrete latent variables represent the composition of these symbols and (iii) the continuous latent variables represent the non-structural quantitites. While our algorithm is suitable for any hybrid generative model, we believe that this particular class of neuro-symbolic models is the most naturally interpretable, which opens up ways to connect to language and symbolic planning.

heading: 4.1 STRUCTURED TIME SERIES
text: We first apply HMWS to the task of finding explainable models for time-series data. We draw inspiration from (Duvenaud et al., 2013), who frame this problem as GP kernel learning (Rasmussen & Williams, 2006). They describe a grammar for building kernels compositionally, and demonstrate that inference in this grammar can produce highly interpretable and generalisable descriptions of the structure in a time series. We follow a similar approach, but depart by learning a set of GP kernels jointly for each in time series in a dataset, rather than individually.
For our model, we define the following simple grammar over kernels
k → k + k | k × k | WN | SE | PER | C, where (11)
• WN is the White Noise kernel, WNσ2(x1, x2) = σ2Ix1=x2 ,
• SE is the Squared Exponential kernel, SEσ2,l2(x1, x2) = σ2 exp(−(x1 − x2)2/2l2), • PER is a Periodic kernel, PERσ2,p,l2(x1, x2) = σ2 exp(−2 sin2(π|x1 − x2|/p)/l2), • C is a Constant, Cσ2(x1, x2) = σ2.
We wish to learn a prior distribution over both the symbolic structure of a kernel and its continuous variables (σ2, l, etc.). To represent the structure of the kernel as zd, we use a symbolic kernel ‘expression’: a string over the characters {(, ),+,×,WN, SE, PER1, . . . , PER4, C}. We provide multiple character types PERi for periodic kernels, corresponding to a learnable coarse-graining of period (short to long). We then define an LSTM prior pθ(zd) over these kernel expressions, alongside a conditional LSTM prior pθ(zc|zd) over continuous kernel parameters. The likelihood is the marginal GP likelihood—a multivariate Gaussian whose covariance matrix is constructed by evaluating the composite kernel on a fixed set of points. Finally, the recognition model qφ(zd, zc|x) mirrors the architecture of the prior except that all the LSTMs additionally take as input an LSTM embedding of the observed signal x. The architecture is summarized in Fig. 5 and described in full in App. B.1.
We construct a synthetic dataset of 100 timeseries of fixed length of 128 drawn from a probabilistic context free grammar which is constructed by assigning production probabilities to our kernel grammar in (11). In Fig. 4 (left), we show that HMWS learns faster than RWS in this domain in terms of the log evidence log pθ(x). We also trained with VIMCO but due to the high variance of gradients, they were unable to learn the model well.
Examples of latent programs discovered by our model are displayed in Fig. 6. For each signal, we infer the latent kernel expression zd by taking the highest scoring memory element zmd according to the memory-based posterior qMEM(zd|, x) and sample the corresponding kernel parameters from the continuous recognition model qφ(zc|zd, x). We show the composite kernel as well as the GP posterior predictive distribution. These programs describe meaningful compositional structure in the time series data, and can also be used to make highly plausible extrapolations.

heading: 4.2 COMPOSITIONAL SCENE UNDERSTANDING
text: Next, we investigate the ability of HMWS to parse images of simple scenes in terms of stacks of toy blocks. Here, towers can be built from a large consortium of blocks, where blocks are drawn from a fixed set of block types. We aim to learn the parameters of the different block types – namely, their size and optionally their color – and jointly infer a scene parse describing which kinds of blocks live
where in each world. Such a symbolic representation of scenes increases interpretability, and has connections to language and symbolic planning.
Our generative and recognition models are illustrated in Fig. 7 (for the full description of the model, see App. B.2). We divide the ground plane into an N ×N (N = 2) grid of square cells and initialize a set of P = 5 learnable block types, which we refer to as the base set of “primitives” with which the model builds scenes. We parameterize each block type by its size and color. Within each cell, we sample the number of blocks in a tower uniformly, from zero to a maximum number of blocks Bmax = 3. For each position in the tower, we sample an integer ∈ [0, Bmax], which we use to index into our set of learnable primitives, and for each such primitive, we also sample its continuous position. The (raw) position of the first block in a tower is sampled from a standard Gaussian and is constrained to lie in a subset of the corresponding cell, using an affine transformation over a sigmoid-transformed value. A similar process is used to sample the positions of subsequent blocks, but now, new blocks are constrained to lie in a subset of the top surface of the blocks below.
Given the absolute spatial locations and identities of all blocks, we render a scene using the PyTorch3D differentiable renderer (Ravi et al., 2020), which permits taking gradients of the generative model probability pθ(zd, zc, x) with respect to the learnable block types θ. All training scenes are rendered from a fixed camera position (a front view); camera position is not explicitly fed into the model. We use an independent Gaussian likelihood with a fixed standard deviation factorized over all pixels – similar to using an L2 loss. The discrete latent variables, zd, comprise both the number of blocks and block type indices, while the continuous latent variables, zc, represent the raw block locations.
The recognition model qφ(zd, zc|x) follows a similar structure to that of the prior p(zd, zc), as shown inside the gray box of Fig. 7. However, the logits of the categorical distributions for qφ(zd|x), as well as the mean and (log) standard deviation of the Gaussian distributions for qφ(zc|zd, x), are obtained by mapping a convolutional neural network (CNN)-based embedding of the image x through small neural networks NN1–NN4; in our case, these are linear layers.
We train two models: one which infers scene parses from colored blocks and another which reasons over unicolor blocks. In both settings, the model must perform a non-trivial task of inferring a scene parse from an exponential space of PBmax·N 2
possible scene parses. Moreover, scenes are replete with uncertainty: towers in the front of the grid often occlude those in the back. Likewise, in the second setting, there is additional uncertainty arising from blocks having the same color. For each model, we generate and train on a dataset of 10k scenes. Both models are trained using HMWS with K = 5,M = 5, N = 5, and using comparison algorithms with S = K(N +M) = 50.
In Fig. 4 (middle and right), we show that in this domain, HMWS learns faster and better models than RWS ad VIMCO, scored according to the log evidence log pθ(x). We directly compare the wall-clock time taken by HMWS compared to RWS in Table 1, highlighting the comparative efficiency of our algorithm. Further, we uncovered two local minima in our domain; we find that a higher portion of HMWS find the superior minima compared to RWS. In Fig. 8, we show posterior samples from our model. Samples are obtained by taking the highest probability scene parses based on the memory-based posterior approximation qMEM(zd|x). These samples illustrate that HMWS-trained models can capture interesting uncertainties over occluded block towers and recover the underlying building blocks of the scenes. For instance, in the colored block domain, we see that the model properly discovers and builds with red, blue, and green cubes of similar sizes to the data.

heading: 5 RELATED WORK
text: Our work builds on wake-sleep algorithms (Hinton et al., 1995) and other approaches that jointly train generative/recognition models, particularly modern versions such variational autoencoders (Kingma & Welling, 2014) and reweighted wake-sleep (Bornschein & Bengio, 2015). While variational autoencoders (VAEs) have been employed for object-based scene understanding (e.g. Eslami et al. (2016); Greff et al. (2019)), they attempt to circumvent the issues raised by gradient estimation with discrete latent variables, either by using continuous relaxation (Maddison et al., 2016) or limited use of control variates (Mnih & Rezende, 2016). The wake-sleep family of algorithms avoids the need for such modification and is better suited to models that involve stochastic branching (see Le et al. (2019) for a discussion), and is thus our primary focus here.
Our most closely related work is Memoised Wake-sleep, which we build directly on top of and extend to handle both discrete and continuous latents. A contemporary of Memoised Wake-sleep, DreamCoder (Ellis et al., 2020), is a closely related program synthesis algorithm following a similar wake-sleep architecture. Like MWS, it lacks principled handling of latent continuous parameters; performing inner-loop gradient descent and heuristically penalizing parameters via the Bayesian Information Criterion (MacKay, 2002). Other instances of wake-sleep models include those that perform likelihood-free inference (Brehmer et al., 2020), or develop more complex training schemes for amortization (Wenliang et al., 2020).
Broadly, our goal of inferring compositionally structured, mixed discrete-continuous objects has strongest parallels in the program synthesis literature. One family of approaches (e.g. HOUDINI (Valkov et al., 2018) and NEAR (Shah et al., 2020)), perform an outer-loop search over discrete structures and an inner-loop optimization over continuous parameters. Others jointly reason over continuous and discrete parameters via exact symbolic methods (e.g. Evans et al. (2021)) or by relaxing the discrete space to allow gradient-guided optimization to run on the whole problem (e.g. DeepProbLog (Manhaeve et al., 2018)). None of these however, learn-to-learn by amortizing the cost of inference, with recent attempts needing quantized continuous parameters (Ganin et al., 2021). What our work contributes to this space is a generic and principled way of applying amortized inference to hybrid discrete-continuous problems. This gets the speed of a neural recognition model–which is valuable for program synthesis–but with Bayesian handling of uncertainty and continuous parameters, rather than relying on heuristic quantization or expensive inner loops.

heading: 6 DISCUSSION
text: Inference in hybrid generative models is important for building interpretable models that generalize. However, such a task is difficult due to the need to perform inference in large, discrete spaces. Unlike deep generative models in which the generative model is less constrained, learning in symbolic models is prone to getting stuck in local optima. While compared to existing algorithms, HMWS improves learning and inference in these models, it does not fully solve the problem. In particular, our algorithm struggles with models in which the continuous latent variables require non-trivial inference, as the quality of continuous inference is directly linked with the quality of the gradient estimators in HMWS. This challenge can potentially be addressed with better gradient estimators. We additionally plan to extend our algorithm to more complex, realistic neuro-symbolic models; for instance, those with more non-cuboidal primitive topologies.

heading: A DERIVATION OF THE GRADIENT ESTIMATOR FOR LEARNING THE GENERATIVE MODEL
text: Here, we derive the gradient estimator for learning the generative model given in (7).
We assume that the memory-induced variational posterior approximates the discrete posterior
pθ(zd|x) ≈ qMEM(zd|x) (12)
= M∑ m=1 ωmδzmd (zd) (13)
and that the continuous posterior expectation Epθ(zc|zmd ,x)[f(zc)] for each m is approximated using a set of weighted samples (zmkc , w̄mk) K k=1
Epθ(zc|zmd ,x)[f(zc)] ≈ K∑ k=1 w̄mkf(z mk c ), (14)
where w̄mk = wmk∑K i=1 wmi .
First, we make use of the Fisher’s identity
∇θ log pθ(x) = Epθ(z|x) [∇θ log pθ(x)] (Integrand is not a function of z) (15) = Epθ(z|x) [∇θ log pθ(x) +∇θ log pθ(z|x)] (Second term vanishes) (16) = Epθ(z|x) [∇θ log pθ(z, x)] (Product rule of probability) (17)
and we continue by using the approximations in (12)–(14) = Epθ(zc|zd,x) [ Epθ(zd|x) [∇θ log pθ(z, x)] ] (Factorize the posterior) (18)
≈ Epθ(zc|zd,x) [ EqMEM(zd|x) [∇θ log pθ(z, x)] ] (Use (12)) (19)
= Epθ(zc|zmd ,x) [ M∑ m=1 ωm∇θ log pθ(zmd , zc, x) ] (Use (13)) (20)
≈ K∑ k=1 w̄mk ( M∑ m=1 ωm∇θ log pθ(zmd , zmkc , x) ) (Use (14)) (21)
= M∑ m=1 K∑ k=1 vmk∇θ log pθ(zmd , zmkc , x), (Combine sums) (22)
where
vmk = w̄mkωm (23)
= wmk∑K i=1 wmi
· ∑K i=1 wmi∑M
i=1 ∑K j=1 wij
(24)
= wmk∑M
i=1 ∑K j=1 wij . (25)

heading: B MODEL ARCHITECTURE AND PARAMETER SETTINGS
text: B.1 STRUCTURED TIME SERIES
Fixed parameters
• Vocabulary V = {WN, SE, PER1, . . . , PER4, C,×,+, (, )} • Vocabulary size vocabulary_size = |V| = 11 • Kernel parametersK = {σ2WN, (σ2SE, `2SE), (σ2PER1 , pPER1 , ` 2 PER1 ), . . . , (σ2PER4 , pPER4 , ` 2 PER4 ), σ2C}
• Kernel parameters size kernel_params_size = |K| = 16 • Hidden size hidden_size = 128 • Observation embedding size obs_embedding_size = 128
Generative model pθ(zd, zc, x)
• Kernel expression LSTM (p) (input size = vocabulary_size, hidden size = hidden_size)
– This module defines a distribution over the kernel expression pθ(zd). – At each time step t, the input is a one-hot encoding of the previous symbol in V (or a
zero vector for the first time step) – At each time step t, extract logit probabilities for each element in V or the end-of-
sequence symbol EOS from the hidden state using a “Kernel expression extractor (p)” (Linear(hidden_size, vocabulary_size + 1)). This defines the conditional distribution pθ(ztd|z 1:t−1 d ).
– The full prior over kernel expression is an autoregressive distribution ∏ t pθ(z t d|z 1:t−1 d )
whose length is determined by EOS. • Kernel expression LSTM embedder (p)
(same as kernel expression LSTM (p)) – The module summarizes the kernel expression zd into an embedding vector e ∈
Rhidden_size – We re-use the kernel expression LSTM (p) above and use the last hidden LSTM state
to be the summary embedding vector. • Kernel parameters LSTM (p)
(input size = kernel_params_size + hidden_size, hidden size = hidden_size) – This module define a distribution over kernel parameters pθ(zc|zd). – At each timestep t, the input is a concatenation of the previous kernel parameters zt−1c ∈ Rkernel_params_size (or zero vector for t = 1) and the embedding vector e ∈ Rhidden_size.
– At each timestep t, extract mean and standard deviations for each kernel parameter in K using a “Kernel parameters extractor (p)” Linear(hidden_size, 2 * kernel_params_size). This defines the conditional distribution pθ(z t c|z1:t−1c , zd).
Recognition model qφ(zd, zc|x)
• Signal LSTM embedder (input size = 1, hidden size = obs_embedding_dim = hidden_size)
– This module summarizes the signal x into an embedding vector ex ∈ Robs_embedding_dim.
– The embedding vector is taken to be the last hidden state of an LSTM where x is fed as the input sequence.
• Kernel expression LSTM (q) (input size = obs_embedding_dim + vocabulary_size, hidden size = hidden_size)
– This module defines a distribution over kernel expression qφ(zd|x). – At each time step t, the input is a concatentation of the one-hot encoding of the previous
symbol in V (or a zero vector for the first time step) and ex. – At each time step t, extract logit probabilities for each element in V or the end-of-
sequence symbol EOS from the hidden state using a “Kernel expression extractor (q)” (Linear(hidden_size, vocabulary_size + 1)). This defines the conditional distribution qφ(ztd|z 1:t−1 d , x).
– The full distribution over kernel expression is an autoregressive distribution∏ t qφ(z t d|z 1:t−1 d , x) whose length is determined by EOS.
• Kernel expression LSTM embedder (q) (input size = vocabulary_size, hidden size = hidden_size)
– This module summarizes the kernel expression zd into an embedding vector ezd ∈ Rhidden_size.
– The embedding vector ezd is taken to be the last hidden state of an LSTM where zd is fed as the input sequence.
• Kernel parameters LSTM (q) (input size = obs_embedding_dim + kernel_params_size + hidden_size, hidden size = hidden_size)
– This module defines a distribution over kernel parameters qφ(zc|zd, x). – At each timestep t, the input is a concatenation of the previous kernel parameters zt−1c ∈ Rkernel_params_size (or zero vector for t = 1), the embedding vector ezd ∈ Rhidden_size and the signal embedding vector ex ∈ Robs_embedding_size. – At each timestep t, extract mean and standard deviations for each kernel parameter in K using a “Kernel parameters extractor (q)” Linear(hidden_size, 2 * kernel_params_size). This defines the conditional distribution qφ(z t c|z1:t−1c , zd, x).
B.2 COMPOSITIONAL SCENE UNDERSTANDING
Fixed parameters
• Number of allowed primitives to learn P = 5 • Maximum number of blocks per cell Bmax = 3 • Number of cells in the x-plane = number of cells in the z-plane = N = 2 • Image resolution I = [3, 128, 128] • Observation embedding size obs_embedding_size = 676
Recognition model qφ(zd, zc|x). Components are detailed to match those shown in Fig. 7.
• CNN-based embedding of the image
– Conv2d(in_channels=4, out_channels=64, kernel_size=3, padding=2)
– ReLU – MaxPool(kernel_size=3, stride=2) – Conv2d(64, 128, 3, 1) – ReLU – MaxPool(3, 2) – Conv2d(128, 128, 3, 0) – ReLU-Conv2d(128, 4, 3, 0) – ReLU – MaxPool(2, 2) – Flatten
Note, the dimensionality of the output is obs_embedding_size (in this case, = 676)
• Neural network for the distribution over blocks (NN1)
Linear(obs_embedding_size, N2 * (1 + Bmax))
• Neural network for the distribution over primitives (NN2)
Linear(obs_embedding_size, N2 * Bmax * P)
• Neural network that outputs the mean for the raw primitive locations (NN3)
Linear(obs_embedding_size, N2 * Bmax)
• Neural network that outputs the standard deviation for the raw primitive locations (NN4)
Linear(obs_embedding_size, N2 * Bmax)

heading: C COMPARATIVE EFFICIENCY: WALL-CLOCK TIME AND GPU MEMORY CONSUMPTION
text: We comprehensively investigate the comparative wall-clock time and GPU memory consumption of HMWS and RWS by comparing the algorithms over a range of matched number of likelihood evaluations. Here, we focus on the scene understanding domain, specifically the model wherein scene parses are inferred with color. We find that HMWS consistently results in more time- and memory-efficient inference, as seen in Tables 1 and 2. The reason for this is that while we match HMWS’s K(N +M) with RWS’s S likelihood evaluations,

heading: D ABLATION EXPERIMENTS
text: D.1 “INTERPOLATING” BETWEEN HMWS AND RWS
We “interpolate” between RWS and HMWS to better understand which components of HMWS are responsible for the performance increase. In particular, we train two more variants of the algorithm,
• “HMWS-” which trains all the components the same way as HMWS except the continuous recognition model qφ(zc|zd, x) which is trained using the RWS loss. This method has a memory.
• “RWS+” which trains all components the same way as RWS except the continuous recognition model qφ(zc|zd, x) which is trained using the HMWS loss in (10).
The learning curves for the time series model in Figure 9 indicate that the memory is primarily responsible for the performance increase since HMWS- almost identically matches the performance of HMWS. Including the loss in (10) in RWS training helps learning but alone doesn’t achieve HMWS’s performance.
D.2 SENSITIVITY TO THE DIFFICULTY OF CONTINUOUS INFERENCE
Since both the marginalization and inference of the continuous latent variable zc is approximated using importance sampling, we expect the performance of HMWS to suffer as the inference task gets more difficult. We empirically confirm this and that RWS and VIMCO suffer equally if not worse (Fig. 10) by arbitrarily adjusting the difficulty of continuous inference by changing the range of the period pPERi of periodic kernels (see App. B.1).
D.3 COMPARING HMWS AND RWS FOR DIFFERENT COMPUTE BUDGETS
We compare HMWS and RWS for a number of different compute budgets (see Fig 11). Here, we keep K ∗ (M + N) = S to match the number of likelihood evaluations of HMWS and RWS, and
keep K = M = N for simplicity. We run a sweep for HMWS with K = M = N in {2, 3, 4, 5, 6} and a corresponding sweep for RWS with S in {8, 18, 32, 50, 72}. We find that HMWS consistently overperforms RWS.

Review summaries:
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.

The paper proposes an extension of memoised wake-sleep which allows applications to hybrid discrete-continuous graphical models. The approach is directed towards Bayesian program synthesis.  Specifically, the memoised wake-sleep algorithm requires computing the joint probabiltiy of a discrete latent variable and an observation, which would require integrating out any continuous latents. This paper proposes to solve that problem using importance sampling.

Memoized wake-sleep is a past method that builds a variational approximate posterior over discrete latent variables by memorising previously drawn samples.  However, memorised wake-sleep can currently be used only on models that are purely discrete.  This paper extends memoized wake-sleep by providing a mechanism for using VI to integrate over continuous latent variables.