Original: A VAE FOR T RANSFORMERS WITH N ONPARAMETRIC V ARIATIONAL I NFORMATION B OTTLENECK
Revised: A VAE FOR T RANSFORMERS WITH N ONPARAMETRIC V ARIATIONAL I NFORMATION B OTTLENECK

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We propose a Variational AutoEncoder (VAE) for Transformers by developing a
Revised: We propose a Variational AutoEncoder (VAE) for Transformers by developing a

Original: Variational Information Bottleneck (VIB) regulariser for Transformer embeddings.
Revised: Variational Information Bottleneck (VIB) regulariser for Transformer embeddings.

Original: We formalise the embedding space of Transformer encoders as mixture distributions, and use
Revised: We formalise such attention-based representations as mixture distributions, and use

Original: Bayesian nonparametrics to develop a Nonparametric VIB (NVIB) for such attentionbased representations.
Revised: Bayesian nonparametrics to develop a Nonparametric VIB (NVIB) for them.

Original: The variable number of mixture components supported by nonparametric methods captures the variable number of vectors supported by attention, and exchangeable distributions from nonparametric methods capture the permutation invariance of attention.
Revised: The variable number of mixture components supported by nonparametrics captures the variable number of vectors supported by attention, and exchangeable distributions from nonparametrics capture the permutation invariance of attention.

Original: We then propose our Transformer VAE (NVAE) using NVIB to regularise the information passing from the Transformer encoder to the Transformer decoder through cross-attention.
Revised: Our Transformer VAE (NVAE) uses NVIB to regularise the information passing from the Transformer encoder to the Transformer decoder.

Original: Evaluations of a NVAE, trained on natural language text, demonstrate that NVIB can regularise the number of mixture components in the induced embedding whilst maintaining generation quality and reconstruction capacity
Revised: Evaluations of a NVAE, trained on natural language text, demonstrate that NVIB can regularise the number of mixture components in the induced embedding whilst maintaining generation quality and reconstruction capacity.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Attention-based deep learning models, such as Transformers (Vaswani et al., 2017; Devlin et al., 2019), have achieved unprecedented empirical success in a wide range of cognitive tasks, in particular in natural language processing.
Revised: Attention-based deep learning models, such as Transformers (Vaswani et al., 2017; Devlin et al., 2019), have achieved unprecedented empirical success in a wide range of cognitive tasks, in particular in natural language processing.

Original: The use of attention allows these models to represent their input with multiple vectors, which is essential for embedding natural language text (Bahdanau et al., 2015).
Revised: The use of attention allows these models to represent their input with multiple vectors, which is essential for embedding natural language text (Bahdanau et al., 2015).

Original: On the other hand, deep variational
Revised: On the other hand, deep variational

Original: Bayesian approaches to representation learning, such as variational autoencoders (VAEs) (Kingma &
Revised: Bayesian approaches to representation learning, such as variational autoencoders (VAEs) (Kingma &

Original: Welling, 2014), have also been shown to have many benefits (Mathieu et al., 2019; Ghosh et al., 2020;
Revised: Welling, 2014) have also been shown to have many benefits (Mathieu et al., 2019; Ghosh et al., 2020;

Original: Vahdat & Kautz, 2020), especially due to their variational information bottleneck (VIB) (Alemi et al., 2017) for regularising the induced latent representations.
Revised: Vahdat & Kautz, 2020), especially due to their variational information bottleneck (VIB) (Alemi et al., 2017) for regularising the induced latent representations.

Original: However, it has not been clear how to combine these two trends, because the latent space induced by Transformers is a set of vectors whose size grows with the size of the input, whereas standard VIB methods only apply to a vector space of a fixed size (Lin et al., 2020;
Revised: However, it has not been clear how to combine these two trends, because the latent space induced by Transformers is a set of vectors whose size grows with the size of the input, whereas standard VIB methods only apply to a vector space of a fixed size (Liu & Liu, 2019;

Original: Fang et al., 2021; Park & Lee, 2021).
Revised: Fang et al., 2021; Park & Lee, 2021).

Original: To define a VIB regulariser for a Transformer’s embedding space, we need to allow the size of a latent representation to vary dynamically depending on the complexity of the individual input, and yet regularise the total amount of information conveyed by the whole representation.
Revised: To define a VIB regulariser for a Transformer’s embedding space, we need to allow the size of a latent representation to vary dynamically depending on the complexity of the individual input, and yet regularise the total amount of information conveyed by the whole representation.

Original: In this paper, we propose such a variational information bottleneck for variable sized latent representations, which we use to regularise the embeddings of a Transformer encoder-decoder, giving us a variational autoencoder for Transformers.
Revised: In this paper, we propose such a variational information bottleneck for variable sized latent representations, which we use to regularise the embeddings of a Transformer encoder-decoder, giving us a variational autoencoder for Transformers.

Original: 1 Like a Transformer’s encoder output and decoder input, the proposed VAE’s sampled encoder output is (a generalisation of) a set of vectors, and the decoder accesses this embedding with (a generalisation of) cross attention.
Revised: 1 Like a Transformer encoder’s embedding space, the proposed VAE’s sampled encoder output is (a generalisation of) a set of vectors, and the decoder accesses this embedding with (a generalisation of) cross attention.

Original: But unlike Transformers, the proposed VIB layer for this VAE regularises the (effective) number of vectors in the set, as well as the information conveyed by each vector.
Revised: But unlike Transformers, the proposed VIB layer for this VAE regularises the (effective) number of vectors in the set, as well as the information conveyed by each vector.

Original: We show that this regularisation improves generative abilities and compresses latent representations.
Revised: We show that this regularisation improves generative abilities and compresses latent representations.

Original: In addition, previous work shows that VAEs are good for disentanglement (Higgins et al., 2017). All these topics are important and active areas of research in NLP, for efficiency of over-parameterised language models (Child et al., 2019),  language generation (Liu & Liu, 2019), and explainability (Mercatali & Freitas, 2021). 
Revised: In addition to the regularisation of over-parameterised language models (Child et al., 2019), previous work shows the efficacy of VAEs for: disentanglement (Higgins et al., 2017), language generation (Liu & Liu, 2019), and explainability (Mercatali & Freitas, 2021). All these topics are important and active areas of research in NLP.

Original: To define this VIB, we need to model distributions over these variable-sized encoder embeddings. We argue that these distributions need to be exchangeable and nonparametric .
Revised: To define this VIB, we need to model distributions over these variable-sized encoder embeddings, as interpreted by cross attention.

Original: Firstly,  the attention function is insensitive to the order of the vectors in the embedding, so it interprets the output of  the encoder as a permutation-invariant set of vectors (Zaheer et al., 2017).
Revised: Firstly, because the attention function returns an interpolation between the vectors output by the encoder, it generalises across the varying number of vectors, which like the input length is theoretically unbounded.

Original: Thus, the distributions over these permutation-invariant embeddings should be exchangeable (Jordan, 2010).
Revised: Thus, to define distributions over these unbounded embeddings, we need to use nonparametric methods (Jordan, 2010).

Original: Secondly,because the attention 
Revised: Secondly, the attention function is insensitive to the

Original: [Figure]
Revised: [Figure]

Original: function returns an interpolation between input vectors, it generalises across the varying size of the sets output by the encoder, which like the input length is theoretically unbounded.
Revised: order of the vectors output by the encoder, so it interprets this embedding as a permutation-invariant set of vectors.

Original: Thus, to define distributions over these unbounded embeddings, we need to use nonparametric methods (Jordan, 2010).
Revised: Thus, the distributions over these permutation-invariant embeddings should be exchangeable (Jordan, 2010).

Revised: Thirdly, the attention function imposes a normalised weighting over the embedding vectors, via the attention weights.

Revised: So we should model an embedding as a distribution rather than a set.

Original: The properties of permutation-invariance and unboundedness are also shared with mixtureprobability distributions, which have been extensively studied using Bayesian nonparametric methods with exchangeable distributions (Blei & Jordan, 2006; Jordan, 2010).
Revised: A normalised weighting over an unbounded permutation-invariant set of fixed-length vectors matches exactly the properties of a nonparametric space of mixture distributions, which have been extensively studied in Bayesian nonparametrics using exchangeable distributions (Blei & Jordan, 2006; Jordan, 2010).

Original: In previous work, Bayesian nonparametrics is typically applied to learning models where the number of parameters grows with the size of the training data (Teh, 2010; Jordan, 2010; Kossen et al., 2021).
Revised: In previous work, Bayesian nonparametrics is typically applied to learning models where the number of parameters grows with the size of the training data (Teh, 2010; Jordan, 2010; Kossen et al., 2021).

Original: In contrast, we apply it to inferring latent representations where the number of parameters grows with the size of the input .
Revised: In contrast, we apply it to inferring latent representations where the number of parameters grows with the size of the input .

Original: To the best of our knowledge, ours is the first work to use nonparametric methods in this way for deep variational Bayesian models.
Revised: We believe this is the first work to use nonparametric methods in this way for deep variational Bayesian models.

Original: To characterise how attention-based representations are related to mixture distributions, we provide an interpretation of attention where the input set of vectors is interpreted as a mixture of impulse distributions, and the attention function uses this as a prior to denoise the observed query vector (depicted in Figure 1b).
Revised: To define a precise equivalence between attention-based representations and mixture distributions, we provide an interpretation of attention where the input set of vectors defines a mixture of impulse distributions, which is used as a prior to denoise the observed query vector (depicted in Figure 1b).

Original: Generalising sets of vectors to mixtures of component distributions and generalising the attention function to query denoising allows us to propose a general deep variational Bayesiancomputational framework for attention-based models using nonparametric methods with exchangeable distributions.
Revised: Generalising sets of vectors to mixture distributions and generalising the attention function to query denoising allows us to propose a general deep variational Bayesian framework for attention-based models using Bayesian nonparametrics.

Original: More specifically, we propose to use Dirichlet processes (DPs) as the exchangeable distributions (Aldous, 1985; Jordan, 2010) to specify distributions over mixtures of impulse distributions, including distributions over the effective number of components in the mixture.
Revised: More specifically, we propose to use Dirichlet processes (DPs) as the exchangeable distributions (Aldous, 1985; Jordan, 2010) to specify distributions over mixtures of impulse distributions, including distributions over the effective number of components in the mixture.

Original: We define a nonparametric VIB (NVIB) layer using a bounded DP prior and posterior to regularise the effective number of vectors in the latent representation.
Revised: We define a nonparametric VIB (NVIB) layer using a bounded DP prior and posterior to regularise the effective size of variable-sized latent representations.

Original: Applying this NVIB regulariser to a Transformer autoencoder gives us our proposed nonparametric variational autoencoder (NVAE), depicted in Figure 1a. The NVIB layer uses exact inference to infer the posterior fromthe prior and a set of pseudo-observations output by the encoder, and we propose efficient approximations to sample from this posterior, with their reparameterisation tricks to allow effective backpropagation.
Revised: This NVIB layer uses exact inference to infer the posterior from a set of pseudo-observations, and uses proposed efficient approximations to sample from this posterior with a reparameterisation trick and to regularise it with the KL divergence with the prior. Applying this NVIB regulariser to a Transformer autoencoder gives us our proposed nonparametric variational autoencoder (NVAE), depicted in Figure 1a.

Original: The noise introduced by sampling from the DP posterior controls the amount of information which flows from the encoder to the decoder, despite the fact that the amount of information required to reconstruct different text inputs varies enormously.
Revised: The noise introduced by sampling from the DP posterior controls the amount of information which flows from the encoder to the decoder, despite the fact that the amount of information required to reconstruct different text inputs varies enormously.

Original: To evaluate the effectiveness of NVIB, we train a NVAE on natural language text and find that it is able to reconstruct, generate and regularise the effective number of vectors in the latent representation, thereby demonstrating that NVAE is a viable VAE.
Revised: To evaluate the effectiveness of NVIB, we train a NVAE on natural language text and find that it is able to reconstruct, generate and regularise the effective number of vectors in the latent representation, thereby demonstrating that NVAE is a viable VAE.

Original: We also find that the regularised latent space is smooth, using a proposed method for interpolating between DP posteriors to generate interpolations between sentences.
Revised: We also find that the regularised latent space is smooth, using a proposed method for interpolating between DP posteriors to generate interpolations between sentences.

Original: Contributions
Revised: Contributions

Original: This paper makes the following contributions: (1) Propose anonparametric variational Bayesian framework for  attention-based representations with denoising attention  (Section 2).
Revised: This paper makes the following contributions: (1) We propose a variational Bayesian framework for modelling attention-based representations using mixture distributions, denoising attention and Bayesian nonparametrics (Section 2).

Original: (2) Propose a nonparametric variational information bottleneck (NVIB) regulariser for  attention-based representations (Section 3).
Revised: (2) We propose a nonparametric variational information bottleneck (NVIB) regulariser for learning attention-based representations (Section 3).

Original: (3) Propose a nonparametric variational autoencoder (NVAE) which is a variational Bayesian reformulation of a Transformer encoder-decoder (Section 4).
Revised: (3) We propose a nonparametric variational autoencoder (NVAE), which is a variational Bayesian extension of a Transformer encoder-decoder (Section 4).

Original: (4) Show that the NVAE model is a competitive VAE which can reconstruct, generate, regularise its latent space and intuitively interpolate between sentences (Section 5).
Revised: (4) We show that the NVAE model is a competitive VAE which can reconstruct, generate, regularise its latent space and intuitively interpolate between sentences (Section 5).

Original: Related work Related work in stochastic attention assume that the keys, queries, values (Martin et al., 2020) or attention weight vectors of the network are treated as latent random variables (Deng et al., 2018;
Revised: Related work Related work in stochastic attention assume that the keys, queries, values (Martin et al., 2020) or attention weight vectors of the network are treated as latent random variables (Deng et al., 2018;

Original: Bah; Fan et al., 2020; Cinquin et al., 2021).
Revised: Bahuleyan et al., 2018; Fan et al., 2020; Cinquin et al., 2022).

Original: Nguyen et al.
Revised: Nguyen et al.

Original: (2022) provides a formulation and interpretation of attention keys as latent mixture distributions, whereas our formulation characterises the whole attention function and is interpreted as Bayesian query denoising.
Revised: (2022) provides a formulation and interpretation of attention keys as latent mixture distributions, whereas our formulation characterises the whole attention function and is interpreted as Bayesian query denoising.

Original: The use of Bayesian nonparametrics to learn a variable sized latent space using a VAE (Nalisnick & Smyth, 2016; Goyal et al., 2017; Echraibi et al., 2020) still assumes a fixed-sized latent representation at test time, unlike our proposal.
Revised: The use of Bayesian nonparametrics to learn a variable sized latent space using a VAE (Nalisnick & Smyth, 2017; Goyal et al., 2017; Echraibi et al., 2020) still assumes a fixed-sized latent representation at test time, unlike our proposal.

Original:  N ONPARAMETRIC B AYESIAN F RAMEWORK FOR T RANSFORMER E MBEDDINGS
Revised: 2 A N ONPARAMETRIC B AYESIAN F RAMEWORK FOR T RANSFORMER E MBEDDINGS

Original: This section proposes a formalisation of attention-based representations as mixture distributions over a vector space, and proposes nonparametric Bayesian methods for modelling information about these mixture distributions.
Revised: This section proposes a formalisation of attention-based representations as mixture distributions over a vector space, and proposes nonparametric Bayesian methods for modelling information about these mixture distributions.

Original: First we show that standard attention functions can be interpreted as implementing Bayesian query denoising, where the set of vectors being accessed specifies a mixture of impulse distributions (Section 2.1).
Revised: First we show that standard attention functions can be interpreted as implementing Bayesian query denoising, where the set of vectors being accessed specifies a mixture of impulse distributions (Section 2.1).

Original: We adopt mixture distributions over vectors as a generalisation of attention-based representations, and adopt this denoising function as a generalisation of attention.
Revised: We adopt mixture distributions over vectors as a generalisation of attention-based representations, and adopt this denoising function as a generalisation of attention.

Original: Then we use Bayesian nonparametrics to propose prior (Section 2.2) and posterior distributions (Section 2.3) over these mixture distributions.
Revised: Then we use Bayesian nonparametrics to propose prior (Section 2.2) and posterior distributions (Section 2.3) over these mixture distributions.

Original: These priors and posteriors form the basis of our nonparametric variational information bottleneck, proposed in Section 3.
Revised: These priors and posteriors form the basis of our nonparametric variational information bottleneck, proposed in Section 3.

