Original: K ANERVA
Revised: K ANERVA

Original: ++:
Revised: ++:

Original: E XTENDING THE K ANERVA M ACHINE
Revised: E XTENDING THE K ANERVA M ACHINE

Original: W ITH
Revised: W ITH

Original: D IFFERENTIABLE ,
Revised: D IFFERENTIABLE ,

Original: L OCALLY
Revised: L OCALLY

Original: B LOCK
Revised: B LOCK

Original: A LLOCATED L ATENT M EMORY
Revised: A LLOCATED L ATENT M EMORY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Episodic and semantic memory are critical components of the human memory model.
Revised: Episodic and semantic memory are critical components of the human memory model.

Original: The theory of complementary learning system (McClelland et al., 1995) suggests that the compressed representation produced by a serial event (episodic memory) is later re-structured to build a more generalized form of re-usable knowledge (semantic memory).
Revised: The theory of complementary learning system (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).

Original: In this work we develop a new principledhierarchical Bayesian memory allocation scheme that bridges the gap between episodic and semantic memory via a heirarchal latent variable model.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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 distribute information within the memory.
Revised: 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.

Original: We demonstrate that this allocation scheme improves performance in conditional image generation, resulting in new state-of-the-art likelihood values on binarized MNIST ( 45.64 nats/image ), binarized Omniglot ( 71.64 nats/image ) as well as presenting competitive performance on CIFAR10, DMLab Mazes, Celeb-A and ImageNet32 × 32.
Revised: We demonstrate that this allocation scheme improves performance in conditional image generation, resulting in new state-of-the-art variational lower bounds on binarized MNIST ( 41.58 nats/image ) , binarized Omniglot ( 66.24 ), as well as presenting competitive performance on CIFAR10, DMLab Mazes, Celeb-A and ImageNet32 × 32.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Memory is a central tenant in the model of human intelligence and is crucial to long term reasoning and planning.
Revised: Memory is a central tenant in the model of human intelligence and is crucial to long-term reasoning and planning.

Revised: Of particular interest is the theory of complementary learning systems

Revised: McClelland et al.

Revised: (1995) which proposes that the brain employs two complementary systems to support the acquisition of complex behaviours: a hippocampal fast learning system that records events as episodic memory, and a neocortical slow learning system that learns statistics across events as semantic memory.

Revised: While the functional dichotomy of the complementary systems are well-established McClelland et al.

Revised: ; Kumaran et al.

Original: Recent work has shown that  memory augmented neural networks  drastically improve the performance of generative models (Wu et al., 2018a;b), language models (Weston et al., 2015), meta-learning (Santoro et al., 2016), long term planning (Graves et al., 2014; 2016) and sample efﬁciency in reinforcement learning (Zhu et al., 2019).
Revised: , 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. Traditional heap allocated memory affords O (1) free / malloc computational complexity and serves as inspiration for K++ which uses differentiable neural proxies. While recent work has shown that using memory augmented neural networks can drastically improve the performance of generative models (Wu et al., 2018a;b), language models (Weston et al., 2015), meta-learning (Santoro et al., 2016), long-term planning (Graves et al., 2014; 2016) and sample efﬁciency in reinforcement learning (Zhu et al., 2019), no model has been proposed to exploit the inherent multi-dimensionality of biological memory.

Original: Inspired by the traditional (computer-science) memory model of heap allocation (Figure 1- Left ), 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 { p 1 , ..., p 4 } in Figure 1, into a latent multi-dimensional memory (Figure 1- Right ).
Revised: Inspired by the traditional (computer-science) memory model of heap allocation (Figure 1- Left ), 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 { p 1 , ..., p 4 } in Figure 1, into a latent multi-dimensional memory (Figure 1- Right ).

Original: The K++ model infers a key distribution as a proxy to the pointers used in Marlow et al.
Revised: The K++ model infers a key distribution as a proxy to the pointers used in Marlow et al.

Original: and is able to embed similar samples to an overlapping latent representation space, thus enabling it to be more efﬁcient on compressing input distributions.
Revised: and is able to embed similar samples to an overlapping latent representation space, thus enabling it to be more efﬁcient on compressing input distributions.

Original: In this work we focus on 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).
Revised: In this work, we focus on 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).

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Variational Autoencoder Variational autoencoders (VAE) (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, 2016), medical applications (Sedai et al., 2017; Zhao et al., 2019) and video analysis (Fan et al., 2020).
Revised: Variational Autoencoder Variational autoencoders (VAE) (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, 2016), medical applications (Sedai et al., 2017; Zhao et al., 2019) and video analysis (Fan et al., 2020).

Original: As a latent variable models, VAEs infer an approximate posterior over a latent representation Z and can be used in downstream tasks such as control in reinforcement learning (Nair et al., 2018; Pritzel et al., 2017).
Revised: As a latent variable models, VAEs infer an approximate posterior over a latent representation Z and can be used in downstream tasks such as control in reinforcement learning (Nair et al., 2018; Pritzel et al., 2017).

Original: VAEs maximize an evidence lower bound (ELBO), L ( X, Z ) , of the log-marginal likelihood, ln p ( X ) >
Revised: VAEs maximize an evidence lower bound (ELBO), L ( X, Z ) , of the log-marginal likelihood, ln p ( X ) >

Original: L ( X, Z ) =
Revised: L ( X, Z ) =

Original: ln p θ ( X | Z ) −
Revised: ln p θ ( X | Z ) −

Original: D KL ( q φ ( Z | X ) || p
Revised: D KL ( q φ ( Z | X ) || p

Original: θ ( Z )) .
Revised: θ ( Z )) .

Original: The produced variational approximation, q φ ( Z | X ) , is typically called the encoder, while p θ ( X | Z ) comes from the decoder.
Revised: The produced variational approximation, q φ ( Z | X ) , is typically called the encoder, while p θ ( X | Z ) comes from the decoder.

Original: Methods that aim to improve these latent variable generative models typically fall into two different paradigms: learning more informative priors or leveraging novel decoders.
Revised: Methods that aim to improve these latent variable generative models typically fall into two different paradigms: learning more informative priors or leveraging novel decoders.

Original: While improved decoder models such as PixelVAE (Gulrajani et al., 2017) and PixelVAE++ (Sadeghi et al., 2019) drastically improve the performance of p θ ( X | Z ) , they suffer from a 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., 2017a).
Revised: While improved decoder models such as PixelVAE (Gulrajani et al., 2017) and PixelVAE++ (Sadeghi et al., 2019) drastically improve the performance of p θ ( X | Z ) , they suffer from a 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., 2017a).

