H YBRID M EMOISED W AKE -S LEEP : A PPROXIMATE I NFERENCE AT THE D ISCRETE -C ONTINUOUS I NTERFACE
Luke Hewitt
A BSTRACT
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.
I NTRODUCTION
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 speciﬁcation 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 efﬁcient inference?
[Equation]
We focus on a particular class of hybrid generative models p
[Equation]
[Equation] ) of observed data x with discrete latent variables z d , continuous latent variables z c and learnable parameters θ with a graphical model structure shown in Figure 1.
In particular, the discrete latent variables z
d represent an underlying structure present in the data, while the remaining continuous latent variables
z c represent non-structural quantities.
For example, in the context of compositional scene understanding, z d can represent a scene graph comprising 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 z c represents the continuous poses of these objects.
Here, 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 inefﬁcient 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 & , 2015) are not applicable to discrete
[Figure]
[Figure]
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., 2017) don’t naturally apply to stochastic control ﬂow 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 z
d
and learn a separate recognition model which is used for importance-sampling based approximate inference and marginalization of the continuous latent variables z c .
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.
B ACKGROUND
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
[Equation]
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)
[Equation]
[Equation]
[Figure]
where z ∼ q φ ( z | x ) .
However, the ﬁrst term is often high-variance which makes learning inefﬁcient.
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., 2017), 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).
[Figure]
- 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 conﬁgurations 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 { z md } Mm =1 for each data point x which induces a variational distribution
[Equation]
consisting of a weighted set of delta masses δ z md centered on the memory elements (see also
Saeedi et al.
This variational distribution is proven to improve the evidence lower bound
ELBO ( x, p θ ( z d , x ) , q MEM ( z d | x ))
(see Sec.
3 of (Hewitt et al., 2020)) by a memory-update phase comprising (i) the proposal of N new values { z (cid:48) nd } Nn =1 ∼ q φ ( z d | x ) , (ii) retaining the best M values from the union of the old memory elements and newly proposed values { z md } Mm =1 ∪ { z (cid:48) nd } Nn =1 scored by p θ ( z d , x ) , and (iii) setting the weights to ω m =
p θ ( z md , x ) / (cid:80) Mi
=1 p θ ( z id , x ) .
M WS , 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. I MPORTANCE 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 = (cid:82) γ ( z ) d z 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 { z k } Kk =1 from a proposal distribution ρ ( z ) , and weight each sample by w k = γ ( z k ) ρ
( z k ) , leading to the estimators
[Equation]
where ¯ w k = w k / ( K ˆ Z ) 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), E p θ ( 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.
H YBRID M EMOISED W AKE -S LEEP
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 θ ( z d , z c , x ) of hybrid discrete ( z d ) and continuous ( z c ) latent variables, and a recognition model q φ ( z d , z c | x ) which factorizes into a discrete recognition model q φ ( z d | x ) and a continuous recognition model q φ ( z c | z
d , x ) .
Like MWS , HMWS maintains a memory of M discrete latent variables { z md } 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 .
3. W AKE
Given a set of M memory elements { z md } Mm =1 , we deﬁne the memory-induced variational posterior q MEM ( z d | x ) = (cid:80) Mm =1 ω m δ z
md ( z d ) similarly as in MWS (eq.
If we knew how to evaluate p θ ( z md , 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
[Equation]
The IS weights {{ w mk } Kk =1 } Mm =1 are used for marginalizing z c here but they will later be reused for estimating expected gradients by approximating the continuous posterior p θ ( z c | z md , x ) .
Given the estimation of p θ ( z md , x ) , the weights { ω m } Mm =1 of the memory-induced posterior q MEM ( z d | x ) are
[Equation]
The resulting wake-phase memory update is identical to that of MWS (Sec. 2.1), except that the evaluation of p θ ( z md , 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 θ ( z d , x ) , q MEM ( z d | x )) like in MWS .
3. S LEEP : R EPLAY
In this phase, we use the memoized discrete latent variables { z md } Mm =1 (“replay from the memory”) to learn the generative model p θ ( z d , z c , x ) , and the discrete and continuous recognition models q φ ( z d | x ) and q φ ( z c | z d , 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.
Input: Generative model p θ ( z d , z c , x ) , recognition model q φ ( z d | x ) q φ ( z c | z
d , x ) , data point x , memory elements { z md } Mm =1 ,
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
θ
[Table]
[Table]
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 ) =
E p θ ( z | x ) [ ∇ θ log p θ ( z, x )] which we estimate using a combination of the memory-based approximation of p θ ( z d | x ) and the IS -based approximation of p θ ( z c | x ) .
The resulting estimator reuses the samples z mkc and weights w mk from the wake phase (see Appendix A for derivation):
[Equation]
where v mk =
[Equation]
Learning the discrete recognition model.
We want to minimize KL ( p θ ( z d | x ) ||
q φ ( z d | x )) whose gradient can be estimated using a memory-based approximation of p θ ( z d | x ) (eq. (3))
[Equation]
where memory elements { z md } 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 θ ( z c | z md , x ) ||
q φ ( z c | z md , x )) over the memory elements z md .
The gradient of this KL is estimated by an IS -based estimation of p θ ( z c |
z md , x )
[Equation]
where ¯ w mk = w mk / (cid:80) Ki =1 w mi .
Both weights and samples z mkc are reused from the wake phase (5).
3. S LEEP : F ANTASY
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 E p θ ( z d ,z c ,x ) [ −∇ φ log q φ ( z d , z c | 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 q MEM ( z d | x ) is a weighted sum of delta masses on { z md } Mm =1 which cannot model the conditional posterior p θ ( z d | x, z c ) as it would require a continuous index space z c .
We therefore use q MEM ( z d | x ) to approximate p θ ( z d | x ) and q φ ( z c | z d , x ) to approximate p θ ( z c | z d , x ) , which allows capturing all dependencies in the posterior p θ ( z c , z d | x ) .
Relaxing the restriction on the factorization of q φ ( z c , z d | x ) is possible, but, we won’t be able to model any dependencies of discrete latent variables on preceding continuous latent variables.
E XPERIMENTS
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
[Figure]
g p ( x )
HMWSRWSVIMCO 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 .
HMWS has a ﬁxed memory cost of O ( M ) per data point which can be prohibitive for large datasets.
However in practice, HMWS is faster and more memory-efﬁcient 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 S test = 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 [Figure]
4. S TRUCTURED TIME SERIES
[Figure]
We ﬁrst apply HMWS to the task of ﬁnding 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 [Figure] 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 deﬁne the following simple grammar over kernels
[Equation]
where
[Equation]
• WN is the White Noise kernel, WN σ 2 ( x 1 , x 2 ) = σ
[Equation]
[Equation]
[Figure]
• SE is the Squared Exponential kernel, SE σ 2 ,l 2 ( x 1 , x 2 ) =
σ 2 exp( − ( x 1 − x 2 ) 2 / 2 l 2 ) , • PER is a Periodic kernel, PER σ 2 ,p,l 2 ( x 1 , x 2 ) =
σ 2 exp( − 2 sin 2 ( π | x 1 − x 2 | /p ) /l 2 ) , • C is a Constant , C σ 2 ( x 1 , x 2 )
=
σ 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 z d , we use a symbolic kernel ‘expression’: a string over the characters { ( , ) , + , × , WN , SE , PER 1 , . .
, PER 4 , C } .
We provide multiple character types PER i for periodic kernels, corresponding to a learnable coarse-graining of period (short to long).
We then deﬁne an LSTM prior p θ ( z d ) over these kernel expressions, alongside a conditional LSTM prior p θ ( z c | z d ) 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 ﬁxed set of points.
Finally, the recognition model q φ ( z d , z c | x ) mirrors the architecture of the prior except that all the LSTM s 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 ﬁxed 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.
For each signal, we infer the latent kernel expression z d by taking the highest scoring memory element z md according to the memory-based posterior q MEM ( z d | , x ) and sample the corresponding kernel parameters from the continuous recognition model q φ ( z c | z d , 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 [Figure] and can also be used to make highly plausible extrapolations [Figure]
4. C OMPOSITIONAL SCENE UNDERSTANDING
[Figure]
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 ﬁxed 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
B max = 3 .
For each position in the tower, we sample an integer ∈ [0 , B max ] , 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 ﬁrst block in a tower is sampled from a standard Gaussian and is constrained to lie in a subset of the corresponding cell, using an afﬁne 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 θ ( z d , z c , x ) with respect to the learnable block types θ .
All training scenes are rendered from a ﬁxed camera position (a front view); camera position is not explicitly fed into the model.
We use an independent Gaussian likelihood with a ﬁxed standard deviation factorized over all pixels – similar to using an L 2 loss.
The discrete latent variables, z d , comprise both the number of blocks and block type indices, while the continuous latent variables, z c , represent the raw block locations.
The recognition model q φ ( z d , z c | x ) follows a similar structure to that of the prior p ( z d , z c ) , as shown inside the gray box of Fig.
However, the logits of the categorical distributions for q φ ( z d | x ) , as well as the mean and (log) standard deviation of the Gaussian distributions for q φ ( z c | z d , 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 P B max · 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 10 k 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 efﬁciency of our algorithm.
Further, we uncovered two local minima in our domain; we ﬁnd that a higher portion of HMWS ﬁnd 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 q MEM ( z d | 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.
R ELATED WORK
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 ( VAE s) have been employed for object-based scene understanding (e.g. Eslami et al.
(2016); Greff et al.
, they attempt to circumvent the issues raised by gradient estimation with discrete latent variables, either by using continuous relaxation (Maddison et al., 2017) or limited use of control variates (Mnih & Rezende, 2016).
The wake-sleep family of algorithms avoids the need for such modiﬁcation 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.
Observations
S a m p l e
Posterior Samples
S a m p l e S a m p l e
S a m p l e S a m p l e
S a m p l e
Observations
S a m p l e
Posterior Samples
S a m p l e S a m p l e
S a m p l e S a m p l e
S a m p l e
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., 2021), 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.
D ISCUSSION
Inference in hybrid generative models is important for building interpretable models that generalize.
However, such a task is difﬁcult 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.
A D ERIVATION OF THE GRADIENT ESTIMATOR FOR LEARNING THE GENERATIVE MODEL
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
[Equation]
) and that the continuous posterior expectation E p θ
[Equation]
z c | z md ,x ) [ f ( z c )] for each m is approximated using a set of weighted samples ( z
[Equation] [Equation]
14) where ¯ w mk =
[Equation]
First, we make use of the Fisher’s identity ∇
[Equation]
] (Integrand is not a function of z ) (15)
[Equation]
(Second term vanishes) (16)
[Equation]
(Product rule of probability) (17) and we continue by using the approximations in (12)–(14)
[Equation]
(Factorize the posterior) (18)
[Equation]
(Use (12)) (19)
[Equation]
(cid:35) (Use (13)
[Equation]
(20) ≈
[Equation]
(cid:33) (Use (14)) (21) =
[Equation]
m
[Equation]
(Combine sums) (22) where
[Equation]
[Equation]
[Equation]
B M ODEL A RCHITECTURE AND P ARAMETER S ETTINGS B. S TRUCTURED TIME SERIES Fixed
parameters • Vocabulary V = { WN , SE , PER 1 , . . .
, PER 4 , C , × , + , ( , ) } • Vocabulary size vocabulary_size = |V| =• Kernel parameters
K = { σ 2 WN , ( σ 2 SE , (cid:96) 2 SE ) , ( σ 2 PER 1 , p PER 1 , (cid:96) 2 PER 1 ) , . . .
, ( σ 2 PER 4 , p PER 4 , (cid:96) 2 PER 4 ) , σ 2 C }
• Kernel parameters size kernel_params_size = |K| =• Hidden size hidden_size =
• Observation
embedding size obs_embedding_size =
[Table] p [Table] [Equation]
x ) • Kernel expression LSTM (p) (input size =
vocabulary_size , hidden size =
hidden_size ) – This module deﬁnes a distribution over the kernel expression p θ ( z d ) .
– At each time step t , the input is a one-hot encoding of the previous symbol in V (or a zero vector for the ﬁrst time step) – At each time step t , extract logit probabilities for each element in V or the end-ofsequence symbol EOS from the hidden state using a “Kernel expression extractor (p)”
( Linear(hidden_size, vocabulary_size + 1) ).
This deﬁnes the conditional distribution p θ ( z td | z 1: t − 1 d ) .
– The full prior over kernel expression is an autoregressive distribution (cid:81) t p θ ( z td | 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 z d into an embedding vector e ∈ R hidden_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 deﬁne a distribution over kernel parameters p θ ( z c | z d ) .
– At each timestep t , the input is a concatenation of the previous kernel parameters z t − 1 c ∈ R kernel_params_size (or zero vector for t = 1 ) and the embedding vector e ∈ R hidden_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 deﬁnes the conditional distribution p θ ( z tc | z 1: t − 1 c , z d ) .
[Table]
• Signal LSTM embedder (input size = 1, hidden size = obs_embedding_dim =
hidden_size ) – This module summarizes the signal x into an embedding vector e x ∈ R obs_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 deﬁnes a distribution over kernel expression q φ ( z d | 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 ﬁrst time step) and e x .
– At each time step t , extract logit probabilities for each element in V or the end-ofsequence symbol EOS from the hidden state using a “Kernel expression extractor (q)” ( Linear(hidden_size, vocabulary_size + 1) ).
This deﬁnes the conditional distribution q φ ( z td | z 1: t − 1 d , x ) .
– The full distribution over kernel expression is an autoregressive distribution (cid:81) 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 z d into an embedding vector e z d ∈ R hidden_size .
– The embedding vector e z d is taken to be the last hidden state of an LSTM where z d 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 deﬁnes a distribution over kernel parameters q φ ( z c | z d , x ) .
– At each timestep t , the input is a concatenation of the previous kernel parameters z t − 1 c ∈ R kernel_params_size (or zero vector for t = 1 ), the embedding vector e z d ∈ R hidden_size and the signal embedding vector e x ∈ R obs_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 deﬁnes the conditional distribution q φ ( z tc | z 1: t − 1 c , z d , x ) .
B. C OMPOSITIONAL SCENE UNDERSTANDING Fixed parameters
• Number of allowed primitives to learn P =• Maximum number of blocks per cell B max =• Number of cells in the x -plane = number of cells in the z -plane = N =• Image resolution I =
[3 , 128 , 128] • Observation embedding size obs_embedding_size =
Recognition model q φ ( z d , z c | x ) .
Components are detailed to match those shown in Fig.
• 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, N 2 * (1 + B max )) • Neural network for the distribution over primitives (NN2) Linear(obs_embedding_size, N 2 * B max * P)
• Neural network that outputs the mean for the raw primitive locations (NN3) Linear(obs_embedding_size, N 2 * B max ) • Neural network that outputs the standard deviation for the raw primitive locations (NN4) Linear(obs_embedding_size, N 2 * B max )
C C OMPARATIVE E FFICIENCY : W ALL -C LOCK T IME AND GPU M EMORY C ONSUMPTION
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, speciﬁcally the model wherein scene parses are inferred with color.
We ﬁnd that HMWS consistently results in more time- and memory-efﬁcient 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,
• The actual number of likelihood evaluations of HMWS is K · L which is typically smaller than K ( N + M ) because we only take L ≤ N + M unique memory elements (step 2 of Algorithm 1).
• The loss terms in lines 7–9 of Algorithm 1 all contain only K · M < K · L ≤ K ( N + M ) log probability terms { log p θ ( z md , z mkc , x ) } M,Km =1 ,k =1 , { log q φ ( z mkc | z md , x ) } M,Km =1 ,k
=1 and M terms { log q φ ( z md | x ) } Mm =1 which means differentiating these loss terms is proportionately cheaper than for RWS which has K ( N + M ) log probability terms.
• We draw only N samples from q φ ( z d | x ) and K · L samples from q φ ( z c | z d , x ) in HMWS in comparison to K · ( N + M ) in RWS .
S or K ( M + N )
HMWS
RWS 67 min 82 min 136 min 196 min 186 min
[Table]
[Table]
D A BLATION EXPERIMENTS D. “I NTERPOLATING ” 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 φ ( z c | z
d , 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 φ ( z c | z d , 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.
[Figure]
D. S ENSITIVITY TO THE DIFFICULTY OF CONTINUOUS INFERENCE
Since both the marginalization and inference of the continuous latent variable z c is approximated using importance sampling, we expect the performance of HMWS to suffer as the inference task gets more difﬁcult.
We empirically conﬁrm this and that RWS and VIMCO suffer equally if not worse (Fig. 10) by arbitrarily adjusting the difﬁculty of continuous inference by changing the range of the period p
PER i of periodic kernels (see App. B.1).
[Figure]
D. C OMPARING 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 ﬁnd that HMWS consistently overperforms RWS .
[Figure]
E U SING VIMCO FOR HYBRID GENERATIVE MODELS
We initially treated the continuous latent variables the same way as the discrete variables.
In the reported results, we used reparameterized sampling of the continuous latent variables for training
VIMCO which decreases the variance of the estimator of the continuous recognition model gradients.
However, this didn’t seem to improve the performance signiﬁcantly.
Since applying VIMCO to hybrid discrete-continuous latent variable settings is not common in the literature, we include a sketch of the derivation of the gradient estimator that we use for completeness.
First, let Q φ ( z 1: K d ) =
(cid:81)
Kk =1 q φ ( z kd ) and Q φ ( z 1: K c |
z 1: K d ) =
(cid:81)
Kk =1 q φ ( z kc | z kd ) be the distributions of all the discrete and continuous randomness.
We also drop the dependence on data x to reduce notational clutter.
Thus, the gradient we’d like to estimate is
[Equation]
where f is the term inside of the IWAE objective f
[Equation]
Next
[Equation]
we apply reparameterized sampling of the continuous variables z 1: K c .
This involves choosing a simple, parameterless random variable (cid:15) ∼ s ( (cid:15) ) and a reparameterization function z c
= r ( (cid:15), z d , φ ) such that its output has exactly the same distribution sampling from the original recognition network z c ∼ q φ ( z c | z d ) .
Hence, we can rewrite the gradient we’d like to estimate as
[Equation]
[Equation]
where f takes in reparameterized samples of the continuous latent variables instead of z 1: K c .
Now
[Equation]
we follow a similar chain of reasoning to the one for deriving a score-function gradient estimator ∇
[Equation]
[Equation]
The VIMCO control variate is applied to the ﬁrst term ∇ φ log Q φ ( z 1: K d ) f ( ... ) .
Importantly, this term does not contain the density of the continuous variables Q φ ( z 1: K c
[Equation]
Hence, this term only provides gradients for learning the discrete recognition network unlike if we treated the continuous latent variables the same as the discrete latent variables (in which case a score-function gradient would be used for learning the continuous recognition network as well).
The gradient for learning the continuous recognition network instead comes from the second term ∇ φ f ( ... ) .