Original: 2. D ENOISING A TTENTION
Revised: 2. D ENOISING A TTENTION

Original: The attention function provides access to a set of vectors by mapping a query vector to the resulting attention vector.
Revised: The attention function provides access to a set of vectors by mapping a query vector to the resulting attention vector.

Original: As the basis of our approach to attention-based representations, we generalise the set of vectors to a probability distribution over vectors, and generalise attention to a function of these probability distributions.
Revised: As the basis of our approach to attention-based representations, we generalise the set of vectors to a probability distribution over vectors, and generalise attention to a function of these probability distributions.

Original: The attention mechanism we assume is scaled dot product attention, standardly used in many attentionbased models including Transformers.
Revised: The attention mechanism we assume is scaled dot product attention, standardly used in many attentionbased models including Transformers.

Original: For simplicity, we consider cross attention, where a single query vector is mapped to a single result vector.
Revised: For simplicity, we consider cross attention, where a single query vector is mapped to a single result vector.

Original: This attention function projects the set of vectors Z ∈ R n × p via  weight matrices W K , W V ∈ R p × d to keys andvalues, respectively, projects the input u ′ ∈ R 1 × p viathe weight matrix W Q ∈ R p ×
Revised: This attention function projects the input vector u ′ ∈ R 1 × p via the weight matrix W Q ∈ R p × d to a query, and projects the set of vectors Z ∈ R n × p via weight matrices W K , W V ∈ R p ×

Original: d to a query, and  uses the keys’ dimensionality d for scaling.
Revised: d to keys and values, respectively. It uses the keys’ dimensionality d for scaling.

Original: We regroup this scaled dot product attention function into a core dot product attention function Attn ( u , Z ) in which all operations are done in the space of Z 
Revised: We regroup this scaled dot product attention function into a core dot product attention function Attn( u , Z ) in which all operations are done in the space of Z .

Original: [Equation]
Revised: [Equation]

Original: where u = ( u ′ W Q ( W K ) T ) ∈ R 1 × p .
Revised: where u = ( u ′ W Q ( W K ) T ) ∈ R 1 × p .

Original: The function Attn ( u , Z ) can then be defined in two equivalent ways (as shown in Appendix G), in terms of a sum over the vectors z i in Z , or in terms of an integral over a distribution which is only nonzero at the z i :
Revised: The function Attn( u , Z ) can then be defined in two equivalent ways (as shown in Appendix G): in terms of a sum over the vectors z i in Z , or in terms of an integral over a distribution which is only nonzero at the z i :

Original: [Equation]
Revised: [Equation]

Original: where δ z i is an impulse distribution at z i , f ( · ) is the probability density function for distribution F , and g ( u ; v , √ d I ) is the multivariate Gaussian function with diagonal variance of √ d .
Revised: where δ z i is an impulse distribution at z i , f ( · ) is the probability density function for distribution F , and g ( u ; v , √ d I ) is the multivariate Gaussian function with diagonal variance of √ d .

Original: As depicted in Figure 1b,
Revised: As depicted in Figure 1b,

Original: DAttn ( u ; F Z ) can be interpreted as query denoising.
Revised: DAttn( u ; F Z ) can be interpreted as query denoising.

Original: The query u is interpreted as an observation of some true vector v which has been corrupted by Gaussian noise, where v was generated from a prior probability distribution F Z specified by Z .
Revised: The query u is interpreted as an observation of some true vector v which has been corrupted by Gaussian noise, where v was generated from a prior probability distribution F Z specified by Z .

Original: The result of Attn ( u , Z ) is the expected value of this true vector v after seeing the noisy observation u , which can be interpreted as a form of denoising.
Revised: The result of Attn( u , Z ) is the expected value of this true vector v after seeing the noisy observation u , which can be interpreted as a form of denoising.

Original: This denoising attention function DAttn ( u ; F ) is actually a generalisation of attention over a set of vectors, in that it is defined for any probability distribution F over a vector space.
Revised: This denoising attention function DAttn( u ; F ) is actually a generalisation of attention over a set of vectors, in that it is defined for any probability distribution F over a vector space.

Original: In the special case where
Revised: In the special case where

Original: F = (cid:80)
Revised: F = (cid:80)

Original: i π
Revised: i π

Original: i
Revised: i

Original: δ
Revised: δ

Original: z i is a finite mixture of impulse distributions, it is the same as Attn ( u , Z ) but with a bias term log( π i ) substituted for 12 √ d
Revised: z i is a finite mixture of impulse distributions, it is the same as Attn( u , Z ) but with a bias term log( π i ) substituted for 12 √ d

Original: || z
Revised: || z

Original: i || 2 .
Revised: i || 2 .

Original: In the rest of this paper we will use DAttn ( u ; F ) as our definition of attention, which allows us to treat the latent space of a Transformer encoder-decoder asa mixture distribution.
Revised: In the rest of this paper, we will use DAttn( u ; F ) as our definition of attention, which allows us to treat the latent space of a Transformer encoder-decoder as mixture distributions.

Original: 2. A P RIOR OVER M IXTURE D ISTRIBUTIONS
Revised: 2. A P RIOR OVER M IXTURE D ISTRIBUTIONS

Original: Given that our attention-based latent representations are formalised as mixture distributions F , a Bayesian approach requires a prior over these distributions.
Revised: Given that our attention-based latent representations are formalised as mixture distributions F , a Bayesian approach requires a prior over these distributions.

Original: Attention-based models place no finite bound on the possible number of vectors in their set of vectors Z , and thus there is no finite bound on the number of parameters needed to specify the equivalent probability distribution F .
Revised: Attention-based models place no finite bound on the possible number of vectors in their set of vectors Z , and thus there is no finite bound on the number of parameters needed to specify the equivalent mixture distribution F .

Original: Nonetheless, we can still specify  distributionsp ( F ) over this infinite space of possible distributions F  using methods from Bayesian nonparametrics.
Revised: Nonetheless, we can still specify probability distributions over this infinite space of possible distributions F , using methods from Bayesian nonparametrics.

Original: These nonparametric Bayesian methods are specifically designed for modelling probability distributions over unboundedly large mixture distributions.
Revised: These nonparametric Bayesian methods, with exchangeable distributions, are specifically designed for modelling probability distributions over unboundedly large mixture distributions.

Original: We base our distributions over probability distributions on Dirichlet processes DP( G 0 ,α 0 ) .
Revised: We base our distributions over mixture distributions on Dirichlet processes DP( G 0 ,α 0 ) .

Original: Dirichlet processes (DPs) are a generalisation of Dirichlet distributions to an infinite support, such as the points in a vector space.
Revised: Dirichlet processes (DPs) are a generalisation of Dirichlet distributions to an infinite support, such as the points in a vector space.

Original: A Dirichlet distribution Dir( α ) is a distribution over probability mass functions π of discrete categories i , 1 ≤ i ≤ κ .
Revised: A Dirichlet distribution Dir( α ) is a distribution over probability mass functions π of discrete categories i , 1 ≤ i ≤ κ .

Original: One useful definition of Dirichlet processes views a DP F ∼ DP( G 0 ,α 0 ) , where Gis the base distribution over vectors and α 0 ∈ R is the concentration parameter, as the limit of a sequence of finite Dirichlet distributions (see Teh (2010)), given in equations 4.
Revised: One useful definition of Dirichlet processes views a DP F ∼ DP( G 0 ,α 0 ) , where G 0 is the base distribution over vectors and α 0 ∈ R is the concentration parameter, as the limit of a sequence of finite Dirichlet distributions (see Teh (2010)), given in equation 4.

Original: Note that the Dirichlet distributions in equations 4 are symmetric, in that all the κ categories i have the same α i
Revised: Note that the Dirichlet distributions in equation 4 are symmetric, in that all the κ categories i have the same α i

Original: = α 0 κ parameter values.
Revised: = α 0 κ parameter values.

Original: However, these categories end up with very different weights π i , due to the most probable categories getting a large proportion of the probability mass and the tail of categories getting an exponentially decreasing amount of probability mass.
Revised: However, these categories end up with very different weights π i , due to the most probable categories getting a large proportion of the probability mass and the tail of categories getting an exponentially decreasing amount of probability mass.

Original: In the infinite limit, this tail is infinitely long with infinitesimal probabilities.
Revised: In the infinite limit, this tail is infinitely long with infinitesimal probabilities.

Original: The number of categories which get nontrivial probabilities is determined by α 0 , and becomes independent of κ as κ gets large.
Revised: The number of categories which get nontrivial probabilities is determined by α 0 , and becomes independent of κ as κ gets large.

Original: As shown in this definition, each sample F from a DP is an infinite mixture of impulse distributions δ z
Revised: As shown in this definition, each sample F from a DP is an infinite mixture of impulse distributions δ z

Original: i , parameterised by an infinite sequence of weight-vector pairs π i , z i .
Revised: i , parameterised by an infinite sequence of weight-vector pairs π i , z i .

Original: This contrasts with the finite Z in attention-based representations.
Revised: This contrasts with the finite Z in attention-based representations.

Original: The infinitenature of F  also causes problems in our variationalBayesian approach.
Revised: Having an infinite F would also cause problems in our variational

Original:  VIB uses a bound on the log-likelihood (the ELBO), which Kingma & Welling (2014) showed has an error of D KL ( q ( F (cid:12)(cid:12) x ) ∥ p ( F (cid:12)(cid:12) x ))
Revised: Bayesian model, because VIB uses a bound on the log-likelihood (see Section 3.1), which Kingma & Welling (2014) showed has an error of D KL ( q ( F (cid:12)(cid:12) x ) ∥ p ( F (cid:12)(cid:12) x ))

Original: (the looseness of the bound), which would be infinite unless both the true posterior p ( F (cid:12)(cid:12) x ) and its approximation q ( F (cid:12)(cid:12) x ) generate a finite mixture F .
Revised: (the looseness of the bound). This would be infinite unless both the true posterior p ( F (cid:12)(cid:12) x ) and its approximation q ( F (cid:12)(cid:12) x ) generate a finite F , so we need a prior which generates finite F .

Original: [Equation]
Revised: [Equation]

Original: The Unbounded Dirichlet Process Prior We do not want a prior which places an apriori bound on the size of Z , so we assume it is finite but unbounded, and propose a prior which is an unbounded sequence of finite approximations to a DP.
Revised: The Unbounded Dirichlet Process Prior We do not want a prior which places an apriori bound on the size of F , so we assume it is finite but unbounded, and propose a prior which is an unbounded sequence of finite approximations to a DP.

Original: We define a bounded DP F ∼ BDP( G 0 ,α 0 ,κ 0 ) as in equations 5.
Revised: We define a bounded DP F ∼ BDP( G 0 ,α 0 ,κ 0 ) as in equation 5.

Original: Our approach to the prior is to use an unbounded but finite κ 0 , so we define a distribution over approximations as κ 0 increases towards infinity, so that every distribution is over a finite number of vectors, but there is no finite bound on the number of vectors in all distributions.
Revised: Our approach to the prior is to use an unbounded but finite κ 0 , so we define a distribution over approximations as κ 0 increases towards infinity. Hence, every distribution is over a finite number of vectors, but there is no finite bound on the number of vectors in all distributions.

Original: Given ϕ is some distribution over positive integers κ ∈ Z
Revised: Given ϕ is some distribution over positive integers κ ∈ Z

Original: + , we define this unbounded DP as UDP( G 0 ,α 0 ,ϕ )
Revised: + , we define this unbounded DP as UDP( G 0 ,α 0 ,ϕ )

Original: =
Revised: =

Original: BDP( G 0 ,α 0 ,κ ) where κ ∼ ϕ .
Revised: BDP( G 0 ,α 0 ,κ ) where κ ∼ ϕ .

Original: We use these definitions both to define a general prior over probability distributions, and to define a conditional prior for each input length.
Revised: We use these definitions both to define a general prior over probability distributions, and to define a conditional prior for each input length.

Original: In both cases, the base distribution G p 0 is assumed to be a unit
Revised: In both cases, the base distribution G p 0 is assumed to be a unit

Original: Gaussian, inspired by Kingma & Welling (2014), and the concentration parameter α p 0 is assumed to be one.
Revised: Gaussian (inspired by Kingma & Welling (2014)) and the concentration parameter α p 0 is assumed to be one.

Original: [Equation]
Revised: [Equation]

Original: The general prior is UDP( G p 0 ,α p 0 ,ϕ p )  where the size distribution ϕ p is determined empirically.
Revised: The general prior is UDP( G p 0 ,α p 0 ,ϕ p ) , where the size distribution ϕ p is determined empirically.

Original: The conditional prior BDP( G p 0 ,α p 0 ,κ 0 ) sets the level of approximation κ 0 as a fixed function of the input length n , in particular κ 0 =( n +1) κ ∆ , where κ ∆ ∈ Z + is a hyperparameter that controls the approximation.
Revised: The conditional prior BDP( G p 0 ,α p 0 ,κ 0 ) sets the level of approximation κ 0 as a fixed function of the input length n , in particular κ 0 =( n +1) κ ∆ , where κ ∆ ∈ Z + is a hyperparameter that controls the approximation.

Original: A Conditional Bounded DP Prior It will be usefulbelow to generalise this conditioning for the level of approximation to any conditional prior which is a fixed function of only the input length.
Revised: A Conditional Bounded DP Prior It will be useful to generalise this conditioning for the level of approximation to any conditional prior which is a fixed function of only the input length.

Original: If we know the input length n , but know nothing about the content of the text, then the distribution of vectors should stay the same as the general prior, G p ′ 

Revised: If we know the input length n , but know nothing about the content of the text, then the distribution of vectors should stay the same as the general prior, G p ′

Original: = G p 0 .
Revised: = G p 0 .

Original: However, the count of observations we expect to have after an input of that length would not be α p 0 , but should include a pseudo-count α ∆ ∈ R ≥ 0 hyperparameter for every token, and thus α p ′ 

Revised: However, the count of observations we expect to have after an input of that length would not be α p 0 , but should include a pseudo-count α ∆ ∈ R ≥ 0 hyperparameter for every token, and thus α p ′

Original: = α p 0 + nα ∆ .
Revised: = α p 0 + nα ∆ .

Original: This then gives us the conditional prior given n of BDP( G p 0 ,α p ′ 0 ,κ 0 ) .
Revised: This then gives us the conditional prior given n of BDP( G p 0 ,α p ′ 0 ,κ 0 ) .

Original: 2. A P OSTERIOR OVER M IXTURE D ISTRIBUTIONS
Revised: 2. A P OSTERIOR OVER M IXTURE D ISTRIBUTIONS

Original: Since a DP is a conjugate prior, we can use exact inference to compute the posterior DP from the prior
Revised: Since a DP is a conjugate prior, we can use exact inference to compute the posterior DP from the prior

Original: DP plus a set of pseudo-observations output by the encoder.
Revised: DP plus a set of pseudo-observations output by the encoder.

Original: Each pseudo-observation is a real-valued pseudo-count α qi ∈ R ≥ 0 and a parametric distribution which represents uncertainty in the observation.
Revised: Each pseudo-observation is a real-valued pseudo-count α qi ∈ R ≥ 0 and a parametric distribution which represents uncertainty in the observation.

Original: We use an isotropic Gaussian, G qi = N ( µ qi , I ( σ qi ) 2 ) , as the parametric distribution, specified by a mean µ qi ∈ R 1 × d and a standard deviation σ qi ∈ R 1 × d > 0 .
Revised: We use an isotropic Gaussian, G qi = N ( µ qi , I ( σ qi ) 2 ) , as the parametric distribution, specified by a mean µ qi ∈ R 1 × d and a standard deviation σ qi ∈ R 1 × d > 0 .

Original: Here we assume that the number of candidate pseudo-observations is the same as the length n of the input, but some of these pseudo-observations may have zero pseudo-counts and thus be effectively removed from the set.
Revised: Here we assume that the number of candidate pseudo-observations is the same as the length n of the input, but some of these pseudo-observations may have zero pseudo-counts and thus be effectively removed from the set.

Original: The formula for exact inference of the posterior F is given in equations 6, where there is an n +1 th component of the base distribution
Revised: The formula for exact inference of the posterior DP is given in equation 6, where there is an n +1 th component of the base distribution

Original: G q 0 which comes from the prior, namely α qn +1 = α p 0 and G qn +1 =
Revised: G q 0 which comes from the prior, namely α qn +1 = α p 0 and G qn +1 =