Original: In contrast, VampPrior (Tomczak & Welling, 2018), Associative Compression Networks (ACN) (Graves et al., 2018), VAE-nCRP (Goyal et al., 2017b) and VLAE (Chen et al., 2017) tighten the variational bound by learning more informed priors.
Revised: In contrast, VampPrior (Tomczak & Welling, 2018), Associative Compression Networks (ACN) (Graves et al., 2018), VAE-nCRP (Goyal et al., 2017b) and VLAE (Chen et al., 2017) tighten the variational bound by learning more informed priors.

Original: VLAE for example, uses a powerful auto-regressive prior; VAE-nCRP learns a non-parametric Chinese restaurant process prior and
Revised: VLAE for example, uses a powerful auto-regressive prior; VAE-nCRP learns a non-parametric Chinese restaurant process prior and

Original: VampPrior learns a Gaussian mixture prior representing prototypical virtual samples.
Revised: VampPrior learns a Gaussian mixture prior representing prototypical virtual samples.

Original: ACN on the other hand takes a two stage approach: by clustering real samples in the space of the posterior and  using these related samples as inputs to a learned prior, ACN provides an information theoretic alternative to improved code transmission.
Revised: On the other hand, ACN takes a 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.

Original: 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.
Revised: 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.

Original: Memory Models Inspired by the associative nature of biological memory, the Hopﬁeld network (Hopﬁeld, 1982) introduced the notion of content-addressable memory, deﬁned by a set of binary neurons coupled with a Hamiltonian and a dynamical update rule.
Revised: Memory Models Inspired by the associative nature of biological memory, the Hopﬁeld network (Hopﬁeld, 1982) introduced the notion of content-addressable memory, deﬁned by a set of binary neurons coupled with a Hamiltonian and a dynamical update rule.

Original: Iterating the update rule minimizes the Hamiltonian, resulting in patterns being stored at different conﬁgurations (Hopﬁeld, 1982; Krotov & Hopﬁeld, 2016).
Revised: Iterating the update rule minimizes the Hamiltonian, resulting in patterns being stored at different conﬁgurations (Hopﬁeld, 1982; Krotov & Hopﬁeld, 2016).

Original: Writing in a Hopﬁeld network thus corresponds to ﬁnding weight conﬁgurations such that stored patterns become attractors via Hebbian rules (Hebb, 1949).
Revised: Writing in a Hopﬁeld network thus corresponds to ﬁnding weight conﬁgurations such that stored patterns become attractors via Hebbian rules (Hebb, 1949).

Original: This concept of memory was later extended to a distributed, continuous setting in Kanerva (1988) and to a complex valued, convolutional binding mechanism by Plate (1995).
Revised: This concept of memory was later extended to a distributed, continuous setting in Kanerva (1988) and to a complex valued, convolutional binding mechanism by Plate (1995).

Original: The central difference between an associative memory model (Hopﬁeld, 1982; Kanerva, 1988) and (Plate, 1995) is that the latter decouples the size of the memory from the input word size.
Revised: The central difference between an associative memory model (Hopﬁeld, 1982; Kanerva, 1988) and (Plate, 1995) is that the latter decouples the size of the memory from the input word size.

Original: 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 a single slot.
Revised: 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 a single slot.

Original: Reading memory traces, z , entails using a vector of addressing weights, r , to attend to the appropriate column of M , z = r T M .
Revised: Reading memory traces, z , entails using a vector of addressing weights, r , to attend to the appropriate column of M , z = r T M .

Original: This paradigm of memory includes models such as the Neural Turing Machine (NTM) (Graves et al., 2014), Differentiable Neural Computer (DNC) (Graves et al., 2016) 1 , Memory Networks (Weston et al., 2015) and Variational Memory Addressing (VMA) (Bornschein et al., 2017).
Revised: This paradigm of memory includes models such as the Neural Turing Machine (NTM) (Graves et al., 2014), Differentiable Neural Computer (DNC) (Graves et al., 2016) 1 , Memory Networks (Weston et al., 2015), Generative Temporal Models with Memory (GTMM) Fraccaro et al. Variational Memory Encoder-Decoder (VMED)

Revised: Le et al.

Revised: (2018), and Variational Memory Addressing (VMA) (Bornschein et al., 2017).

Original: VMA differs from  DNC, NTM and Memory Networks by taking a stochastic approach to discrete key-addressing instead of the deterministic approach of the latter models.
Revised: 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.

Original: Recently, the Kanerva Machine (KM) (Wu et al., 2018a) and its extension, the Dynamic Kanerva
Revised: Recently, the Kanerva Machine (KM) (Wu et al., 2018a) and its extension, the Dynamic Kanerva

Original: Machine (DKM) (Wu et al., 2018b), interpreted memory writes and reads as inference in a generative model, wherein memory is now treated as a distribution, p ( M ) .
Revised: Machine (DKM) (Wu et al., 2018b), interpreted memory writes and reads as inference in a generative model, wherein memory is now treated as a distribution, p ( M ) .

Original: Under this framework, memory reads and writes are recast as sampling or updating the memory posterior.
Revised: Under this framework, memory reads and writes are recast as sampling or updating the memory posterior.

Original: The DKM model differs from the KM model by introducing a dynamical addressing rule that could be used throughout training.
Revised: The DKM model differs from the KM model by introducing a dynamical addressing rule that could be used throughout training.

Original: While providing an intuitive and theoretically sound bound on the data likelihood, the
Revised: While providing an intuitive and theoretically sound bound on the data likelihood, the

Original: DKM model requires an inner optimization loop which entails solving an ordinary least squares (OLS) problem.
Revised: DKM model requires an inner optimization loop which entails solving an ordinary least squares (OLS) problem.

Original: Typical OLS solutions require a matrix inversion ( O ( n 3 ) ), preventing the model from scaling to large memory sizes.
Revised: Typical OLS solutions require a matrix inversion ( O ( n 3 ) ), preventing the model from scaling to large memory sizes.

Original: More recent work has focused on employing a product of smaller Kanerva memories (Marblestone et al., 2020) in an effort to minimize the computational cost of the matrix inversion.
Revised: More recent work has focused on employing a product of smaller Kanerva memories (Marblestone et al., 2020) in an effort to minimize the computational cost of the matrix inversion.

