Kanerva++: Extending the Kanerva Machine With Differentiable, Locally Block Allocated Latent Memory
Authors: Jason Ramapuram
Yan Wu
Alexandros Kalousis
Venue: ICLR 2021 Poster (2021)
Keywords: memory
generative model
latent variable
heap allocation

Abstract: Episodic and semantic memory are critical components of the human memory model. The theory of complementary learning systems (McClelland et al., 1995) suggests that the compressed representation produced by a serial event (episodic memory) is later restructured to build a more generalized form of reusable knowledge (semantic memory). In this work, we develop a new principled Bayesian memory allocation scheme that bridges the gap between episodic and semantic memory via a hierarchical latent variable model. We take inspiration from traditional heap allocation and extend the idea of locally contiguous memory to the Kanerva Machine, enabling a novel differentiable block allocated latent memory. In contrast to the Kanerva Machine, we simplify the process of memory writing by treating it as a fully feed forward deterministic process, relying on the stochasticity of the read key distribution to disperse information within the memory. We demonstrate that this allocation scheme improves performance in memory conditional image generation, resulting in new state-of-the-art conditional likelihood values on binarized MNIST (≤41.58 nats/image) , binarized Omniglot (≤66.24 nats/image), as well as presenting competitive performance on CIFAR10, DMLab Mazes, Celeb-A and ImageNet32×32.

Published as a conference paper at ICLR 2021
EXTENDING THE KANERVA MACHINE WITH DIFFERENTIABLE, LOCALLY BLOCK ALLOCATED LATENT MEMORY
KANERVA++:
Jason Ramapuram 1.2, Yan Wu', Alexandros Kalousis
1 University of Geneva & Geneva School of Business Administration, HES-SO
2 Currently at Apple
3 Deepmind rranaprrmmaapple.. com, yanwu@google.com alexandros www.sskissssssss..coooom
ABSTRACT
Episodic and semantic memory are critical components of the human memory model. The theory of complementary learning systems (McClelland et al.) 1995) suggests that the compressed representation produced by serial event (episodic memory) is later restructured build more generalized form reusable knowledge (semantic memory). In this work we develop new principled Bayesian memory allocation scheme that bridges the gap between episodic and semantic memory via hierarchical latent variable model. We take inspiration from traditional heap allocation and extend the idea of locally contiguous memory to the Kanerva Machine, enabling novel differentiable block allocated latent memory. In contrast to the Kanerva Machine, we simplify the process of memory writing by treating it fully feed forward deterministic process, relying the stochasticity of the read key distribution to disperse information within the memory. We demonstrate that this allocation scheme improves performance in memory conditional image generation, resulting in new state-of-the-art conditional likelihood values on binarized MNIST (541.58 nats/image) binarized Omniglot (566.24 nats/image), as well as presenting competitive performance on CIFARIO, DMLab Mazes, Celeb-A and ImageNet32x32.
INTRODUCTION
Memory is a central tenet in the model of human intelligence and is crucial to long-term reasoning and planning. Of particular interest is the theory of complementary learning systems McClelland et al.) (1995) which proposes that the brain employs two complementary systems to support the acquisition of complex behaviours: hippocampal fast-learning system that records events as episodic memory, and neocortical slow learning system that learns statistics across events as semantic memory. While the functional dichotomy of the complementary systems are well-established McClelland et al.) (1995); Kumaran et al.) (2016), it remains unclear whether they are bounded by different computational principles. In this work we introduce a model that bridges this gap by showing that the same statistical learning principle can be applied to the fast learning system through the construction of a hierarchical Bayesian memory.
pl malloc(16)
p2 p4
Sequential Operations
p2 = malloc(9)
p3 malloc(4)
p4 malloc(4)
Figure 1: Example final state of traditional heap allocator (Marlww et al. 2008) (Left) vs. K++
(Right); final state created by sequential operations listed on the left. K++ uses key distribution
to stochastically point to memory sub-region while Marlow al. (2008)
direct pointer.
Traditional heap allocated memory affords ((1) free malloc computational complexity and serves inspiration for K+ which uses differentiable neural proxies.

---

Published as a conference paper at ICLR 2021
While recent work has shown that using memory augmented neural networks can drastically improve the performance of generative models (Wu et al.
[2118a] b). language models (Weston et al.) 2015),
meta-learning (Santoro et al.] 2016), long-term planning (Graves et al. 2014 2016) and sample efficiency in reinforcement learning (Zhu et al.
2019), no model has been proposed to exploit
the inherent multi-dimmenionnlity of biological memory Reimann et al. (2017). Inspired by the traditional (computer-science) memory model of heap allocation (Figure LLLeft), we propose a novel differentiable memory allocation scheme called Kanerva ++ (K++), that learns to compress an episode of samples, referred to by the set of pointers {pl.... p4} in Figure [ into a latent multi- multi-diimnnsioal memory (Figure [Right). The K++ model infers a key distribution as a proxy to the pointers (Marlow et al.] 2008) and is able to embed similar samples to an overlapping latent representation space, thus enabling to be more efficient on compressing input distributions. In this work, we focus applying this novel memory allocation scheme to latent variable generative models, where we improve the memory model in the Kanerva Machine (Wu et al. 2018a b).
2 RELATED WORK
Variational Autoencoders: Variational autoencoders (VAEs) Kingma & Welling, 2014) are a fundamental part of the modern machine learning toolbox and have wide ranging applications from generative modeling (Kingma & Welling 2014; Kingma et al.] 2016, Burda et al. 2016).
learning graphs (Kipf & Welling (016), medical applications (Sedai et 2l.] Zhao et al. 2019) and video analysis (Fan et al., 2020). As a latent variable model, VAËS infer an approximate and can be used in downstream tasks such as control in reinforcement learning (Nair et al. 2018 Pritzel et al. 2017). VAEs maximize an evidence lower bound (ELBO), L(X, 2). of the log- aar-giaaa likelihood, Im)(X) > L(X, Z) = In pe ( (X|Z) Z)) (g (Z) (Z)())). The produced variational approximation. go (Z)))),, is typically called the encoder. while Pa Z comes from the decoder. Methods that aim to improve these latent variable generative models typically fall into two different paradigms: learning more informative priors or leveraging novel decoders. While improved decoder models such as PixeIVAE (Gulrajani et al.) 2017) and PixeIVAE++ (Sadeghi et al. 2019) drastically improve the performance of ( ( they suffer from phenomenon called posterior collapse (Lucas et al. (2019), where the decoder can become almost independent of the posterior sample, but still retains the ability to reconstruct the original sample by relying on its auto-regressive property (Goyal et al. (017a))
posterior over a latent representation
In contrast, VampPrior (Tomezak & Welling, 2018), Associative Compression Networks (ACN) (Graves et al. 2018). VAE-nCRP (Goyal et al. (2017)) and VLAAE (Chen et al. 2017) tighten the variational bound by learning more informed priors. VLAE for example, uses a powerful auto-regressive prior; VAE-nCRP learns a non-parametric Chinese restaurant process prior and VampPrior learns a Gaussian mixture prior representing prototypical virtual samples. On the other hand, ACN takes two-stage approach: by clustering real samples in the space of the posterior; and by using these related samples as inputs to a learned prior, ACN provides an information theoretic alternative to improved code transmission. Our work falls into this latter paradigm: we parameterize a learned prior by reading from a common memory, built through a transformation of an episode of input samples.
Memory Models: Inspired by the associative nature of biological memory, the Hopfield network (Hopfield) 1982) introduced the notion of content-addressable memory, defined by a set of binary neurons coupled with a Hamiltonian and a dynamical update rule. Iterating the update rule minimizes the Hamiltonian, resulting in patterns being stored at different configurations (Hopfield) 1982 Krotov& & Hopfield 2016). Writing in Hopfield network, thus corresponds to finding weight configurations such that stored patterns become attractors via Hebbian rules (Hebb) 1949.. This concept of memory was extended to a distributed, continuous setting in Kanerva (1988) and to a complex valued, holographic convolutional binding mechanism by Plate (1995). The central difference between associative memory models HHoffild (1982); Kanervaj (1988) and holographic memory Plate| 1995) is that the latter decouples the size of the memory from the input word size.
Most recent work with memory augmented neural networks treat memory in a slot-based manner (closer to the associative memory paradigm), where each column of a memory matrix, M, represents single slot. Reading memory traces, 2, entails using vector of addressing weights, r, to attend to the appropriate column of M, 2 = 77 M. This paradigm of memory includes models such as the Neural Turing Machine (NTM) (Graves et al. 2014), Differentiable Neural Computer (DNC)
2