Original: G p 0 .
Revised: G p 0 .

Original: [Equation]
Revised: [Equation]

Original: We derive an alternative factorisation of the posterior DP (Appendix H) which helps with the sampling method in Section 3.2.
Revised: We derive an alternative factorisation of the posterior DP (Appendix H) which helps with the sampling method in Section 3.2.

Original: We then bound this factorised DP so that it generates the same number κ 0 =( n +1) κ ∆ of weighted vectors as the prior BDP( G p 0 ,α p 0 ,κ 0 ) .
Revised: We then bound this factorised DP so that it generates the same number κ 0 =( n +1) κ ∆ of weighted vectors as the prior BDP( G p 0 ,α p 0 ,κ 0 ) .

Original: The resulting bounded posterior
Revised: The resulting bounded posterior

Original: F ∼ BFDP( G q , α q ,κ ∆ ) is given in equations 7, which define our posterior distribution.
Revised: F ∼ BFDP( G q , α q ,κ ∆ ) is given in equation 7, which defines our posterior distribution q ( F (cid:12)(cid:12) x ) .

Original: This posterior simplifies to a mixture F = (cid:80) n +1 i =1 (cid:80) κ ∆ j =1 ρ
Revised: This posterior simplifies to a mixture F = (cid:80) n +1 i =1 (cid:80) κ ∆ j =1 ρ

Original: i π ′ ij δ z ij of impulse distributions, with ρ ∼ Dir( α q 1 ,...,α qn +1 ) , π ′ i ∼ Dir( α qi κ ∆ , κ ∆ ... , α qi κ ∆ ) , and z ij ∼ G qi .
Revised: i π ′ ij δ z ij of impulse distributions, with ρ ∼ Dir( α q 1 ,...,α qn +1 ) , π ′ i ∼ Dir( α qi κ ∆ , κ ∆ ... , α qi κ ∆ ) , and z ij ∼ G qi .

Original: T HE N ONPARAMETRIC V ARIATIONAL I NFORMATION B OTTLENECK
Revised: The Mean Posterior Mixture Distribution A VAE is trained on samples from the posterior, but at test time VAEs typically use the mean of this distribution. Generalising the latent space to mixture distributions makes this straightforward, since the mean of our BFDP posterior is its base distribution G q 0 . This base distribution is a continuous distribution, whereas at training time all samples are discrete distributions. Nonetheless, when accessed via denoising attention, the base distribution looks like a typical sample from the posterior.

Revised: This is visualised in Figure 1 by comparing the vector returned by denoising attention (in green) given the continuous mean distribution (Figure 1c) and a typical sample from this distribution (Figure 1b).

Revised: Thus, the function defined by applying denoising attention to a sampled distribution can be seen as a noisy version of the function defined by applying denoising attention to the mean distribution.

Revised: For our Gaussian mixture base distribution G q 0 , there is a closed-form solution to computing denoising attention, given in Appendix F.

Revised:  T HE N ONPARAMETRIC V ARIATIONAL I NFORMATION B OTTLENECK

Original: By generalising attention-based representations to mixture distributions and generalising the attention function to denoising attention, we can define a VIB regulariser for attention-based interfaces.
Revised: By generalising attention-based representations to mixture distributions and generalising the attention function to denoising attention, we can define a VIB regulariser for attention-based interfaces.

Revised: Such encoder-decoder interfaces take a set-of-vectors representation and return a function from query to result vectors.

Original: We map its input set-of-vectors to a set of pseudo-observations, and we replace its output attention function with denoising attention.
Revised: We map the input set-of-vectors to a set of pseudo-observations, and define the returned function with denoising attention.

Original: Then, given the nonparametric prior and posterior from Section 2, we can define our nonparametric VIB regulariser by specifying: how to compute the KL divergence between  prior and posterior; how to effectively sample from the posterior for training; and the mean of these samples to be used at test time.
Revised: Then, given the nonparametric prior and posterior from Section 2, we can define our nonparametric VIB regulariser by specifying how to compute the KL divergence between the prior and posterior, and how to effectively sample from the posterior for training.

Original: As far as we are aware, this proposal is the first VIB model for attention-based representations like Transformer embeddings.
Revised: As far as we are aware, this proposal is the first VIB model for attention-based representations like Transformer embeddings.

Original: The VIB layer of a VAE controls the amount of information passing through it by introducing noise according to a posterior output by the encoder, and regularises this information by minimising the KL divergence between this posterior and a prior.
Revised: The VIB layer in a VAE controls the amount of information passing through it by introducing noise according to a posterior output by the encoder, and regularises this information by minimising the KL divergence between this posterior and an uninformative prior.

Original: One of the known difficulties with VAEs is that they can be difficult to train due to the posterior collapsing to the prior.
Revised: One of the known difficulties with VAEs is that they can be difficult to train due to the posterior collapsing to the prior (Bowman et al., 2016).

Original: Similarly to the free-bits objective proposed to address this problem for vector-space VAEs (Kingma et al., 2016), instead of regularising towards the prior, which gets no pseudo-counts from the input, we regularise towards the conditional prior BDP( G p 0 ,α p ′ 0 ,κ 0 ) , which gets nα ∆ pseudo-counts from the input but knows nothing about the information they convey.
Revised: Similarly to the freebits objective proposed to address this problem for vector-space VAEs (Kingma et al., 2016), instead of regularising towards the prior, which gets no pseudo-counts from the input, we regularise towards the conditional prior BDP( G p 0 ,α p ′ 0 ,κ 0 ) , which gets nα ∆ pseudo-counts from the input but knows nothing about the information they convey.

Original: We found that this helps with the stability of training and avoiding posterior collapse.
Revised: We found that this helps with the stability of training and avoiding posterior collapse.

Original: 3. T HE V ARIATIONAL I NFORMATION B OTTLENECK L OSS
Revised: 3. T HE V ARIATIONAL I NFORMATION B OTTLENECK L OSS

Original: The ELBO is commonly used in variational Bayesian methods as an objective which approximately maximises the log-likelihood of the observation x , wherehere x is the input text.
Revised: The evidence lower bound (ELBO) is commonly used in variational Bayesian methods as an objective which approximately maximises the log-likelihood of the observation x , where x is the input text.

Original: [Equation]
Revised: [Equation]

Original: The first term of the bound is the reconstruction loss( L R ) computed from samples F from the approximate posterior q ( F (cid:12)(cid:12) x ) , and the second term is the KL divergence between this posterior and the prior p ( F ) .
Revised: The first term of the bound is the reconstruction loss L R , computed using samples F from the approximate posterior q ( F (cid:12)(cid:12) x ) , and the second term is the KL divergence between this posterior and the prior p ( F ) .

Original: For the KL divergence term, since both the prior and the posterior are conditioned on the same bound κon the number of vectors they generate, we can compute a meaningful finite KL divergence between the conditional prior p ( F )=
Revised: For the KL divergence term, since both the prior and the posterior are conditioned on the same bound κon the number of vectors they generate, we can compute a meaningful finite KL divergence between the conditional prior p ( F )=

Original: BDP( G p 0 ,α p ′ 0 ,κ 0 ) and the posterior q ( F (cid:12)(cid:12) x )
Revised: BDP( G p 0 ,α p ′ 0 ,κ 0 ) and the posterior q ( F (cid:12)(cid:12) x )

Original: =BFDP( G q , α q ,κ ∆ ) , given κ ∆ and κ 0 =( n +1) κ ∆ .
Revised: =BFDP( G q , α q ,κ ∆ ) , given κ ∆ and κ 0 =( n +1) κ ∆ .

Original: The derivation of the KL divergence is given in Appendix I.
Revised: The derivation of the KL divergence is given in Appendix I.

Original: It has two terms, one( L D) for the distribution of weights π generated by the Dirichlet distributions, and one( L G) for the distribution of vectors Z generated by the component Gaussians.
Revised: It has two terms, one L D for the distribution of weights π generated by the Dirichlet distributions, and one L G for the distribution of vectors Z generated by the component Gaussians.

Original: Using the exact KL divergence for our bounded DPs would regularise each component equally, even for components which have zero α qi and thus have no impact on the posterior.
Revised: Using the exact KL divergence for our bounded DPs would regularise each component equally, even for components which have zero α qi and thus have no impact on the posterior.

Original: We instead approximate a KL divergence where only samples with a nontrivial weight are regularised.
Revised: We instead approximate a KL divergence where only samples with a nontrivial weight are regularised.

Original: Marginalising over the number of nontrivial weights for each component does not appear to be tractable for L D , but since the relationship is approximately linear (see Appendix I), we simply substitute the expected number α qi α q 0 κ 0 for the actual number in the equation for KL divergence.
Revised: Marginalising over the number of nontrivial weights for each component does not appear to be tractable for L D , but since the relationship is approximately linear (see Appendix I), we simply substitute the expected number α qi α q 0 κ 0 for the actual number in the equation for KL divergence.

Original: This gives us the following loss terms for the KL divergence.
Revised: This gives us the following loss terms for the KL divergence, where Γ is the gamma function and ψ is the digamma function.

Original: [Equation]
Revised: [Equation]

Original: where Γ is the gamma function and ψ is the digamma function.

Original: Both L D and L G scale approximately linearly with κ 0 .
Revised: Both L D and L G scale approximately linearly with κ 0 .

Original: The view of the ELBO  as a way to regularise the amount of information which passes through a latent representationallows it to be generalised beyond autoencoders (Alemi et al., 2017).
Revised: To generalise the ELBO beyond autoencoders, it can be viewed as a way to regularise the amount of information which passes through the latent representation (Alemi et al., 2017).

Original: Under this motivation, there is no need for the different parts of the objective to have the same weight.
Revised: This VIB interpretation allows the different parts of the objective to have different weights.

Original: In this paper, we split the ELBO into three parts and introduce two hyperparameters to control the relative weight of these parts, which defines our VIB loss L .
Revised: We introduce two hyperparameters to control the relative weight of the above three parts of the ELBO, which defines our VIB loss L .

Original: [Equation]
Revised: [Equation]

Original: 3. S AMPLING A M IXTURE D ISTRIBUTION FROM THE P OSTERIOR
Revised: 3. S AMPLING A M IXTURE D ISTRIBUTION FROM THE P OSTERIOR

Original: To control the amount of information which passes from the encoder to the decoder, at training time a
Revised: To control the amount of information which passes from the encoder to the decoder, at training time a

Original: VAE (Kingma & Welling, 2014) samples from the encoder’s posterior distribution and uses this sample to reconstruct the input.
Revised: VAE (Kingma & Welling, 2014) samples from the encoder’s posterior distribution and uses this sample to reconstruct the input.

Original: The “reparameterisation trick” is used to ensure that backpropagation of the reconstruction error through this sampling step can be done effectively.
Revised: The “reparameterisation trick” is used to ensure that backpropagation of the reconstruction error through this sampling step can be done effectively.

Original: We propose a novel reparameterisation trick for bounded Dirichlet processes which allows sampling without any categorical choices, and propose specific sampling methods which result in effective backpropagation through the sampling step.
Revised: We propose a novel reparameterisation trick for bounded Dirichlet processes which allows sampling without any categorical choices, and propose specific sampling methods which result in effective backpropagation through the sampling step.

Original: For our NVIB model, we sample the parameters ⟨ π , Z ⟩ of a mixture distribution F generated by our bounded Dirichlet process posterior BFDP( G q , α q ,κ ∆ ) , where F consists of a set of impulse distributions δ z k each with a weight π k .
Revised: For our NVIB model, we sample the parameters ⟨ π , Z ⟩ of a mixture distribution F generated by our bounded Dirichlet process posterior BFDP( G q , α q ,κ ∆ ) , where F consists of a set of impulse distributions δ z i each with a weight π i .

Original: A straightforward approach to sampling from a Dirichlet process would independently sample weights π from a (theoretically infinite) Dirichlet distribution and sample vectors
Revised: A straightforward approach to sampling from a Dirichlet process would independently sample weights π from a (theoretically infinite) Dirichlet distribution and sample vectors

Original: Z from the base distribution of the DP, where sampling from the base distribution involves first sampling a component of the base distribution and then sampling a vector from that component’s Gaussian.
Revised: Z from the base distribution of the DP, where sampling from the base distribution involves first sampling a component of the base distribution and then sampling a vector from that component’s Gaussian.

Original: A Factorised Sampling Method We propose a factorised sampling method where every component of the base distribution has vectors sampled from it an equal (theoretically infinite) number of times, and the only difference between components is the distribution of weights for these different sets of vectors.
Revised: A Factorised Sampling Method The problem is that sampling a component is a discrete choice, for which there is no exact reparameterisation trick. Instead, we note that the components do not differ in the number of vectors sampled from each one (always theoretically infinite for a DP), but only differ in the distribution of weights for those vectors.

Original: These weight distributions are specified in the factorised DP in equations 7. Rather than doing a sequence of discrete choices between components, this factorisation of a DP does a continuous choice of a set of weights ρ , each one corresponding to the total weight  for one component. Each total weightρ i is then combined with individual weights π ′ i and vectors
Revised: As specified in the factorised DP in Section 2.3, we characterise this distribution over weights by factorising it into two steps: first choosing how the total weight is distributed across components ( ρ ), and then for each component choosing how its weight is distributed across its vectors ( π ′ i ). These are both continuous choices.

Original: Z i sampled for each component independently.
Revised: The vectors can then be sampled independently from each component.

Original: Reparameterisation trick The factorised sampling allows for independent sampling of weights and components. Since our components use Gaussian distributions over vectors,  we can use the same reparameterisation trick as Kingma & Welling (2014) for sampling vectors
Revised: Reparameterisation tricks Each individual component specifies a Gaussian distribution over vectors, so we can use the same reparameterisation trick as Kingma & Welling (2014) for sampling vectors

Original: Z i from an individual component i .
Revised: Z i from an individual component i .

Original: The  bounded nature of our posterior BFDP( G q , α q ,κ ∆ ) means that the total weights ρ and the individual weights π
Revised: The factorised and bounded nature of our posterior BFDP( G q , α q ,κ ∆ ) means that the total weights ρ and the individual weights π

Original: ′ i are all sampled from Dirichlet distributions.
Revised: ′ i are all sampled from Dirichlet distributions.

Original: A Dirichlet distribution over a set of category weights can be sampled by sampling from a Gamma distribution for each category and then normalising.
Revised: A Dirichlet distribution over a set of category weights can be sampled by sampling from a Gamma distribution for each category and then normalising.

Original: There is no closed-form, explicit reparameterisation trick for the exact Gamma distribution, but there are for approximations.
Revised: There is no closed-form, explicit reparameterisation trick for the exact Gamma distribution, but there are for approximations.

Original: Knowles (2015) proposes two such approximations which we combine, one which is more accurate for small α and one which is more accurate for larger α .
Revised: Knowles (2015) proposes two such approximations which we combine, one which is more accurate for small α and one which is more accurate for larger α .

Original: We leave the investigation of implicit reparameterisation gradients (Figurnov et al., 2018), an alternative approach to explicit reparameterisation, to future work.
Revised: We leave the investigation of implicit reparameterisation gradients (Figurnov et al., 2018), an alternative approach to explicit reparameterisation, to future work.

Original: More specifics about the sampling methods and their reparameterisation tricks is given in Appendix J. We provide an ablation (Appendix B.3) showing that κ ∆ =1 is an efficient and effective sampling method.
Revised: More specifics about the sampling methods and their reparameterisation tricks is given in Appendix J. We provide an evaluation (Appendix B.3) showing that κ ∆ =1 is an efficient and effective sampling method.

Original: In this case, there is no need to sample from the Dirichlet for each individual component, but we still sample the weights ρ across components.
Revised: In this case, there is no need to sample from the Dirichlet for each individual component, but we still sample the weights ρ across components.

Original: 3. T HE M EAN M IXTURE D ISTRIBUTION AT T EST T IME

Original: During test time evaluation, instead of sampling, VAEs typically use the mean of the sampling distribution, which is the base distribution G q 0 for our BFDP posterior.

Original: Despite the base distribution being a continuous distribution, whereas at training time all samples are discrete distributions, when accessed via denoising attention the base distribution looks like a typical sample from the posterior.

Original: This is visualised in Figureby comparing the vector returned by denoising attention (in green) given the continuous mean distribution (Figure 1c) and a typical sample from this distribution (Figure 1b).