Original: In contrast, we propose a simpliﬁed view of memory creation by treating memory writes as a deterministic process in a fully feed-forward setting.
Revised: In contrast, we propose a simpliﬁed view of memory creation by treating memory writes as a deterministic process in a fully feed-forward setting.

Original: 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.
Revised: 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.

Revised: While the removal of memory stochasticity might be interpreted as reducing the representation power of the model, we demonstrate through our experiments that our empirically model performs better, trains quicker and is simpler to optimize.

Revised: The choice of a deterministic memory is further reinforced by research in psychology, where human visual memory has been shown to change deterministically Gold et al.

Revised: (2005); Spencer & Hund (2002);

Revised: Hollingworth et al.

Original:  M ODEL
Revised:  M ODEL

Original: To better understand the K++ model, we examine each of the individual components to understand their role within the complete generative model.
Revised: To better understand the K++ model, we examine each of the individual components to understand their role within the complete generative model.

Original: We begin by ﬁrst deriving a variational lower bound (Section 3.1), describing the optimization objective and probabilistic assumptions.
Revised: We begin by ﬁrst deriving a variational lower bound (Section 3.1), describing the optimization objective and probabilistic assumptions.

Original: We then describe the write operand (Section 3.3), the generative process (Section 3.4) and ﬁnally the read operand (Section 3.5).
Revised: We then describe the write operand (Section 3.3), the generative process (Section 3.4) and ﬁnally the read and iterative read operands (Section 3.5).

Original: 3. P RELIMINARIES
Revised: 3. P RELIMINARIES

Original: K++ operates over an exchangeable episode (Aldous, 1985) of samples, X = { x t } T t =1 ∈ D , drawn from a dataset D , as in the Kanerva Machine.
Revised: K++ operates over an exchangeable episode (Aldous, 1985) of samples, X = { x t } Tt =1 ∈ D , drawn from a dataset D , as in the Kanerva Machine.

Original: Therefore, the ordering of the samples within the episode does not matter.
Revised: Therefore, the ordering of the samples within the episode does not matter.

Original: This enables factoring the conditional, p ( X | M ) , over each of the individual samples: (cid:81)
Revised: This enables factoring the conditional, p ( X | M ) , over each of the individual samples: (cid:81)

Original: Tt =1 p ( x t | M ) , given the memory, M ∼ p ( M ) , M ∈ R ˆ C × ˆ W × ˆ H .
Revised: Tt =1 p ( x t | M ) , given the memory, M ∼ p ( M ) , M ∈ R ˆ C × ˆ W × ˆ H .

Original: Our objective in this work is to maximize the expected conditional log-likelihood as in (Bornschein et al., 2017; Wu et al., 2018a):
Revised: Our objective in this work is to maximize the expected conditional log-likelihood as in (Bornschein et al., 2017; Wu et al., 2018a):

Original: [Equation]
Revised: [Equation]

Original: As alluded to in Barber & Agakov (2004) and Wu et al.
Revised: As alluded to in Barber & Agakov (2004) and Wu et al.

Original: (2018a), this objective can be interpreted as maximizing the mutual information, I ( X ; M ) , between the memory, M , and the episode, X , since
Revised: (2018a), this objective can be interpreted as maximizing the mutual information, I ( X ; M ) , between the memory, M , and the episode, X , since

Original: I ( X ; M ) =
Revised: I ( X ; M ) =

Original: H ( X ) + J = H ( X )
Revised: H ( X ) + J = H ( X )

Original: − H ( X ; M ) and given that the entropy of the data, H ( X ) , is constant.
Revised: − H ( X ; M ) and given that the entropy of the data, H ( X ) , is constant.

Original: In order to actualize Equation 1 we rely on a variational bound which we derive in the following section.
Revised: In order to actualize Equation 1 we rely on a variational bound which we derive in the following section.

Original: 3. V ARIATIONAL L OWER B OUND
Revised: 3. V ARIATIONAL L OWER B OUND

Original: To efﬁciently read from the memory, M , we introduce a set of latent variables corresponding to the K addressing read heads, Y =
Revised: To efﬁciently read from the memory, M , we introduce a set of latent variables corresponding to the K addressing read heads, Y =

Original: { Y t = { y tk } Kk =1 } Tt =1 , y tk ∈ R 3 , and a set of latent variables corresponding to the readout from the memory, Z = { z t } Tt =1 , z t ∈ R L .
Revised: { Y t = { y tk } Kk =1 } Tt =1 , y tk ∈ R 3 , and a set of latent variables corresponding to the readout from the memory, Z = { z t } Tt =1 , z t ∈ R L .

Original: Given these latent variables, we can decompose the conditional, ln p ( X | M ) , using the product rule and introduce variational approximations q φ ( Z | X )  and q φ ( Y | X ) via a multipy-by-one trick:
Revised: Given these latent variables, we can decompose the conditional, ln p ( X | M ) , using the product rule and introduce variational approximations q φ ( Z | X ) 2 and q φ ( Y | X ) via a multipy-by-one trick:

Original: [Equation]
Revised: [Equation]

Original: To simplify the objective, Equation 2 assumes that Y is independent from M : p ( Y | M )
Revised: To simplify the objective [Equation] Equation 2 assumes that Y is independent from M : p ( Y | M )

Original: = p ( Y ) .
Revised: = p ( Y ) .

Original: This decomposition results in Equation 3, which includes two KL-divergences against true  posteriors, p ( Z | M, Y, X ) and p ( Y | M, X ).
Revised: This decomposition results in Equation [Equation] which includes two KL-divergences against true (unknown) posteriors, p ( Z | M, Y, X ) and p ( Y | M, X )

Original: We can then train the model by maximizing the evidence lower bound (ELBO), L T , to the true conditional, ln p ( X | M ) >
Revised: We can then train the model by maximizing the evidence lower bound (ELBO), L T , to the true conditional, ln p ( X | M ) >

Original: L T :
Revised: L T :

Original: [Equation]
Revised: [Equation]

Original: [Table] 
Revised: [Table] [Equation] [Table] [Equation]

Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original:  The bound in Equation 4 is tight if q φ ( Z | X )
Revised: [Equation] The bound in Equation 4 is tight if q φ ( Z | X )

Original: =
Revised: =