---

Published as a conference paper at ICLR 2021
Graves al. 2016 Memory Networks /Weston et al. 2015). Generative Temporal Models with Memory (TMM) Fraccaro et al. (2018). Variational Memory Encoder- Decoder (VMED) Le et al. (0188), and Variational Memory Addressing (VMA) (Bornschei et al.] 2017). VMA differs from GTMM, VMED, DNC, NTM and Memory Networks by taking a stochastic approach to discrete key-addressing, instead of the deterministic approach of the latter models.
Recently, the Kanerva Machine (KM) (Wu et al., 201 and its extension, the Dynamic Kanerva Machine (DKM) (Wuet al. (018b), interpreted memory writes and reads as inference in generative model, wherein memory is now treated as a distribution, p(M). Under this framework, memory reads and writes are recast as sampling or updating the memory posterior. The DKM model differs from the KM model by introducing dynamical addressing rule that could be used throughout training. While providing an intuitive and theoretically sound bound on the data likelihood, the DKM model requires an inner optimization loop which entails solving an ordinary least squares (OLS) problem. Typical OLS solutions require a matrix inversion (O(n")) preventing the model from scaling to large memory sizes. More recent work has focused on employing product of smaller Kanerva memories (Marblestone et al. 2020) in an effort to minimize the computational cost of the matrix inversion. In contrast, we propose simplified view of memory creation by treating memory writes as a deterministic process in a fully feed-forward setting. Crucially, we also modify the read operand such that it uses localized sub-regions of the memory, providing an extra dimension of operation in comparison with the KM and DKM models. While the removal of memory stochasticity might be interpreted as reducing the representation power of the model, we empirically demonstrate through our experiments that our model performs better, trains quicker and simpler optimize. The choice of deterministic memory is further reinforced by research in psychology, where human visual memory has been shown to change deterministically Gold et al. (2005): Spencer & Hund (2002): Hollingworth et al. (2013).
MODEL
To better understand the K++ model, we examine each of the individual components to understand their role within the complete generative model. We begin by first deriving a conditional variational lower bound (Section 3.1).. describing the optimization objective and probabilistic assumptions. We then describe the write operand (Section 3,3), , the generative process (Section 3.4 and finally the read and iterative read operands (Section 3.5
3.1 PRELIMINARIES
K++ operates over an exchangeable episode (Aldous 1985) of samples, X [xx)7 E D. drawn from dataset D, as in the Kanerva Machine. Therefore, the ordering of the samples within the episode does not matter. This enables factoring the conditional, p( M)M), over each of the individual samples: IT= MM), given the memory, M p(M), M E RCx * Our objective in this
work is to maximize the expected conditional log- likeliiol as (Bornschein et al. 2017 at, 2018ab:
J ==E((x), p(MIX) In In ( / / ((X))(M)X) Σ In (=
As alluded to in Barber & Agakov (0044) and Wu et al. 2018a), this objective can be interpreted as maximizing the mutual information, ((X; M,), between the memory, M, and the episode, X, since T(X; = HXX)= + J H(X) H(X; .)) and given that the entropy of the data, H(X), is constant. In order actualize Equation we rely on variational bound which we derive in the following section.
3.2 VARIATIONAL LOWER BOUND
To efficiently read from the memory, M, we introduce a set of latent variables corresponding to the K addressing read heads, Y = {Y, = (vex ) k= 1T Vsk E R3 and set of latent variables corresponding to the readout from the memory, Z = { 24 R'_ Given these latent variables, we can decompose the conditional, In (MXMM, using the product rule and introduce variational approximations 4o (Z Fndd (YIx) via a multipy-by-one trick:
While DNC is slot based, it should be noted that DNC reads rows rather than columns. We use qo (Z) (Z) X) our variational approximation instead of (Z)X, Y, M) in DKM it presents more stable objective. We discuss this in detail Section 3.5
3

---

Published as a conference paper at ICLR 2021
ln)(X = In p(X, 2, (M
p =In MI P(
((X
M) p(Y) 96(2|X) (x())
R (21)
(2)
p(Z X2]M,Y,X))((Y/MM, (x(Z (Y
= LI + (O () 2000:///DDDD0000////////22/21X
(3)
20
DKL (X)LIX(MXX)
Equation assumes that Y is independent from M: p p This decomposition results in Equation which includes two KL-divergences against true (unknown) posteriors, p X)) and p Y M, X). We can then train the model by maximizing the evidence lower bound (ELBO), Cr, to the true conditional, In (M) > fr:
Lr Bys(Zix), (x In po po (X|Z,
Decoder
,,, D Dxzlx (ttpr///PPPreeee//xxp/p///////a
(4)
Amortized latont variable posterior vs. memory readout prior
(X
Amortized key posterior vs. key prior
The bound in Equation is tight if 9s (Z = D(Z),,Y, and (Y M. however, it involves inferring the entire memory M (M This prevents us from decoupling the size of the memory from inference and scales the computation complexity based on the size of the memory. To alleviate this constraint, we assume purely deterministic memory, M 6[fmem (Xen(())) built by transforming the input episode, X, via a deterministic encoder and memory transformation model, fmem fene- We also assume that the regions of memory which are useful in reconstructing a sample, It. can be summarized by a set of K localized contiguous memory sub-blocks described in Equation Sa below. The intuition here is that similar samples, I, = 2r, might occupy a disjoint part of the representation space and the decoder, pa (X) [.), would need read multiple regions properly handle sample reconstruction. For example, the digit "3" might share part of the representation space with a "2" and another part with "5".
M (m+ = fsT fmem (Sene(X) (X)), Bea H
(MIxX)
(5)
M in equation 5 represents set of K dirac-delta memory sub-regions, determined by the addressing
key, Yek. and a spatial transformer (ST) network Jaderberg et al. (015)). fsr Our final
optimization objective, Cr, is attained by approximating M/ ~ (M X, Y from Equation 4 with
M (Equation 5) and is summarized by the graphical model in below.
M
M
M
M
773
Uk
m2
Uk
mi
yk
Zt
St
St
(d)
Figure 2: (a): Generative model (3.4)..
(b): Read inference model (3.5) (c): Iterative read
inference model (3.5). (d): Write inference model (33.3). Dashed lines represent approximate inference, while solid lines represent computing of conditional distribution. Double sided arrow in (c) represents the KL divergence between qs (z(2 x and pe ( (Z),, (M,Y) Y') from Equation 4 Squares
represent deterministic nodes. Standard plate notation is used depict a repetitive operation.
We provide brief review of spatial transformers in Appendix
4

---

Published as a conference paper at ICLR 2021
3.3
WRITE MODEL
sample episode:
Memory
X₂ fx1,..., ]T D;
Embedding
compute embedding:
Encoder
Writer
Memony
E (X.))
infer keys: V
W(Hozey (E), ( (E));
Test Episode
write memory:
M ~ 5(f0mem (E))
Figure 3: Left: Write model. Right: Write operation.
Writing to memory in the K++ model (Figure 3) entails encoding the input episode, X; = {z,}/= | through the encoder, the pooled representation with the memory writer, M fome. (E). In this work, we employ Temporal Shift Module (TSM) (Lin et al., 2019), applied on a ResNet18 (He et al.) 2016). TSM works by shifting feature maps of two-dimensional vision model in the temporal dimension in order to build richer representations of contextual features. the case K++, this allows the encoder to build better representation of the memory by leveraging intermediary episode specific features. Using a TSM encoder over a standard convolutional stack improves the performance of both K++ and DKM, where the latter observes improvement of 6.32 nats/image test conditional variational lower bound of 77.2 nats/image (Wu et al.] (0188) for the binarized Omniglot dataset. As far not been explored and is contribution of this work.
fo (Xt), pooling the representation over the episode and encoding
the reported
are aware, the application of TSM encoder memory models has
The memory writer model, fomem. in Figure 3 allows K++ non-linearly transform the pooled embedding to better summarize the episode. In addition to inferring the deterministic memory, M, we also project the non-poote non-pooled embedding. E, through key model. fer..
(Y)) = pocay (E) + oZ (E) O Yek E TR., N(0
(6)
The reparameterized keys will be used to read sample specific memory traces, M, from the full memory, M. The memory traces, M, are used in training through the learned prior, Ps M, (,Y) Y) N(H,, from Equation 4 via the KL divergence. (r DxI.)xxxll.22222 (Z) Y)). This KI. divergence constrains the optimization objective to keep the representation of the amortized approximate posterior, qo (7lX). (probabilistically) close to the memory readout representation of the learned prior, po (o M Y). routed from the learned prior, po Z (Z), Y), to the decoder, po (X)),, in a similar manner to standard VAEs. We detail this process in the following section.
generative setting, this constraint enables memory traces to be
3.4 SAMPLE GENERATION
given memory: M
Spatial
Merrry
Transformer
sample keys: Y, P(Y);
extract regions: M
Decoder
(fosr (M, yea ]k
Traddc
infer latent; Zo
Single Generation
N (M), ob (M))
Episode
decode: &,
N N(uae (22z),
Figure 4: Left: Generative model. Right: Generative operation.
5

---

Published as a conference paper at ICLR 2021
The Kanerva++ model, like the original KM and DKM models, enables sample generation given an existing memory or set of memories. K samples from the prior key distribution, {yr, ly- ~ = (0, 1), M/ R3, are used to parameterize the spatial transformer, fsr, which indexes the deterministic memory, M. The result this differentiable indexing is a set of memory sub-regions, M, which are used in the decoder, (x()),), to generate synthetic samples. Reading samples in this manner forces the encoder utilize memory sub-regions that are useful for reconstruction, as non-read memory regions receive zero gradients during backpropagation. This insight allows us to use the simple fewa-foward write process described in Section 3.3 while still retaining the ability to produce locally contiguous block allocated memory.
3.5
READ/ ITERATIVE READ MODEL
given embedding: E: if training then infer latent: Zt ~
Memory
Triiiin
N(Ho, (K), (b))
(A
infer prior: Ze
Zi
Neo, (M). o. (a1))
heraoive rew reeding prior evaluation)
clse
rocontrr.cton ction
Eplande
infer latent; 2, ~ Nipo (M), 0ā. (MT)
end
decude: X, No ((zz)).²))
Figure 5: Left: Read model: bottom branch from embedding used during iterative reading and prior evaluation. Stable top branch used to infer go (Z X) during training. Right: Read operation.
K++ involves two forms of reading (Figure 5) iterative reading and a simpler and more stable read model used for training. During training we actualize go (2)x) from Equation 4 using an amortized isotropic- isotropic-gaussian gaussian posterior that directly transforms the embedding of the episode, E, using learned neural network (Figure b). As mentioned in Section 3.5 the readout, Z, of the memory traces, M, are encouraged learn a meaningful structured representation through memory read-out KL divergence, (x)) DKr.
P&()), Y)), which attempts to minimize the (probabilisticc distance between go (Z and Po(Z)
Kanerva memory models also possess the ability to gradually improve a sample through interative inference (Figure (c),, whereby noisy samples can be improved by leveraging contextual information stored in memory. This can be interpreted as approximating the posterior, g( X, M), by marginalizing the approximate key distribution:
d(]|X - -f qs (o(Y))X)) MYOY po = YY,,
where Y (Y X) in Equation 7 uses single sample Monte Carlo estimate, evaluated by re-infering the previous reconstruction, X (X))), through the approximate key posterior. Each subsequent memory readout, Z, improves upon its previous representation by absorbing additional information from the memory.
4 EXPERIMENTS
We contrast K++ against state-of-the-art memory conditional vision models and present empirical results in Table Binarized datasets assume Bernoulli output distributions, while continuous values are modeled by a discretized mixture logistics [Salimans et al. (2017). As is standard in literature logisticsSalimans
Burda al. (2016); [Sadeghi al. (2019); Ma al. (2018): 8); Chen et al. (201)). we provide results for binarized MNIST and binarized Omniglot in nats/image and rescale the corresponding results to bits/dim for all other datasets. We deseribe the model architecture, the optimization procedure and the memory creation protocol in Appendix E and E.I Finally, extra Celeb-A generations and test image reconstructions for all experiments are provided in Appendix Bnd Appendix respectivvl.

---

Published as a conference paper at ICLR 2021
Binarized
Binarized Omniglot
Fashion
CIFARIO
DMLab
Method
MNIST
MNIST
Mazes
(nats/ / image)
(nats/ / image)
(bits/ dim)
(bits /im)
(iiss / dim)
VAE Kingma & Welling ( 2014
87.86
104.75
5.84
6.3
IWAF Burda et al. 2016
85.32
103.38
Improved decoders
PixeIVAE++ Sadeghi et al. 2019
78.00
2.90
Ma al. 2018
77.98
89.09
2.95
DRAW Gregor et 2015
87.4
96.5
3,58
MatNet Bachman 2016
78.5
89.5
3.24
Richer priors
Ordered Graves et al. 2018
737
$3.07
VI.AF Chen al. 2017
78.53
102.11
2.95
VampProo Tomezak Welling 2018
78.45
89.76
Memory conditioned models
VMA Borschhin et 2017
103.6
KM al. 20188
$68.3
$4.37*
DNC Graves al. 2001
$100
DKM Wu et al. 20188
$75.3*
$77.2
$4,79
$2.75t
DKM w/TSM (our impl)
<51.84
$70.88
<4.15
<4.31
<2.92
Kanerva++ (ours)
<41.58
$66.24
$3.40
<3.28
$2.881
Table 1: Negative test likelihood and conditional test likelihood values (lower is better). was
graciously provided by original authors. + estimated from Wu et al. (2018a) Appendix Figure 12. variadic performance due online generation of DMLab samples.
K++ presents state-of-the-arr results for memory conditioned binarized MNIST and binarized Omniglot, and presents competitive performance for Fashion MNIST, CIFARIO and DMLab mazes. The performance gap on the continuous valued datasets can be explained by our use of a simple convolutional decoder, rather than the autoregressive decoders used in models such as PixeIVAE Sadeghi et al. (2019). We leave the exploration of more powerful decoder models to future work and note that our model can integrated with autoregressive decoders.
4.1 ITERATIVE INFERENCE
Figure 6: Left: First column to left visualizes first random key generation. Following columns created by inferring previous sample through K++. Right: denoising of salt & pepper (top), speckle (middle) and Poisson noise (bottom).
One of the benefits of K++ is that uses the memory to learn a more informed prior by condensing the information from an episode of samples. One might suspect that based on the dimensionality of the memory and the size the read traces, the memory might only learn prototypical patterns, rather than full mamatioaiion of the input episode. This presents problem for generation, as described Section 3.4. and can be observed in the first column of Figure o Left where the first generation from random key appears blurry. To overcome this limitation, we rely the iterative inference of Kanerva memory models (Wu et al.) (2088) 8a b). By holding the memory, M, fixed and repeatedly inferring the latents, are able to clcan-up the pattern by leveraging the contextual information contained within the memory ($ 3.5). This is visualized in the proceeding columns Figure Left, where observe slow but clear improvement in generation quality. This property of iterative inference one of the central benefits of using a memory model over a tradition solution like a VAE. We also present results of iterative inference more classical image noise distributions such as halt-and-pepper, speckle and Poisson noise in Figure For each original noisy pattern (top rows) provide the resultant final reconstruction after ten steps of clean-up. The proposed K++ is robust to input noise and is able to clean-up most of the patterns in semantically meaningful way.

---

Published as a conference paper at ICLR 2021
4.1.1 IMAGE GENERATIONS
m
H
H
60
d0
8
Figure 7: Key perturbed generations. Left: DML.ab mazes. Center: Omniglot, Right: Celeb-A 64x64.
R
%
H
L
O
/
G
2
7
3
x
C
18
2
R
oo
s
b
x
U
J K
Figure Random key generations. Lefr: DML.ah mazes. Center: Omniglot. Right: Celeb-A 64x64.
Typical VAEs use high dimensional isotropic Gaussian latent variables (2 R Burda et al. (2016); Kingma & Welling 2014). A well known property of high dimensional Gaussian distributions
that most their mass is concentrated on the surface area of a high dimensional ball. Perturbations of valid density easily move it ct al. 2018 White 2016), causing blurry or irregular generations. In the case of K++, since the key distribution, 1/2 p( pYY), R3, is within low dimensional space, local perturbations, Yt + e 0.1,) are likely in regions with high probability density. We visualize this form of generation in Figure for DMLab Mazes, Omniglot and Celeb-A 64x64, as well the more traditional random key generations, ~ p(Y), in Figure g
to sample in
invalid density region (Arvanitidis)
Interestingly, local key perturbations of trained DMLab Maze K++ model induces resultant generations that provide natural traversal of the maze observed by scanming Figure 7 Left, row by row, from left right. In contrast, the random generations of the same task (Figure 8 Left) present a more discontinuous set of generations. We see similar effect for the Omniglot and Celeb-A datasets, but observe that the locality is instead tied to character or facial structure as shown in Figure Center and Figure Right. Finally, in contrast VAE generations, K++ is able to generate sharper images ImageNet32x32 as shown in Appendix C Future work will investigate this form of locally perturbed generation through an MCMC lens.
4.2 ABLATION: IS BLOCK ALLOCATED SPATTAL MEMORY USEFUL.?
Embedding
Encoder
Prior
Test Episode
Figure 9: Left: Simplified write model directly produces readout prior from Equation 4by projecting embedding E via learned network. Right: Test negative variational lower bound (mean + 1std).
8

---

Published as a conference paper at ICLR 2021
While Figure 7 demonstrates the advantage of having low dimensional sampling distribution and Figure d the benefit of iterative inference, is unclear whether the performance benefit in Table is achieved from the episodic training, model structure. optimization procedure memory allocation scheme. To isolate the cause of the performance benefit, we simplify the write architecture from Section 3.3 shown in Figure Leftt In this scenario, we produce the learned memory readout, 2, via an equivalently sized dense model that projects the embedding, E, while keeping all other aspects the same. We train both models five times with the exact same TSM-ResNet18 encoder, decoder, optimizer and learning rate scheduler. As shown in Figure Right, the test conditional variational lower bound of the K++ model 20.6 nats/image better than the baseline model for the evaluated binarized Omniglot dataset. This confirms that the spatial, block allocated latent memory model proposed in this work useful when working with image distributions. Future work will explore this dimension for other modalities such as audio and text.
4.3
ABLATION: EPISODE LENGTH (T) AND MEMORY READ STEPS (K).
vafiational owe bound (natshmege)
44
42-
40-
30
0.28
0.24
36
16
52
Number of reads 0C)
30 40
Figure 10: Binarized MNIST. Left: Episode length (T) ablation showing negative test conditional variational lower bound (meanttstd) Right: Memory read steps (K) ablation showing test KL divergence.
To further explore K++, we evaluate the sensitivity of the model to varying episode lengths (T) in using the binarized MNIST dataset. We train K++ five times (each) for episode lengths ranging from 5 to 64 and observe that the model performs within margin of error for increasing episode lengths, producing negative test conditional variational bounds within a 1-std of +0.625 nats/image. This suggests that for the specific dimensionality
Figure 10-left and memory read steps (K) in Figure [ight I
memory (IR 64x6444 used in this experiment, K++ was able to successfully capture the semantics of the binarized MNIST distribution. We suspect that for larger datasets this relationship might not necessary hold and that the dimensionality of the memory should scale with the size of the dataset, but leave the prospect of such capacity analysis for future research.
While ablating the number of memory reads (K) in Figure 10ght, we observe that the total test KL-divergence varies by 1-std of +-0041 nats/image for range of memory reads from 64. A lower KL divergence implies that the model is able to better fit the approximate posteriors To (Z and 9s (X)) to their correspondings priors in Equation 4 should however be noted that lower KL-divergence does not necessary imply better generative model [Theis et al. (2016). While qualitatively inspecting the generated samples, we observed that K++ generated more semantically sound generations at lower memory read steps. We suspect that the difficulty of generating realistic samples increases with the number of disjoint reads and found that K = produces high quality results. We use this value for all experiments in this work.
CONCLUSION
In this work, we propose a novel block allocated memory in a generative model framework and demonstrate its state-of-the-art performance on several memory conditional image generation tasks. We also show that stochasticity in low-dimensional spaces produces higher quality samples in comparison to high-dimensional latents typically used VAEs. Furthermore, perturbations to the low-dimensional key generate samples with high variations. Nonetheless, there are still many unanswered questions: would hard attention based solution to differentiable indexing prove to be better than a spatial transformer? What is the optimal upper bound of window read regions based on the input distribution? Future work will hopefully be able to address these lingering issues and further improve generative memory models.

---

Published as a conference paper at ICLR 2021
REFERENCES
David Aldous. Exchangcability and related topics. In École d' d'ééé de Probabilités de Saint-Flow XIII-193,, pp. 1-198. Springer, 1985.
Georgios Arvanitidis, Lars Kai Hansen, and Søren Hauberg. Latent space oddity: On the curvature of deep generative models. In 6th International Conference on Learning Representations, ICLR 2018, 2018
Philip Bachman. An architecture for deep. hierarchical generative models. In Daniel D. Lee, Masashi Sugiyama, Ulrike Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Conference Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 4826 4834, 2016. URL
nips. cc/paper/
David Barber and Felix V Agakov. Information maximization in noisy channels: A variational approach. In Advances in Neural Information Processing Systems, pp. 201-208, 2004.
2004.
Jörg Bomschein, Andriy Mnih, Daniel Zoran, and Danilo Jimenez Rezende. Variational memory addressing in generative models. In Advances in Neural Information Processing Systems, pp- 3920-3929, 2017.
Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL
httpg//bs/150... 00519
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. ICLR, 2017.
Yaxiang Fan, Gongjian Wen, Deren Li, Shaohua Qiu, Martin D Levine, and Fei Xiao. Video anomaly detection and localization via gaussian mixture fully convolutional variational autoencoder. Computer Vision and Image Understanding, pp. 102920, 2020.
Marco Fraccaro, Danilo Jimenez Rezende, Yori Zwols, Alexander Pritzel, S. M. Ali Eslami, and Fabio Viola. Generative temporal models with spatial memory for partially observed environments. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssunn Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1544-1553, PMLR, 2018. URL httt:://proceed.nnt
mlr. poress/vvU fraccarol html
Jason M Gold, Richard F Murray, Allison B Sekuler, Patrick J Bennett, and Robert Sekuler. Visual memory decay deterministic. Psychological Science, 16(10):799-77, 2005.
Anirudh Goyal Alias Parth Goyal, Alessandro Sordoni, Marc-Alexandre Côté, Nan Rosemary Ke, and Yoshua Bengio. Z-forcing: Training stochastic recurrent networks. In Advances in neural information processing systems, pp. 6713-6723, 2017a
Prasoon Goyal, Zhiting Hu, Xiaodan Liang, Chenyu Wang, and Eric P Xing. Nonparametric variational auto-encoders for hierarchical representation learning. In Proceedings of the IEEE International Conference Computer Vision, pp. 5094-5102, 2017b.
Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: training imagenet in hour. CoRR, abs/1706.02677,, 2017c. URL htttp::/aaxii..orr/bbb////
02677
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint 1rX1v.1400.4401, 2014.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwifske Barwinska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626) 2016.
10

---

Published as a conference paper at ICLR 2021
Alex Graves, Jacob Menick, and Aaron den Oord. Associative compression networks for representation learning. arXiv preprint 1804. 02476, 2018.
Karol Gregor, Ivo Danihelka, Alex Graves, Danilo Rezende, and Daan Wierstra. Draw: A recurrent neural network for image generation. In International Conference on Machine Learning, pp. 1462-1471, 2015.
Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taiga, Francesco Visin, David Vázquez, and Aaron C. Courville. Pixelvae: A latent variable model for natural images. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. httts:::: /ttt:ss://openeeieess net/Eon Beet//oouunddd3JJKV..
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 770-778, 2016.
Donald Olding Hebb. The organization of behavior: neuropsychological theory. Wiley; Chapman & Hall, 1949.
Andrew Hollingworth, Michi Matsukura, and Steven Luck. Visual working memory modulates rapid eye movements to simple onset targets. Psychological science, 24(5):790-796, 2013.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8)22554 1982.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis R. Bach and David M. Blei (eds.), Proceedings the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 448 456. JMLR.org, 2015. URL httt::/pprcceedi ngs ml press/ wttpss/////offffffoooo html
Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. In Advances in neural information processing systems, pp. 2017-2025, 2015.
Pentti Kanerva. Sparse distributed memory. MIT press, 1988.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of the 3rd International Conference on Learning Representations (ICLR), 2014.
Durk P. Kingma and Max Welling. Auto-encoding variational bayes. ICLR, 2014.
Durk Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling- Improved variational inference with inverse autoregressive flow. In Advances in neural information processing systems, pp. 4743-4751, 2016.
Thomas N. Kipf and Max Welling, Variational graph auto-encoders. CoRR, abs/161 1.07308, 2016. URL http:://axxxi. wtt.://orgg/bbb///./oo// /tt.httpp:///aaii..oorraaab//bbs
07308
Dmitry Krotov and John J Hopfield. Dense associative memory for pattern recognition. In Advances in neural information processing systems, pp. 1172-1180, 2016.
Dharshan Kumaran, Demis Hassabis, and James L McClelland. What learning systems do intelligent agents need? complementary learning systems theory updated. Trends in cognitive sciences, 20 (7):512-534, 2016.
Hung Le, Truyen Tran, Thin Nguyen, and Svetha Venkatesh. Variational memory encoder-decoder. In Advances in Neural Information Processing Systems, pp. 1508-1518, 2018.
Ji Lin, Chuang Gan, and Song Han. Tsm: Temporal shift module for efficient video understanding. In Proceedings of the IEEE International Conference on Computer Vision, pp. 7083-7093, 2019.
Hanxiao Liu, Andrew Brock, Karen Simonyan, and Quoc V. Le. Evolving nomaaiinnnnmmmliizttionnaaatiit layers. CoRR, abs/2004.02967, 2020. Utt.tttpss//arrs...r////////...0

---

Published as a conference paper at ICLR 2021
Iya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference on Learning Representations. ICLR 2017. Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. htLL:ttts///ppeeeeee
net/forum? satrrummid=skk89.... d= 9Scxx
James Lucas, George Tucker, Roger B. Grosse, and Mohammad Norouzi. Understanding posterior collapse in generative latent variable models. In Deep Generative Models for Highly Structured Data, ICLR 2019 Workshop, New Orleans, Louisiana, United States, May 6, 2019. OpenReview.net, 2019. URL
Xuezhe Ma, Chunting Zhou, and Eduard Hovy. Mae: Mutual pnsterior-divergence regularization for variational autoencoders. In International Conference on Learning Representations, 2018,
Adam Marblestone, Yan Wu, and Greg Wayne. Product kanerva machines: Factorized bayesian memory. Bridging Al and Cognitive Science Workshop. ICLR, 2020.
Simon Marlow, Tim Harris, Roshan James, and Simon Peyton Jones. Parallel generationallcopying garbage collection with a block-structured heap. In Proceedings of the 7th international symposium on Memory management, pp. 1-220. ACM, 2008.
James L McClelland, Bruce L McNaughton, and Randall C O'Reilly. Why there are complementary learning systems the hippocampus and neocortex: insights from the successes and failures of connectionist models of learning and memory. Psychological review, 102(3):419, 1995.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. net/ formm?? d=B1 QRgziT
OpenReview.net, 2018. URL httpt:://///
openrevieeeeeee /
Ashvin Nair, Vitchyr Pong. Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa- Besa-Bianchi, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, 3-8 December 2018, Montréal, Canada, pp. 9209-9220, 2018. URL http::// nips. wtppscccpappe. 1022-visaal-aa reinfrreeee learning with imagined goals
Vinod Nair and Geoffrey F. Hinton. Rectified linear units improve restricted boltzmann machines. In Johannes Fürnkranz and Thorsten Joachims (eds.), Proceedings of the 27th International Conference on Machine Learming (ICML-10), June 21-24, 2010, Haifa, Israel, pp. 807-814. Omnipress, 2010. URL https: /ttp:///cmmma..mm cc (ecnfferencee/200//ppppeeee
pdf,
Tony A Plate. Holographic reduced representations. IEEE Transactions on Neural networks, 6(3): 623-641, 1995.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomènech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In Doina Precup and Yee whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 2827-2833 PMLR, 2017. URL /tttpe:/proccedingg.. mlr. Dxress/v0/pp. te 7a html
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- May 3, 2018, Workshop Track Proceedings. OpenReview.net, 2018. URL https: //openreview. net wttt:/oorrm...-Huu222
Michael W. Reimann, Max Nolte, Martina Scolamiero, Katharine Turner, Rodrigo Perin, Giuseppe Chindemi, Pawel Diotko, Ran Levi, Kathryn Hess, and Henry Markram. Cliques of neurons bound into cavities provide missing link between structure and function. Frontiers in
Computational Neuroscience, 1 11:48, 2017. ISSN 1662-5188. doi: 10.3389/fncom.2217.0002 0.0048. 3389/ frccom 2017.0000
htl.ttt:: //ww. front ieponiierseee www.srssiinnngrrttiiiiiooooo
12

---

Published as a conference paper at ICLR 2021
Hossein Sadeghi, Evgeny Andriyash, Walter Vinci, Lorenzo Buffoni, and Mohammad H. Amin. Pixelvae++: Improved pixelvae with discrete prior. CoRR, Ses/1998.09948,, 2019. URL http:
//axxxv.crg/abs///// a9tp//////a//9...... 09948
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. Pixelen+++: Improving the
pixelenn with discretized logistic mixture likelihood and other modifications. In 5th International
Conference on Learning Representations. ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL
net/
Forum? FFuum?dd=BJrFFGcey
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy P. Lillicrap. Meta-leaming with memory-augmentee neural networks. In Maria-Florins Balcan and Kilian Q. Weinberger (eds.). Proceedings of the 33nd International Conference on Machine Learning, ICMI. 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp- 1842-1850. JMIR.org, 2016. URL mlr. wtt.ppee/////astorooooo..... html
Suman Sedai, Dwarikanath Mahapatra, Sajini Hewavitharanage, Stefan Maetschke, and Rahil Garnavi. Semi-supervised segmentation of optic cup in retinal fundus images using variational autoencoder. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 75-82. Springer, 2017.
Leslie N. Smith and Nicholay Topin. Super-convergence: Very fast training of residual networks using large learning rates. CoRR, Abs/1708.00122 2017. URL httt::///rxi...rg/////
1708.077 07120
John P Spencer and Alycia M Hund. Prototypes and particulars: Geometric and snperience-dependenn spatial categories. Journal of Experimental Psychology: General, 131(1): 16, 2002.
Lucas Theis, Aäron van den Oord, and Matthias Bethge. A note on the evaluation of generative models. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URLptttp://axxx....
abs 01844
Jakub Tomezak and Max Welling. Vae with a vampprior. In International Conference on Artificial Intelligence and Statistics, pp. 1214-1223, 2018.
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. In Yoshua Bengio and Yann LeCun (eds.) ), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http://arxiv.orrt
abs/ 3916
Tom White. Sampling generative networks. arXiv preprint 1rXiv:1609.04466 2016.
Yan Wu, Greg Wayne, Alex Graves, and Timothy Lillicrap. The kanerva machine: A generative distributed memory. ICLR, 2018a.
Yan Wu, Gregory Wayne, Karol Gregor, and Timothy Lillicrap. Learning attractor dynamics for generative memory. In Advances in Neural Information Processing Systems, pp. 9379-9388, 2018b.
Yuxin Wu and Kaiming He. Group normalization. Inf. J. Comput. Vis., 128(3):742-755, 2020. doi: 10.1077/412233-----11199 s11263- 019-01198-111
URL. httpps://di http:////ooooorg////o/ 1ttpe/iioro1100...0..0.100....
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arxiv: 03888, 2017.
Qingyu Zhao, Nicolas Honnorat, Ehsan Adeli, Adolf Pfefferbaum, Edith V Sullivan, and Kilian M Pohl. Variational autoencoder with truncated mixture of gaussians for functional connectivity analysis. In International Conference on Information Processing in Medical Imaging, pp- 867-879. Springer, 2019.
13

---

Published as a conference paper at ICLR 2021
Guangxiang Zhu, Zichuan Lin, Guangwen Yang, and Chongjie Zhang. Episodic reinforcement learning with associative memory. In International Conference on Learning Representations, 2019.
14

---

Published as a conference paper at ICLR 2021
A SPATIAL TRANSFORMER REVIEW
Indexing matrix, M + Az, y + Ay), is typically non-differentiable operation since involves hard cropping around an index. Spatial transformers (Jaderberg al., 2015) provide solution to this problem by decoupling the problem into two differentiable operands:
1. Learn an afline transformation of coordinates.
2. Use differntiable bilinear transformation.
The affine transformation of source coordinates,
to target coordinates,
is defined as:
2
0
2
y/u
0
ja
(8)
9
y/o
y/2
Here, the affine transform, 0
0 S has three learnable scalars: {s, 25, y) which define
scaling and translation the components of the key sample, {yo. 31, 322 I2 shown in Equation g After transforming the co-ordinates (not to be confused with the actual data), spatial transformers learn differentiable bilinear transform which can be interpreted learning differentiable mask that is element-wise multiplied by the original data, M:
and j respectively. In the case of K++, these three scalars represent
Σ
M max(0, 1 1 r|) max(0, 1 -15h1 - nl)
0.5
0
0.3
Consider the following example where 0
this parameterization differntiably
0
0.5
0.5
extracts the region shown in Figure
Right from Figure
Left:
100
20
300
30
400
500
700
10
20
30
10
Figure 11: Spatial transformer example. Left: original image with region inlaid. Right: extracted grid.
The range of values for (8, is bound between [-1, 1], where the center of the image is (0, 0).
15

---

Published as a conference paper at ICLR 2021
B CELEB-A A GENERATIONS
Figure 12: Random key Celeb-A generations.
We present random key generations of Celeb-A 64x64, trained without center cropping in Figure
C VAE VS. K++ IMAGENET32X32 GENERATIONS
Figure 13: ImageNet32x32 generations. Left: VAE; Right: K++.
Figure 13 shows the difference in generations of standard VAE K++. In contrast to the
standard VAE generation (Figure Left, the K++ generations (Figure 13 Right) appear much sharper, avoiding the blurry generations observed with standard VAEs.
16

---

Published as a conference paper at ICLR 2021
D
TEST IMAGE RECONSTRUCTIONS
P
P
::
6
7
3
4
7
0
3 1
G T h 2 2 R 4 L 2
d
4
2
4
/
0
{
8
5
R
J
9
8
4
4
7
4
/ 3 1 3 3
$ 6 4 0
5
8
9
3
6
8
6
U
O
2
O
0
L
9
41
HL
4
5
8
3
P d
:4
LA
7
3
1
9
8
0
T
C
2
4
O
8
5 5
20
2
P
2
6
12
9
9
4
4
7
4
/
D
5
18
9
#
6
8
6
3 \
6 4
O
2
0
0
M
O
4
9
5
8
3 3
O
Figure 14: Binarized test reconstructions; top row true samples. Left: Omniglot; Right: MNIST.
Figure 15; Test set reconstructions; top row are true samples. Left: ImageNe164x64, Right: DMLab Mazes.
17

---

Published as a conference paper at ICLR 2021
ing
Figure 16: Test set reconstructions; top row are true samples. Celeb-A 64x64.
E MODEL ARCHITECTURE & TRAINING
Encoder: As mentioned in Section 3
TSM-Resnetll 18 encoder with Batchnorm (loff &
Szegedy] 2015) and ReLU activations (Nir & Hinton 2010) for all tasks. We apply a fractional shift of the feature maps by 0. 0.125 as suggested by the authors.
18

---

Published as a conference paper at ICLR 2021
Decoder: Our decoder is simple conv-transpose network with EvoNormSO Liu et al. 2020
inter-spliced between each layer. EvonormsO is similar in stride to Groupnorm combined with the swish activation function (Ramachandran al. 2018
Wu & He 2020
Optimizer & LR scheule: We use LARS YYu et al. 2017) coupled with ADAM (Kingma & Ba
2014) and one-cycle (Smith & Topin 2017) cosine learning rate schedule (Loshchilov & Hutter 2017 A linear warm-up of 10 epochs (Goyal et al. 2017c) is also used for the schedule. weight decay of le is used on every parameter barring biases and the affine terms of batchnorm. Each task is trained for 500 or 1000 epochs depending on the size of the dataset.
Dense models: All dense models such our key network simple three layer deep linear dense models with a latent dimension of 512 coupled with spectral normalization (Miyato et al. 20188).
femen uses a deep linear conv-transpose decoder on the pooled embedding, E with a base feature map projection size of 256 with division by 2 per layer. We use memory size of 3x364444 for all the experiments in this work.
Memory writer:
Learned Prior: PM Y) uses convolutional encoder that stacks the K read traces, (M, bt)) along the channel dimension and projects it to the dimensionality of Z.
In practice, we observed that K++ is about 2x as fast (wall clock) compared to our re-implementation of DKM. We mainly attribute this to not having to solve an inner OLS optimization loop for memory inference.
E. MEMORY CREATION PROTOCOL
The memory creation protocol of K++ is similar in stride to that of the DKM model, given the deterministic relaxations and addressing mechanism described in Sections 3.3 and 3.4 Each memory, M m[ [een(())) is function of an episode of samples, X { D. To efficiently optimize the conditional lower bound in Equation we parallelize the learning objective using set of minibatches, as typical with the optimization of neural networks. As with the DKM model, K++ computes the train and test conditional evidence lower bounds in Table by first inferring the memory, M ~ 8[fmem [fenc(X] (X))] from the input episode, followed by the read out procedure as described in Section 3.5
19


Reviews:
[{"review_id": "mW5wlSBj5bS", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a new memory mechanism based on the Kanerva Machine inspired by computer heap allocation. It stores the knowledge with several segments, which can be shared between data. The machine stores the knowledge more efficiently through the sharable part-based system and the authors simplified the writing mechanism by designing the memory deterministic.\n\nStrengths:\n1. They designed a new Kanerva Machine having a simplified writing mechanism and sharable part-based memory.\n2. They evaluated their method on several datasets (e.g., DMLab Mazes or CIFAR10) and showed better performance than previous works.\n3. They analyzed local key permutation and ablation study with the version without the memory.\n\nWeaknesses:\n1. I think that section 3.3 requires the additional description. As I understand, the memory is updated directly from the input episode, X through f_enc and f_mem, then where is the inferred key used? in the paper, the authors mentioned the key model is described in Section 3.4, but I can't find the description in the section. I think that it is used in the Read inference model, but it is not clearly described. \n2. The lack of analysis: this model uses part-based memory, then what happens when the K is decreasing or increasing? I think that the ratio between memory size and the K can be important, when the ratio (K/memory size) is increasing, what happens? What happens with a simplified write mechanism? (I saw you mentioned that you observed K++ is about 2 times faster than DKM, but it is because of the writing mechanism?)\n\nThe correctness of their claim and Clarity:\n\nThis paper is well-written and correct I think, but if the section 3.3 and 3.5 was written more kindly, then it would be better.\n\nAdditional feedback:\n\nThank you for submitting, I enjoyed reading. Basically, I think the idea of this paper (simplified write mechanism and sharable part-based memory) is good enough and you evaluated it well. However, as I mentioned, for me, some sections are hard to follow and some interesting analyses are not included. And, the discussions about the down-stream tasks like how can we use it for other tasks like model-based RL? also make this paper more concrete.\n\nMinor things are\n\nOn page5, \nequation 6, R^{T x 3} is right? not T x K x 3?\nfigure 4, infer latent part, \\mu_{\\theta_{\\hat{M}}}(E) and \\sigma^2_{\\theta_{\\hat{M}}}(E) are right? there is no E in that part.\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "gy3wUQnkBC", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe authors develop a hierarchical Bayesian memory allocation scheme to bridge the gap between episodic and semantic memory via a hierarchical latent variable model. They take inspiration from traditional heap allocation and extend the idea of locally contiguous memory to the Kanerva Machine, enabling a novel differentiable block allocated latent memory. In contrast to the Kanerva Machine, the authors simplify the process of memory writing by treating it as a fully feed forward deterministic process, relying on the stochasticity of the read key distribution to distribute information within the memory.\n\nPros:\n- The authors combine the idea of differentiable indexing in Spatial Transformer (Jaderberg et al., 2015) into the memory of Kanerva Machine (Wu et al, 2018a;b) and prove by experiments that this allocation scheme on the memory helps improve the test negative likelihood. Also, its speed is about 2x faster than the Dynamic Kanerva Machine. \n- The authors show the efficiency of Temporal Shift Module (TSM) (Lin et al., 2019) in the encoder of memory models. Replacing a standard convolutional stack by TSM improves the ELBO in Dynamic Kanerva Machine by 6.32 nats/image for the Omniglot dataset.\n- The experiments are well reported for different tasks, such as reconstruction and generation, on various datasets.\n\nCons:\n- The whole article is just like a mechanical mixture of old ideas. To be specific, the K++ model is Kanerva Machine + Spatial Transformer + a powerful encoder (namely, Temporal Shift Module). The authors do not introduce any significant improvement or novel insight for old models and techniques.\n- Is there any theory support for the idea that we should use an allocated deterministic memory instead of a full variational memory? The authors mention the theory of complementary learning system in the Abstract and heap allocation at the beginning of Section 1, but there is no further analysis for these two theoretical intuitions. \n\nComments and questions:\n- The basic idea is clear and reasonable, but the authors should provide deeper analysis as well as deeper insights for their new model (and for old models that they use, if possible).\n- The authors use q_phi(Z|X) in the ELBO (eq. 4) instead of q_phi(Z|X, Y, M) as in the Kanerva Machine. How is the memory used in the read model?\n- Where do the results in Table 1 come from? For example, in Kanerva Machine and Dynamic Kanerva Machine paper (Wu et al, 2018a;b), the authors did not report the negative likelihood for CIFAR10 dataset.\n- Is there any explanation for the significantly low negative likelihood (-2344.5 bits/dim) of K++ for CIFAR10 dataset?\n- The authors should include a brief introduction section about Kanerva Machine and Dynamic Kanerva Machine. Moreover, the function δ in eq. 5 is not defined beforehand.\n\n\nREFERENCES\nJames L McClelland, Bruce L McNaughton, and Randall C O’Reilly. Why there are complementary learning systems in the hippocampus and neocortex: insights from the successes and failures of connectionist models of learning and memory. Psychological review, 102(3):419, 1995.\nMax Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. In Advances in neural information processing systems, pp. 2017–2025, 2015.\nYan Wu, Greg Wayne, Alex Graves, and Timothy Lillicrap. The Kanerva machine: A generative distributed memory. ICLR, 2018a.\nYan Wu, Gregory Wayne, Karol Gregor, and Timothy Lillicrap. Learning attractor dynamics for generative memory. In Advances in Neural Information Processing Systems, pp. 9379–9388, 2018b.\nJi Lin, Chuang Gan, and Song Han. TSM: Temporal shift module for efficient video understanding. In Proceedings of the IEEE International Conference on Computer Vision, pp. 7083–7093, 2019.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "a_aUHGkdIsh", "reviewer": "AnonReviewer2", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a generative memory modeling method.  Specifically, this paper proposed a novel memory allocation scheme,  replacing the stochastic memory writing process in prior works KM[1]  and  DKM[2] with a set of deterministic operations. This deterministic process is implemented as the spatial transformation[3] on the pooled embedding space of input data.  The authors showed the proposed model achieved state-of-the-art results on the binarized MNIST data set and the binarized Omniglot data set.  The authors also showed that introducing the work in TSM[4] helps learn a richer context representation compared with 2D convolution layers.\n\nSummary of positive:\nThe method integrated with TSM and showed improvement on DKM.\nThe quantitative experiment result shows a great improvement on previous work.\n\n\nSummary of negative:\nThe method section is not very clear to me, e.g. how is the key used in the writing process. The reading inference process seems not involved with the memory (section3.5 and figure 5), which is not consistent with the graphic model in figure 2. \nTable 1 shows the result of one of the key experiments, however, it is not well-formatted, and some measurements are not well aligned or not completed.  For example, the result on CIFAR10 is reported as nats/img while the baseline models report bits/dim. \nThe baseline model in the ablation study is too weak. Is it possible to compare with KM with multiple addressing keys? The description of the baseline model is not clear, from the text, it says the learned memory prior (p(z|y,M)?) was replaced with a dense layer. While compared with the figure in section 3.3, it seems the memory writing (f_{mem}) was replaced.\n\n\nOthers:  \nHow is the value K set, and how is it affecting the performance?  \nIs there any analysis on how the possibly overlapped sub-regions provide a better-compressed representation than the sparse distributed memory?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "S3Fl-PP3Dy", "reviewer": "AnonReviewer3", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a generative memory (K++) that takes inspiration from Kanerva Machine and heap memory allocation. An episode of images is encoded into keys and the memory using Temporal Shift Module and transposed convolution. Given the memory and the keys,  a Spatial Transformer reads multiple contiguous memory blocks, which are then used to generate samples using a transposed convolutional decoder.  The model is trained by maximizing the ELBO of the conditional lnP(X|M). The experiments show that K++ can achieve competitive performance on various image datasets.\n\nPros:\n- The idea of using Spatial Transformer to simulate heap memory allocation is interesting\n- The performance of the proposed model is very  good for binarized image datasets\n\nCons:\n- The writing is sometimes confusing, hard to follow\n- Some details  are missing\n\nDetailed comments and questions:\n- The abstract mentions complementary learning system and hierarchical Bayesian memory allocation. However, in the main text, no detail is given to elaborate on these points. How do these concepts embody in the proposed model?\n- K++ architecture is related to generative memory models with deterministic memory (especially those that have multiple keys [1,2]). Please make a clear distinction between K++ and those works. \n- Eq. 4, How is $q_\\phi(Z|X)$ implemented? Fig. 5, How does the memory trace contribute to the read model as mentioned in Sec. 3.4? During training, how are the generative/read models used?\n- Eq.6 Writing  $R^{T\\times3}$ is inappropriate since $Y_t$ contains $y_{tk}$\n- Table 1, reporting two metrics in the same column is misleading\n- Table 1, any explanation for the reported values of K++ for CIFAR10? The numbers look strange and use different metric (nats/img), which makes it hard to compare with other baselines.\n- Fig. 7, it seems that your model can resist against local perturbation of keys. It is better if you can show that this property is unique to K++. E.g., other models such as Kanerva Machine fails to do so.\n- Sec. 4.2 only validates the important role of the memory. To verify the benefit of block memory allocation, you should compare it with other memory indexing schemes. The simplest baseline is to use the whole memory: $\\hat{M} = M$ or to read from M using  multi-head slot-based attention (see [1,2]).\n- What are the values of T and K used in the experiments?\n- Possible writing/syntax error: \nin abstract \"heirarchar\" -->\"hierarchical \";\nFig. 4's equations, $Y_t\\sim P(Y_t)$ --> $y_t\\sim P(Y_t)$;\nFig. 4's equations, given only the memory and keys, where does E come from? It is better to draw Z in the generative model picture; \npage 4, Jaderberg et al. (2015) --> (Jaderberg et al., 2015);\npage 7,  Burda et al. (2016) --> (Burda et al., 2016);\n\nI like the idea of the paper. However, I feel that the writing is a bit rushed and need revision. Hence, given the current manuscript, I can only give a marginal score. If the authors can address my concerns, improve the writing, explain unclear parts, I will raise my score. \n\n[1]  Mevlana Gemici, Chia-Chun Hung, Adam Santoro, Greg Wayne, Shakir Mohamed, Danilo J Rezende, David Amos, and Timothy Lillicrap. Generative temporal models with memory. arXiv preprint arXiv:1702.04649, 2017.\n\n[2] Hung Le, Truyen Tran, Thin Nguyen, and Svetha Venkatesh. \"Variational memory encoder-decoder.\" In Advances in Neural Information Processing Systems, pp. 1508-1518. 2018.\n\n\n\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a variant of Kanerva Machine Wu et al. (2018) by introducing a spatial transformer to index the memory storage and Temporal Shift Module Lin et al., (2019). The KM++ model learns to encode an exchangeable sequence locally via the spatial transformer. The proposed method is evaluated on conditional image generation tasks. The empirical results demonstrated the nearby keys in the memory encoded related and similar images. Several issues of clarity and the correctness of the main theoretical result were addressed during the rebuttal period in a way that satisfied the reviewers. The basic ideas in the paper are interesting to both the machine learning and the wider cognitive science communities. However, additional experiments should be included in Table 1 to complete the "DKM w/TSM (our impl)" row on Fashion MNIST, CIFAR-10, and DMLab in the final revision for completeness. 
Decision: Accept (Poster)