Original: Thus, the function defined by applying denoising attention to a sampled distribution can be seen as a noisy version of the function defined by applying denoising attention to the mean distribution.

Original: For our Gaussian mixture base distribution G q 0 , there is a closed-form solution to computing denoising attention, given in Appendix F .

Original:  T HE N ONPARAMETRIC V ARIATIONAL A UTOENCODER
Revised:  T HE N ONPARAMETRIC V ARIATIONAL A UTOENCODER

Original: We define a VAE for Transformers by using the nonparametric VIB defined in Section 3 to regularise the attention-based representation between the encoder and decoder of a Transformer autoencoder, as depicted in Figure 1a.
Revised: We define a VAE for Transformers by using the nonparametric VIB defined in Section 3 to regularise the attention-based representation between the encoder and decoder of a Transformer autoencoder, as depicted in Figure 1a.

Original: In this NVAE model, the Transformer encoder is used to estimate the parameters ⟨ α q , µ q , σ q ⟩ of the posterior given the input text x .
Revised: In this NVAE model, the Transformer encoder is used to estimate the parameters ⟨ α q , µ q , σ q ⟩ of the posterior given the input text x .

Original: The Transformer decoder is used to reconstruct the input text x using denoising attention over a sample F from this posterior.
Revised: The Transformer decoder is used to reconstruct the input text x using denoising attention over a sample F from this posterior.

Original: The NVIB Regulariser Our NVIB layer regularises the amount of information which passes from the encoder to the decoder through this posterior.
Revised: The NVIB Regulariser Our NVIB layer regularises the amount of information which passes from the encoder to the decoder through this posterior.

Original: As with VIB for vector spaces, the KL divergence encourages the encoder to output component Gaussians with smaller µ qi and larger σ qi .
Revised: As with VIB for vector spaces, the KL divergence encourages the encoder to output component Gaussians with smaller µ qi and larger σ qi .

Original: With NVIB, the KL divergence also encourages the encoder to output smaller and sparser α qi , which regularises the effective number of components in the posterior as well as the noisiness of their weights.
Revised: With NVIB, the KL divergence also encourages the encoder to output smaller and sparser α qi , which regularises the effective number of components in the posterior as well as the noisiness of their weights.

Original: The Transformer Encoder The Transformer encoder q ( F | x ) with a text x of n number of tokens is used to compute a vector for each token i of the input.
Revised: The Transformer Encoder The Transformer encoder q ( F | x ) with a text x of n number of tokens is used to compute a vector for each token i of the input.

Original: From each of these n individual token embeddings, the encoder then linearly projects to three parameters, α qi ∈ R , µ qi ∈ R 1 × d and log( σ qi ) ∈ R 1 × d .
Revised: From each of these n individual token embeddings, the encoder then linearly projects to three parameters, α qi ∈ R , µ qi ∈ R 1 × p and log( σ qi ) ∈ R 1 × p .

Original: The variance parameters are exponentiated to be strictly positive, whereas the pseudo-count parameters α qi are estimated using a Rectified Linear Unit (ReLU) activation (Nair & Hinton, 2010), which results in masking the vector during cross-attention when it is exactly zero.
Revised: The variance parameters are exponentiated to be strictly positive, whereas the pseudo-count parameters α qi are estimated using a Rectified Linear Unit (ReLU) activation (Nair & Hinton, 2010), which results in masking the vector during cross-attention when it is exactly zero.

Original: Thus, the DP posterior has one component ⟨ α qi , µ qi , σ qi ⟩ of its base distribution for each token of the input (plus one for the prior).
Revised: Thus, the DP posterior has one component ⟨ α qi , µ qi , σ qi ⟩ of its base distribution for each token of the input (plus one for the prior).

Original: The Transformer Decoder The Transformer decoder q ( x | F ) receives a distribution F over vectors and reconstructs the input text x .
Revised: The Transformer Decoder The Transformer decoder q ( x | F ) receives a distribution F over vectors and reconstructs the input text x .

Original: During training, F is specified by the sampled vectors Z ∈ R n × d and the sampled weights π ∈ R n × 1 , and at test time F is specified by the output of the encoder α q ∈ R n × 1 , µ q ∈ R n × d and log ( σ ) q ∈ R n × d .
Revised: During training, F is specified by the sampled vectors Z ∈ R κ 0 × p and the sampled weights π ∈ R κ 0 × 1 , and at test time F is specified by the output of the encoder α q ∈ R κ 0 × 1 , µ q ∈ R κ 0 × p and log ( σ ) q ∈ R κ 0 × p .

Original: In both cases, the decoder accesses F using denoising attention in the same way that standard Transformer decoders use cross attention.
Revised: In both cases, the decoder accesses F using denoising attention in the same way that standard Transformer decoders use cross attention.

Original: We include the exact equations used for a deep learning implementation of denoising attention in Appendix K.
Revised: We include the exact equations used for a deep learning implementation of denoising attention in Appendix K.

Original: During training, the text is predicted using teacher forcing, and during test time the text is predicted autoregressively using greedy decoding until the end of sequence token is generated or the sentence generated is 50 tokens larger than the target length.
Revised: During training, the text is predicted using teacher forcing, and during test time the text is predicted autoregressively using greedy decoding until the end-of-sequence token is generated or the sentence generated is 50 tokens larger than the target length.

Original: The Generative Model To use our NVAE model as a generative model, we sample from the prior and use the trained Transformer decoder to generate a sentence.
Revised: The Generative Model To use our NVAE model as a generative model, we sample from the prior and use the trained Transformer decoder to generate a sentence.

Original: As discussed in Section 2.2, to sample from the same prior as used for training, we need to first sample a sentence length, and then sample from the conditional prior given that (approximate) sentence length.
Revised: As discussed in Section 2.2, to sample from the same prior as used for training, we need to first sample a sentence length, and then sample from the conditional prior given that (approximate) sentence length.

Original: For simplicity, we sample the sentence length from the empirical distribution of sentence lengths in the training data.
Revised: For simplicity, we sample the sentence length from the empirical distribution of sentence lengths in the training data.

Original:  I NTRINSIC E VALUATIONS OF NVIB IN NVAE
Revised:  I NTRINSIC E VALUATIONS OF NVIB IN NVAE

Original: To support our theoretical contributions, we  demonstrate that our proposed NVIB regulariser performs as predicted.
Revised: To support our theoretical contributions, we provide proof-of-concept experiments which demonstrate that our proposed NVIB regulariser performs as claimed.

Original: We evaluate it in our proposed NVAE model by training NVAEs on natural language text and evaluating the resulting models.
Revised: We evaluate it in our proposed NVAE model by training NVAEs on natural language text and evaluating the resulting models.

Original: We show that the NVAE is a viable VAE model as it exhibits a competitive reconstruction versus generation trade-off (Section 5.1).
Revised: We show that the NVAE is a viable VAE model as it exhibits a competitive reconstruction versus generation trade-off (Section 5.1).

Original: We show that the NVIB layer is able to dynamically choose the number of components it needs in its embeddings (Section 5.2).
Revised: We show that the NVIB layer is able to dynamically choose the number of components it needs in its embeddings (Section 5.2).

Original: Additionally, NVIB provides an intuitive way to interpolate between sentence embeddings, which provides an evaluation of the smoothness of the latent space (Section 5.3).
Revised: Additionally, NVIB provides an intuitive way to interpolate between sentence embeddings, which provides an evaluation of the smoothness of the latent space (Section 5.3).

Original: Data The Wikitext-2 and Wikitext-103 (Merity et al., 2017) encyclopedia datasets were selected as they are general English language corpora of a small and large scale containing high quality Wikipedia articles.
Revised: Data The Wikitext-2 and Wikitext-103 (Merity et al., 2017) encyclopedia datasets were selected as they are general English language corpora of a small and large scale containing high quality Wikipedia articles.

Original: Baselines We compare to various alternative ways to define a VAE from a Transformer autoencoder.
Revised: Baselines We compare to various alternative ways to define a VAE from a Transformer autoencoder.

Original: As representative of a standard fixed-length-vector VAE, the Variational Transformer Pooled (VTP) baseline pools its vectors across the sequence length dimension, and then applies a Gaussian VIB layer (Kingma & Welling, 2014).
Revised: As representative of a standard fixed-length-vector VAE, the Variational Transformer Pooled (VTP) baseline pools its vectors across the sequence length dimension, and then applies a Gaussian VIB layer (Kingma & Welling, 2014).

Original: At the other extreme, the Variational Transformer (VT) baseline keeps all its vectors and applies a Gaussian VIB layer to each one.
Revised: At the other extreme, the Variational Transformer (VT) baseline keeps all its vectors and applies a Gaussian VIB layer to each one.

Original: In between these baselines, as a hand-coded solution to constraining the quantity of latent vectors, Variational Transformer Stride (VTS) baselines, with parameter
Revised: In between these baselines, as a hand-coded solution to constraining the quantity of latent vectors, Variational Transformer Stride (VTS) baselines, with parameter

Original: S , masks 1 − S proportion of the embedding vectors based on their position.
Revised: S , masks 1 − S proportion of the embedding vectors based on their position.

Original: For comparability, all our baselines only differ from the NVAE model in the latent representation between the encoder and decoder.
Revised: For comparability, all our baselines only differ from the NVAE model in the latent representation between the encoder and decoder, with the same Transformer encoder and Transformer decoder architectures.

Original: 5. R ECONSTRUCTION VERSUS G ENERATION
Revised: 5. R ECONSTRUCTION VERSUS G ENERATION

Original: This section shows that the NVAE Transformer model is a competitive VAE in both reconstruction of input sentences and generation by sampling from the prior.
Revised: This section shows that the NVAE Transformer model is a competitive VAE in both reconstruction of input sentences and generation by sampling from the prior.

Original: All models undergo hyperparameter tuning on the validation set (Appendix B) across 5 seeds, to select the best models and then report results on the Wikitext-2 test set.
Revised: All models undergo hyperparameter tuning on the validation set (Appendix B) across 5 seeds, to select the best models and then report results on the Wikitext-2 test set.

Original: For reconstruction, we report the Sacre BLEU metric (Papineni et al., 2002; Post, 2018).
Revised: For reconstruction, we report the SacreBLEU metric (Papineni et al., 2002; Post, 2018).

Original: For generation, we report forward perplexity (F-PPL) and reverse perplexity (R-PPL)
Revised: For generation, we report forward perplexity (F-PPL) and reverse perplexity (R-PPL)

Original: (Zhao et al., 2017; C ´ ıfka et al., 2018), which trains an external language model on the gold training text and evaluates it on the generated text (F-PPL) or vice versa (R-PPL).
Revised: (Zhao et al., 2018; C´ıfka et al., 2018), which trains an external language model on the gold training text and evaluates it on the generated text (F-PPL) or vice versa (R-PPL).

Original: Training details are provided in Appendix A.
Revised: Training details are provided in Appendix A.

Original: Figures 2a and 2b show the reconstruction versus generation trade-off, where lower right is better.
Revised: Figures 2a and 2b show the reconstruction versus generation trade-off on the Wikitext-2 test set, where lower right is better.

Original: The single vector baseline VTP is unable to reconstruct well (low BLEU) or generate diverse sentences (high R-PPL).
Revised: The single vector baseline VTP is unable to reconstruct well (low BLEU) or generate diverse sentences (high R-PPL).

Original: Even with larger capacity  the model performs poorly in generation (AppendixSection B.4).
Revised: Even with larger capacity and more training data, the model performs poorly in generation (Appendix B.4).

Original: The full vector baseline VT is unable to  generate fluent sentences (high F-PPL) whereas the location-based dropout baseline VTS shows that regularisation of the space is beneficial for both reconstruction and generation quality.
Revised: The full vector baseline VT is unable to consistently generate fluent sentences (high F-PPL), whereas the position-based dropout baseline VTS shows that regularisation of the space is beneficial for both reconstruction and generation quality.

Original: The best NVAE models are competitive with the best VTS baselines in being able to both reconstruct and generate, even slightly better at generation with high reconstruction accuracy.
Revised: The best NVAE models are competitive with the best VTS baselines in being able to both reconstruct and generate, even slightly better at generation with high reconstruction accuracy.

Original: In particular, reverse perplexity shows that NVAEs are able to generate a diverse collection of sentences from the prior, and forward perplexity shows that these are fluent sentences (samples in Appendix C).
Revised: In particular, reverse perplexity shows that NVAEs are able to generate a diverse collection of sentences from the prior, and forward perplexity shows that these are fluent sentences (samples in Appendix C).

Original: We believe this advantage is the result of learning what vectors to keep instead of a hand-coded location-based dropout. 
Revised: We believe this advantage is the result of learning what vectors to keep instead of a hand-coded position-based dropout. [Figure] (a)

Original: [Figure]
Revised: [Figure]

Original: 5. R EGULARISATION
Revised: 5. R EGULARISATION

Original: This section shows that the NVIB layer is able to regularise the number of vectors in the latent representation of a NVAE.
Revised: This section shows that the NVIB layer is able to regularise the number of vectors in the latent representation of a NVAE.

Original: Without regularisation the NVIB layer becomes a standard transformer without any noise and retaining all latent vectors (Ablation in AppendixSection B.1).
Revised: Without regularisation the NVIB layer becomes a standard transformer without any noise and retaining all latent vectors (ablation in Appendix B.1).

Original: The proportion of vectors retained is controllable by the conditional prior hyperparameter α ∆ .
Revised: The proportion of vectors retained is controllable by the conditional prior hyperparameter α ∆ .

Original: Figure 2c shows that there exist α ∆ where the NVAE is able to remove a significant proportion (1 − ν ) of vectors whilst maintaining a high reconstruction performance.
Revised: Figure 2c shows that there exist α ∆ where the NVAE is able to remove a significant proportion (1 − ν ) of vectors whilst maintaining a high reconstruction performance.

Original: Figure 2d plots the number of latent vectors retained during evaluation against the number of input tokens.
Revised: Figure 2d plots the number of latent vectors retained during evaluation against the number of input tokens.

Original: The NVAE models are able to learn to dynamically regularise the number of vectors based on the information within the text, without hand-coding a function of length as in VTS.
Revised: The NVAE models are able to learn to dynamically regularise the number of vectors based on the information within the text, without hand-coding a function of length as in VTS.

Original: The different NVAE lines show the effect of α ∆ on vector retention, with larger values allowing more vectors to be used (also see the ablation in AppendixSection B.2).
Revised: The different NVAE lines show the effect of α ∆ on vector retention, with larger values allowing more vectors to be used (also see the evaluation in Appendix B.2).

Original: 5. I NTERPOLATION
Revised: 5. I NTERPOLATION

Original: The NVIB framework provides an intuitive interpolation between latent sets of vectors that overcomes the challenges with the baselines of latent vector alignment (Ablation in AppendixSection D) and variable set sizes.
Revised: The NVIB framework provides an intuitive interpolation between latent sets of vectors that overcomes the challenges with the baselines of latent vector alignment (analysis in Appendix D) and variable set sizes.

Original: We simply interpolate the probability assigned to each point in vector space.
Revised: We simply interpolate the probability assigned to each point in vector space.

Original: Given two latent mixture distributions F 1 and F 2 , we decode from the combined mixture distribution ( τF 1 +(1 − τ ) F 2 ) , for varying interpolation rates 0 ≤ τ ≤ 1 .
Revised: Given two latent mixture distributions F 1 and F 2 , we decode from the combined mixture distribution ( τF 1 +(1 − τ ) F 2 ) , for varying interpolation rates 0 ≤ τ ≤ 1 .

Original: For the baselines we use τ Z 1 +(1 − τ )
Revised: For the baselines we use τ Z 1 +(1 − τ )

Original: Z 2 such that the interpolation is over the content of the latent vectors
Revised: Z 2 such that the interpolation is over the content of the latent vectors

Original: Z i .
Revised: Z i .

Original: We align latent sets of vectors by input position and pad smaller sets with zero vectors, which is the mean of the Gaussian prior.
Revised: We align latent sets of vectors by input position and pad smaller sets with zero vectors, which is the mean of the Gaussian prior.

Original: We evaluate the interpolation with selected larger scale models (Ablation in AppendixSection B.4) and use the Wikitext-103 validation set for Sand its reverse order for S 2 .
Revised: We evaluate the interpolation with selected larger scale models (see Appendix B.4) and use the Wikitext-103 validation set for S 1 and its reverse order for S 2 .