Original: p ( Z | M, Y, X ) and q φ ( Y | X ) =
Revised: p ( Z | M, Y, X ) and q φ ( Y | X ) =

Original: p ( Y | M, X ) .
Revised: p ( Y | M, X ) .

Original: While the ELBO in Equation 4 is principled, it involves inferring the entire memory M .
Revised: While the ELBO in Equation 4 is principled, it involves inferring the entire memory M .

Original: This prevents us from decoupling the size of the memory from inference and scaling the computation complexity based on the size of the memory.
Revised: This prevents us from decoupling the size of the memory from inference and scales the computation complexity based on the size of the memory.

Original: To alleviate this constraint, we assume a purely deterministic memory, M ∼ δ [ f mem ( f enc ( X ))] , built by transforming the input episode, X , via an encoder and memory transformation model, f mem ◦ f enc .
Revised: To alleviate this constraint, we assume a purely deterministic memory, M ∼ δ [ f mem ( f enc ( X ))] , built by transforming the input episode, X , via an encoder and memory transformation model, f mem ◦ f enc .

Original: We also assume that the regions of memory which are useful in reconstructing a sample, x t , can be summarized by a set of K localized contiguous memory sub-blocks:
Revised: We also assume that the regions of memory which are useful in reconstructing a sample, x t , can be summarized by a set of K localized contiguous memory sub-blocks as described in Equation 5 below.

Revised: The intuition here is that similar samples, x t ≈ x r , might occupy a disjoint part of the representation space and the decoder, p θ ( X |· ) , would need to read multiple regions to properly handle sample reconstruction.

Revised: For example, the digit “3” might share part of the representation space with a “2” and another part with a “5”.

Original: [Equation]
Revised: [Equation]

Original: (5) ˆ M in equation 5 represents a set of K  memory sub-regions, determined by the addressing key, y tk , and a spatial transformer ( ST ) network Jaderberg et al.
Revised: ˆ M in equation 5 represents a set of K dirac-delta memory sub-regions, determined by the addressing key, y tk , and a spatial transformer ( ST ) network Jaderberg et al.

Original:  f ST 2 .
Revised: (2015), f ST 3 .

Original: Our ﬁnal optimization objective, L T ≈ L T | ˆ M ≈ M , is attained by replacing M in Equation 4 with ˆ M and is summarized by the graphical model in 2 below.3. W RITE M ODEL
Revised: Our ﬁnal optimization objective, L T ≈ L T | ˆ M ≈ M , is attained by replacing M in Equation 4 with ˆ M and is summarized by the graphical model in 2 below.

Original: [Figure]
Revised: [Figure]

Original: Writing to memory in the K++ model (Figure 3) entails encoding the input episode, X t
Revised: Writing to memory in the K++ model (Figure 3) entails encoding the input episode, X t

Original: = { x t } Tt =1 , through the encoder, E = f
Revised: = { x t } Tt =1 , through the encoder, E = f

Original: θ enc ( X t ) , pooling the representation over the episode and encoding the pooled representation with the memory writer, M = f θ mem ( E ) .
Revised: θ enc ( X t ) , pooling the representation over the episode and encoding the pooled representation with the memory writer, M = f θ mem ( E ) .

Original: In this work we employ a
Revised: In this work we employ a

Original: Temporal Shift Module (TSM) (Lin et al., 2019), applied on a ResNet18 (He et al., 2016).
Revised: Temporal Shift Module (TSM) (Lin et al., 2019), applied on a ResNet18 (He et al., 2016).

Original: TSM works by shifting feature maps of a 2d vision model in the temporal dimension.
Revised: TSM works by shifting feature maps of a 2d vision model in the temporal dimension.

Original: This allows the model to build richer representations of contextual features.
Revised: This allows the model to build richer representations of contextual features.

Original: In the case of K++, this allows the encoder to build a better representation of the memory.
Revised: In the case of K++, this allows the encoder to build a better representation of the memory by leveraging intermediary episode speciﬁc features.

Original: Using a TSM encoder over a standard convolutional stack improves the performance of both K++ and DKM, where the latter observes an improvement of 6.32 nats/image over the reported test ELBO of 77.2 nats/image (Wu et al., 2018b) for the Omniglot dataset.
Revised: Using a TSM encoder over a standard convolutional stack improves the performance of both K++ and DKM, where the latter observes an improvement of 6.32 nats/image over the reported test variational lower bound of 77.2 nats/image (Wu et al., 2018b) for the Omniglot dataset.

Original: As far as we are aware, the application of a TSM encoder to memory models has not been explored and is a contribution of this work.
Revised: As far as we are aware, the application of a TSM encoder to memory models has not been explored and is a contribution of this work.

Original: The memory writer model, f θ mem , in Figure 3, allows K++ to non-linearly transform the pooled embedding to better summarize the episode.
Revised: The memory writer model, f θ mem , in Figure 3, allows K++ to non-linearly transform the pooled embedding to better summarize the episode.

Original: In addition to inferring the deterministic memory, M , we also project the non-pooled embedding, E , through a key model, f θ key , described in Section 3.4:
Revised: In addition to inferring the deterministic memory, M , we also project the non-pooled embedding, E , through a key model, f θ key :

Original: [Equation]
Revised: [Equation]

Revised: ) (cid:12) (cid:15), y

Original: [Equation]
Revised: [Equation]

Original: given memory: M ; sample keys: Y t ∼ P
Revised: ∈ R 3 , (cid:15) ∼ N (0 , 1) . The reparameterized keys, y tk , will be used to read sample speciﬁc memory traces, ˆ M , from the full memory, M .

Original: [Equation]
Revised: The memory traces, ˆ M , are used in training through the learned prior, p θ ( Z | ˆ M, Y ) = N ( µ z , σ 2 z ) , from Equation 4 via the KL divergence, E q φ ( Y | X ) D KL [ q φ ( Z | X ) || P

Original: t ) ; extract regions: ˆ M =
Revised: θ ( Z | ˆ M, Y )] .

Original: [Equation]
Revised: This

Original: M
Revised: KL divergence constrains the optimization objective to keep the representation of the amortized approximate posterior, q φ ( Z | X ) , (probabilistically) close to the memory readout representation of the learned prior, p θ ( Z | ˆ M, Y ) .

Original: [Equation] ; infer latent: Z [Equation] ∼
Revised: While in the generative setting this constraint enables memory traces to be routed from the learned prior, p θ ( Z | ˆ M, Y ) , to the decoder, p θ ( X |· ) , in a similar manner to standard VAEs.

