Hybrid Memoised Wake-Sleep: Approximate Inference at the Discrete-Continuous Interface
Authors: Tuan Anh Le
Katherine M. Collins
Luke Hewitt
Kevin Ellis
Siddharth N
Samuel Gershman
Joshua B. Tenenbaum
Venue: ICLR 2022 Poster (2022)
Keywords: wake-sleep
variational inference
neuro-symbolic generative models

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 importance-sampling 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.

Published as a conference paper at ICLR 2022
HYBRID MEMOISED WAKE-SLEEP. APPROXIMATE IN- FERENCE AT THE DISCRETE-CONTINUOUS INTERFACE
Tuan Anh Le
Katherine M. Collins'
Luke Hewitt'
Kevin Ellis?
N. Siddharth
Samuel Gershman
Joshua B. Tenenbaum'
ABSTRACT
Modeling complex phenomena typically involves the use of both discrete and continuous variables. Such setting applies across wide range of problems, from identifying trends in time-series data to performing effective compositional scene understanding 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 recent approach, Memoised Wake-Sleep (MWS), which alleviates part of the problem by memoising discrete variables, and extend it to allow for principled and effective way handle continuous variables by learning a separate recognition model used for importance- sampling based approximate inference and marginalization. We evaluate HMWS in the GP-kernel learning and 3D) scene understanding domains, and show that it outperforms current fate-of-thhearr inference methods.
INTRODUCTION
We naturally understand the world around us in terms of discrete symbols. When looking scene, we automatically parse what is where, and understand the relationships between objects in the scene. We understand that there book and a table, and that the book is on the table. Such symbolic representations often necessary for planning, communication and abstract reasoning. They allow the specification 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 re forced to describe such plans, among many other things we say, in discrete words. And further, abstract reasoning requires creating symbols by composing old ones which allows us to generalize to completely new situations. A "tower" structure made out of books still understood "tower" when built out of Jenga blocks. How do represent and learn models that support such symbolic reasoning while supporting efficient inference?
Zd
We focus on particular class of hybrid generative models (o že, x) of observed data x with discrete latent variables žd, continuous latent variables &e and learnable parameters f with a graphical model structure shown in Figure [] In particular, the discrete latent variables represent an underlying structure present the data, while the remaining continuous latent variables 2, represent non-structural quantities. For example, in the context compositional scene understanding, 2a can represent a scene graph comprising object identities and the relationships between them, like "a small green pyramid is on top yellow cube; a blue doughnut leans Figure I: Hybrid generative on the yellow eube; the large pyramid is next to the yellow cube" while model (= ze. ar). represents the continuous poses of these objects. Here, we assume that object identities are discrete, symbolic variables indexing into generative model parameters 0. The idea is to make these primitives learn to represent concrete objects like "yellow cube" "large green pyramid" from data in an unsupervised fashion.
Z
T
of learnable primitives, parameterized by subset of the
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 reparameteri- zation trick (Kingma & Welling 2014 Rezende Mohamed, 2015) are not applicable to discrete
MTT, Comell University, UUnivessity of Edinhurgh & The Alan Turing Institute, "Harvard University

---