Original: The results in Table 1 and Figure 3 show that the NVIB regulariser in NVAE provides smoother interpolations, and improvements in fluency of interpolations over the baselines.
Revised: The results in Table 1 and Figure 3 show that the NVIB regulariser in NVAE provides smoother interpolations, and improvements in fluency of interpolations over the baselines.

Original: 3 These finding are qualitatively confirmed through different examples (AppendixSection E).
Revised: These findings are qualitatively confirmed through different examples (Appendix E).

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: In this work, we generalise the latent representations of attention-based models to mixture distributions over a vector space, and propose a nonparametric variational information bottleneck (NVIB) to regularise these latent representations.
Revised: In this work, we generalise the latent representations of attention-based models to mixture distributions over a vector space, and propose a nonparametric variational information bottleneck (NVIB) to regularise these latent representations.

Original: Using this NVIB model, we propose a nonparametric variational autoencoder (NVAE), which uses a Transformer encoder to embed text in a nonparametric space of distributions over mixture distributions, and uses a Transformer decoder to generate text given a sampled mixture distribution.
Revised: Using this NVIB model, we propose a nonparametric variational autoencoder (NVAE), which uses a Transformer encoder to embed text in a nonparametric space of distributions over mixture distributions, and uses a Transformer decoder to generate text given a sampled mixture distribution.

Original: This nonparametric Bayesian formalisation of attention-based representations captures two key properties of the attention function, namely its invariance to permutations of its input vectors, and that this input can vary widely in size.
Revised: This nonparametric Bayesian formalisation of attention-based representations captures two key properties of the attention function, namely its invariance to permutations of its input vectors, and that this input can vary widely in size.

Original: Our NVIB model adds the ability to regularise attention-based representations so that the size of the representation is appropriate for the complexity of the input being encoded.
Revised: Our NVIB model adds the ability to regularise attention-based representations so that the size of the representation is appropriate for the complexity of the input being encoded.

Original: This is a crucial ability for encoding text, where the size of a text can vary enormously.
Revised: This is a crucial ability for encoding text, where the size of a text can vary enormously.

Original: Empirical evaluations indicate that this model: is a competitive VAE in that it is able to reconstruct input sentences and generate a good distribution over sentences from the prior; regularises the size of the induced latent representations as desired; and is able to intuitively interpolate  between latent mixture distributions.
Revised: Empirical evaluations indicate that this model: is a competitive VAE in that it is able to reconstruct input sentences and generate a good distribution over sentences from the prior; regularises the size of the induced latent representations as desired; and is able to intuitively interpolate smoothly between latent mixture distributions.

Original: Current work is looking at larger scale experiments. Future work will evaluate the effectiveness of NVIB  when applied to multi-head attention and to self-attention layers, and evaluate the usefulness of these pretrained representations  in downstream tasks. 
Revised: Future work will evaluate the effectiveness of NVIB at a larger scale, when applied to multi-head attention and to self-attention layers, and when its pretrained representations are used in downstream tasks. A UTHOR C ONTRIBUTIONS

Revised: James Henderson is responsible for the high-level vision and the theoretical derivations and proofs related to Bayesian nonparametrics.

Revised: Fabio Fehr is responsible for the sampling method, evaluations, implementations and running of experiments.

Original: R EPRODUCIBILITY STATEMENT
Revised: A CKNOWLEDGEMENTS

Revised: We would like to thank Florian Mai, Andrei Catalin Coman, Melika Behjati, Andreas Marfurt, and other members of the Idiap NLU group for helpful comments and suggestions.

Revised: This work was funded in part by the Swiss National Science Foundation under the NCCR grant Evolving Language, Swiss National

Revised: Science Foundation Agreement #51NF40 180888.

Revised: R EPRODUCIBILITY STATEMENT

Original: We faithfully describe the details of method in the text and provide detailed derivations for theoretical grounding (Sections  G, H,I).
Revised: We faithfully describe the details of method in the text and provide detailed derivations for theoretical grounding (Sections F, G, H, I).

Original: We mention all relevant hardware, hyperparameters and datasets to reproduce our experiments (Section A).
Revised: We mention all relevant hardware, hyperparameters and datasets to reproduce our experiments (Section A).

Original: We provide variance estimations across seeds and ablations to justify design choices (Section B).
Revised: We provide variance estimations across seeds and ablations to justify design choices (Section B).

Original: Finally, we provide exact equations for denoising attention used to make implementation easier (Section K).
Revised: Finally, we provide exact equations for denoising attention used to make implementation easier (Section K).

Original: We shall release the open source code for this research before publication
Revised: The open source code for this research has been released at https://github.com/idiap/nvib for the NVIB layer in PyTorch and https://github.com/idiap/nvib transformers for the experiments.

Original: A E XPERIMENTAL S ETUP
Revised: A E XPERIMENTAL S ETUP

Original: Training details All models are trained, without pretraining using the same encoder and decoder configuration for comparability.
Revised: Training details All models are trained, without pretraining, using the same encoder and decoder configuration for comparability.

Original: We use a two layer Transformer encoder and decoder with a single attention-head.
Revised: We use a two layer Transformer encoder and decoder with a single attention-head.

Original: The size for the word embedding vectors and model projections are 256 , feed forward dimensions 1024 , which leads to models of approximately 19 million trainable parameters.
Revised: The size for the word embedding vectors and model projections are 256 , feed forward dimensions 1024 , which leads to models of approximately 19 million trainable parameters.

Original: The BERT base-uncased tokeniser is used for tokenisation with a vocabulary of approximately 30 K. During training we use: a constant learning rate of 1 e − 4 , Adam optimiser (Kingma & Ba, 2015), a batch size of 256 , gradient norm clipping 0 . 1 and trained for 50 epochs ( ≈ 15 K steps).
Revised: The BERT base-uncased tokeniser is used for tokenisation with a vocabulary of approximately 30 K. During training we use: a constant learning rate of 1 e − 4 , Adam optimiser (Kingma & Ba, 2015), a batch size of 256 , gradient norm clipping 0 . 1 and trained for 50 epochs ( ≈ 15 K steps).

Original: The number  epochs were selected considering model convergence and minimising computation time.
Revised: The number of epochs were selected considering model convergence and minimising computation time.

Original: As a form of regularisation we use a dropout rate of 0 . 
Revised: As a form of regularisation we use a dropout rate of 0 . 1 and the

Original: 1 and the VIB parameters λ G , λ D , α ∆ and κ ∆ are selected througha hyperparameter tuning.
Revised: VIB parameters λ G , λ D , α ∆ and κ ∆ are selected through hyperparameter tuning.

Original: Learning rate schedules, KL annealing strategies and free-bits KL loss objectives were considered but unnecessary for convergence.
Revised: Learning rate schedules, KL annealing strategies and free-bits KL loss objectives were considered but unnecessary for convergence.

Original: Each model experiment takes approximately 2hrs to  on  single NVIDIA GeForce RTX 3090.
Revised: Each model experiment takes approximately 2hrs to run on a single NVIDIA GeForce RTX 3090.

Original: Generation metrics
Revised: Generation metrics

Original: The two automatic generation metrics forward/reverse perplexity (Zhao et al., 2017;
Revised: The two automatic generation metrics forward/reverse perplexity (Zhao et al., 2018;

Original: C´ıfka et al., 2018)
Revised: C´ıfka et al., 2018)

Original: both involve training an external language model, for which we use a Transformer language  with the same configuration, without VIB regularisation.
Revised: both involve training an external language model, for which we use a Transformer language model with the same configuration as in the previous paragraph, but without VIB regularisation.

Original: First we generate 100k sample sentences from the model we wish to evaluate.
Revised: First we generate 100k sample sentences from the model we wish to evaluate.

Original: The forward perplexity (F-PPL) measures the perplexity of the external language model trained on training data and evaluated on the generated text .
Revised: The forward perplexity (F-PPL) measures the perplexity of the external language model trained on training data and evaluated on the generated text.

Original: This measures fluency of the text.
Revised: This measures fluency of the text.

Original: The reverse perplexity (R-PPL) measures the perplexity of the external language model trained on generated samples and evaluated on  validation or test data.
Revised: The reverse perplexity (R-PPL) measures the perplexity of the external language model trained on generated samples and evaluated on the validation or test data.

Original: This captures the word frequency and overall proximity of our generated text distribution to the true text distribution.
Revised: This captures the word frequency and overall proximity of our generated text distribution to the true text distribution.

Original: Data In general the Wikitext-2 dataset, which is a small subset of encyclipedia Wikitext-103 dataset (Merity et al., 2017), is used and we reserve the larger scale data for the larger model experiments (AppendixSection B.4) and interpolations (Section 5.3).
Revised: Data In general the Wikitext-2 dataset, which is a small subset of the encyclopedia Wikitext-103 dataset (Merity et al., 2017), is used, and we reserve the larger scale data for the larger model experiments (Appendix B.4) and interpolations (Section 5.3).

Original: The datasets are cleaned and segmented at the sentence level using the NLTK toolkit (Bird et al., 2009), keeping only inputs with length from 5 towordpiece tokens using the BERT tokeniser (Devlin et al., 2019).
Revised: The datasets are cleaned and segmented at the sentence level using the NLTK toolkit (Bird et al., 2009), keeping only inputs with length from 5 to 50 wordpiece tokens using the BERT tokeniser (Devlin et al., 2019).

Original: Dataset statistics can be found in Table 2.
Revised: Dataset statistics can be found in Table 2.

Original: [Table]
Revised: [Table]

Original: B H YPERPARAMETER T UNING 
Revised: B H YPERPARAMETER T UNING AND A BLATIONS

Original: The data is trained on the Wikitext-2 training dataset and tuned on the validation  with the aim to be able to both reconstruct and generate output.
Revised: The models are trained on the Wikitext-2 training dataset using the loss from equation 10. They are tuned on the validation dataset with the aim to be able to both reconstruct and generate output.

Revised: Because both L D and L G scale approximately linearly with κ 0 , and κ 0 is linear in the sentence length, the D KL divergence losses L D and L G grow linearly with the sentence length n .

Original: Preliminary experiments suggest that training converges better without this linear dependence on sentence length, so we set the weights on the Gaussian and Dirichlet KL divergences to be linear in 1 n .
Revised: Preliminary experiments suggest that training converges better without this linear dependence on sentence length, so we set the weights on the Gaussian and Dirichlet KL divergences to be linear in 1 n .

Original: In addition we scale the weights on the Gaussian KL divergence byd , removing the dependence on the dimensionality d of vectors.
Revised: In addition we scale the weights on the Gaussian KL divergence by 1 d , removing the dependence on the dimensionality d of vectors.

Original: [Equation]
Revised: [Equation]

Revised: where λ ′ D and λ ′

Revised: G are fixed hyperparameters.

Original: All combinations of the following hyperparameters were considered in a grid search for the respective models:
Revised: All combinations of the following hyperparameters were considered in a grid search for the respective models:

Original: • λ ′ G = { 1 , 1 e − 1 , 1 e − 2 , 1 e − 3 , 1 e − 4 , 1 e − 5 , 0 } • λ ′ D
Revised: • λ ′ G = { 1 , 1 e − 1 , 1 e − 2 , 1 e − 3 , 1 e − 4 , 1 e − 5 , 0 } • λ ′ D

Original: = { 10 , 1 , 1 e − 1 , 1 e − 2 , 1 e − 3 , 1 e − 4 , 1 e − 5 , 0 } • α ∆ = { 1 , 0 .
Revised: = { 10 , 1 , 1 e − 1 , 1 e − 2 , 1 e − 3 , 1 e − 4 , 1 e − 5 , 0 } • α ∆ = { 1 , 0 .

Original: 1 , 0 } • κ ∆ = { 1 , 2 , 5 } • S = { 0 . 9 , 0 . 8 , 0 .
Revised: 1 , 0 } • κ ∆ = { 1 , 2 , 5 } • S = { 0 . 9 , 0 . 8 , 0 .

Original: } • P = { mean , max , one }
Revised: } • P = { mean , max , one }

Original: where λ ′ G and λ ′ D are the weights on the Gaussian and Dirichlet KL divergences for all variational models, respectively.
Revised: where λ ′ G and λ ′ D are the weights on the Gaussian and Dirichlet KL divergences for all variational models, respectively.

Original: The α ∆ and κ ∆ are NVAE specific parameters and represent the conditional prior parameter and number of samples per component.
Revised: The α ∆ and κ ∆ are NVAE specific parameters and represent the conditional prior parameter and number of samples per component.

Original: The stride parameter S for the VTS model results in 1 − S proportion of vectors.
Revised: The stride parameter S for the VTS model results in 1 − S proportion of vectors being kept.

Original: Finally, P is the pooling parameter for the single vector model VTP.
Revised: Finally, P is the pooling method for the single vector model VTP.

Original: Baselines Empirically we found the best KL divergence parameter for VT, VTP and VTS is λ ′ G =1 e −and using max pooling for VTP.
Revised: Baselines Empirically we found the best KL divergence parameter for VT, VTP and VTS is λ ′ G =1 e −and using max pooling for VTP.

Original: All stride parameters are considered to adjust the number of vectors.
Revised: All stride parameters are considered to adjust the number of vectors.

Original: This provides the best trade off of reconstruction accuracy with high BLEU score and generative sampling ability achieved by low F-PPL and R-PPL scores.
Revised: This provides the best trade-off of reconstruction accuracy with high BLEU score versus generative sampling ability achieved by low F-PPL and R-PPL scores.

Original: NVAE
Revised: NVAE

Original: The hyperparameter tuning for NVAE aims to discover models which do not collapse to a single vector or use all vectors, reconstruct  and  able to sample effectively from the prior by achieving low FPPL and RPPL scores.
Revised: The hyperparameter tuning for NVAE aims to discover models which: neither collapse to a single vector nor use all vectors, reconstruct accurately, and are able to sample effectively from the prior by achieving low F-PPL and R-PPL scores.

Original: Empirically we find the parameters λ ′ G = 1 e − 3 , λ ′ D
Revised: Empirically we find the parameters λ ′ G =1 e − 3 , λ ′ D

Original: = 1 and κ ∆ = 1 to produce the best trade off between reconstruction accuracy and generative ability.
Revised: =1 and κ ∆ = 1 to produce the best trade off between reconstruction accuracy and generative ability.

Original: The α ∆ parameter is able to control the proportion of vectors (AppendixSection B.2) retained and κ ∆ =1 provides an efficient sampling of the model (AppendixSection B.3).
Revised: The α ∆ parameter is able to control the proportion of vectors (Appendix B.2) retained and κ ∆ =1 provides an efficient sampling of the model (Appendix B.3).

Original: Validation results Table 3 display the validation reconstruction and generation results across 5 seeds for the best performing parameters.
Revised: Validation results Table 3 displays the validation reconstruction and generation results across 5 seeds for the best performing parameters.

Original: The VT model is able reconstruct well, however a high F-PPL score suggest a poor fluency of generated text and large variation.
Revised: The VT model is able reconstruct well, but a high F-PPL score suggest a poor fluency of generated text and large variation.

Original: The VTP models show that a single vector bottleneck is insufficient to reconstruct.
Revised: The VTP models show that a single-vector bottleneck is insufficient to reconstruct.

Original: Low F-PPL and high R-PPL suggest the model has collapsed toa just sampling a few fluent sentences.
Revised: Moreover, low F-PPL and high R-PPL suggest the model has collapsed to just sampling a few fluent sentences.

Revised: The VTS models show that some fixed proportions ν of vectors retained result in good overall performance.

Original: We notice that NVAE is able to findcomparable models with respects to average proportion  of vectors retainedν to  the VTS models.
Revised: NVAE is able to find models with comparable average proportion ν of vectors retained to those hand-coded in the VTS models.

Original: Their performance is comparable  with respects to reconstruction and generation.
Revised: These NVAE and VTS models have comparable performance with respect to reconstruction and generation.

Original: However, the NVAE models have notably more variance of metrics across seeds.
Revised: However, the NVAE models have notably more variance of metrics across seeds

Original: [Table]
Revised: [Table]

Original: Figures 4c 4a and 4b  visually display the reconstruction versus generation trade-off for the validation data across seeds.
Revised: Figures 4a, 4b and 4c visually display the reconstruction versus generation trade-off for the validation data across seeds.

Original: The best models have both low generation perplexities and high BLEU scores whilst dropping a large proportion of vectors.
Revised: The best models have both low generation perplexities and high BLEU scores whilst dropping a large proportion of vectors.