Original: [Equation] [Figure] [Equation] [Figure]
Revised: We detail this process in the following section.

Original: [Equation]
Revised: given memory: M ;

Original: [Figure]
Revised: [Figure]

Original: The Kanerva++ model, like the original KM and DKM models, enables sample generation given an existing memory or set of memories.
Revised: The Kanerva++ model, like the original KM and DKM models, enables sample generation given an existing memory or set of memories.

Original: K samples from the prior key distribution, { y k } Kk =1 ∼ p ( Y ) =
Revised: K samples from the prior key distribution, { y k } Kk =1 ∼ p ( Y ) =

Original: N (0 , 1) , y k ∈ R 3 , are used to parameterize the spatial transformer, f ST , which indexes the deterministic memory, M .
Revised: N (0 , 1) , y k ∈ R 3 , are used to parameterize the spatial transformer, f ST , which indexes the deterministic memory, M .

Original: The result of this differentiable indexing is a set of memory sub-regions, ˆ
Revised: The result of this differentiable indexing is a set of memory sub-regions, ˆ

Original: M , which are used in the decoder, p θ ( X |· ) , to generate synthetic samples.
Revised: M , which are used in the decoder, p θ ( X |· ) , to generate synthetic samples.

Original: Reading samples in this manner forces the encoder to utilize memory sub-regions that are useful for reconstruction, as non-read memory regions receive zero gradients during backpropagation.
Revised: Reading samples in this manner forces the encoder to utilize memory sub-regions that are useful for reconstruction, as non-read memory regions receive zero gradients during backpropagation.

Original: The memory traces, ˆ M , are used in  the decoder training (Section 3.5) through the learned prior, p θ ( Z | ˆ M, Y ) =
Revised: This insight allows us to use the simple feed-foward write process described in Section 3.3, while still retaining the ability to produce locally contiguous block allocated memory. [Figure] given embedding: E ; if training then

Original: N ( µ z , σ [Equation]
Revised: [Figure] latent [Figure]

Original: z ) , from Equation 4. The KL divergence,
Revised: else infer latent:

Original: [Equation]
Revised: [Figure]

Original: ] constrains the optimization objective to keep the representation of the amortized approximate posterior, q φ ( Z | X )  , close to  the memory readout representation of the learned prior, p θ ( Z | ˆ M, Y ) .
Revised: 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 q φ ( Z | X ) from Equation 4 using an amortized isotropic-gaussian posterior that directly transforms the embedding of the episode, E , using a learned neural network (Figure 2- Center-left ). As mentioned in Section 3.5, the readout, Z , of the memory traces, ˆ M , are encouraged to learn a meaningful structured representation through the memory read-out KL divergence, E q φ ( Y | X ) D KL [ q φ ( Z | X ) || P θ ( Z | ˆ M, Y )] , which attempts to minimize the (probabilistic) distance between q φ ( Z | X ) and P θ ( Z | ˆ M, Y ) .

Original: 3. R EAD M ODEL
Revised: Kanerva memory models also possess the ability to gradually improve a sample through interative inference (Figure 2- Center-right ), whereby noisy samples can be improved by leveraging contextual information stored in memory.

Original: given embedding: E ;
Revised: This can be interpreted as approximating the posterior, q ( Z | X, M ) , by marginalizing the approximate key distribution:

Original: [Equation]
Revised: [Equation]

Original: The K++ read model (Figure 5) is similar in structure to a standard VAE with the exception that there exist T samples per episode.
Revised: where Y ∗ ∼ q φ ( Y | X =

Original: As in KM and DKM, we use the mean of the latent variable, µ Zt , as input to the decoder. In contrast to (Tomczak & Welling, 2018; Graves et al., 2018; Salimans et al., 2017; van den Oord et al., 2016), we use a simple convolutional decoder for all our experiments.
Revised: ˆ X ) in Equation 7 uses a single sample monte-carlo estimate, evaluated by re-infering the previous reconstruction, ˆ X ∼ p θ ( X |· ) , through the approximate key posterior. Each subsequent memory readout, Z , improves upon its previous representation by absorbing additional information from the memory.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We contrast K++ against state of the art vision models for likelihood estimation and present empirical results in Table 1.
Revised: We contrast K++ against state of the art vision models for likelihood estimation and present empirical results in Table 1.

Original: A 1000 sample importance weighted estimate of the likelihood (Burda et [-al., 2016), ln p θ ( x ) ≈ (cid:80) 1000 i =
 L t i , (for each sample in the empirical test distribution)-] is used as a Monte Carlo estimate to integrate out latent variables.
Revised: Binarized datasets assume Bernoulli output distributions, while continuous values are modeled by a discretized mixture of logistics Salimans et [-al., 2016), ln p θ ( x ) ≈ (cid:80) 1000 i =
 L t i , (for each sample in the empirical test distribution)-]al. As is standard in literature Burda et al.

Original: This allows us to fairly compare various models with different latent variable structures.
Revised: ; Sadeghi et al. ; Ma et al.

Original: Binarized datasets assume Bernoulli output distributions, while continuous values are modeled by a Gaussian output distribution with a single learnable variance term (Dai & Wipf, 2019).
Revised: ; Chen et al. , we provide results for binarized MNIST and Omniglot in nats/image and rescale the corresponding results to bits/dim for all other datasets.

Original: Model architecture and optimization details are provided in Appendix E as well as the included code.
Revised: Model architecture and optimization details are provided in Appendix E as well as the included code .

Original: Finally, extra Celeb-A generations and test image reconstructions for all experiments are provided in Appendix B and Appendix D respectively.
Revised: Finally, extra Celeb-A generations and test image reconstructions for all experiments are provided in Appendix B and Appendix D respectively.

Original: [Table]
Revised: [Table]

Revised: estimated from Wu et al.

Revised: (2018a)

Revised: Appendix Figure 12.

Revised: K++ presents state-of-the-art results for binarized MNIST and binarized Omniglot and presents competitive performance for Fashion MNIST, CIFAR10 and DMLab mazes.

Revised: 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 PixelVAE Sadeghi et al.

Revised: We leave the exploration of powerful decoder models to future work and note that our model is not mutually exclusive with autoregressive decoders.

Original: 4. I TERATIVE INFERENCE
Revised: 4. I TERATIVE INFERENCE