Published as a conference paper at ICLR 2022
Sermel:
WN((0.99) (040.49, 0.50)
Sa(0.50, 0.51) (0.50 1.00, 0.50)
(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).
slock("Hlue"; Poo(" Pos(-112 1.09
Bleck("Green) "peeee"PPcs(..... Pcs(1.0. 1.0)
"ON" Block("Red", Pos(-LL. 1.0)) "ow' Block("slus", Poa(-1O.
"oN' (nd(Pos(O, 0)),
BLock("Red", ""eo"""""","o( Pos(-Lo,,, 0.00)
'ON' Block("Red". Pos(-L6,
1.0])
0.0))
"OW' Gad(Pos(O, 1))
'63' Block("Blue", P8
Block("Red" Poa(D.8, 0.0))
0.0))
"ow Gnd(Poa(), 9))
'0N' Snd(Pos(1, 1))
Block(" Pcs(-L11,00)
"oN' Und(Hos(L, 1))
(b) Compositional scene understanding over primitives, attributes, and spatial arrangements.
Figure 2: Examples highlighting the hybrid discrete-continuond nature of data. Colours indicate samples from discrete (violct) continuous (cyan) random variables respectively.
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 flow models (Le et al. 2019). Second, wake-sleep methods (Hinton et al.
1995 Dayan al. 1995) like reweighted wake-sleep (RWS) (Bornschein & Bengio 2015] et al. 2019) require inferring discrete latent variables at every learning iteration, without saving previously performed inferences. Memoised wake-sleep (MWS) (Hewitt al. 2020) addresses this issue, but is only applicable to purely discrete latent variable models.
We propose hybrid memoised wake-sleep (IMWW) method for learning and amortized inference in probabilistic generative models with hybrid discrete-coniinuuus latent variables. HMWS combines the strengths of MWS in memoising the discrete latent variables and RWS handling continuous latent variables. The core idea in HMWS is to memoise discrete latent variables z and learn separate recognition model which used for aepoortacc-sampllin based approximate inference and marginalization of the continuous latent variables
We empirically compare HMWS with late-of-the-ar 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 two domains: structured time series and compositional 3D scene understanding, respectively.
BACKGROUND
Our goal to learn the parameters of generative model p0(z, æ) of latent variables z and data æ, and parameters & of recognition model go (z(xr) which acts as an approximation to the posterior po(=)x)) This can achieved by maximizing the evidence lower bound
ELBBO (x. (2,, (=(=a)) lo
(1)
which maximizes the evidence Po(p) (æ) while minimizing the Kullback- (KL) divergence, thereby encouraging the recognition model to approximate the posterior.
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 [992 Schulman al. 2015)
ELBO (z, po go(2)))) ~ log (2
V log logg)(z)z) + Vo Po (o
(2)
(=(x)x)
g
2

---

Published as a conference paper at ICLR 2022
Memory
1-2 Fant forevery trainin
Wake
Pylo
Sleep: Replay
KLindz al (o)) 9, KLKL (l
1= Rg(Ceel 83
Generative model Pa
Sleep: Fantasy
Recognition model 4ytia (4
RO Top to from (:5") (eG) bassd on aylllig)))
V lo (2(2
Figure 3: Learning phases of Hybrid Memoised wake-sleep. The memory stores a set of discrete latents for each training data point z. The Wake phase updates the memory using samples from the recognition model. The Replay phase updates the generative model and the recognition model using the memory. The Sleep: Fantasy phase updates the recognition model using samples from the generative model.
Sleep:
where 2 ~ 4s (2).. 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 jet 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).
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 (=|xz) at every iteration. This repeated sampling can be wasteful, especially when only a few latent configurations explain the data well. Memoised wake- sleep (MWS) (Hewitt et al. 2020) extends wake-sleep by introducing memory- a set of M unique M
discrete latent variables [2d Jm=1 for each data point a which induces a variational distribution
M
ONEM(ZZ (x) =
Wm Wmddeg (za).
(3)
consisting a weighted set of delta masses 8pm centered on the memory elements (see also Saeedi et al. (2017)). This variational distribution is proven to improve the evidence lower bound ELBO( z, Po &d. x). (MEx((a(Iz) (see Sec. 3 of (Hewitt et al. 2020 by a memory-update phase comprising (i) the proposal of N new values {2/"T..
Jn= go (0(Za|zz) z). (ii) retaining the best M values from the union of the old memory elements and newly proposed values {2) IM
Jm=1 (222 J==1 scored (=)2222) 2).
by (o(a,, v)., and (iii) setting the weights to wm = po(" x)/ ΣM
M
Mws, however, only works on models with purely discrete latent variables. If we try to use the same approach for hybrid discrete-continune 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 (18) to perform approximate inference and marginalization. In general. given an unnormalized density (2), and its corresponding normaliz- ing constant Z rz)) dz and normalized density T2) = /Z, we want to estimate Z and the expectation of an arbitrary function f. [f)s)](())]. To do this, we sample K values {z ) from
2(2x) leading to the estimators proposal distribution p(z), and weight each sample by Wk p(a))
K
Z
Σ k=1
wk = 2.
Ex())f(2))
Σ @xf((zk =:
(4)
where Wx - ww/(k) 2) is the normalized weight. The estimator 2 is often used to estimate marginal
distributions, for example p(x), with 7(2) being the joint distribution p(2, x). It is unbiased and its variance decreases linearly with K. The estimator is often used to estimate the posterior expectation of gradients, for example the "wake-ø" gradient of RWS (Bornschein & Bengio 2015
Le et al. 2019). Epo(=]=)]-VV - & log go go (z|==] with >2) pe (=(z, x). ((2) = (=(2) and zz)
~v log go (z)x). This estimator is asymptotically unbiased and its asymptotic variance decreases linearly with K (Owen. 2013 Eg. 9.8) which means increasing K improves the estimator.
3

---

Published as a conference paper at ICLR 2022
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. HMWs, we learn a generative model pa (zd: Ze. 2) of hybrid discrete (za) and continuous (ze) latent variables, and recognition model go (2d, 2 zz)z) which factorizes into a discrete recognition model go (za(z) and a continuous recognition model go (=<]/=a; æ). Like MWS, HMWS maintains memory of M discrete latent variables Im=1 12 per data point 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 data generated from the generative model as well. We summarize these learning phases in Fig, [3] the full algorithm in Alg. [ and describe each learning phase in detail below. For notational clarity, we present the algorithm for a single training data point x.
3.1
WAKE
Given set of M memory elements (= 1m=1- define the memory-induced variational posterior 127
(Zal)) = Σm= Wm (za) similarly as in MWS (eq. (3). If we knew how to evaluate po(" .7), the wake-phase for updating the memory of HMWS would be identical to MWS. Here, we an Is-based estimator of this quantity based K samples from the continuous recognition model
K
Po(= 26 It)
p x) = K
Σ
Wmk::
Wmnk
x).
(5)
4o (2 1 xz)
The weights {[um 1K m== are used for marginalizing ze here but they will later be reused for estimating expected gradients by approximating the continuous posterior po (=el==" " z²) Given the estimation of (= x)), the weights {wm ] of the memory-induced posterior Tursa (zal)( are
pT x)
# Σ Wimk
E 1 Wmk
Win
(6)
En Pa x= Σ R 1 Way
ΣS Wag
The resulting wake-phase memory update is identical to that of MWS (Sec. 2.1p. except that the evaluation of pe (2z", 2) and wm is done by (5) and (6) instead. For large K, these estimators are more accurate, and if exact, guaranteed to improve EI,OO po (zd, x)), TMEM (za (z) like in MWs.
3.2 SLEEP: REPLAY
In this phase, we use the memoized discrete latent variables {>(] Im=1 ('replay from the memory") ") to learn the generative model po(ad, žey at), and the discrete and continuous recognition models go (zale (x) and go (==] ka, x). We now derive the gradient estimators used for learning each of these distributions turn and summarize how they are used within HMWS in the "Sleep: Replay" part of Alg. !
{am)M M
Algorithm 1 Hybrid Memoised Wake-Sleep (a single learning iteration)
Input: Generative model ps (zd, Ze, .7), recognition model (o(aa) (=)]==) =)), data point zr memory elements {2T M
Jm=1, replay factor À E [0, 1), # of importance samples K, # of proposals N.
Output: Gradient estimators updating and di go and ga: updated memory elements {=/ (2 mm=1-
1: Propose (=' ~ (=(((x))
>
Wake
2: Select [z/] 4- unique( (z," M.r.MM M U (= where I S SM+N
3: Sample {2(hh fh=1 qo(zelzd 2) for =1i===... 1,...,L L
4: Compute [[wx) 1K ] (Po #))= (eq. 55) and (wi) (eq. (6)
5: Select {2/7 (2
K " 1k=1, (win 1k=1 (m ) + best M values from (=))/.. H Me and (wi) according to (x) M.K
( Man
[[wa 18
6: Compute generative model gradient go + go (Wmn, 2 J Jm=ik (eq.
Sleep: Replay
7: Compute Ld gu t af (1=17 ((=/"""mlii (m) (oq.
8: Compute RE, + RE. (Wine 2ith M.K mm=1k=1) (eq. 10
9: Compute PENNTASY f- - ΣK Vo log (o (xx) from JN JA (o
Sleep: Fantasy
10: Aggregate recognition model gradient In + X(0 + gemuna)) ) + (1 A)g
11: Return: go. 9á. [2]7] ] mm=1
4

---

Published as a conference paper at ICLR 2022
Learning the generative model. We want to minimize the negative evidence log (a We express its gradient using Fisher's identity Ve log pe = (z|xz) (Vo log (e po x)] which we estimate using combination of the memory- approximation of pp((elar) (a)). The resulting estimator reuses the samples ze and weights Wok from the wake phase (see Appendix A for derivation):
approximation of po(2al and the Is-based
M
Ve logpo(x)
Σ
mmk Vo log po (2d 2mk 22) =: go (f 26 2i 2mnkk Im= 1m=1, (7)
k=1
Wonk
where Vmnk
M K
Σラ
Wij
Learning the discrete recognition model. We want to minimize KL(ps ((al=) (alla)llllaaaaa (zalx)) whose gradient can be estimated using memory-based approximation of po (=d(zr (eq.
( ppo(zal)) |- V log go (za[2)]
M
Wm V log ( =: PdRREPLAY (12 Wm 1. Iw=)).
(9)
M==1
( where memory elements {= I and weights {win
} mm= are reused from the wake phase (
Learning the continuous recognition model. We want to minimize the average of (L(po(2cl2") ;=) 1x)1lg0(2e1=?? ,))) over the memory elements 2Ø The gradient of this KL is esti- mated by an IS- 1s-based estimation po po(ell""
M
M
1
V VoKI(po(Zzlz) ( " =
Σ m=
Eoo [-V Vo log do (zcl=") x)]
M
M
m=1
M
K
Wonk V log (= (r),.))) x))
Fc mk M.K RERRE ([wmk
Jm=1,k=1)
(10)
M
where Tonk wmk// Ex Wmis Both weights and samples an are reused from the wake phase (5)
K
3.3 SLEEP: FANTASY
This phase is identical "sleep" the original wake-sleep algorithm (Hinton al.) 1995 Dayan
et al. 1995) which minimizes the Monte Carlo estimation of (po(sa.ay.2) 1- V log Qø 2dd 2e [)
equivalent 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 model-dependent replay factor À E [0, 1] which modulates the relative importance of the Sleep: Replay versus the Sleep: Fantasy gradients for the recognition model. This similar to modulating between the "wake-" and "sleep-" updates of the recognition model in RWS (Bornschei & Bengio 2013). We provide a high level overview of HMWS in Fig, and its full description in Alg.
Factorization of the recognition model. The memory-induced variational posterior QNMM (z)())) is weighted sum of delta masses on {> m.m... JM Jm=1 which cannot model the conditional posterior po(2alæ; 2e) as it would require a continuous index space 2g- We therefore use (MEM((2,) (z) to approximate po(zal (2alaz) and (= x) to approximate po p= æ), which allows capturing all dependencies in the posterior Po (2,, 2a 2z) Relaxing the restriction on the factorization of 9s Zes 2 2a(z)
possible, but, won't i be able model any dependencies of discrete latent variables preceding continuous latent variables.
EXPERIMENTS
In experiments, compare HMWS hybrid generative models against oftte-of-the-at wake- sleep-based RWS and variational inference based VIMCO. To ensure fair comparison, we match the number of likelihood evaluations typically the most time-consuming part of training across
5

---

Published as a conference paper at ICLR 2022
Vieets
-0.00
-00
13-
Touale
10000
Figure 4: Hybrid memoised wake-sleep (HMWS) learns faster than the baselines: reweighted wake-sleep (RWS) and VIMCO based on the marginal likelihood. in both the time series model (left), and the scene understanding models with learning shape and color (middle) and learning shape only (right). HMWs also learns better scene understanding models. The gradient estimator of VIMCO was too noisy and failed learn the time series model. (Median with the shaded interquartile ranges over 20 runs is shown.)
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 fixed memory cost of ((M) per data point which can be prohibitive for large datasets. However 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 Sie 100 sample importance weighted autoencoder (TWAE) estimation of the log marginal likelihood po()). We also tried the variatiomal-inferenae based REINFORCE but found that it underperforms VIMCO.
We foeus 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.
4.1
STRUCTURED TIME SERIES
Prior PolZat 2p)
Likellhood p(x Eig 2))
Data
Recognition model gglEgh Z (x)
Kemci eap cosionsion L5
Construct keme!
Kerell repasszze LSMM
R WNN
LBTM ar-beddier
"C Pare Vit
"D Per WiN
Compute covoriance
Karmal suprrsaiion LSI arr ambeedeer Jcl
Korne expression LSTM embedder
Kornel parameners LSTM joi
Koree parometers LSTM
Concolerade and leed to 5TM
Sample autoregneor why from ISTM
LSTM state the
Figure 5: The generative model structured time series data consists of prior that uses LSTMS to first sample the discrete Gaussian process (GP) kernel expression and then the continuous kernel parameters, and (ii) a likelihood which constructs the final GP kernel expression. The recognition model (on gray background) mirrors the architecture of the prior but additionally conditions an I.STM embedding of data.
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 GP kemee learning (Rasmussen) Williams 2006). They describe 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 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 x WN SE PER C, where
(11)
WN is the White Noise kernel, WN,2 a =ag-

---

Published as a conference paper at ICLR 2022
W MW
Figure 6: Learning model structured time series data with Gaussian processes (GPS) first inferring the kemel expression (shown as text in top-left corner) and the kernel parameters. The blue curve the 128-dimensional observed signal, and the orange curve is a GP extrapolation based on the inferred kernel. We show the mean (dark orange), the +2 standard deviation range (shaded) and one sample (light orange).
SE is the Squared Exponential kernel, S (2 x2) = o² exp(- (01 - ==2)2
PER is Periodic kernel, p,, (#1, 02) o² exp(- 1op(-2sin) sin |m x22l/)/12),
C is Constant, ( x2) = =²
We wish learn prior distribution over both the symbolic structure of kernel and its continuous variables (o², I, etc.). To represent the structure of the kernel as Ed, we use a symbolic kernel 'ex- pression': string over the characters {(); + x, WN, SE, PERI:..... PER4; C). We provide multiple character types PER; for periodic kernels, corresponding to learnable coarse-grainine geassenraiinig of period (short to long). We then define an LSTM prior po(zd) over these kernel expressions, alongside conditional LSTM prior po (z 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 [o(za, 2e(2r) 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.
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 (TT.. In Fig. (left), we show that HMWS learns faster than RWS in this domain in terms of the log evidence log (2:). 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, infer the latent kernel expression 2, by taking the highest scoring memory element =/" according to the memory-based posterior QUEM ( x)) and sample the corresponding kernel parameters from the continuous recognition model to (2, I.). We show the composite kernel as well as the GP 2d,
posterior predictive distribution. These programs describe meaningful compositional structure in the time series data, and can also be used to make highly plausible extrapolations.
4.2 COMPOSITIONAL SCENE UNDERSTANDING
Prior play +
Likelihood py(x 1ig 2
Data
Recognition model dz(ا,t (x)
Far: nell vangen can
Compute spatal
For mill in renge (run, anl 1631
niw blooke samplefe agorttoals
for btock cunpainum blockals
for bleck in ruspetous blocke prinisive sumpán (Cateporical joalo,
loytter mm
tay location exeple peoomat Joew scalemasp [pwill bluck] 1)
Ml aestemapp (II) 1)
Learnalle primitives
Render
CNN
Figure 7: Generative model of compositional scenes. The prior places stochastic number of blocks into each cell, where cells form imaginary grid the ground plane. For each cell, stack of blocks is built by both: i) sampling blocks from learnable set of primitives and ii) sampling their relative location to the object below (i.e.. either the ground or the most recently stacked block). The likelihood uses a differentiable renderer to produce an image. The recognition model (on gray background) mirrors the structure of the prior, but parametrizes distributions as learnable functions
of a CNN- based embedding of the image.
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 large consortium of blocks, where blocks are drawn from 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 scene parse describing which kinds of blocks live
7

---

Published as a conference paper at ICLR 2022
where in each world. Such symbolic representation of scenes increases interpretability, and has connections to language and symbolic planning.
Our generative and recognition models are illustrated in Fig. (for the full description of the model, see App. B.2). We divide the ground plane into an N x N (N 2) grid of square cells and initialize set of p=5 learnable block types, which refer 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. 3. For each position in the tower, we sample an integer E [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 constrained to lie in subset of the corresponding cell, using an affine transformation over a iigmoid-trassfrme 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 Po(=d, =e, -r) with respect to the learnable block types 0. 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 fixed standard deviation factorized over all pixels similar to using an L2 loss. The discrete latent variables, žd, comprise both the number of blocks and block type indices, while the continuous latent variables, Ze, represent the raw block locations.
The recognition model qs (2d, [z) follows a similar structure to that of the prior p(za, ze), as shown inside the gray box of Fig. However. the logits of the categorical distributions for go (zalz)). as well as the mean and (log) standard deviation of the Gaussian distributions for go (z(zel (Zelza, 4)), are obtained by mapping convolutional neural network (CNN)- embedding of the image x through small neural networks NN1-NN4 4; 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 pB possible scene parses. Moreover, scenes 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, 5, N = 5, and using comparison algorithms with 5 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 po (2)). We directly compare the wall-clock time taken by HMWS compared to RWS in Table highlighting the comparative efficiency of our algorithm. Further, we uncovered two local minima in our domain: we find that higher portion of HMWS find the superior minima compared to RWS. In Fig. 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 (=alz)) (z). These samples illustrate that HMWsetrained 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.
5 RELATED 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 (VAEs) have been employed for object-based scene understanding (e.g. [slami 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. (017)) or limited use of control variates (Mnih & Rezende, 2016). The wake- 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.
8

---

Published as a conference paper at ICLR 2022
Inferring scene parse with color
Inferring scene parse without color
Posterior Samples
Posterior Samples
Observations
Observations
Sample
Sempay
Samole
Sample
Sempe
Sample
Figure 8: Samples from the posterior when inferring scene parses with color (left) and seene parses without color diversity (right). Conditioned on a single observation of the front view of scene (left column), HMWS infers posterior over blocks that make up the scene. Three samples from the posterior shown per scene, sorted by log probahility under the model; e-g., the first sample most probable under HMWs. Sampled scenes are rendered from three different camera angles; position of camera depicted emphasize that the model has never seen the 3/ /-vvis... The sampled scenes highlight that we are able to handle occlusion, capturing distribution over possible worlds that are largely consistent with the observation.
figure insets.
Our most closely related work is Memoised Wake-sleep. which we build directly on top of and extend handle both discrete and continuous latents. contemporary of Memoised Wake-sleep, DreamCoder (Ellis et al. 2021). I, is a closely related program synthesis algorithm following a similar wake-slcep 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-ffi -rreie inference (Brehmer et al. 2020), or develop more complex training schemes amortization (Wenliang al. 2020).
Broadly, our goal of inferring compositionally structured, mixed discrete objects has strongest parallels in the program synthesis literature. One family of approaches (e.g. HOU- DINI (Valkov et al. 2018) and NEAR (Shah et al. 20200), perform 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)) by relaxing the discrete space to allow mondient-guided optimization to run on the whole problem (e.g. DeepProbl (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. (02T). What our work contributes this space is generic and principled way of applying amortized inference to hrbiid
problems. This gets the speed of a neural recognition model-which is valuable for program synthesis- with Bayesian handling of uncertainty and continuous parameters, rather than relying on heuristic quantization or expensive inner loops.
DISCUSSION
Inference hybrid generative models is important for building interpretable models that generalize. However, such task is difficult due to the need to perform inference in large, discrete spaces. Unlike deep generative models in which the generative model 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, does not fully solve the problem. particular, our algorithm struggles with models in which the continuous latent variables require non-trivial inference, as the quality of continuous inference directly linked with the quality of the gradient estimators in HMWS. This challenge potentially be addressed with better gradient estimators. We additionally plan to extend our algorithm to more complex, realistic neuro-symbe symbolic models; for instance, those with more non-cuboidal primitive topologies.

---

Published as a conference paper at ICLR 2022
REFERENCES
Jörg Bornschein and Yoshua Bengio. Reweighted wake-sleep. In International Conference Learning Representations, 2015.
Johann Brehmer, Gilles Louppe, Juan Pavez, and Kyle Cranmer. Mining gold from implicit models to improve likelihood- free inference. Proceedings of the National Academy of Sciences, 117(10): 5249, 2020.
Peter Dayan, Geoffrey Hinton, Radford M Neal, and Richard Zemel. The Helmholtz machine. Neural computation, 7(5):889-99) 904, 1995.
David Duvenaud, James Lloyd, Roger Grosse, Joshua Tenenbaum, and Ghahramani Zoubin. Structure discovery in nonparametric regression through compositional kernel search. In International Conference on Machine Learning, 2013.
Kevin Ellis,, Catherine Wong, Maxwell Nye, Mathias Sablé-Meyer, Lucas Morales, Luke Hewitt, Luc Cary, Armando Solar-Lezama, and Joshua Tenenbaum. Dreamcoder: Bootstrapping inductive program synthesis with wake-sleep library learning. In Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation, pp. 835-850, 2021.
SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances Neural Information Processing Systems, 2016.
Richard Evans, Matko Bošnjak, Lars Buesing, Kevin Ellis, David Pfau, Pushmeet Kohli, and Marek Sergot. Making sense of raw input. Artificial Intelligence, 299:103521, 2021.
Yaroslav Ganin, Sergey Bartunov, Yujia Li, Ethan Keller, and Stefano Saliceti. Computer-aiie aided design as language. Advances in Neural Information Processing Systems, 34, 2021.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoff Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In International Conference on Learning Representations, 2018
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In International Conference on Machine Learning, 2019.
Luke Hewitt, Tuan Anh Le, and Joshua Tenenbaum. Learning to learn generative programs with memoised wake-sleep. In Uncertainty in Artificial Intelligence, 2020.
Geoffrey E Hinton, Peter Dayan, Brendan J Frey, and Radford M Neal. The "wake-sleep" algorithm for unsupervised neural networks. Science, 1995.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-softmax. In International Conference on Learning Representations, 2017.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014.
Tuan Anh Le, Adam R. Kosiorek, N. Siddharth, Yee Whye Teh, and Frank Wood. Revisiting reweighted wake-sleep for models with stochastic control flow. In Uncertainty in Artificial Intelligence, 2019.
David J. C. MacKay. Information Theory, Inference & Learning Algorithms. Cambridge University Press, USA, 2002. ISBN 0521642981.
Maddison, Mnih, and Y Teh. The Concrete distribution: A continuous relaxation of discrete random variables. In International Conference Learning Representations, 2017.
10

---

Published as a conference paper at ICLR 2022
Robin Manhaeve, Sebastijan Dumancic, Angelika Kimmig, Thomas Demeester, and Lue De Raedt. Deepproblog: Neural probabilistic logic programming. Advances Neural Information Processing Systems, 2018.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In International Conference on Machine Learning, 2014.
Andriy Mnih and Danilo Rezende. Variational inference for Monte Carlo objectives. In International Conference on Machine Learning. 2016.
Art Owen. Monte Carlo theory, methods and examples. 2013.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian processes for machine learning. Adaptive computation and machine learning. MIT Press, 2006. ISBN 026218253X.
Nikhila Ravi, Jeremy Reizenstein, David Novotny, Taylor Gordon, Wan- Yen Lo. Justin Johnson, and Georgia Gkioxari. Accelerating 3d deep learning with pytorch3d. arXiv:2007.08501, 2020.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning, pp. 1530-1538. PMLR, 2015.
Ardavan Saeedi, Tejas D) Kulkarni, Vikash K Mansinghka, and Samuel J Gershman. Variational particle approximations. The Journal of Machine Learning Research, (8(1):2328-23555 2017.
John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. In Advances in Neural Information Processing Systems, 2015.
Ameesh Shah, Eric Zhan, Jennifer Sun, Abhinav Verma, Yisong Yue, and Swarat Chaudhuri. Learn- ing differentiable programs with admissible neural heuristics. Advances in neural information processing systems, 2020.
George Tucker, Andriy Mnih, Chris Maddison, John Lawson, and Jascha Sohl-Dickstein. REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models. In Advances in Neural Information Processing Systems, 2017.
Lazar Valkov, Dipak Chaudhari, Akash Srivastava, Charles Sutton, and Swarat Chaudhuri. Houdini: Lifelong learning as program synthesis. Advances in Neural Information Processing Systems, 2018.
Li Wenliang, Theodore Moskovitz, Heishiro Kanagawa, and Maneesh Sahani. Amortised learning by wake-sleep. In International Conference Machine Learning. 2020.
Ronald J Williams.
Simple statistical follwing algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-256, 1992.

---

Published as a conference paper at ICLR 2022
A DERIVATION OF THE GRADIENT ESTIMATOR FOR LEARNING THE GENERATIVE MODEL
Here, we derive gradient estimator for learning the generative model given in 7
We assume that the memory-induced variational posterior approximates the discrete posterior
Po (2 =2 QMEM (za(z>) (z)
(12)
M
Wm mm (Zd)
(13)
m=
and that the continuous posterior expectation Ep (2cl=3 .n) (2c)) for each m is approximated using of weighted samples (zith, Wmk)k k=
K
Bpo (ff((())))))))))) Σ
wmkf (z (zmk))
(14)
where mik Σk Moni"
Wmrk
First, make of the Fisher's identity
Vo log (z( = [pp(tla)) V log Pa (z:)
(Integrand is not a function of 2)
(15)
(=)()) V log pa (x) + V log pa(z]æ)]
(Second term vanishes) (16)
= (p)()) V log Po (2,
(Product rule of probability) (17)
and we continue by using the approximations in 12
(=.laa
(=al=) V log
(Factorize the posterior)
(18)
B (z V log (z. x)]
(Use (12)
(19)
Wm Vo log (20 Zer x)
(Use (3)
(20)
1==1
K
M
Wmnk
Wm V log zith ,))
(Use (4]
(21)
m=1
K
mk Umkk Vo logpo =z mm x),
(Combine sums)
(22)
where
Umk WmkWWm
(23)
Wmnk
Wms
(24)
1=1 Wms
Σ
E1=1
M
#=1
K
Wij
Wink
(25)
M
Σj=10 Wig
B
MODEL ARCHITECTURE AND PARAMETER SETTINGS
B.1
STRUCTURED TIME SERIES
Fixed parameters
Vocabulary = {WN, SE, PER PER,, C, X, +, ())
Vocabulary size poccablaary_szeee =1=11
Kernel parameters K. {o ( ( (o PrERi: 22 Pex,) ....; (0 PEERA e Pex.),
12

---

+ 1)). This defines the Published as a conference paper at ICLR 2022
Kernel parameters size Kernel_paranee size |K] 16
Hidden size hi hidder_sieee saiz==a = 128
Observation embedding size obs obs_embedding_size = 128
Generative model po (zd; Ro; 2)
Kernel expression LSTM (p)
(input size = vocabulary_siee hidden size = hidden_size)
This module defines a distribution over the kernel expression po(za).
At each time step 4, the input is one-hot encoding of the previous symbol in (or zero vector for the first time step)
At each time step I, extract logit probabilities for each element in or the end-of- sequence symbol EOS from the hidden state using "Kernel expression extractor (p)" (Linear (hidden size, vocabulary_ size 1)). This defines the conditional distribution pe (2 1
The full prior over kernel expression is an autoregressive distribution II, pe (2 whose length determined by EOS.
Kernel expression LSTM embedder (p)
(same as kernel expression LSTM (p))
The module summarizes the kernel expression 2d into an embedding vector e e
ghhidden_aize
We re-use the kernel expression LSTM (p) above and the last hidden LSTM state to be the summary embedding vector.
Kernel parameters LSTM (p)
(input size kernel laaeneel_paaamm..mm ze hidden sisisiiieee hidden size hidden_siieee
This module define distribution over kernel parameters po(ze|za).
At each timestep I, the input is concatenation of the previous kernel parameters z E [ieerne__arrmm__sie (or zero vector for t = 1) and the embedding vector E ghidden_size
At each timestep I, extract mean and standard deviations for each kernel param- eter in I. using "Kernel parameters extractor (p)" Linear (hi dhee_asizz,, Laaeeel__rrass______ pe(2?122 za).
This defines the conditional distribution
Recognition model (z(za, ze(x)
Signal LSTM embedder
(input size 1, hidden size obs_ bobo__mmeddinn_diee dim hidddhi_diisee
This module summarizes the signal x into embedding vector
gobo_cmbedding_.di.
The embedding vector 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 abs__mmeedding dim + vocabulary_size _ize,, hidden size hidden_siite size)
This module defines a distribution over kernel expression 4o (=a(tr))
At each time step f, the input is concatentation of the one-hot encoding of the previous symbol in V (or zero vector for the first time step) and cg-
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 tbiden_ssize conditional distribution go (=<l=> 1:4-1 2x).
vocabulary_size
The full distribution over kernel expression is an autoregressive distribution Hl qa 1:t-1 1:t-1 whose length is determined by EOS.
13

---

Published as a conference paper at ICLR 2022
Kernel expression LSTM embedder (q)
(input size = vocabulary_sicee.. hidden size hidden_size,
This module summarizes the kernel expression &d into an embedding vector esa R Bhidden_s117,
-The embedding vector Cs2 is taken to be the last hidden state of an LSTTM where =d is fed as the input sequence.
Kernel parameters LSTM (q) (input size obs less__ebeednnggiie + @ermel_params_siee hidden_size, hidden size aide_sizze)
This module defines a distribution over kernel parameters (a)
At each timestep f, the input is concatenation of the previous kernel parameters (or zero vector for / = 1), the embedding vector ez2 E Rhidden_size and the signal embedding vector < [gobs_embediing
.t11 E [Bkernei_pazam____0
At each timestep extract mean and standard deviations for each kernel param- eter in I. using "Kernel parameters extractor (q)" Linear (hidden_ste sin_ssieeee __iesssssee_____ This defines the conditional distribution 10(=01=0 x),
B.2 COMPOSITIONAL SCENE UNDERSTANDING
Fixed parameters
Number of allowed primitives learn
Maximum number of blocks per cell Bmax
Number of cells in the z-plane number of cells in the z-plane = N
Image resolution [3, 128, 128]
Observation embedding size cbs www_eeeeednn_iiiieeee size 676
Recognition model (o(=a, =x)2r). Components are detailed to match those shown in Fig.
CNN-based embedding of the image
Conv2d [iin_channels=t out_channelsste 8=64, kernel_sh _innel_ssiee, padding-2)
ReLU
MaxPool (kernel_size-) stride=2)
Conv2d (64, 128, 3, 1)
ReIU
MaxPool (3, 2)
Conv2d (128, 128, 3, 0)
ReIU-Conv2d ( (128, 4, 3, 0)
ReLU
MaxPool (2, 2)
Flatten
Note, the dimensionality of the output is obs abs_emmeddingg_iie (in this case, = 676)
Neural network for the distribution over blocks (NNI)
Linear (obs_embeedingg_ize,, N +
Bmas)
Neural network for the distribution over primitives (NN2)
Linear (obs bobb__ebbdddigg.cmm size, N Bmaz * P)
14

---

Published as a conference paper at ICLR 2022
Neural network that outputs the mean for the raw primitive locations (NN3)
Linear (obs_ weweeeedingsiiieeeeeeedeeeeeeeee N Bmax)
Neural network that outputs the standard deviation for the raw primitive locations (NN4)
Linear (obs_ember (obs_embedding_sizzee N -Bmaa)
C COMPARATIVE EFFICIENCY: WALL-CLOCK TIME AND GPU MEMORY CONSUMPTION
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, focus the scene understanding domain, specifically the model wherein scene parses are inferred with color. We find that HMWS consistently results more time- and memory-efficienn inference, HMWS's ((N + M) with RWs's S likelihood evaluations,
in Tables T a 2
The reason for this is that while match
The actual number of likelihood evaluations of HMWS K L which is typically smaller than K(N + M) because we only take I. S N + M unique memory elements (step 2 of Algorithm
The loss terms in lines 7-9 of Algorithm 1 all contain only K M < K L M)
log probability terms {log Po (22". 2. . ,x) .x)) M,K
I=)mm= flog (=(= 22 x)) M,K m= and M
terms {log go x (r)] m== 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 (=(zal=) and K samples from (s(=elza X) in HMWS in comparison to K (N + M) in RWS.
S or K(M + N) HMWS
RWS
8
67 min
82 min
136 min
196 min
32
186 min
327 min
50
390 min
582 min
72
499 min
N/A
Table 1: Wall-clock comparison (time to 5000 iterations). HMWS is faster in terms of absolute time to reach a fixed number of iterations regardless of the setting of the number of likelihond evaluations. Note that HMWS is even able with higher number of samples (e.g, K(M + N) 72). unlike RWS which fails due computationally prohibitive memory requirements.
S or K(M + N)
HMWS
RWS
8
2277 MB
3467 MB
18
4645 MB
7646 MB
32
8028 MB
13556 MB
50
12485 MB
21155 MB
72
17963 MB
N/A
Table 2: Comparison of GPU memory consumption. HMWS is more memory-efficient than RWS over a range of hyperparameter settings. Note that HMWS is even able with higher number of samples (e.g. K (M + N) 72), unlike which fails due computationally prohibitive memory requirements.
15

---

Published as a conference paper at ICLR 2022
D ABLATION EXPERIMENTS
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 IIWWS except the continuous recognition model qo (z) which is trained using the RWS loss. This method has memory.
"RWS+" which trains all components the same way as RWS except the continuous recognition model Ea, x) which is trained using the HMWS loss in (TO).
The learning curves for the time series model in Figure indicate that the memory is primarily responsible for the performance increase since HMWS- almost identically matches the performance of HMWS. Including the loss in (IO) in RWS training helps learning but alone doesn't achieve HMWs's performance.
35
30
25
10
VNNG
HMWS-
KWS
Eoration
10000
Figure 9; Augmenting the learning curves HMWS, RWS, VIMCO on the time series domain Figure (left) by HMWS- and RWS+ highlights the importance of HMWS's memory.
D.2 SENSITIVITY THE DIFFICULTY OF CONTINUOUS INFERENCE
Since both the marginalization and inference of the continuous latent variable ze 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 worse (Fig. [0) by arbitrarily adjusting the difficulty of continuous inference by changing the range of the period PeEs, of periodic kernels (see App. B.T
HMWS
R195
20
063)
-10
-30
Range
Figure 10: Inereasing the difficulty of the continuous inference task results in worse performance in HMWS. Baseline algorithms RWS and VIMCO suffer equally
D.3 COMPARING HMWS AND RWS FOR DIFFERENT COMPUTE BUDGETS
We compare HMWS and RWS for a number of different compute budgets (see Fig II Here, we keep K * (M N) S to match the number of likelihood evaluations of HMWS and RWS, and
16

---

Published as a conference paper at ICLR 2022
keep K M N for simplicity. We run sweep for HMWs with K = M = N (2,3,4,5,8 and a corresponding sweep for RWS with S in {8, 18, 30, 72). We find that HMWS consistently overperforms RWS.
TIMESERIES
SCENE UNDERSTANDING
-4300
20
log pld
-65400
109
-55600
-30
70
10
50
Number likethood evallvttion 00M4 M) for HSWSS S RIS)
Number likelnood evalustions (KIM NI for HMNS, f RW55
Figure 11: Comparing HMWS and RWS other compute budgets. Increasing overall compute by increasing K and keeping K = M = N improves performance while HMWS consistently outperforms RWS with equivalent compute. Note that RWS didn't 't run for S = 72 due memory constraints.
E USING VIMCO FOR HYBRID GENERATIVE MODELS
We initially treated the continuous latent variables the same way 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 't seem to improve the performance significantly.
Since applying VIMCO hybrid discrete-continuos 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 (=(= I do (=1)) and Qa (= |=]:k) = I (0 (=)) be the distributions of all the discrete and continuous randomness. We also drop the dependence data reduce notational clutter. Thus, the gradient we'd like to estimate is
V Bos [<] ( [[(( Φ. zd 1:K =1:K)]
=0
(26)
where f is the term inside of the IWAE objective ((...) log k Ek= po po(=
(=(
Next, we apply reparameterized sampling of the continuous variables 2 This involves choosing simple, parameterless random variable s(e) and reparameterization function &e r(e, za; o) such that its output has exactly the same distribution sampling from the original recognition network ze (z(l)=a)) Hence, we can rewrite the gradient we'd d like to estimate as
VolQoste ()(()()))
(r(e", = {=){ [J Φ, =d
(27)
where / takes in reparameterized samples of the continuous latent variables instead of 21:K
Now, we follow a similar chain of reasoning the one for deriving score-function gradient estimator
V Vo VoBoo(e)** (S(Aa)) [f(0, o, 2d 2 {r(ck. 2 K=)]
(28)
/ (o(Qu(z) SS f(...) + (2 1. S((e (
l:
(29)
/ Qa(=!i liK SS(
log Qa () Qs(z) 1 )S( 1: Vo)) de 1:K
del
(30)
E V a log Qo ((...) VoS(...))) V..(..
(31)
The VIMCO control variate is applied to the first term V log Qs (21k (f..))(....)
]f...). Importantly, this term does not contain the density of the continuous variables Q+(2?: provides gradients for learning the discrete recognition network unlike we treated the continuous latent variables the same the discrete latent variables (in which would used for learning the continuous recognition network as well). The gradient for learning the continuous recognition network instead comes from the second term V V.....)
1: 1= Hence, this term only
a score-function gradient
17


Reviews:
[{"review_id": "zWq4OUNjSI", "reviewer": "Reviewer_G5jQ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**\n\n* The paper deals with a relevant issue - how to extend an existing approach to hybrid discrete-continuous graphical models.\n* The exposition is clear. Related work and background are explained well. It is made easy to see where the authors' work fits in to previous work.\n* The proposed method is principled and based on a well-established method from approximate inference (importance sampling). There are no ad-hoc terms or deviations from theory.\n* Experiments are favorable towards the approach and show that it is better suited to program synthesis than other applicable methods.\n\n**Weaknesses**\n\n* Somewhat weak in novelty.\n\n**Other questions/comments**\n\n* In 2.1 you say:\n\n> 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.\n\nI don't immediately see why this is the case. Is this an empirical observation or does it follow from theory? Also there is a typo in that sentence.\n\n* In theory we could approximately integrate out the continuous variables using other methods. Did you consider comparing to a simple MC estimate instead of importance sampling (maybe even with a single sample, as done in VAEs)? Or is there a reason why such a comparison is not possible?\n\n", "summary_of_the_paper": "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.\n\nSpecifically, 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.", "summary_of_the_review": "I think the paper is good and recommend accepting.\n\nThe main weakness I see is that the contribution seems to be limited to using importance sampling for marginalizing out continuous variables. The actual presence of continuous variables can be seen in the Gaussian mixture models experiment of the memoised wake-sleep paper [1], where it was possible to treat them analytically. Of course, the general case doesn't allow for that and therefore the paper is still making a valuable contribution by investigating what to do in the general setting.\n\nFurther, the paper is specifically aimed at settings where discrete variables arise in a programmatic context and impact control flow. The experiments show that the proposed approach has an edge over other general methods in this setting.\n\n[1] Hewitt et al., \"Learning to learn generative programs with Memoised Wake-Sleep\", https://arxiv.org/pdf/2007.03132.pdf", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Ta1Fx21YDAx", "reviewer": "Reviewer_FY4y", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The paper is clearly written with good motivation and methods sections. The rigorous and concise notations convey a lot of information. The diagram in Fig 5 is less clear but many details are given in the appendix. \n2. The method is technically sound and involves a non-trivial application of importance sampling, which is of value.\n3. The experiments are reasonably complicated (but see below) and serve a good purpose demonstrating the potential of this method.\n\nWeaknesses:\n1. My main concern is the significance of this paper. The applications are relevant to comp. cogsci community, but the authors did not discuss much of the cognitive link or compare it with any cognitive/behavioural data. As a machine learning paper, I do not find the comparison with VIMCO and RWS to be sufficient. Further, although the two problems in the experiments involve reasonably complicated generative models, these models are not the main contributions of the paper, and they are arguably low-dimensional in latent space. For example, in the second experiment, there are only 4 discrete locations among other latent variables. Again, as a machine learning paper, the question of scalability is not well addressed. As a simple test, is it possible to learn a flexible parametric likelihood (Gaussian with NN mean) for experiment 2, rather than using a differentiable renderer?\n2. The fact that the problems are relatively low-D raises two further questions: how does this method scale to larger latent dimensions? How does this method compare with ablations of this method? As far as I understand, the gap between RWS and the proposed method include a) memory; b) importance sampling to maintain extra dependency between $z_d$ and $z_c$. Either of these can be augmented to RWS for improvement. How do these two contribute to performance? I am not asking the authors to \"just compare more baselines\", but to conduct a more careful comparison between contributions of various components in their model.\n3. Importance sampling has a proveably large variance. How does this method not suffer from this problem? Is it because of the sleep-phase updates? I hope the authors could address this problem. \n\nOther suggestions\n1. On the algorithmic side, I'd like the authors to expand a bit more on the paragraph before experiment 4, which I think touches on important upside (mentioned dependence between $z_c$ and $z_d$) and downside: for model where the induced true posterior dependency factories the other way as $p(z_d|z_c,x)p(z_c|x)$, the proposed method may suffer.\n2. In the explanation before equation (10), the authors mentioned that \"We want to minimize the average of KL\". Why does this make sense? The authors could include a derivataion starting from the joint KL\n$$\nKL[p(z_c,z_d|x) || q(z_c, z_d|x)] = \\int  p(z_c,z_d|x) \\log \\frac{p(z_c,z_d|x)}{q(z_c,z_d|x)} = \\int p(z_c|z_d,x)  p(z_d|x) \\log \\frac{p(z_d|x)}{q(z_d|x)}\\frac{p(z_c|z_d,x)}{q(z_c|z_d,x)}\n$$\nThis will then expose another assumption that the aforementioned expectation is in fact taken over the approximate $q(z_d|x)$. This is also related to the expectation propagation method. \n3. For the runtime comparison, I'd like the authors to also show the memory cost comparison and be explicit about the tradeoff.", "summary_of_the_paper": "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.", "summary_of_the_review": "In general, I think the method is an interesting and valuable contribution, I also enjoyed the precise notations the authors adopted. But either the impact/scalability of the method or the machine learning power needs to be more clearly demonstrated for acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Hff_60hzBX1", "reviewer": "Reviewer_pAhD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Good paper, giving an interesting new method.\n\nMy main concerns are around the empirical evaluations and wallclock times.  In particular, I can easily see that memoization in MWS is going to significantly help runtime.  This isn't so clear in the present algorithm, because there's an awful lot more going on (all the steps involving continuous latent variables don't seem to benefit much from memoization).  This concern is further reinforced by comparing algorithms by iterations (it would also be helpful to look at algorithms by runtime).  Could the authors clarify why their method should help runtime (or whether it won't, and its primarily intended to help performance), and include wallclock time comparisons?\n\nI am also surprised by how much better the algorithm performs relative to RWS.  Intuitively, it would seem that HMWS is doing something very similar to RWS, but memoizing the best samples from the approximate posterior over latent variables.  That would seem to help, but it isn't clear to me that it would help _that_ much, as I wouldn't expect the best samples from the approximate posterior to be that different from typical samples.  Can the authors comment? This is especially concerning, given that:\n* There is no hyper parameter optimisation (Adam with default hyper parameters).\n* There is no discussion of \"matching\" the RWS and HMWS architectures to ensure a fair comparison (that I could see).", "summary_of_the_paper": "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.", "summary_of_the_review": "Good paper, presenting an interesting new method.  A few concerns as regards evaluations + performance tradeoffs.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose a method called Hybrid Memoised Wake-Sleep (HMWS) for training models with both discrete and continuous latent variables efficiently using amortized inference. They extend Memoised Wake-Sleep (MWS), which can only handle discrete latent variables, to discrete-continuous systems by using importance sampling to approximately marginalize out the continuous variables and then applying MWS to the discrete variables.

This is well motivated and well written paper. The method is novel, clearly described, and evaluated in two fairly different interesting settings. However, while the empirical evaluation was considerably strengthened by the ablation studies and other experiments included in response to the reviewers, it is still on the weak side. The main issues are the relatively low-dimensional latent spaces and insufficiently tuned baselines. As one reviewer pointed out, importance sampling is unlikely to scale well to high dimensions, so exploring this aspect experimentally would strengthen the paper. The use of default Adam parameters for all methods and models substantially undermines the results and might at least in part explain the underwhelming baseline performance. For example, VIMCO has been successfully used to train a discrete/continuous model of seemingly comparable complexity in [1] and yet here it seems to fail completely. To shed some light on this, the authors might want to explain how they used VIMCO and whether their approach was substantially different from the one from [1].

Finally, it at least somewhat misleading to claim, as is done e.g. in Sec. 4, that HMWS is more memory efficient than the baselines, when unlike them, it needs to store M discrete latent configurations per training case. Please make this claim more precise to avoid potential confusion.

[1] Variational Memory Addressing in Generative Models, Bornschein et al., NIPS 2017
Decision: Accept (Poster)