Original: Figure 4c shows that the NVAE model is able to dynamically reduce the average proportion of vectors and reconstruct comparatively to the VTS model.
Revised: Figure 4c shows that the NVAE model is able to dynamically reduce the number of vectors and still reconstruct, comparably to the VTS models.

Original: We notice that their exist some NVAE models that have a good reconstruction-generation trade-off but are less clustered than the VTS models.
Revised: We notice that their exist some NVAE models that have a good reconstruction-generation trade-off but are less clustered than the VTS models.

Original: Note that the R-PPL and F-PPL plot limits are cropped at 1.2 to focus on the higher performing models.
Revised: Note that the R-PPL and F-PPL plot limits are cropped at 1.2 to focus on the higher performing models.

Original: [Figure]
Revised: [Figure]

Original: BLEU [Figure] Testresults
Revised: Test

Original:  The best seed models from Table 3 are selected for the baselines and NVAE and then evaluated on the test set and shown in Table 4 and a subset of this α ∆ = { 0 . 75 , 0 . 4 , 0 . 3 , 0 . 2 } is plotted in Figure 2. 
Revised: results The best seed models from Table 3 are selected for the baselines and NVAE and then evaluated on the test set and shown in Table 4 and a subset of this α ∆ = { 0 . 75 , 0 . 4 , 0 . 3 , 0 . 2 } is plotted in Figure 2. Reconstruction Generation

Revised: Model ν

Original: [Table]
Revised: BLEU ( ↑ ) PPL ( ↓ ) F-PPL ( ↓ ) R-PPL ( ↓ )

Revised: VT 1.

Revised: VTS

Revised: S =0 .0.

Revised: VTS

Revised: S =0 .0.

Revised: VTS

Revised: S =0 .0.

Revised: VTS

Revised: S =0 .0.

Revised: VTP

Revised: P = max 0.05*

Revised: [Table]

Original: B. N O REGULARISATION ABLATION
Revised: B. N O R EGULARISATION A BLATION

Original: In this ablation we consider the behaviour of the NVAE without any regularisation.
Revised: In this ablation we consider the behaviour of the NVAE without any regularisation.

Original: Table 5 shows that the model is able to ignore the noise and revert to a standard Transformer model.
Revised: Table 5 shows that the model is able to ignore the noise and revert to a standard Transformer model.

Original: We can see this by the model retaining all its vectors as ν =1 and being able to perfectly reconstruct the input.
Revised: We can see this by the model retaining all its vectors as ν =1 and being able to perfectly reconstruct the input.

Original: Reconstruction
Revised: Reconstruction

Original: [Table]
Revised: [Table]

Original: B. C ONDITIONAL PRIOR ABLATION
Revised: B. C ONDITIONAL P RIOR E XPERIMENT

Original: In this ablation we consider the affect of the conditional prior α ∆ .
Revised: In this experiment we consider the effect of the conditional prior α ∆ .

Original: We select a subset of hyperparameters that managed to achieve good reconstruction and generation performance: λ ′ G = { 1 e − 3 , 1 e − 4 } , λ ′ D
Revised: We select a subset of hyperparameters that managed to achieve good reconstruction and generation performance

Original: = { 1 , 10 } , κ ∆ = { 1 , 2 , 5 } each across 5 random seed samples.

Original: [Figure]
Revised: [Figure]

Original: The figures shows validation reconstruction BLEU against the average proportion  of vectors ν is controlled by the conditional prior α ∆ across these parameters. The parameter α ∆ roughly corresponds to the proportion of  vectors you wish to retain making themodels regularisation controllable.
Revised: Figure 5 displays the validation reconstruction and validation generation metrics versus the average proportion ν of vectors retained. We see that the conditional prior hyperparameter α ∆ roughly corresponds to the proportion of retained vectors ν , across all the selected hyperparameter subsets. Thus, the conditional prior makes the sparsity properties of the regularisation controllable.

Original: B. N UMBER OF SAMPLES ABLATION
Revised: B. N UMBER OF S AMPLES E XPERIMENT

Original: In this ablation we consider the affect of the number of samples κ ∆  per component.
Revised: In this experiment we consider the effect of the number of samples κ ∆ = { 1 , 2 , 5 } per component.

Original: We consider a subset of hyperparameters λ ′ G = { 1 e − 3 , 1 e − 4 } , λ ′ D
Revised: We select a subset of hyperparameters λ ′ G = { 1 e − 3 , 1 e − 4 } , λ ′ D

Original: = { 1 , 10 } , α ∆ = { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 , 0 .
Revised: = { 1 , 10 } , α ∆ = { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 , 0 .

Original: 75 , 1 } each across 5 random seed samples.
Revised: 75 , 1 } each across 5 random seed samples.

Original: [Figure]
Revised: [Figure]

Original: The figures shows validation generation metrics F-PPL and R-PPL against reconstruction BLEU fora different number of samples κ ∆ per component does not have a significant improvement in the generation or reconstruction ability of the models.
Revised: Figure 6 shows validation generation metrics F-PPL and R-PPL against reconstruction BLEU for different numbers of samples κ ∆ per component. Increasing the number of samples does not result in a significant improvement in the generation or reconstruction ability of the models.

Original: Hence, we use a single sample κ ∆ =1 per component.
Revised: Hence, we use a single sample κ ∆ =1 per component because it is more efficient.

Original: B. L ARGER SCALE ABLATION
Revised: B. L ARGER S CALE M ODELS

Original: We conduct a larger scale experiment to if NVIB is able to generalise to models with more parameters and larger datasets such as Wikitext-103.
Revised: We conduct a larger scale experiment to test whether NVAE is able to generalise to models with more parameters and larger datasets such as Wikitext-103.

Original: Due to restricted resources we needed to be selective with the larger scale experimentation.
Revised: Due to restricted resources we need to be selective with the larger scale experimentation.

Original: We scale up the models (inspired by the Transformer base size (Vaswani et al., 2017)) by using  six layer Transformer encoder and decoder with the word embedding vectors and model projections are 512, feed forward  dimensions  2048, which leads to models of approximatelymillion trainable parameters.
Revised: We scale up the models (inspired by the Transformer base size (Vaswani et al., 2017)) by using a six layer Transformer encoder and decoder with the word embedding vectors and model projections of size 512. The feed forward layers’ dimensions are 2048, which leads to models of approximately 76 million trainable parameters.

Original: We use a constant learning rate of 1 e − 5 and trained for 11 epochs ( ≈ 150 K steps).
Revised: We use a constant learning rate of 1 e − 5 and trained for 11 epochs ( ≈ 150 K steps).

Original: Each model experiment takes approximately 24hrs to  on a single NVIDIA Tesla v100 which was the largest compute within budget.
Revised: Each model experiment takes approximately 24hrs to run on a single NVIDIA Tesla v100, which was the largest compute within budget.

Original: Otherwise the same training details of AppendixSection A are used.
Revised: Otherwise the same training details as Appendix A are used.

Original: We initially trained the NVAE model with the best parameters (based on lowest R-PPL score) from Table 4 thereafter we trained the VTS model with the closest proportion of vectors retained for comparability.
Revised: We initially trained the NVAE model with the best parameters (based on lowest R-PPL score) from Table 4. Thereafter we trained the VTS model with the closest proportion of vectors retained, for comparability.

Original: [Table] In Table B.4 we see thatwith larger model capacity allows the VTP models to be more competitive in reconstruction however, still struggles in generation.
Revised: In Table 6 we see that larger model capacity allows the VTP model to be more competitive in reconstruction, but it still struggles in generation.

Original: The NVAE model is able to dynamically regularise more vectors and is competitive with the baselines at scale with respects to both reconstruction and generation. 
Revised: The NVAE model is able to dynamically regularise the number of vectors and is competitive with all baselines at scale with respects to both reconstruction and generation. [Table] [Table] [Table]

Original: C G ENERATED S AMPLE E XAMPLES
Revised: C G ENERATED S AMPLE E XAMPLES

Original: [Table] [Table]
Revised: Tables 7 and 8 give examples of text generated by sampling from the prior.

Original: D A LIGNMENT A BLATION
Revised: D A LIGNMENT A NALYSIS

Original: To highlight the problem of latent alignment  and argue for alignment based on position we conduct the following experiment.
Revised: This experiment highlights the problem of latent alignment in non-NVIB models [Figure] and evaluates alignment based on position, which we use in the interpolation experiments.

Original: We consider  VTS models from Table 4 and use them to encode a sentence into their latent space.
Revised: We consider the VTS models from Table 4 and use them to encode a sentence into their latent space.

Original: We then iteratively perturb the non-masked latent components with gaussian noise and consider the  autoregressively decoded output.
Revised: For each latent component retained by the VTS model, we perturb it with Gaussian noise and consider the resulting autoregressively decoded output.

Original: We record the location of words differing through these perturbations and discard sentences of differing length  (onlyaccount for 2% over 100 samples)
Revised: We plot the percentage of the time a given position in the output is changed by perturbing a given latent component, discarding sentences where the length is changed (only 2% over 100 samples)

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: The Figure 7 show that the latent space for the VTS S =0 .
Revised: Figure 7 shows that the latent space for the VTS S =0 .

Original: 5 model are approximately aligned by location.
Revised: 5 model are approximately aligned by location.

Original: However, in Figure 8 for more condensed representations where S =0 .
Revised: However in Figure 8, for more condensed representations where S =0 .

Original: 75  there is  unclear alignment of latent vectors to their position.
Revised: 75 , there is an unclear alignment of latent vectors to their position.

Revised: As discussed in Section 5.3, the use of mixture distributions instead of sets of vectors allows NVIB representations to avoid this problem of alignment.

Original: E I NTERPOLATION E XAMPLES
Revised: E I NTERPOLATION E XAMPLES

Revised: Tables 9 through 12 give examples of the text generated by interpolations in the latent space.

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original: F T HE M EAN OF THE P OSTERIOROVER M IXTURE D ISTRIBUTIONS
Revised: F D ENOISING A TTENTION FOR THE M EAN OF THE P OSTERIOR

Original: As introduced in Section 3.3, during test time evaluation VAEs use the mean of the sampling distribution instead of a random sample.
Revised: As introduced in Section 2.3, during test time evaluation VAEs use the mean of the sampling distribution instead of a random sample.

Original: The mean of our BFDP posterior distribution over mixture distributions is its base distribution G q 0 .
Revised: The mean of our BFDP posterior distribution over mixture distributions is its base distribution G q 0 .

Original: However, the base distribution is not a discrete distribution, whereas at training time all samples are discrete distributions.
Revised: However, the base distribution is not a discrete distribution, whereas at training time all samples are discrete distributions.

Original: This is one of the main advantages of generalising the attention function to denoising attention, proposed in Section 2.1.
Revised: This is one of the main advantages of generalising the attention function to denoising attention, proposed in Section 2.1.

Original: Denoising attention can equally well be applied to the mixture of Gaussians of the base distribution as to the mixture of impulse distributions from sampling, as depicted in Figure 1c.
Revised: Denoising attention can equally well be applied to the mixture of Gaussians of the base distribution as to the mixture of impulse distributions from sampling, as depicted in Figure 1c.

Original: On the other hand, the standard attention function is only defined for a finite set of vectors, so even if we convert the sampled discrete distributions to finite sets of vectors, at testing time there is no finite set of vectors which is the mean of the sampled sets of vectors.

Original: In particular, the set of mean vectors is not the mean of the set of vectors, since it underestimates the variance of the mean distribution.

Original: To understand why the continuous base distribution looks to the model like a typical sample of a discrete distribution, we need to consider how this representation is being interpreted by the attention function.
Revised: To understand why the continuous base distribution looks to the model like a typical sample of a discrete distribution, we need to consider how this representation is being interpreted by the attention function.

Original:  For any given query, we want a test-time representation such that attention will return a good approximation to the mean of the vectors returned for the sampled representations.
Revised: Given this representation and a query vector, attention returns a result vector. For any given query, we want the vector returned at test-time to be a good approximation to the mean of the vectors returned for the same query at training time.

Original: This is what the base distribution is giving us.
Revised: This is what the base distribution is giving us.

Original: Thus, the function defined by applying denoising attention to a sampled distribution can be seen as a noisy version of  the function defined by applying denoising attention to the mean distribution.
Revised: In contrast, with standard attention there is no finite set of vectors we can use at test time which will have this same property. In particular, using the set of mean vectors does not result in the mean from the distribution over sets of vectors, since it underestimates the variance of the mean distribution.

Original: More precisely, our latent representation is a parameterisation of an attention function from query vectors to attention vectors.
Revised: More precisely, we can consider the latent representation as a parameterisation of an attention function from query vectors to result vectors.

Original: The encoder outputs two things, a mean function and a noise level, parameterised by  G q 0 and  α 0 respectively. At training time we sample noisy versions of the mean function where the amount of noise is determined by α 0 .
Revised: Using denoising attention, the function parameterised by a sample from the posterior is a noisy version of the function parameterised by the posterior’s base distribution, in that for any query the former function returns a noisy version of the vector returned by the latter function. When the encoder specifies the posterior’s base distribution G q 0 and total pseudo-count α 0 , the base distribution G q 0 specifies the mean of a distribution over such query-result functions, and the pseudo-count α 0 specifies the amount of noise (i.e. larger values mean less noise). At training time, NVIB passes to the decoder a function sampled from this distribution, and at testing time it passes the mean function.

Original: At test timewe just use the mean function parameterised by G q 0 .
Revised: The decoder then accesses this function by repeatedly running queries through it, receiving at test time the mean of the vectors it receives at training time for the same query.

Original: To efficiently compute denoising attention applied to the mixture of Gaussians G q 0 , we take advantage of the fact that a multiplication of Gaussians is a Gaussian, giving us:
Revised: To efficiently compute denoising attention applied to the mixture of Gaussians G q 0 , we take advantage of the fact that the multiplication of two Gaussian distributions is a Gaussian distribution, giving us:

Original: [Equation]
Revised: [Equation]

Original: where all algebraic calculations over vectors σ qi are done componentwise
Revised: where all algebraic calculations over vectors σ qi are done componentwise.

Original: [Equation]
Revised: See Appendix K for a version of this formula which is useful for implementation.

Original: G E QUIVALENCE OF S ET - OF -V ECTOR A TTENTION AND D ENOISING A TTENTION
Revised: G E QUIVALENCE OF S ET - OF -V ECTOR A TTENTION AND D ENOISING A TTENTION

Original: In this section we show the equivalence of standard attention, expressed as a sum over vectors
Revised: In this section we show the equivalence of standard attention, expressed as a sum over vectors

Original: z i in Z , and denoising attention, expressed as an integral over a distribution F Z which is only nonzero at the z i .
Revised: z i in Z , and denoising attention, expressed as an integral over a distribution F Z which is only nonzero at the z i .

Original: The attention mechanism we assume is scaled dot product attention,  standardly used in many attention-based models, including Transformers.
Revised: The attention mechanism we assume is scaled dot product attention, which is standardly used in many attention-based models, including Transformers.

Original: For simplicity, we consider cross attention, where a single query vector is mapped to a single result vector.
Revised: For simplicity, we consider cross attention, where a single query vector is mapped to a single result vector.

Revised: In this work we view cross attention as an interface between an encoder, which provides a set of vectors Z , and a decoder, which receives a function from query vectors u ′ to result vectors.

Revised: The decoder can then apply this function to as many queries as it likes.

Original: This attention function uses keys  K  ∈ R n × d and values  V  ∈ R n × dwhich are both projections from the same set of vectors Z ∈ R n × p via weight matrices W K , W V ∈ R p × d respectively. The query q ∈ R 1 × d is a projection from the input u ′  ∈ R 1 × p via the weight matrix W Q ∈ R p × d .
Revised: But for this equivalence we only need consider the composed function from both u ′ and Z to a result vector. Scaled dot product attention first maps the set of vectors Z ∈ R n × p into keys ( ZW K ) ∈ R n × d and values ( ZW V ) ∈ R n × d via weight matrices W K , W V ∈ R p × d , respectively, and maps the query u ′ ∈ R 1 × p into key space ( u ′ W Q ) ∈ R 1 × d via the weight matrix W Q ∈ R p × d .

Original: The keys’ dimensionality d is used for scaling.
Revised: The key space’s dimensionality d is used for scaling.