Original: One of the beneﬁts of K++ is that it uses the memory to learn a more informed prior by condensing the information from an episode of samples.
Revised: One of the beneﬁts of K++ is that it uses the memory to learn a more informed prior by condensing the information from an episode of samples.

Original: One might suspect that, based on the dimensionality of the memory and the size of the read traces, the memory might only learn prototypical patterns, rather than a full amalgamation of the input episode.
Revised: One might suspect that, based on the dimensionality of the memory and the size of the read traces, the memory might only learn prototypical patterns, rather than a full amalgamation of the input episode.

Original: This presents a problem for generation, as described in Section 3.4 and can be observed in the left most column of Figure
Revised: This presents a problem for generation, as described in Section 3.4 and can be observed in the left most column of Figure

Original: 6- Left where the ﬁrst generation from a random key appears blurry.
Revised: 6- Left where the ﬁrst generation from a random key appears blurry.

Original: To overcome this limitation we rely on the iterative inference of Kanerva memory models (Wu et al., 2018a;b).
Revised: To overcome this limitation we rely on the iterative inference of Kanerva memory models (Wu et al., 2018a;b).

Original: By holding the memory, M , ﬁxed and repeatedly inferring the latents, we are able to cleanup the pattern by leveraging the contextual information contained within the memory.
Revised: By holding the memory, M , ﬁxed and repeatedly inferring the latents, we are able to cleanup the pattern by leveraging the contextual information contained within the memory ( § 3.5).

Original: This is visualized in the proceeding columns of Figure 6, where we observe a slow but clear improvement in generation quality.
Revised: This is visualized in the proceeding columns of Figure 6, where we observe a slow but clear improvement in generation quality.

Original: This property of iterative inference is one of the central beneﬁts of using a memory model over a tradition solution like a VAE.
Revised: This property of iterative inference is one of the central beneﬁts of using a memory model over a tradition solution like a VAE.

Original: We also present results of iterative inference on more classical image noise distributions such as salt-and-pepper, speckle and poisson noise in Figure
Revised: We also present results of iterative inference on more classical image noise distributions such as salt-and-pepper, speckle and poisson noise in Figure

Original: 6- Right .
Revised: 6- Right .

Original: For each original noisy pattern (top rows) we provide the resultant ﬁnal reconstruction after ten steps of cleanup.
Revised: For each original noisy pattern (top rows) we provide the resultant ﬁnal reconstruction after ten steps of cleanup.

Original: The proposed K++ ispretty robust to input noisedistribution and is able to cleanup most of the patters in a semantically meaningful way.
Revised: The proposed K++ is robust to input noise and is able to cleanup most of the patterns in a semantically meaningful way.

Original: 4.1. I MAGE G ENERATIONS
Revised: [Figure]

Original: Typical VAEs use high dimensional isotropic Gaussian latent variables ( ≥ R 16 ) Burda et al.
Revised: Typical VAEs use high dimensional isotropic Gaussian latent variables ( ≥ R 16 ) Burda et al.

Original: Kingma & Welling (2014).
Revised: Kingma & Welling (2014).

Original: A well known property of high dimensional Gaussian distributions is that most of their mass is concentrated on the surface area of a high dimensional ball.
Revised: A well known property of high dimensional Gaussian distributions is that most of their mass is concentrated on the surface area of a high dimensional ball.

Original: Perturbations to a sample in an area of valid density can easily move it to an invalid density region (Arvanitidis et al., 2018; White, 2016), causing blurry or irregular generations.
Revised: Perturbations to a sample in an area of valid density can easily move it to an invalid density region (Arvanitidis et al., 2018; White, 2016), causing blurry or irregular generations.

Original: In the case of K++, since the key distribution, y t ∼ p ( Y ) , y t ∈ R 3 , is within a low dimensional space, local perturbations, y t + (cid:15), (cid:15) ∼
Revised: In the case of K++, since the key distribution, y t ∼ p ( Y ) , y t ∈ R 3 , is within a low dimensional space, local perturbations, y t + (cid:15), (cid:15) ∼

Original: N (0 , 0 . 1) , are likely in regions with high probability density.
Revised: N (0 , 0 . 1) , are likely in regions with high probability density.

Original: We visualize this form of generation in Figure 7 for DMLab Mazes, Omniglot and Celeb-A 64x64 as well as the more traditional random key generations, y t ∼ p ( Y ) , in Figure 8.
Revised: We visualize this form of generation in Figure 7 for DMLab Mazes, Omniglot and Celeb-A 64x64 as well as the more traditional random key generations, y t ∼ p ( Y ) , in Figure 8.

Original: Interestingly, local key perturbations of a trained DMLab Maze K++ model induces resultant generations that provide a natural traversal of the maze as observed by scanning Figure
Revised: Interestingly, local key perturbations of a trained DMLab Maze K++ model induces resultant generations that provide a natural traversal of the maze as observed by scanning Figure

Original: 7- Left , row by row, from left to right.
Revised: 7- Left , row by row, from left to right.

Original: In contrast, the random generations of the same task (Figure 8- Left ) present a more discontinuous set of generations.
Revised: In contrast, the random generations of the same task (Figure 8- Left ) present a more discontinuous set of generations.

Original: We see a similar effect for the Omniglot and
Revised: We see a similar effect for the Omniglot and

Original: Celeb-A datasets, but observe that the locality is instead tied to character or facial structure as shown in Figure 7- Center and Figure
Revised: Celeb-A datasets, but observe that the locality is instead tied to character or facial structure as shown in Figure 7- Center and Figure

Original: Right .
Revised: Right .

Original: Finally, in contrast to VAE generations, K++ is able to generate sharper images of ImageNet32x32 as shown in Appendix C. Future work will investigate this form of locally perturbed generation through an MCMC lens. [Figure]
Revised: Finally, in contrast to VAE generations, K++ is able to generate sharper images of ImageNet32x32 as shown in Appendix C. Future work will investigate this form of locally perturbed generation through an MCMC lens. 4. A BLATION : I S B LOCK A LLOCATED S PATIAL M EMORY U SEFUL ?

Original: While Figure 7 demonstrates the advantage of having a low dimensional sampling distribution and
Revised: While Figure 7 demonstrates the advantage of having a low dimensional sampling distribution and

Original: Figure 6 demonstrates the beneﬁt of iterative inference, it is unclear whether the performance beneﬁt in Table 1 is achieved from the episodic training, model structure, optimization procedure or memory allocation scheme.
Revised: Figure 6 demonstrates the beneﬁt of iterative inference, it is unclear whether the performance beneﬁt in Table 1 is achieved from the episodic training, model structure, optimization procedure or memory allocation scheme.

Original: To isolate the cause of the performance beneﬁt, we simplify the write architecture from Section 3.3 as shown in Figure
Revised: To isolate the cause of the performance beneﬁt, we simplify the write architecture from Section 3.3 as shown in Figure

Original: 9- Left .
Revised: 9- Left .

Original: In this scenario, we replace the learned memory prior, with an equivalently sized dense model that projects the embedding, E , while keeping all other aspects the same.
Revised: In this scenario, we produce the learned memory readout, Z , via an equivalently sized dense model that projects the embedding, E , while keeping all other aspects the same.

Original: We train both models ﬁve times with the exact same TSM-ResNet18 encoder, decoder, optimizer and learning rate scheduler.
Revised: We train both models ﬁve times with the exact same TSM-ResNet18 encoder, decoder, optimizer and learning rate scheduler.

Original: As shown in Figure
Revised: As shown in Figure

Original: 9- Right , the likelihood of the K++ model is 20.6 nats/image better than the baseline model for the evaluated  Omniglot dataset.
Revised: 9- Right , the test variational lower bound of the K++ model is 20.6 nats/image better than the baseline model for the evaluated binarized Omniglot dataset.

Original: This conﬁrms that the spatial, block allocated latent memory model proposed in this work is useful when working with image distributions.
Revised: This conﬁrms that the spatial, block allocated latent memory model proposed in this work is useful when working with image distributions.

Original: Future work will explore this dimension for other modalities such as audio and text.
Revised: Future work will explore this dimension for other modalities such as audio and text.

Original: C ONCLUSION
Revised: 4. A BLATION : EPISODE LENGTH (T) AND MEMORY READ STEPS (K).

Revised: To further explore K++, we evaluate the sensitivity of the model to varying episode lengths (T) in Figure 10- left and memory read steps (K) in Figure

Revised: 10- right using the binarized MNIST dataset.

Revised: K++ performs within margin of error for increasing episode lengths, producing negative test variational bounds within a 1-std of ± 0 .

Revised: 625 nats/image for episode lengths ranging from 5 to 64.

Revised: This suggests that for the speciﬁc dimensionality of memory ( R 64 × 64 ) used in this experiment, K++ was able to successfully capture the semantics of the binarized MNIST distribution.

Revised: 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.

Revised: While ablating the number of memory reads (K) in Figure

Revised: 10- right , we observe that the total test

Revised: KL-divergence varies by 1-std of ± 0.041 nats/image for a range of memory reads from 1 to 64.

Revised: A lower KL divergence implies that the model is able to better ﬁt the approximate posteriors q φ ( Z | X ) and q φ ( Y | X ) to their correspondings priors in Equation 4.

Revised: It should however be noted that a lower KL-divergence does not necessary imply a better generative model Theis et al.

Revised: While qualitatively inspecting the generated samples, we observed that K++ generated more semantically sound generations at lower memory read steps.

Revised: We suspect that the difﬁculty of generating realistic samples increases with the number of disjoint reads and found that K = 2 produces high quality results.

Original:  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 conditional image generation tasks.
Revised: We use this value for all experiments in this work.  C ONCLUSION 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 conditional image generation tasks.

Original: We also show that stochasticity in low-dimensional spaces produces higher quality samples in comparison to high dimensional latents typically used in VAEs.
Revised: We also show that stochasticity in low-dimensional spaces produces higher quality samples in comparison to high dimensional latents typically used in VAEs.

Original: Furthermore, perturbations to the low-dimensional key generate samples with high variations.
Revised: Furthermore, perturbations to the low-dimensional key generate samples with high variations.

Original: There are still however many unanswered questions: would a hard attention based solution to differentiable indexing prove to be better than a spatial transformer?
Revised: There are still however many unanswered questions: would a hard attention based solution to differentiable indexing prove to be better than a spatial transformer?

Original: What is the optimal upper bound of window read regions based on the input distribution?
Revised: What is the optimal upper bound of window read regions based on the input distribution?

Original: Future work will hopefully be able to address these lingering issues and further improve generative memory models.
Revised: Future work will hopefully be able to address these lingering issues and further improve generative memory models.

Original: A S PATIAL T RANSFORMER R EVIEW
Revised: A S PATIAL T RANSFORMER R EVIEW

Original: Indexing a matrix, M [ x : x + ∆ x, y : y + ∆ y ] , is typically a non-differentiable operation since it involves hard cropping around an index.
Revised: Indexing a matrix, M [ x : x + ∆ x, y : y + ∆ y ] , is typically a non-differentiable operation since it involves hard cropping around an index.

Original: Spatial transformers (Jaderberg et al., 2015) provide a solution to this problem by decoupling the problem into two differentiable operands:
Revised: Spatial transformers (Jaderberg et al., 2015) provide a solution to this problem by decoupling the problem into two differentiable operands:

Original: Learn an afﬁne transformation of coordinates.
Revised: Learn an afﬁne transformation of coordinates.

Original: Use a differntiable bilinear transformation.
Revised: Use a differntiable bilinear transformation.

Original: The afﬁne transformation of source coordinates,
Revised: The afﬁne transformation of source coordinates,

Original: [Equation]
Revised: [Equation]

Original: , to target coordinates,
Revised: , to target coordinates,

Original: [Equation]
Revised: [Equation]

Original: is deﬁned as:
Revised: is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: Here, the afﬁne transform, θ = (cid:20) s
Revised: Here, the afﬁne transform, θ = (cid:20) s

Original: [Equation]
Revised: [Equation]

Original: x s y (cid:21) has three learnable scalars: { s, x, y } which deﬁne a scaling and translation in i and j respectively.
Revised: x s y (cid:21) has three learnable scalars: { s, x, y } which deﬁne a scaling and translation in i and j respectively.

Original: In the case of K++, these three scalars represent the components of the key sample, { y 0 , y 1 , y 2 } ∈ R 3 as shown in Equation 7.
Revised: In the case of K++, these three scalars represent the components of the key sample, { y 0 , y 1 , y 2 } ∈ R 3 as shown in Equation 8.