Original: Scaled dot product attention is then defined as:
Revised: Scaled dot product attention is then defined as:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original:  V , where u = u ′ W Q ( W K ) T =
Revised: ZW V , where u = u ′ W Q ( W K ) T =

Original: Attn ( u , Z ) W V where u ∈ R 1 × p is the input query u ′ projected into Z space.
Revised: Attn( u , Z [Equation] W V where u ∈ R 1 × p is the input query u ′ projected into Z space.

Original: In the last line we rewrite scaled dot product attention in terms of a core dot product attention function Attn ( u , Z ) where all operations are done in the space of Z :
Revised: In the last line we rewrite scaled dot product attention in terms of a core dot product attention function Attn( u , Z ) where all operations are done in the space of Z :

Original: [Equation]
Revised: [Equation]

Original: where z i is the i th row of Z .
Revised: where z i is the i th row of Z .

Original: In models that use cross attention, an encoder computes the set of vectors Z , which is the attention-based representation of the input. A decoder then accesses this representation by issuing queries u (or equivalently u ′ ). We interpret Z as specifying a probability distribution over a vector space, and want to define a function over such probability distributions which, when given any vector u , always returns the vector Attn ( u , Z ) .
Revised: We interpret Z as specifying a probability distribution over a vector space, and we interpret the function

Original: We interpret the function Attn ( u , Z ) as  denoising  u with a generative model specified by Z , as depicted in Figure 1b.
Revised: Attn( u , Z ) as Bayesian denoising of u using this distribution, as depicted in Figure 1b.

Original: The vector u is interpreted as an observation of some true vector v ∈ R 1 × p which has been corrupted by Gaussian noise.
Revised: The vector u is interpreted as an observation of some true vector v ∈ R 1 × p which has been corrupted by Gaussian noise.

Original: The true vector v was generated from a prior probability distribution specified by Z .
Revised: The true vector v was generated from a prior probability distribution specified by Z .

Original: The result of Attn ( u , Z ) is the expected value of this true vector v after seeing the observation u , which can be considered the best guess of the true vector given the noisy observation, and thus is a form of  denoising.
Revised: The result of Attn( u , Z ) is the expected value of this true vector v after seeing the observation u , which can be considered the best guess of the true vector given the noisy observation, and thus is a form of Bayesian denoising.

Original: To derive this interpretation, we interpret the set Z of vectors
Revised: To derive this interpretation of attention as Bayesian query denoising, we interpret the set Z of vectors

Original: z i as specifying a mixture distribution F Z over vectors v which consists of one impulse distribution δ z i at each vector z
Revised: z i as specifying a mixture distribution F Z over vectors v which consists of one impulse distribution δ z i at each vector z

Original: i weighted by the softmax over their scaled L 22 norms:
Revised: i weighted by the softmax over their scaled L 22 norms:

Original: [Equation]
Revised: [Equation]

Original: Then we can derive this interpretation by replacing attention’s sum over i with an integration over v.
Revised: Then we can derive this interpretation by replacing attention’s sum over i with an integration over v

Original: [Table]

Original: [Equation]

Original: where f Z ( · ) is the probability density function for distribution F Z , and g ( u ;

Original: [Equation]
Revised: [Equation]

Original:  , √ [Equation] =[Equation]
Revised: Attn( u , Z ) =

Original: √

Original: [Equation]
Revised: [Equation]

Original: exp( −

Original: [Equation]
Revised: [Equation]

Original: √ d

Original: [Equation] pk = [Equation] u k − v k ) 2 ) is the multivariate Gaussian function with diagonal variance of √ d .
Revised: where f Z ( · ) is the probability density function for distribution F Z , and g ( u ; v , √ d I ) = 1 √ 2 π √ d exp( − 12 √ d (cid:80) pk =1 ( u k − v k ) 2 ) is the multivariate Gaussian function with diagonal variance of √ d .

Original: The first step just adds terms which don’t effect the value.
Revised: The first step just adds terms which don’t effect the value.

Original: The second step changes some instances of z i into an integral over v which is only nonzero when v =
Revised: The second step changes some instances of z i into an integral over v which is only nonzero when v =

Original: z i (i.e. δ z i ).
Revised: z i (i.e. δ z i ).

Original: Thereafter, the terms are rearranged such that the formula reduces to an expected value over v with weights proportional to the probability of generating v with the distribution F and generating the query u with Gaussian noise
Revised: Thereafter, the terms are rearranged such that the formula reduces to an expected value over v with weights proportional to the probability of generating v with the distribution F and generating the query u with Gaussian noise

Original: N ( 0 , √ d I ) added to v .
Revised: N ( 0 , √ d I ) added to v .

Original: Scaling the variance of the multi-dimensional Gaussian noise by √ d reduces the impact of the dimensionality d on the similarity g ( u ; v , √ d I ) between u and v .
Revised: Scaling the variance of the multi-dimensional Gaussian noise by √ d reduces the impact of the dimensionality d on the similarity g ( u ; v , √ d I ) between u and v .

Original: The above derivation is inspired by the interpretation of softmax as Bayesian classification with normally distributed classes (Bishop, 1995), and it is similar to the interpretation of attention keys as latent mixture distributions provided by Nguyen et al.
Revised: The above derivation is inspired by the interpretation of softmax as Bayesian classification with normally distributed classes (Bishop, 1995), and it is similar to the interpretation of attention keys as latent mixture distributions provided by Nguyen et al.

Original: However, here the Gaussian represents uncertainty about the observation or query vector instead of uncertainty about the class or key vectors, exploiting the fact that a Gaussian function is symmetric in its argument and mean.
Revised: However, here the Gaussian represents uncertainty about the observation or query vector instead of uncertainty about the class or key vectors, exploiting the fact that a Gaussian function is symmetric in its argument and mean.

Original: This allows us to incorporate the value part of the attention function into a Bayesian denoising interpretation, so we have a Bayesian interpretation of the entire attention function.
Revised: This allows us to incorporate the value part of the attention function into a Bayesian denoising interpretation, so we have a Bayesian interpretation of the entire attention function.

Original: To the best of our knowledge, this interpretation of attention is novel.
Revised: To the best of our knowledge, this interpretation of attention is novel.

Original: The function from equation 12 is a special case of the definition of denoising attention DAttn ( u ; F ) given above in equation 3, where F = F Z .
Revised: The function from equation 14 is a special case of the definition of denoising attention DAttn( u ; F ) given above in equation 3, where F = F Z .

Original: The construction above indicates that any scaled dot product attention function is an example of the DAttn ( u ; F ) function.
Revised: The construction above indicates that any scaled dot product attention function is an example of the DAttn( u ; F ) function.

Original: However, while attention is only defined over sets-of-vectors Z , denoising-attention is defined over any probability distribution F over a vector space, not just finite sets of impulse distributions.
Revised: However, while attention is only defined over sets-of-vectors Z , denoising-attention is defined over any probability distribution F over a vector space, not just finite sets of impulse distributions.

Original: H D ERIVING THE F ACTORISED D IRICHLET P ROCESS
Revised: H D ERIVING THE F ACTORISED D IRICHLET P ROCESS

Original: In this section we derive an alternative factorisation of a DP which helps with the sampling method in
Revised: In this section we derive an alternative factorisation of a DP which helps with the sampling method in

Original: Section 3.2.
Revised: Section 3.2.

Original: For notational convenience, in this section we use c as the number of components for the base distribution instead of n +1 .
Revised: For notational convenience, in this section we use c as the number of components for the base distribution instead of n +1 .

Original: This is still intended to include both the output of the encoder and the prior component.
Revised: This is still intended to include both the output of the encoder and the prior component.

Original: Here we provide the proof that
Revised: Here we provide the proof that

Original: [Equation]
Revised: [Equation]

Original: where G q =( G q 1 ,...,G qc ) , α q =( α q 1 ,...,α qc ) , G q 0 = (cid:80) ci =
Revised: where G q =( G q 1 ,...,G qc ) , α q =( α q 1 ,...,α qc ) , G q 0 = (cid:80) ci =

Original: 1 α qi α q 0 G qi , α q 0 = (cid:80) ci
Revised: 1 α qi α q 0 G qi , α q 0 = (cid:80) ci

Original: =1 α qi , and F ∼ FDP( G q , α q ) is defined as
Revised: =1 α qi , and F ∼ FDP( G q , α q ) is defined as

Original: [Equation]
Revised: [Equation]

Original: We start with the definition of a DP as an infinite symmetric Dirichlet distribution.
Revised: We start with the definition of a DP as an infinite symmetric Dirichlet distribution.

Original: A Dirichlet process
Revised: A Dirichlet process

Original: F ∼ DP( G 0 ,α 0 ) can be defined as the limit of a sequence of finite Dirichlet distributions (see Teh (2010)):
Revised: F ∼ DP( G 0 ,α 0 ) can be defined as the limit of a sequence of finite Dirichlet distributions (see Teh (2010)):

Original: [Equation]
Revised: [Equation]

Original: Note that the weights π and the vectors z k are independent of each other, so we can treat these two issues separately.
Revised: Note that the weights π and the vectors z k are independent of each other, so we can treat these two issues separately.

Original: For the vectors, we know that after generating an infinite number of z k from G 0 , a proportion of exactly α qi α q 0 of them will be generated from G qi .
Revised: For the vectors, we know that after generating an infinite number of z k from G 0 , a proportion of exactly α qi α q 0 of them will be generated from G qi .

Original: For a finite number of vectors κ 0 , let κ i be the number of z k generated from G qi , for each i .
Revised: For a finite number of vectors κ 0 , let κ i be the number of z k generated from G qi , for each i .

Original: So we have
Revised: So we have

Original: [Equation]
Revised: [Equation]

Original: Given the exchangeability of Dirichlet distributions, we can renumber the κ 0 categories of Dir( α 0 κ 0 , κ 0 ..., α 0 κ 0 )
Revised: Given the exchangeability of Dirichlet distributions, we can renumber the κ 0 categories of Dir( α 0 κ 0 , κ 0 ..., α 0 κ 0 )

Original: so that π = ( π 11 ,...,π 1 κ 1 , c ..., π c 1 ,...,π cκ c ) and the π i 1 ,...,π iκ i are all weights for vectors z ij generated from component i .z ij
Revised: so that π = ( π 11 ,...,π 1 κ 1 , c ..., π c 1 ,...,π cκ c ) and the π i 1 ,...,π iκ i are all weights for vectors z ij generated from component i .

Original: [Equation]
Revised: [Equation]

Original: For the weights, we again consider the case of finite κ 0 before taking the limit as κ 0 goes to infinity, using the above indexing where categories ij are partitioned according to their vector’s base distribution component i .
Revised: For the weights, we again consider the case of finite κ 0 before taking the limit as κ 0 goes to infinity, using the above indexing where categories ij are partitioned according to their vector’s base distribution component i .

Original: We define ( ρ 1 , c ...,ρ c ) to be the vector of total weights ρ i
Revised: We define ( ρ 1 , c ...,ρ c ) to be the vector of total weights ρ i

Original: = (cid:80) κ i j =1 π ij for each of these partitions i .
Revised: = (cid:80) κ i j =1 π ij for each of these partitions i .

Original: By the rule for merging categories in a Dirichlet distribution, we know that these total weights are themselves distributed according to a Dirichlet distribution.
Revised: By the rule for merging categories in a Dirichlet distribution, we know that these total weights are themselves distributed according to a Dirichlet distribution.

Original: [Equation]
Revised: [Equation]

Original: Now we take advantage of the neutrality property of Dirichlet distributions.
Revised: Now we take advantage of the neutrality property of Dirichlet distributions.

Original: It states that this vector ( ρ 1 , c ...,ρ c ) of partition weights and all of the vectors ( π i 1 ρ i , κ i ..., π iκi ρ i ) of normalised weights inside each partition are independent.
Revised: It states that this vector ( ρ 1 , c ...,ρ c ) of partition weights and all of the vectors ( π i 1 ρ i , κ i ..., π iκi ρ i ) of normalised weights inside each partition are independent.

Original: In essence, this means that the only way that the weights inside each partition constrain each other is through normalisation, so when normalisation is factored out they become independent.
Revised: In essence, this means that the only way that the weights inside each partition constrain each other is through normalisation, so when normalisation is factored out they become independent.

Original: This independence allows us to compute the distribution over ( π i 1 ρ i , κ i ..., π iκi ρ i ) by simply marginalising out all the other categories.
Revised: This independence allows us to compute the distribution over ( π i 1 ρ i , κ i ..., π iκi ρ i ) by simply marginalising out all the other categories.

Original: We first merge all the categories outside partition i into a single category, whose weight is thus 1 − ρ i .
Revised: We first merge all the categories outside partition i into a single category, whose weight is thus 1 − ρ i .

Original: This gives us the marginalised Dirichlet distribution ( π i 1 , κ i ...,π iκ i , (1 − ρ i )) ∼ Dir(α q 0 κ
Revised: This gives us the marginalised Dirichlet distribution ( π i 1 , κ i ...,π iκ i , (1 − ρ i )) ∼ Dir(

Original: [Equation]
Revised: [Equation]

Original:  κ i ...
Revised: 0 κ 0 ,

Original: [Equation]
Revised: [Equation]

Original:  α
Revised: i ..., α

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

Original:  , αq 0 (1 −
Revised: 0 , α

Original: [Equation]
Revised: [Equation]

Original:  i κ 0 )) .
Revised: 0 (1 − κ i κ 0 )) .

Original: Let d ( π ; α ) be the probability density function for the distribution Dir ( α ) :
Revised: Let d ( π ; α ) be the probability density function for the distribution Dir ( α ) :

Original: [Equation]
Revised: [Equation]

Original: Now we can marginalise out the weight of the outside category by integrating over ρ i .
Revised: Now we can marginalise out the weight of the outside category by integrating over ρ i [Equation]

Original: (cid:90)ρ i=
Revised: (cid:90)ρ i

Original: [Equation]
Revised: [Equation]

Original: where in the last step we note that the integral (assuming it is well defined) is simply part of the normalisation constant, which we know from the definition of the Dirichlet distribution must be
Revised: where in the last step we note that the integral (assuming it is well defined) is simply part of the normalisation constant, which we know from the definition of the Dirichlet distribution must be

Original: B ( α q 0 κ 0 , κ i ..., α q 0 κ 0 ) .
Revised: B ( α q 0 κ 0 , κ i ..., α q 0 κ 0 ) .

Original: This gives us
Revised: This gives us

Original: [Equation]
Revised: [Equation]

Original: Now that we have all the individual distributions, we can put them together to get the factorised distribution for the case of finite κ 0 . 
Revised: Now that we have all the individual distributions, we can put them together to get the factorised distribution for the case of finite κ 0 . [Equation] ij

Original: [Equation]
Revised: [Equation]

Original: Noting that lim κ 0 →∞ α q 0 κ 0 =
Revised: Noting that lim κ 0 →∞ α q 0 κ 0 =

Original: α qi κ i , we can then take the limit as κ 0 goes to infinity to get our definition of the weights for the factorised Dirichlet distribution.
Revised: α qi κ i , we can then take the limit as κ 0 goes to infinity to get our definition of the weights for the factorised Dirichlet distribution.

Original: [Equation]
Revised: [Equation]

Original: Thus the weights of a DP can be rewritten as the weights of an equivalent FDP using the above construction.
Revised: Thus the weights of a DP can be rewritten as the weights of an equivalent FDP using the above construction.

Original: Putting the vectors and weights together, we get the distribution F i over the weighted vectors in each partition i .
Revised: Putting the vectors and weights together, we get the distribution F i over the weighted vectors in each partition i .

Original: [Equation]
Revised: [Equation]

Original: and thus
Revised: and thus

Original: [Equation]
Revised: [Equation]

Original: This concludes our proof that, if F ∼ DP( G q 0 ,α q 0 ) , then:
Revised: This concludes our proof that, if F ∼ DP( G q 0 ,α q 0 ) , then:

Original: [Equation]
Revised: [Equation]

Original: and thus FDP( G q , α q )=
Revised: and thus FDP( G q , α q )=

Original: DP( G q 0 ,α q 0 ) .
Revised: DP( G q 0 ,α q 0 ) .

Original: I D ERIVING THE KL D IVERGENCE
Revised: I D ERIVING THE KL D IVERGENCE

Original: In this section we derive the KL divergence between the prior and the posterior.
Revised: In this section we derive the KL divergence between the prior and the posterior.

Original: We also argue that this function is approximately linear in the number of sampled vectors for each component.
Revised: We also argue that this function is approximately linear in the number of sampled vectors for each component.

Original: To directly compare the posterior with the prior, we first reformulate the prior as a factorised DP with the same form as the posterior.
Revised: To directly compare the posterior with the prior, we first reformulate the prior as a factorised DP with the same form as the posterior.

Original: We can do this without changing the distribution specified by the prior, simply by making n + 1 copies of the base distribution G p 0 and weighting those copies proportionately to the weights α qi α q 0 of the components of the posterior base distribution.
Revised: We can do this without changing the distribution specified by the prior, simply by making n + 1 copies of the base distribution G p 0 and weighting those copies proportionately to the weights α qi α q 0 of the components of the posterior base distribution.

Original: This gives us the prior BFDP( G p , α p ′ ,κ ∆ ) where G p =( G p 0 , n {++1 ... ,G p 0 ) and α p ′ = α q α p ′ 
 α q 0 =( α p [Equation]+}
Revised: This gives us the prior BFDP( G p , α p ′ ,κ ∆ ) where G p =( G p 0 , n {++1 ... ,G p 0 ) and α p ′ = α q α p ′ 
 α q 0 =( α p [Equation]+}

Original: [Equation]
Revised: [Equation]

Original: 1 ... ,G p
Revised: n

Original: [Equation]
Revised: [Equation]

Original: and
Revised: 1 ... ,α

Original: [Equation]
Revised: [Equation]

Original:  The formulation of both the prior and posterior as bounded factorised DPs of the same form simplifies the computation of the KL divergence, because the KL divergence for each respective pair of factors can be computed separately, and then combined.
Revised: α qn +1 α q 0 ) . The formulation of both the prior and posterior as bounded factorised DPs of the same form simplifies the computation of the KL divergence, because the KL divergence for each respective pair of factors can be computed separately, and then combined.

Original: First consider the factors for the Dirichlet distributions over the partitions for the different components i .
Revised: First consider the factors for the Dirichlet distributions over the partitions for the different components i .

Original: There is a closed-form solution to the KL divergence between two Dirichlet distributions.
Revised: There is a closed-form solution to the KL divergence between two Dirichlet distributions.

Original: [Equation]
Revised: [Equation]

Original:   where Γ is the gamma function and ψ is the digamma function.
Revised:   where Γ is the gamma function and ψ is the digamma function.

Original: For the bounded DPs for each individual component i , there are two factors, a symmetric Dirichlet distribution over the weights and a Gaussian distribution over each vector.
Revised: For the bounded DPs for each individual component i , there are two factors, a symmetric Dirichlet distribution over the weights and a Gaussian distribution over each vector.

Original: For the symmetric Dirichlet distribution, in the case where κ i =1 , then the KL for this term is zero, since there is no choice to make for this weight.
Revised: For the symmetric Dirichlet distribution, in the case where κ i =1 , then the KL for this term is zero, since there is no choice to make for this weight.

Original: In the case where κ i > 1 , the KL divergence between the posterior and prior versions of these weight distributions again has a closed-form solution.
Revised: In the case where κ i > 1 , the KL divergence between the posterior and prior versions of these weight distributions again has a closed-form solution.

Original: [Equation]
Revised: [Equation]

Original: This term then needs to be summed across components 1 ≤ i ≤ c +1 .
Revised: This term then needs to be summed across components 1 ≤ i ≤ n +1 .

Original: For the factors for generating vectors from each individual component of the base distribution, because the different components have been factorised, there is also a closed-form solution to computing these KL divergences.
Revised: For the factors for generating vectors from each individual component of the base distribution, because the different components have been factorised, there is also a closed-form solution to computing these KL divergences.

Original: Each Gaussian component of the posterior’s base distribution is compared independently to the Gaussian of the prior’s base distribution.
Revised: Each Gaussian component of the posterior’s base distribution is compared independently to the Gaussian of the prior’s base distribution.

Original: TheD KL divergence between two Gaussians (with diagonal covariance with values σ ) is:
Revised: The KL divergence between two Gaussians (with diagonal covariance with values σ ) is:

Original: [Equation]
Revised: [Equation]

Original: where the last step assumes µ p = 0 , ( σ p ) 2 = 1 .
Revised: where the last step assumes µ p = 0 , ( σ p ) 2 = 1 .

Original: This term then needs to be multiplied by the number κ i of vectors for this component, and summed across components 1 ≤ i ≤ c +1 .
Revised: This term then needs to be multiplied by the number κ i of vectors for this component [Equation] and summed across components 1 ≤ i ≤ n + [Equation]

Original: Given these exact closed-form solutions for each pair of factors, we can compute the full KL divergence.
Revised: Given these exact closed-form solutions for each pair of factors [Equation] we can compute the full KL divergence.

Original: We start by combining the formulas for the weight factors, where some terms cancel:
Revised: We start by combining the formulas for the weight factors [Equation] where some terms cancel:

Original: [Equation]
Revised: [Equation]

Original: Now we can put all these pieces together
Revised: Now we can put all these pieces together.

Original: [Equation]
Revised: [Equation]

Original: Equation equation 13 gives the KL portion of the loss when we are given the full set κ of numbers of vectors κ i generated for each component i .
Revised: Equation 15 gives the KL portion of the loss when we are given the full set κ of numbers of vectors κ i generated for each component i .

Original: Equation equation 13 is approximately linear in κ i , when the variation in κ i is fairly small relative to the values of κ i .
Revised: If we assume that the κ i are chosen stochastically, then we can take advantage of the fact that equation 15 is approximately linear in κ i , when the variation in κ i is fairly small relative to the values of κ i .

Original: The Gaussian term is exactly linear in κ i , and the terms ψ ( α qi κ i ) and κ i (cid:16) logΓ( α qi α [Equation] 0 κ i ) − logΓ( α qi κ i ) (cid:17) are both approximately linear in κ i .
Revised: The Gaussian term is exactly linear in κ i , and the terms ψ ( α qi κ i ) and κ i (cid:16) logΓ( α qi α q 0 κ i ) − logΓ( α qi κ i ) (cid:17) are both approximately linear in κ i .

Original: This allows us to approximate the expectation over κ i of this loss as this loss of the expectation over κ i , as discussed in Section 3.1.
Revised: This allows us to approximate the expectation over κ i of this loss as this loss of the expectation over κ i , as discussed in Section 3.1.

Original: In this case, this approximation of the full KL divergence is: 
Revised: In this case, this approximation of the full KL divergence is: [Equation] [Equation] n

Original: [Equation]
Revised: [Equation]

Original: J R EPARAMETERISATION T RICK AND S AMPLING
Revised: J R EPARAMETERISATION T RICK AND S AMPLING

Original: In this section we consider the reparameterisation trick to allow back propagation through the sampling step.
Revised: In this section we consider the reparameterisation trick to allow backpropagation through the sampling step.

Original: We consider the Gaussian base distribution and DP weights seperately.
Revised: We consider the component Gaussians of the base distribution and the weights generated by the DP separately.

Original: J. S AMPLING V ECTORS FROM A C OMPONENT OF THE B ASE D ISTRIBUTION
Revised: J. S AMPLING V ECTORS FROM A C OMPONENT OF THE B ASE D ISTRIBUTION

Original: Each vector z k is sampled independently from some specific component G qi of the base distribution G q 0 .
Revised: Each vector z k is sampled independently from some specific component G qi of the base distribution G q 0 .

Original: Since we assume that all these components are distributed according to a Gaussian G qi =
Revised: Since we assume that all these components are distributed according to a Gaussian G qi =

Original: N ( µ qi , I ( σ qi ) 2 ) , we can sample from this distribution using location-scale shifting (Kingma & Welling, 2014):
Revised: N ( µ qi , I ( σ qi ) 2 ) , we can sample from this distribution using location-scale shifting (Kingma & Welling, 2014):

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

Original: Since the random sampling comes from an unparameterised unit Gaussian, there is no need to backpropagate error into this sampling step, but the error can be backpropagated into µ qi and σ qi given a specific sample.
Revised: Since the random sampling comes from an unparameterised unit Gaussian, there is no need to backpropagate error into this sampling step, but the error can be backpropagated into µ qi and σ qi given a specific sample.

Original: This is the reparameterisation trick for Gaussian distributions.
Revised: This is the reparameterisation trick for Gaussian distributions.

Original: J. S AMPLING FROM A D IRICHLET D ISTRIBUTION
Revised: J. S AMPLING FROM A D IRICHLET D ISTRIBUTION

Original: A Dirichlet distribution over category weights can be sampled by sampling from a Gamma distribution for each category and then normalising.
Revised: A Dirichlet distribution over category weights can be sampled by sampling from a Gamma distribution for each category and then normalising.

Original: A sum-normalised set of κ random variables π 1 ,...,π κ ∼ Dir ( α 1 ,...,α κ ) follows a Dirichlet distribution if the unnormalised random variables γ i each follow a Gamma distribution.
Revised: A sum-normalised set of κ random variables π 1 ,...,π κ ∼ Dir ( α 1 ,...,α κ ) follows a Dirichlet distribution if the unnormalised random variables γ i each follow a Gamma distribution.

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where Γ( α i ,β =1) is the Gamma distribution with β =1 , whose PDF is f ( x )
Revised: where Γ( α i ,β =1) is the Gamma distribution with β =1 , whose PDF is f ( x )

Original: = 1 Γ( α i ) exp(( α i − 1)log( x ) − x ) .
Revised: = 1 Γ( α i ) exp(( α i − 1)log( x ) − x ) .

Original: There is no closed-form explicit reparameterisation trick for theexact Gamma distribution, but there are for approximations.
Revised: There is no closed-form explicit reparameterisation trick for the Gamma distribution, but there are for approximations.

Original: We propose to use a combination of two approximations for the Gamma distribution which have a reparameterisation trick, one for small values of α i and one for larger values of α i .
Revised: We propose to use a combination of two approximations for the Gamma distribution which have a reparameterisation trick, one for small values of α i and one for larger values of α i .

Original: Inverse CDF approximation to Gamma distributions The Gamma distribution cannot use location-scale shifting for sampling due to its asymmetry, nor can the parameters and noise components be decoupled in the inverse CDF.
Revised: Inverse CDF approximation to Gamma distributions The Gamma distribution cannot use location-scale shifting for sampling due to its asymmetry, nor can the parameters and noise components be decoupled in the inverse CDF.

Original: Hence, (Knowles, 2015) suggests sampling using an approximation to the inverse CDF of the Gamma distribution of the following form:
Revised: Hence, (Knowles, 2015) suggests sampling using an approximation to the inverse CDF of the Gamma distribution of the following form:

Original: [Equation]
Revised: [Equation]

Original: This approximation allows the inverse CDF of the Gamma distribution to be a function of the parameters and independent noise from a uniform distribution U (0 , 1) .
Revised: This approximation allows the inverse CDF of the Gamma distribution to be a function of the parameters and independent noise from a uniform distribution U (0 , 1) .

Original: However, this approximation is only recommended when the value of α i < 1 and β =1 .
Revised: However, this approximation is only recommended when the value of α i < 1 and β =1 .

Original: In our case β =

Revised: In our case β =

Original: but we sometimes have large α i .
Revised: but we sometimes have large α i .

Original: Gaussian approximation to Gamma distributions Knowles (2015) further mentions that the Gaussian distribution can be used to approximate a Gamma distribution for larger α .
Revised: Gaussian approximation to Gamma distributions Knowles (2015) further mentions that the Gaussian distribution can be used to approximate a Gamma distribution for larger α .

Original: Bahuleyan et al.
Revised: Bahuleyan et al.

Original: (2017) uses a similar approach to approximate variational attention weights.
Revised: (2018) uses a similar approach to approximate variational attention weights.

Original: The Gaussian distribution is a symmetric distribution which can be sampled by location-scale shifting, as discussed above.
Revised: The Gaussian distribution is a symmetric distribution which can be sampled by location-scale shifting, as discussed above.

Original: The Gamma distribution Γ( α,β =1) can be approximated with a Gaussian of the form γ ∼N ( α, √ α ) (Knowles, 2015), which gives us:[Equation]
Revised: The Gamma distribution Γ( α,β =1) can be approximated with a Gaussian of the form γ ∼N ( α, √ α ) (Knowles, 2015), which gives us:

Original: [Equation]
Revised: [Equation]

Original: The Gaussian distribution is symmetric and can take on negative values.
Revised: The Gaussian distribution is symmetric and can take on negative values.

Original: Hence, this approximation is inappropriate for the Gamma distribution unless the α i parameter is sufficiently large, otherwise the sample will need to be truncated to a value greater than zero.
Revised: Hence, this approximation is inappropriate for the Gamma distribution unless the α i parameter is sufficiently large, otherwise the sample will need to be truncated to a value greater than zero.

Original: The combined reparameterisation of Gamma distributions To visualise the error for these two approximations, their average L 1 distance from the true Gamma inverse CDF is plotted in Figure 9 for different values of α .
Revised: The combined reparameterisation of Gamma distributions To visualise the error for these two approximations, their average L 1 distance from the true Gamma inverse CDF is plotted in Figure 9 for different values of α .

Original: The plot shows that the approximation error is equal when α =0 .
Revised: The plot shows that the approximation error is equal when α =0 .

Original: To take advantage of the strengths of both these approximations, we propose to reparameterise the Gamma distribution as a blend of these two approximations with a switch at α =0 .
Revised: To take advantage of the strengths of both these approximations, we propose to reparameterise the Gamma distribution as a blend of these two approximations with a switch at α =0 .

Original: 6363 and truncate negative samples to zero.
Revised: 6363 and truncate negative samples to zero.

Original: [Figure]
Revised: [Figure]

Original: K P RACTICAL I MPLEMENTATION 
Revised: K P RACTICAL I MPLEMENTATION OF D ENOISING A TTENTION

Original: In this section we provide the equations used to allow denoising attention to be implemented in a deep learning framework at training time and test time.
Revised: In this section we provide the equations used to allow denoising attention to be implemented in a deep learning framework at training time and test time.

Original: Denoising Attention atTraining Time During training  the set of vectors Z ∈ R n × p  are  sampled and outputted by the NVIB layer.
Revised: Denoising attention at training time During training, the set of vectors Z ∈ R n × p and their log-probability values log ( π ) ∈ R 1 × n are both sampled and output by the NVIB layer, thereby specifying the sampled mixture distribution F .

Original: The query u ′ ∈ R 1 × p is projected by the grouped matrices W Q , W K ∈ R p ×
Revised: For each use of denoising attention, the query u ′ ∈ R 1 × p is projected by the grouped matrices W Q , W K ∈ R p ×

Original: d tosuch that u =( u ′ W Q ( W K ) T ) .
Revised: d to u =( u ′ W Q ( W K ) T ) .

Original: The probability values log ( π ) ∈ R 1 × p are sampled from the NVIB layer.

Original: The keys’ dimensionality d is  for scaling.
Revised: The keys’ dimensionality d is used for scaling.

Revised: Denoising attention can then be computed as:

Original: [Equation]
Revised: [Equation]

Original: We define this for u ∈ R 1 × p however, this can easily extend to multiple queries.
Revised: We define this for u ∈ R 1 × p , but this can easily be extended to multiple queries.

Original: Denoising Attention atTest Time During test time  we do not sample Z ∈ R n × p but rather use the mean  posterior distribution.
Revised: Denoising attention at test time During test time, we do not sample F , but instead use the mean of the posterior distribution, which is its base distribution G q 0 .

Revised: The NVIB layer takes its input from the encoder and maps it to the parameters ( µ q , σ q , α q α q 0 ) of this base distribution G q 0 = (cid:80) i α qi α q 0 N ( µ qi , I ( σ qi ) 2 ) .

Original: We use the posterior mean µ q ∈ R n × p and posterior variance parameters ( σ q ) 2 ∈ R n × p
Revised: For convenience let ( σ ri ) 2 =( √ d +( σ qi ) 2 ) . Test-time denoising attention can then be computed as:

Original: DAttn
Revised: DAttn

Revised: [Equation]

Revised: G

Original: [Equation]
Revised: [Equation]

Revised: where 1 p is a row vector of p ones.

Original: A caveat of this derivation is that  applies simply for single-head attention and is not trivial to apply for multi-head attention.
Revised: A caveat of this derivation is that it applies only for single-head attention and is not trivial to extend for multi-head attention.

Original: We leave this for future work.
Revised: We leave this for future work.