Original: After transforming the co-ordinates (not to be confused with the actual data), spatial transformers learn a differentiable bilinear transform which can be interpreted as learning a differentiable mask that is element-wise multiplied by the original data, M :
Revised: After transforming the co-ordinates (not to be confused with the actual data), spatial transformers learn a differentiable bilinear transform which can be interpreted as learning a differentiable mask that is element-wise multiplied by the original data, M :

Original: [Equation]
Revised: [Equation]

Original: (cid:19) 
Revised: (cid:19) (9)

Original: Consider the following example where θ = (cid:20) 0 . 0 . 0 .0 .(cid:21) ; this parameterization differntiably extracts the region shown in Figure
Revised: Consider the following example where θ = (cid:20) 0 . 0 . 0 .0 .(cid:21) ; this parameterization differntiably extracts the region shown in Figure

Original: 10- Right from Figure
Revised: 11- Right from Figure

Original: 10- Left :
Revised: 11- Left :

Original: The range of values for { s, x, y } is bound between [ − 1 , 1] , where the center of the image is [0 , 0] .
Revised: The range of values for { s, x, y } is bound between [ − 1 , 1] , where the center of the image is [0 , 0] .

Original: B C ELEB -A G ENERATIONS
Revised: B C ELEB -A G ENERATIONS

Original: We present random key generations of Celeb-A 64x64, trained without center cropping in Figure 11.
Revised: We present random key generations of Celeb-A 64x64, trained without center cropping in Figure 12.

Original: C VAE VS . K++ I MAGE N ET 32 X 32 G ENERATIONS
Revised: C VAE VS . K++ I MAGE N ET 32 X 32 G ENERATIONS

Original: Figure 12 shows the difference in generations of a standard VAE vs. K++.
Revised: Figure 13 shows the difference in generations of a standard VAE vs. K++.

Original: In contrast to the standard VAE generation (Figure 12- Left ), the K++ generations (Figure 12- Right ) appear much sharper, avoiding the blurry generations observed with standard VAEs.
Revised: In contrast to the standard VAE generation (Figure 13- Left ), the K++ generations (Figure 13- Right ) appear much sharper, avoiding the blurry generations observed with standard VAEs.

Original: D T EST I MAGE R ECONSTRUCTIONS
Revised: D T EST I MAGE R ECONSTRUCTIONS

Original: E M ODEL A RCHITECTURE & T RAINING
Revised: E M ODEL A RCHITECTURE & T RAINING

Original: Encoder :
Revised: Encoder :

Original: As mentioned in Section 3.3, we use a TSM-Resnet18 encoder with Batchnorm (Ioffe &
Revised: As mentioned in Section 3.3, we use a TSM-Resnet18 encoder with Batchnorm (Ioffe &

Original: Szegedy, 2015) and ReLU activations (Nair & Hinton, 2010) for all tasks.
Revised: Szegedy, 2015) and ReLU activations (Nair & Hinton, 2010) for all tasks.

Original: We apply a fractional shift of the feature maps by 0.125 as suggested by the authors.
Revised: We apply a fractional shift of the feature maps by 0.125 as suggested by the authors.

Original: Decoder : Our decoder is a simple conv-transpose network with EvoNormS0 (Liu et al., 2020) inter-spliced between each layer.
Revised: Decoder : Our decoder is a simple conv-transpose network with EvoNormS0 (Liu et al., 2020) inter-spliced between each layer.

Original: Evonorms0 is similar in stride to Groupnorm (Wu & He, 2020) combined with the swish activation function (Ramachandran et al., 2018).
Revised: Evonorms0 is similar in stride to Groupnorm (Wu & He, 2020) combined with the swish activation function (Ramachandran et al., 2018).

Original: Optimizer & LR scheule :
Revised: Optimizer & LR scheule :

Original: We use LARS (You et al., 2017) coupled with ADAM (Kingma & Ba, 2014) and a one-cycle (Smith & Topin, 2017) cosine learning rate schedule (Loshchilov & Hutter, 2017).
Revised: We use LARS (You et al., 2017) coupled with ADAM (Kingma & Ba, 2014) and a one-cycle (Smith & Topin, 2017) cosine learning rate schedule (Loshchilov & Hutter, 2017).

Original: A linear warm-up of 10 epochs (Goyal et al., 2017c) is also used for the schedule.
Revised: A linear warm-up of 10 epochs (Goyal et al., 2017c) is also used for the schedule.

Original: A weight decay of 1 e − 3 is used on every parameter barring biases and the afﬁne terms of batchnorm.
Revised: A weight decay of 1 e − 3 is used on every parameter barring biases and the afﬁne terms of batchnorm.

Original: Each task is trained for 500 or 1000 epochs depending on the size of the dataset.
Revised: Each task is trained for 500 or 1000 epochs depending on the size of the dataset.

Original: Dense models : All dense models such as our key network are simple three layer deep linear dense models with a latent dimension of 512 coupled with spectral normalization (Miyato et al., 2018).
Revised: Dense models : All dense models such as our key network are simple three layer deep linear dense models with a latent dimension of 512 coupled with spectral normalization (Miyato et al., 2018).

Original: Memory writer :
Revised: Memory writer :

Original: f θ mem uses a deep linear conv-transpose decoder on the pooled embedding, E with a base feature map projection size of 256 with a division by 2 per layer.
Revised: f θ mem uses a deep linear conv-transpose decoder on the pooled embedding, E with a base feature map projection size of 256 with a division by 2 per layer.

Original: We use a memory size of R 3 × 64 × 64 for all the experiments in this work.
Revised: We use a memory size of R 3 × 64 × 64 for all the experiments in this work.

Original: Learned Prior : p θ ( Z | ˆ M, Y ) uses a convolutional encoder that stacks the K read traces, { f ST ( M, y tk ) } Kk =1 , along the channel dimension and projects it to the dimensionality of Z .
Revised: Learned Prior : p θ ( Z | ˆ M, Y ) uses a convolutional encoder that stacks the K read traces, { f ST ( M, y tk ) } Kk =1 , along the channel dimension and projects it to the dimensionality of Z .

Original: In practice, we observed that K++ is about 2x as fast (wall clock) compared to  DKM.
Revised: In practice, we observed that K++ is about 2x as fast (wall clock) compared to our re-implementation of DKM.

Revised: We mainly attribute this to not having to solve an inner OLS optimization loop for memory inference.