A VAE for Transformers with Nonparametric Variational Information Bottleneck
Authors: James Henderson
Fabio James Fehr
Venue: ICLR 2023 poster (2023)
Keywords: VAE
VIB
Bayesian nonparametrics
Transformers
natural language

Abstract: We propose a Variational AutoEncoder (VAE) for Transformers by developing a Variational Information Bottleneck (VIB) regulariser for Transformer embeddings.  We formalise such attention-based representations as mixture distributions, and use Bayesian nonparametrics to develop a Nonparametric VIB (NVIB) for them.  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.  Our Transformer VAE (NVAE) uses NVIB to regularise the information passing from the Transformer encoder to the Transformer decoder.  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.

Published as a conference paper at ICLR 2023
A VAE FOR TRANSFORMERS WITH NONPARAMETRIC VARIATIONAL INFORMATION BOTTLENECK
James Henderson
Fabio Fehr
Idiap Research Institute, Switzerland james henderson@ di.p..h
Idiap Research Institute and EPFL, Switzerland fabio. fehr@ i di ap.ch
ABSTRACT
We propose a Variational AutoEncoder (VAE) for Transformers by developing a Variational Information Bottleneck (VIB) regulariser for Transformer embeddings. We formalise such attention- representations as mixture distributions, and use Bayesian nonparametries to develop Nonparametric VIB (NVIB) for them. 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. Our Transformer VAE (NVAE) uses NVIB to regularise the information passing from the Transformer encoder to the Transformer decoder. 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.
INTRODUCTION
Attention-bast based deep learning models, such as Transformers (Vaswani al. 2017 Devlin et al. 2019). have achieved unprecedented empirical success in a wide range of cognitive tasks, in particular in natural language processing. The of attention allows these models to represent their input with multiple vectors, which essential for embedding natural language text (Bahdanau et al. 2015). On the other hand, deep variational Bayesian approaches to representation learning, such as variational autoencoders (VAEs) (Kingma &
Welling, (2014), have also been shown to have many benefits (Mathicu et al. 2019 Ghosh et al. 2020
Vahdat & Kautz 2020), especially due to their variational information bottleneck (VIB) (Alemi al. 2017)
for regularising the induced latent representations. However, it has not been clear how to combine these two trends, because the latent space induced by Transformers 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) Fang et al. 2021 Park & Lee 202T). To define a VIB regulariser for a Transformer's 's embedding space, we need to allow the size of a latent representation vary dynamically depending on the complexity of the individual imput, and yet regularise the total amount of information conveyed by the whole representation.
In this paper, propose such variational information bottleneck for variable sized latent representations, which we use to regularise the embeddings of a Transformer encoder-decoder, giving us a variational au- toencoder for Transformers Like Transformer encoder's '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. 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: We show that this regularisation improves generative abilities and compresses latent representations. In addition to the regularisation of over-paramen parameterised language models (Child et al., 2019), previous work shows the efficacy of VAEs for: disentanglement (Higgins al., 2017). language generation (Liu & Liu, 2019). and explainability (Mercatali & Freitas 2021). All these topics are important and active areas of research in NLP.
To define this VIB, we need to model distributions over these variable-sized encoder embeddings, as interpreted by cross attention. Firstly, because the attention function returns an interpolation between the vectors output by the encoder, generalises across the varying number of vectors, which like the input length is theoretically unbounded. Thus, to define distributions over these unbounded embeddings, we need to use nonparametric methods (Jordan 2010). Secondly, the attention function is insensitive to the
'he code is available https:/ggthaah.cm///daaapiiib and https//gthabbccm/ddaaiiihh transformers

---

Published as a conference paper at ICLR 2023
NVIB
Cp. Go
sampling
F:
denoising
intersnce
attention
(b) Denoising attention.
|set of ireeudo-bsservalion
query, denoising function
Transformer encoder
Transformer decoder
Denoising you need
Denoising is you need
(a) The NVAE model and its NVIB layer.
Test-time denoising attention.
Figure 1: (a) Illustration of the NVAE model, with its NVIB layer. (b) Query denoising attention at training time, with the sampled distribution the query prior, noisy query observation, and the expected value of the denoised query. (c) Query denoising attention at test time using the mean distribution.
order of the vectors output by the encoder, so it interprets this embedding as a permutation -nvvariant set of vectors. Thus, the distributions over these cormutatiooninvariann embeddings should be exchangeable [Jordan 2010). Thirdly, the attention function imposes a normalised weighting over the embedding vectors, via the attention weights. So we should model an embedding as distribution rather than a set.
A normalised weighting over an unbounded permutation-invarian set of fixed-length vectors matches ex- actly the properties a nonparametric space of mixture distributions, which have been extensively studied in Bayesian nonparametrics using exchangeable distributions (Blei & Jordan 2006 Jordan 2010). In previ- ous 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). In contrast, apply it to inferring latent representations where the number of parameters grows with the size of the input. We believe this the first work to use nonparametric methods this way for deep variational Bayesian models.
To define a precise equivalence between attention-based representations and mixture distributions, provide an interpretation of attention where the input set of vectors defines mixture of impulse distributions, which is used as a prior to denoise the observed query vector (depicted in Figure [b]. Generalising sets of vectors to mixture distributions and generalising the attention function to query denoising allows us to propose general deep variational Bayesian framework for iicention-base models using Bayesian nonparametrics. More specifically, we propose to use Dirichlet processes (DPs) as the exchangeable distributions (Aldous) 1985 Jordan (2000) to specify distributions over mixtures of impulse distributions, including distributions over the effective number of components in the mixture.
We define a nonparametric VIB (NVIB) layer using a bounded DP prior and posterior to regularise the effective size of variable-sized latent representations. This NVIB layer uses exact inference to infer the posterior from 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 Ia 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.
To evaluate the effectiveness of NVIB, we train NVAE 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 viable VAE. We also find that the regularised latent space smooth, using a proposed method for interpolating between DP posteriors to generate interpolations between sentences.
This paper makes the following contributions: (1) We propose variational Bayesian framework for modelling attention- based representations using mixture distributions, denoising attention and Bayesian nonparametrics (Section 2 (2) We propose nonparametric variational information bottleneck (NVIB) regulariser for learning attention-basee
Contributions
representations (Section 3). (3) We propose a
nonparametric variational autoencoder (NVAE), which is variational Bayesian extension of a Transformer decoder-decoder (Section 4.. (4) We show that the NVAE model a competitive VAE which can reconstruct, generate, regularise its latent space and intuitively interpolate between sentences (Section 5).
2

---

Published as a conference paper at ICLR 2023
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 Bahuleyan et al.) 2018: Fan et al. 2020 Cinquin al. 2022). Nguyen et al. (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. The use of Bayesian nonparametries to learn variable sized latent space using a VAE (Nalisnick & Smyth [0117] |ooyal et al.
2017
Echraibi et al. 2020) still assumes fixed-sized latent representation at test time, unlike our proposal.
2 A NONPARAMETRIC BAYESIAN FRAMEWORK FOR TRANSFORMER EMBEDDINGS
This section proposes formalisation of attention- -based representations mixture distributions over vector space, and proposes nonparametric Bayesian methods for modelling information about these mixture distri- butions. 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 We adopt mixture distributions over vectors as generalisation of attention based representations, and adopt this denoising function as a generalisation of attention. Then we use Bayesian nonparametrics to propose prior (Section 2.2 and posterior distributions (Section 2.3 over these mixture distributions. These priors and posteriors form the basis of our nonparametric variational information bottleneck, proposed in Section 3
2.1 DENOISING ATTENTION
The attention function provides access to a set of vectors by mapping a query vector to the resulting attention vector. As the basis of our approach to attention-based representations, we generalise the set of vectors to probability distribution over vectors, and generalise attention to function of these probability distributions.
The attention mechanism we assume is scaled dot product attention, standardly used in many attention- based models including Transformers. For simplicity, we consider cross attention, where a single query vector is mapped to a single result vector. This attention function projects the imput vector w e <R' Xx via the weight matrix 8º e RPxd to query, and projects the set of vectors Z E R"XP via weight matrices WKK WV xR to keys and values, respectively. It uses the keys' dimensionality d for scaling. 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.
Attention(u' u/ Z: 2 W W = Attn(u'W' (W Z) W = Attn(u, WV
where u - (u' Q (
IIxp_ The function Attn(u, Z) can then be defined in two equivalent ways (as shown in Appendix Gp in terms of a sum over the vectors z; in Z, or in terms of an integral over distribution which is only nonzero at the Z;:
Attn(u.Z)
softmax
1 u
Z
DAttn(u; Fz)
(1)
Va
exp( 2va
Fz
8z.
(2)
exp( 2vd
J. f(v) g(u; ddI)
DAttn(u; F)
dv
(3)
(o.(() g(u; v. v.vdr) do
where 52, is impulse distribution z,, f(-) is the probability density function for distribution F, and glu; u, (T) is the multivariate Gaussian function with diagonal variance of Vd. As depicted in Figure 1b DAtin(u; Fz) can be interpreted as query denoising. The query u is interpreted as an observation of some true vector v which has been corrupted by Gaussian noise, where was generated from a prior probability distribution I'z specified by Z. The result of Attn(u. (u, Z) is the expected value of this true vector e after seeing the noisy observation u, which can be interpreted as a form of denoising.
This denoising attention function DAttn(u; F) is actually a generalisation of attention over set of vectors, that is defined for any probability distribution F over vector space. In the special case where T,,d, is finite mixture of impulse distributions, it is the same as Atin(u.2) but with a bias term log( (T)) substituted for 1 2 ||2. 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 Transformer encoder-decoder mixture distributions.
3

---

Published as a conference paper at ICLR 2023
2.2 A PRIOR OVER MIXTURE DISTRIBUTIONS
Given that our attention-base latent representations are formalised as mixture distributions F, a Bayesian approach requires a prior over these distributions. Attention-based models place 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. Nonetheless, we can still specify probability distributions over this infinite space of possible distributions F, using methods from Bayesian nonparametrics. These nonparametric Bayesian methods, with exchangeable distributions, are specifically designed for modelling probability distributions over unboundedly large mixture distributions.
We base our distributions over mixture distributions on Dirichlet processes Oo)) Dirichlet processes (DPs) are a generalisation of Dirichlet distributions to an infinite support, such as the points in vector space. Dirichlet distribution Dir(ox) is a distribution over probability mass functions π of discrete categories i, 1 <<ig.. K. One useful definition of Dirichlet processes views DP F ~ (oo,)),)) where Go is the base distribution over vectors and oop E R is the concentration parameter, as the limit of sequence of finite Dirich- let distributions (see Teh 2010)), ). given in equation 4 Note that the Dirichlet distributions in equation 4 Jrr symmetric, in that all the A categories have the same a; = cxo parameter values. However, these categories end up with very different weights T, 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. In the infinite limit, this tail is infinitely long with infinitesimal probabilities. The number of categories which get nontrivial probabilities is determined by og. and becomes independent of K as λ gets large.
As shown in this definition, each sample F' from a DP is an infinite mixture of impulse distributions Se, parameterised by an infinite sequence of weight-vector pairs T This contrasts with the finite Z in attention- based representations. Having an infinite F would also cause problems in our variational Bavesian model, because VIB uses bound on the log- ligliiilioood (see Section (.1),
(3.I), which Kingma &
Welling (014) showed has an error of D [a) II p(F [z))) (the looseness of the bound). This would
be infinite unless both the true posterior p(F) x) and its approximation g(F) 2)) generate a finite F, so we need a prior which generates finite F.
F
(4)
F =
Σ T=OZ.
(5)
COD K
Oo
000
lim Dir
Dir
KO
KO
Go for i=
Go for I......0
The Unbounded Dirichlet Process Prior We do not want a prior which places an apriori bound the size of F. so we assume it is finite but unbounded, and propose a prior which is an unbounded seguence of finite approximations
a DP. We define a bounded DP F (Gg, as in equation 5 Our approach to the prior is to use an unbounded but finite Kg. so we define a distribution over approximations as Kg 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. Given & is some distribution over positive integers & E Z we define this unbounded DP as (D) (G). where gー
We use these definitions both to define a general prior over probability distributions, and to define a conditional prior for each input length. In both cases, the base distribution Gj is assumed to be a unit Gaussian (inspired by Kingma& & Welling 2014h) and the concentration parameter a is assumed be one
Go = =N(He ((o")) of 1; 1²² = O: =1=1
The general prior UDP oD og,))
o)), where the size distribution of is determined empirically. The conditional prior BDP o Ku) sets the level of approximation Ki as a fixed function of the input length n, in particular KO (n= 1),4, where A E E a hyperparameter that controls the approximation.
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. If we know the imput length n, but know nothing about the content of the text, then the distribution of vectors should stay the same as the general prior, Go = G. However, the count of observations we expect to have after an input
A Conditional Bounded DP Prior
We will use "p" and "q" superscripts to designate variables for the prior and posterior, respectively. Similarly, zero subscript is part of the name of the variable, contrast positive integer subscripts which indices.
4

---

Published as a conference paper at ICLR 2023
of that length would not be ob. but should include pseudo-count a R2o hyperparameter for every token, and thus ag o + +na This then gives us the conditional prior given n of BDP .o))-
2.3 A POSTERIOR OVER MIXTURE DISTRIBUTIONS
Since a DP is a conjugate prior, we can use exact inference to compute the posterior DP from the prior DP plus set of pseudo-observations output by the encoder. Each pseudo-observation is a real-valued pseudo-count a, E R2 and a parametric distribution which represents uncertainty in the observation. We use isotropic Gaussian, G = Ni [(o)))) as the parametric distribution, specified by mean μ gxxd and a standard deviation of R. 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-observation may have zero pseudo-counts and thus be effectively removed from the set. The formula for exact inference of the posterior DP is given in equation 6
where there is an n-11"h component of the base distribution G3 which comes from the prior, namely ax+1 =o and n+ Go
F ~ DP og
(6)
F
Σ P.P.
n+1
a+1
i=1
G =
Dir(a)
=
F
ADP(C,aff.t for fori=1,,m 11,...,,,
We derive an alternative factorisation of the posterior DP (Appendix HD which helps with the sampling method in Section 3.2 We then bound this factorised DP so that it generates the same number (n= of weighted vectors as the prior BDP ag. Ku). The resulting bounded posterior F ~ BFDP(G ad, A²) is given in equation 7 which defines our posterior distribution q(F)) x). This posterior simplifies to a mixture F = E Σj= Pet of impulse distributions, with p~ DrrrDaaf........ Din of (. and Ziy Gg.
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 Gg. This base dis- tribution is a continuous distribution, whereas at training time all samples are discrete distributions. Nonethe- less, when accessed via denoising attention, the base distribution looks like typical sample from the poste- rior. This is visualised in Figure by comparing the vector returned by denoising attention (in green) given the continuous mean distribution (Figure Ic} and a typical sample from this distribution (Figure lb). 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 the mean distribution. For our Gaussian mixture base distribution Go there is a closed-form solution to computing denoising attention, given in Appendix F
THE NONPARAMETRIC VARIATIONAL INFORMATION BOTTLENECK
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. Such encoder-decoder interfaces take set-of-vectors representation and return function from query result vectors. We map the input set-of-vecton to set of oos.udddobser,innnss, and define the returned function with denoising attention. Then, given the nonparametric prior and posterior from Section 2 we can define our nonparametric VIB regulariser by specifying how to compute the KI. divergence between the prior and posterior, and how to effectively sample from the posterior for training. As far as we are aware, this proposal is the first VIB model for attention-baan based representations like Transformer embeddings.
The VIB layer a VAE controls the amount of information passing through it by introducing noise accord- ing to a posterior output by the encoder, and regularises this information by minimising the KL divergence between this posterior and an uninformative prior. 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). Similarly to the free- bits objective proposed to address this problem for vector-space VAES (Kingma et al. 2016), instead of regu- larising towards the prior, which gets no pseudo-counts from the input, we regularise towards the conditional
5

---

Published as a conference paper at ICLR 2023
prior BD of .Ko), which gets no pseudo-counts from the input but knows nothing about the infor- mation they convey. We found that this helps with the stability of training and avoiding posterior collapse.
3.1 THE VARIATIONAL INFORMATION BOTTLENECK LOss
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 is the input text. log(p(x))
2 E
((F log(p(æ>
F)) (KL(q(F)) [2) p(F))
(8)
Ln -EE E of x) log(n(>) | F))
The first term of the bound is the reconstruction loss T.R, computed using samples F from the approximate posterior q( (F) [x)), and the second term is the KI. divergence between this posterior and the prior p( ()).
For the KL divergence term, since both the prior and the posterior conditioned on the same bound Ko on the number of vectors they generate, we can compute a meaningful finite KL divergence between the conditional prior p(F)=BDP(C oy .Ko) and the posterior q( =)=BFDP , ,.),, given x² and (o-=n+1)x=- The derivation of the KL divergence is given in Appendix It has two terms, one Ln for the distribution of weights π generated by the Dirichlet distributions, and one LG for the distribution of vectors Z generated by the component Gaussians. Using the exact KL divergence for our bounded DPs would regularise each component equally, even for components which have zero of and thus have no impact on the posterior. We instead approximate a KI. divergence where only samples with a nontrivial weight are regularised. Marginalising over the number of nontrivial weights for each component does not appear to be tractable for LD. but since the relationship is approximately linear (see Appendix []. we simply substitute the expected number [ Ko for the actual number in the equation for KL divergence. This gives the og
following loss terms for the KL divergence, where is the gamma function and da is the digamma function.
LD x DKL(g(F p(F))
LD = (ogI ++aaa-oat +10 Ki
(-va0)0 ) two
logI" logo'
CO
Ko
( -,, (0 - 1o (0
LG
No i=1
(a))² (o
(o
Both LD and LG scale approximately linearly with Kg.
To generalise the ELBO beyond autoencoders, it can be viewed as way regularise the amount of information which passes through the latent representation (Alemi et al. 2017). This VIB interpretation allows the different parts of the objective to have different weights. We introduce two hyperparameters to control the relative weight the above three parts of the ELBO, which defines our VIB loss L. Ln AnLD AgLa
L
(10)
3.2 SAMPLING A MIXTURE DISTRIBUTION FROM THE POSTERIOR
To control the amount of information which passes from the encoder to the decoder, at training time VAE (Kingma & Welling 2014) samples from the encoder's 's posterior distribution and uses this sample reconstruct the input. The "reparameterisation trick" is used to ensure that backpropagation of the recon- struction error through this sampling step can be done effectively. 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.
For our NVIB model, we sample the parameters (TT, (Tr,Z) of mixture distribution F generated by our bounded Dirichlet process posterior BFDP (G",""",a" , AAA,, where F consists of a set of impulse distributions dz. each with weight ,.. A straightforward approach sampling from a Dirichlet process would independently sample weights л from a (theoretically infinite) Dirichlet distribution and sample vectors 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 vector from that component's 's Gaussian.
Factorised Sampling Method The problem is that sampling a component is a discrete choice, for which there exact reparameterisation trick. Instead, we note that the components do not differ in the

---

Published as a conference paper at ICLR 2023
number of vectors sampled from each one (always theoretically infinite for DP), but only differ in the distri- bution of weights for those vectors. As specified in the factorised DP in Section 2.3] we characterise this dis- tribution over weights by factorising into two steps: first choosing how the total weight is distributed across components (p), and then for each component choosing how its weight is distributed across its vectors (T(). These are both continuous choices. The vectors then he sampled independently from each component.
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 Z; from an individual component i.
The factorised and bounded nature of our posterior BFDP means that the total weights p and the individual weights T.. are all sampled from Dirichlet distributions. A Dirichlet distribution over a set of category weights can be sampled by sampling from a Gamma distribution for each category and then normalising. There is no dosed-forms explicit reparameterisatioo trick for the exact Gamma distribution, but there are for approximations. Knowles (2015) proposes two such approximations which we combine, one which is more accurate for small and one which is more accurate for larger a. We leave the investigation of implicit reparameterisation gradients (Figurnov et al., 2018), an alternative approach to explicit reparameterisation, to future work. More specifies about the sampling methods and their reparametcrisation tricks is given in Appendix ] We provide an evaluation (Appendix B.3) showing that =A=1 is an efficient and effective sampling method. In this case. there is no need sample from the Dirichlet for each individual component, but we still sample the weights across components.
4 THE NONPARAMETRIC VARIATIONAL. AUTOENCODER
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 Ta In this NVAE model, the Transformer encoder is used to estimate the parameters (od ,rt o²)) of the posterior given the imput text x. The Transformer decoder is used to reconstruct the input text using denoising attention over sample F from this posterior.
The NVIB Regulariser Our NVIB layer regularises the amount of information which passes from the encoder to the decoder through this posterior. As with VIB for vector spaces, the KL divergence encourages the encoder to output component Gaussians with smaller lu and larger o¹. With NVIB, the KI. divergence also encourages the encoder to output smaller and sparser a², which regularises the effective number of components in the posterior as well as the noisiness of their weights.
The Transformer Encoder The Transformer encoder (F((z)) with a text x of n number of tokens is used to compute vector for each token i of the input. From each of these individual token embeddings, the encoder then linearly projects to three parameters, of E RR, p4? E RRXx and log(a?) ) E RIxp, The variance parameters are exponentiated to be strictly positive, whereas the pseudo-count parameters of are estimated using Rectified Linear Unit (ReLU) activation (Nair & Hinton (2100), which results in masking the vector during cross-attention when it is exactly zero. Thus, the DP posterior has one component (a? 1gf 44.00 . of its base distribution for each token of the input (plus one for the prior).
The Transformer Decoder The Transformer decoder q(x(F) receives a distribution F over vectors and reconstructs the input text During training, F is specified by the sampled vectors Ze 1o and the sampled weights π E Ro and at test time F is specified by the output of the encoder a 1 2e2 <o and log(o ) R In both cases, the decoder accesses F using denoising attention in the same way that standard Transformer decoders use cross attention. We include the exact equations used for deep learning implementation of denoising attention in Appendix K During training, the text predicted using teacher forcing. and during test time the text is predicted autoregressively using greedy decoding until the ond-of-sequence token is generated or the sentence generated is 50 tokens larger than the target length.
The Generative Model
To use our NVAE model generative model, we sample from the prior and use the trained Transformer decoder to generate a sentence. As discussed in Section 2.2 to sample from the same prior used for training, we need to first sample sentence length, and then sample from the conditional prior given that (approximate) sentence length. For simplicity, we sample the sentence length from the empirical distribution of sentence lengths in the training data.
7

---

Published as a conference paper at ICLR 2023
5 INTRINSIC EVALUATIONS OF NVIB IN NVAE
To support our theoretical contributions, we provide proof-of-concepl experiments which demonstrate that our proposed NVIB regulariser performs as claimed. We evaluate it in our proposed NVAE model by training NVAES on natural language text and evaluating the resulting models. We show that the NVAE a viable VAE model it exhibits a competitive reconstruction versus generation trade-off (Section 5.1). We show that the NVIB layer is able to dynamically choose the number of components it needs its embeddings (Section 5.2). Additionally, NVIB provides an intuitive way interpolate between sentence embeddings, which provides an evaluation of the smoothness of the latent space (Section (3.
Data The Wikitext-2 and Wikitext- 103 (Meriy 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.
Baselines We compare to various alternative ways to define a VAE from a Transformer autoencoder. As representative of a standard seed-leengh--vector VAF, the Variational Transformer Pooled (VTP) baseline pools its vectors across the sequence length dimension, and then applies a Gaussian VIB layer (Kingmal & Welling, 2014). At the other extreme, the Variational Transformer (VT) baseline keeps all its vectors and applies a Gaussian VIB layer to each one. In between these baselines, as a hand-coded solution to constraining the quantity of latent vectors, Variational Transformer Stride (VTS) baselines, with parameter S, masks 1-S proportion of the embedding vectors based on their position. 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.
5.1 RECONSTRUCTION VERSUS GENERATION
This section shows that the NVAE Transformer model is competitive VAE in both reconstruction of input sentences and generation by sampling from the prior. All models undergo hyperparameter tuning on the validation set (Appendix [B across seeds, to select the best models and then report results on the Wikitext-2 test set. For reconstruction, we report the SacreBLEUU metric (Papineni et al.) 2002, Post
2018).
For generation, we report forward perplexity (F-PPL) and reverse perplexity (R-PPL) Zhao et al.
2018 Cilka et al. 2018), which trains an external language model on the gold training text and evaluates 1t on the generated text (F-PPL) or vice versa (R-PPL). Training details are provided in Appendix
Figures 2a and 25 show the reconstruction versus generation trade-off on the Wikitext-2 test set, where lower right is better. The single vector baseline VTP is unable to reconstruct well (low BLEU) or generate diverse sentences (high R-PPL) Even with larger capacity and more training data, the model performs poorly in generation (Appendix [B.4). 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.
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. In particular, reverse perplexity shows that NVAEs are able to generate diverse collection of sentences from the prior, and forward perplexity shows that these are fluent sentences (samples in Appendix (.. We believe this advantage the result of learning what vectors to keep instead of hand-coded position- -based dropout.
1.20
100
30
Nak VT
VIS.
Number of Viccoos
25
1.3
VIP
20
BLEU(
80
1.2
R-PPL
1.10
35
1.1
1.05
X
70
WAKE VTS
20
1.00
50
80
20
(
(1)
Input Tokens
(a) F-PPL vs test BLEU.
(b) R-PPL vs test BLEU
(c) Test BLEU vs proportion (d)
Latent quantity of
of retained vectors v.
vectors vs input tokens.
Figure 2: Reconstruction versus generation trade-off (a), (b) and regularisation analysis (c), (d).
8

---

Published as a conference paper at ICLR 2023
5.2 REGULARISATION
This section shows that the NVIB layer is able to regularise the number of vectors in the latent representation of a NVAE. Without regularisation the NVIB layer becomes a standard transformer without any noise and retaining all latent vectors (ablation in Appendix B.I.. The proportion of vectors retained controllable by the conditional prior hyperparameter a Figure 2c shows that there exist a where the NVAE is able to remove a significant proportion (1 (1) of vectors whilst maintaining a high reconstruction performance. Figure 2d plots the number of latent vectors retained during evaluation against the number of input tokens. 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, The different NVAE lines show the effect of a on vector retention, with larger values allowing more vectors to be used (also see the evaluation in Appendix B.2)
5.3 INTERPOLATION
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. We simply interpolate the probability assigned to each point in vector space. Given two latent mixture distributions Fi and F2, we decode from the combined mixture distribution interpolation rates S 1. For the baselines we use +Z + (1-T)ZZ such that the interpolation is over the content of the latent vectors Z, We align latent sets of vectors by input position and pad smaller sets with zero vectors, which is the mean of the Gaussian prior. We evaluate the interpolation with selected larger scale models (see Appendix (B.4) and use the Wikitext- 103 validation set for S, and its reverse order for S2. The results in Table and Figure show that the NVIB regulariser in NVAE provides smoother interpolations, and improvements in fluency of interpolations over the baselines These findings are qualitatively confirmed through different examples (Appendix E
for varying
100
NWW
VTS
Model
0.25
0.5
0.75
F-PPL (L)
80
VT?
VI
VT
0.04
1.00
0.04
1+1.6e
60
VTP
P=moa
0.57
0.99
0.57
143.3e
40
VTS
0.8
0.04
1.00
0.04
1 1+2.8e
NVAE
a =044
0.88
0.89
0.88
1+1.3e
20
75
100
Table I: The proportion of interpolations different from S and S2 by varying the interpolation rate T. Fluency metric F-PPL of interpolations when T =.55.
S, BLEU
Figure 3: BLEU with S versus S2 for varying interpolations T.
CONCLUSIONS
In this work, we generalise the latent representations of attention- based models to mixture distributions over a vector space, and propose nonparametric variational information bottleneck (NVIB) to regularise these latent representations. Using this NVIB model, we propose nonparametric variational autoencoder (NVAE), which uses a Transformer encoder to embed text in nonparametric space of distributions over mixture distributions, and uses Transforner decoder to generate text given sampled mixture distribution. 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. Our NVIB model adds the ability to regularise so that the size of the representation is appropriate for the complexity of the input being encoded. This is a crucial ability for encoding text, where the size of text can vary enormously, Empirical evaluations indicate that this model: a competitive VAE in that it is able to reconstruct input sentences and generate 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.
pveetion-based representations
Future work will evaluate the effectiveness of NVIB at larger scale, when applied to multi-heal attention and to self-attention layers, and when its pretrained representations are used in downstream tasks.
3p-PPI. is calculated across interpolations using Transfomer language model trained on Wikitext- 103 at the same scale the larger models. We remove any collapses exactly S, or S2 this will bias the F-PPL. metric favourably.

---

Published as a conference paper at ICLR 2023
AUTHOR CONTRIBUTIONS
James Henderson is responsible for the high-level vision and the theoretical derivations and proofs related to Bayesian nonparametrics. Fabio Fehr is responsible for the sampling method, evaluations, implementations and rumning of experiments.
ACKNOWLEDGEMENTS
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, This work was funded in part by the Swiss National Science Foundation under the NCCR grant Evolving Language, Swiss National Science Foundation Agreement #5 INF401180888
REPRODUCIBIL STATEMENT
We faithfully describe the details of method in the text and provide detailed derivations for theoretical
grounding (Sections F G D. We mention all relevant hardware, hyperparameters and datasets to reproduce our experiments (Section AA). We provide variance estimations acToss seeds and ablations to justify design choices (Section B). Finally, we provide exact equations for denoising attention used to make implementation easier (Section Kj. The open source code for this research has been released at https://gthub..om/diaap for the NVIB layer in PyTorch and httpss//gihuu..cm//daa//////// transformers for the experiments.
REFERENCES
David Aldous. Exchangeability and related topics. In P. L. Hennequin (ed.), École d' d'Éíé Probabilités de Saint-Flour XIII - 1983, pp. 1-198, Berlin, Heidelberg, 1985. Springer Berlin Heidelberg. ISBN 978-3-540-39316-0.
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck. In 5th International Conference on Learning Representations, ICLR 2017, Conference Track Proceedings, Toulon, France, 2017. OpenReview.net. URL net/ wtatt//frumiii-Hyooo
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, Conference Track Proceedings. San Diego, CA, USA, May 7-9 2015. URL htRLLhttt////xttee www/abs//aaeeeeeeeeooeooeee 409. 0473
Hareesh Bahuleyan, Lili Mou, Olga Vechtomova, and Pascal Poupart. Variational attention for eequence-to-seqquecce models. In Proceedings of the 27th International Conference Computational Linguistics, pp. 1672-1682, Santa Fe, New Mexico, USA, August 2018. Association for Computational Linguistics. URL https:://///mm https/////aaaatttoooogeee org/c18 1142
Steven Bird, Ewan Klein, and Edward Loper. Natural language processing with Python: analyzing text with the natural language toolkit. O' O'Reilly Media, Inc., 2009.
Christopher M. Bishop. Neural Networks for Pattern Recognition. Oxford University Press, Oxford, UK 1995.
David M. Blei and Michael I. Jordan. Variational inference for Dirichlet process mix- doi: 10.1214/06-BAII
tures. Bayesian Analysis, 1(1):121 143, 2006.
URL
https://doi. http://oo.orgg///// .org/ 1214//06 -8166-811444
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from continuous space. In 20th SIGNLL Conference on Computational Natural Language Learning, CoNLL 2016, pp. 10-21, Association for Computational Linguistics (ACL), 2016.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv: Machine Learning, abs/ 1904, 10509, 2019. doi: htthttss: .rrg wwwpg///////aaa /alos/19044 10509
10

---

http::///oeererreeeesn///oouuu Published as a conference paper at ICLR 2023
Ondrej Cifka, Aliaksei Severyn, Enrique Alfonseca, and Katja Filippova. Eval all, trust few, do
wrong to none: Comparing sentence generation models. CoRR, abs/1804.07972,, 2018. URL
http://arxiv wttpe/iiivrgggb/s//eeeooa wwt///gg//s....o.o http//aasss1c.... 07972)
Tristan Cinquin, Alexander Immer, Max Hom, and Vincent Fortuin. Pathologies in priors and inference for bayesian transformers. In Fourth Symposium on Advances in Approximate Bayesian Inference, 2022.
wwtp//veeeeviieeeeeee ew net. fort//ooummm d= -e0-T-3hw0C9999
Yuntian Deng, Yoon Kim, Justin Chiu, Demi Guo, and Alexander Rush. Latent alignment and variational attention. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https: h ptttpp:://procceednnggs newr ips. httpsc/a/////////EEL... 16691344cef0ddaabe 4406-14447722 445544477777/555--PP-ee
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, Volume (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/11 1423. URL
Amine Echraibi, Joachim Flocon-Choles Stéphane Gosselin, and Sandrine Vaton. On the variational poste- rior of dirichlet process deep latent gaussian mixture models. ArXiv: Computer Science, abs/2006.08999, 2020. doi: 0044850//RXX..2206.008
URL htl.tttt::///// /tt///axiivv..////// wwwivorgg/ae abs /2006.2206..00 08993
Xinjie Fan, Shujian Zhang. Bo Chen, and Mingyuan Zhou. Bayesian attention modules. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 16362-16376. Curran Asso- ciates, Inc., 2020. URL https:://proceedingg www.r.pps.ccppppee//////e. SrF.56552225510003030022222-2-02
.ddf
Le Fang, Tao Zeng. Chaochun Liu, Liefeng Bo, Wen Dong, and Changyou Chen. Transformer-base conditional variational autoencoder for controllable story generation. CoRR, abs/2101.00828, 2021.
Uttttt / arxiv. orgg httrg/ab//211111 00828
Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. In Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018. URL
www.:eewewdinin.nenenneeeeeeeeee neur httprrppp..c/aaeee/////////
Partha Ghosh, Mehdi S. M. Sajjadi, Antonio Vergari, Michael Black, and Bernhard Scholkopf. From variational to deterministic autoencoders. In 8th International Confer- ence on Learning Representations, ICIR 2020, Conference Track Proceedings, 2020. URL
Prasoon Goyal, Zhiting Hu, Xiaodan Liang, Chenyu Wang, Eric P. Xing, and Camegic Mel-
lon. Nonparametric variational auto-encoders for hicrarchical representation learning. 2017
IEEE International Conference on Computer Vision (ICCV), 2017/papers/Goya theevf www.rf..ooccoooo contectntett....mcm ICCV 2012 Nonparametrric_avraatttt Variational Auto-Encoders Variiattinoa_Attt--EEnddder ICCV 2017 11017_ppper_ppee
5104-5112, 2017. URL
https://
Irina Higgins, Loie Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Mathew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. In 5th International Conference on Learning Representations, ICLR 2017, Conference Track Proceedings, Toulon, France, April 24-26 2017. OpenReview.net. URL
M.I. Jordan. Bayesian nonparametric learning: Expressive priors for intelligent systems. In R. Dechter, H. Geffner, and Halpem (eds.). Heuristics, Probability and Causality: Tribute to Judea Pearl, chapter 10. College Publications, 2010.
11

---

Published as a conference paper at ICLR 2023
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd Internationaa Conference on Learning Representations. ICLR 2015, Conference Track Proceedings, 2015. httthht:/aaaae org/ abs 6980
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference Leurning Representations, ICLR 2014, Conference Truck Proceedings, Banff, AB, Canada, 2014, URL.
/112. 6114
Diederik P. Kingma, Tim Salimans, Rafal Józefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improving variational autoencoders with inverse autoregressive flow. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, pp. 4736-4744, Barcelona, Spain, 2016. URL httpt:://////........ httptt/pss///proceeeeeiiioo http:/2201//aaahaeeeeeeeaaa http:///eeeeeataa.....22222226a3
neurips CC/ html
David A. Knowles. Stochastic gradient variationaa bayes for gamma approximating distributions. arXiv: Machine Learning, 2015. ttl.httts://aaxxie org.
Jannik Kossen, Neil Band, Clare Lyle, Aidan N Gomez, Thomas Rainforth, and Yarin Gal. Self-attention between datapoints: Going beyond individual input-output pairs in deep learn ing. In M. Ranzato, A. Beygelzimer. Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 28742-28756. Curran As- sociates, Inc., 2021. URL httptttt////ppceeeeiiiig wwwiinng....uureeeeeeeeeeeeee neur cteurr/ps.cccppper//00//// FIL0.abbSSEs22ftaacc1111131aaaa. Pappre
Danyang Liu and Gongshen Liu. A transformer-hase variational autoencoder for sentence generation. In 2019 International Joins Conference Neural Networks (IJCNN) pp. 1-7. IEEE, 2019. doi: 10.1 10.1109//CNN222119888222 /R/tttss:///axxx http://////a.ii...rr////////////
12738
Alice Martin, Charles Ollion, Florian Strub, Sylvain Le Corff, and Olivier Pietquin. The monte carlo transformer: a stochastic self-attention model for sequence prediction. ArXiv: Mathematics, Computer Science, abs/2007 006600, 2020.
http://aaavvvvvv. org abs 2007 08620
Emile Mathieu, Tom Rainforth, N Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autoencoders. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings Machine Learning Research, pp. 4402-4412, Long Beach, California, USA, 09-15 Jun 2019. PMLR. URL https://arxiv.oom
alss 1812. 02833
Giangiacomo Mercatali and André Freitas. Disentangling generative factors in natural language with discrete variational autoencoders. In Findings of the Association for Computational Lin- guistics: EMNLP 2021, pp. 3547-3556, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653w122222.ffninn.---- findiings--mnl..001... URL
https://aclanne http://acaanthhoooooooooo otw.aaatooooooooo//00220000
findings-ennnnnnn emmn-mmmllp p.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher: Pointer sentinel
mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Conference Track Proceedings, Toulon, France, 2017. OpenReview.net. ht https://openreviees www.nett/orommm.ddyyyyyauuuu
URL
Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Johannes Fürnkranz and Thorsten Joachims (eds.), Proceedings of the 27th International Conference on Machine Learning (1) pp. 807-814, Haifa, Israel, 21-24 June 2010. Omnipress. URL
https: /ttess///mmmmmm.mmm/ weee/Cooneeeeeeeeeeeeeeeee 201 http://00000pppp////ea 44f..pffffffff
Erie T. Nalisnick and Padhraic Smyth. Stick-breaking variational autoencoders. In 5th International Con- ference on Learning Representations, ICLR 2017, Conference Track Proceedings. Toulon, France, April 24-26 2017. OpenReview.net. URL https: http:/::////eneeeiiiee http://wwwwwwweeeeeeeeeeeeeeeee wttp://wwww..eeeoooedoeee..eae jmjotxx
Tam Minh Nguyen, Tan Minh Nguyen, Dung D. D. Le, Duy Khuong Nguyen, Viet-Anh Tran, Richard Baraniuk, Nhat Ho, and Stanley Osher. Improving transformers with probabilistic attention keys. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato
12

---

Published as a conference paper at ICLR 2023
(eds.). Proceedings of the 39th International Conference Machine Learning, volume 162 of
Proceedings of Machine Learning Research, pp. 16595-16621. PMLR, 17-23 Jul 2022. URL
https:://procedinggs.ttt wttpp//essw.aaaaagaeeeeo...ttt. mttp://2yynnn22.tmmm..
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: method for automatic evaluation of machine translation. In Proceedings of the 40th Annual Meeting of the Association for Compulational Linguistics, pp. 311-318, Philadelphia, Pennsylvania, USA. July 2002. Association for Computational URL
Linguistics. doi: 10.3115/107330831177111 URUptttsss//cccnnntool... -ttp://ooooooooooo/00-0-111111
Seongmin Park and Jihwa Lee. Finetuning pretrained transformers into variational autoencoders. In Proceedings of the Second Workshop on Insights from Negative Results in NLP, pp. 29-35, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguisties. doi: 10. 10.18653/11202111 Linsights- 1.5.
wwt.llooy.rgg//2021.11 inghhtss 1.5
Matt Post. A call for clarity in reporting BLEU scores. In Proceedings of the Third Con- ference on Machine Translation: Research Papers, pp. 186-191, Brussels, Belgium, Octo- ber 2018. Association for Computational Linguistics. doi: 10.18653/1/W18-631199 URL
https: hatpss//aaclaaten /ww./aaannnntoy........ wwq//w/--... 6319
Yee Whye Teh. Dirichlet processes. Encyclopedia of Machine Learning, volume 1063, pp. 280-287.
Springer, 2010.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 19667-19679. Curran Associates, URL https:://rrooeeddggs.o www.ies....... http:///rorrrrrroooo//////////// 13221961BccfF2201166666633311333
Inc.. 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett (eds.), Advances in Neurul Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL httt::::/pproeeeeee
newrips
Junbo Zhao, Yoon Kim, Kelly Zhang, Alexander Rush, and Yann LeCun. Adversarially regularized autoen- coders. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 5902-5911, PMLR, 10-15 Jul 2018. URL https: /ttp///epccccccseeeeeeeeeoe...se mlr www.s/wwwsssspp 2haol
13

---

Published as a conference paper at ICLR 2023
A EXPERIMENTAL SETUP
Training details All models are trained, without pretraining, using the same encoder and decoder config- uration for comparability. We two layer Transformer encoder and decoder with single attention-head. The size for the word embedding vectors and model projections are 256, feed forward dimensions 1024, which leads models of approximately 19 million trainable parameters. The BERT base-uncased tokeniser is used for tokenisation with a vocabulary of approximately 30K. During training we use: constant learning rate of le Adam optimiser (Kingma & Ba 2015). a batch size of 256, gradient norm clipping 0.1 and trained for 50 epochs (== 15K steps). The number of epochs were selected considering model con- vergence and minimising computation time. As form of regularisation we use dropout rate 0.1 and the VIB parameters λg. ÅD. a and & are selected through hyperparameter tuning. Learning rate schedules, KL annealing strategies and free-bits KL loss objectives were considered but unnecessary for convergence. Each model experiment takes approximately 2hrs to run on a single NVIDIA GeForce RTX 3090.
Generation metrics The two automatic generation metrics forward/reverse perplexity (Zhao et al. 2018 Cifka et al. 2018) both involve training external language model, for which we Transformer lan- guage model with the same configuration as in the previous paragraph, but without VIB regularisation. First we generate 100k sample sentences from the model we wish to evaluate. The forward perplexity (F-PPL) measures the perplexity of the external language model trained on training data and evaluated on the gener- ated text. This measures fluency of the text. The reverse perplexity (R- (R-PPL) measures the perplexity of the external language model trained on generated samples and evaluated on the validation or test data. This cap- tures the word frequency and overall proximity of our generated text distribution to the true text distribution.
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 [.4) and interpolations (Section 5.3). The datasets cleaned and segmented 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 al. 2019). Dataset statistics can be found in Table
Train/ Val/Tet Tokens
Wikitext-2 Wikitext- 103
77K/8K/900 /9K
212
3578K /9K/8K
+10
Table 2: Dataset statistics. Number of sentence examples in the train, validation and test sets. Number of word piece tokens per sentence example.
B HYPERPARAMETER TUNING AND ABLATIONS
The models are trained on the Wikitext-2 training dataset using the loss from equation I0 They are tuned on the validation dataset with the aim to be able both reconstruct and generate output.
Because both LD and LG scale approximately linearly with Ko, and Kg is linear in the sentence length, the DKL divergence losses Lp and LG grow linearly with the sentence length n. 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 4. addition we scale the weights on the Gaussian KL divergence by 3. removing the dependence on the dimensionality d of vectors.
1
11
Ap
-Npi AG
dn X
n
where XD and Xa are fixed hyperparameters.
All combinations of the following hyperparameters were considered in grid search for the respective models:
X (1, le 1, 1e 14-4, le 5.03
Xp -10, 1e 33 le 1e 5.0}
4 =1,, 0., 0.3, 0.2, 011,0)
4 2,,
14

---

Published as conference paper at ICLR 2023
{0.9, 0.8, 0.75, 0.5, 0.25}
. P mean, max, one}
where Ng and Ny are the weights the Gaussian and Dirichlet KL divergences for all variational models, respectively. The a and R NVAE specific parameters and represent the conditional prior parameter and number of samples per component. The stride parameter for the VTS model results in 1 proportion of vectors being kept. Finally, P is the pooling method for the single vector model VTP.
Baselines Empirically we found the best KL divergence parameter for VT, VTP and VTS is X/ -1e and using max pooling for VTP. All stride parameters are considered adjust the number of vectors. 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.
NVAE The hyperparameter tuning for NVAE aims to discover models which: neither collapse to 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. Empirically we find the parameters X/: =111 1e- XD = and A = to produce the best trade off between reconstruction accuracy and generative ability. The a parameter able to control the proportion vectors (Appendix B.2 retained and N provides an efficient sampling of the model (Appendix B.3).
Validation results 3 displays the validation reconstruction and generation results across seeds for the best performing parameters. The VT model is able reconstruct well, but high F-PPL score suggest poor fluency of generated text and large variation. The VTP models show that single-vector bottleneck insufficient to reconstruct. Moreover, low F-PPI. and high R-PPI. suggest the model has collapsed to just sampling a few fluent sentences. The VTS models show that some fixed proportions // of vectors retained result in good overall performance. NVAE is able to find models with comparable average proportion / of vectors retained to those hand-coded the VTS models. These NVAE and VTS models have comparable performance with respect reconstruction and generation. However, the NVAE models have notably more variance of metrics across seeds.
Reconstruction
Generation
Model
BLEU (1)
PPL. (})
F-PPL
R-PPL (L)
VT
1.00
99.63+00.00
1.00
+0.00
1.96 =0.99
1.06 00.02
VTS
5=0.
0.50
99.59 #0.01
1.00
#0.00
1.03 =0.01
1.06 +0.01
VTS
0.7.5
0.25
88.92 10.73
3.74
+0.37
1.00
1.08 +0.01
VTS
0
0.20
77.99 +0063
18.78
+1.41
1..00:0000
1.08 +0.01
VTS
509.9
0.10
65.20 #102
66.71
#11.66
1.00
1.09 10.01
VTP
P mar
0,05
46.36 +0.40
1659.24
=70.28
1.01 10.00
10.03
VTP
P mean
0.05*
42.94 +0.49
2425.25
197.65
1.09 10.03
1.35 10,0
VTP
P one a =1
0.05
38.33 11.08
2902.39
445..97
1.00 #0.00
1.29 +003
NVAE
0.50 10.15
98.90 +0.67
1.07
10.07
1.03 +0.02
1,50 +0.65
NVAE
a =0.75
0.40 10.10
98.19
1.15
10.21
1.02 #0.02
1.15 10.10
NVAE
a 005
0.27 10.66
92.78 14.61
2.14
#1.06
1.02 10.00
1.15 #0.09
NVAE
2 =0.4
0.23 #006
33.93-144
29.75
+60.42
1.02 #0.01
1.11 =0.5
NVAE
Q =0.3
0.18 #0.02
67.01 100.30
119.97
+151.33
1.01 #0.01
1.12 #0.09
NVAE
a 0.2
0.13 +0.03
54.02 14467
1145.56
1490082
1.01 #0.01
1.26 +0.27
NVAE
a 001
0.11 0.02
44.46 +6.59
1635.51
=1122.01
1.01 +0.00
1.21 +0.10
Table 3: Results for regularisation and generation validation Wikitext-2 averaged over 5 seeds. The
average proportion of latent vectors retained during evaluation is reported by v. #The VTP models only use a single vector.
Figures 4al 4b and 4 visually display the reconstruction versus generation trade-off for the validation data across seeds. The best models have both low generation perplexities and high BLEU scores whilst dropping a large proportion of vectors. Figure 4e shows that the NVAE model able to dynamically reduce the num- ber of vectors and still reconstruct, comparably to the VTS models. We notice that their exist some NVAE models that have a good reconstruction. generatiennn trade-off but are less clustered than the VTS models. Note that the R-PPI. and F-PPI. plot limits are cropped 1. to focus the higher performing models.
15

---

Published as a conference paper at ICLR 2023
1,000
MIRE
1.00
NYME
100-
1175-
1.755
90-
BLEU (11
70
80
1825-
1,025
Noke
VTS
1000
50
100
1,000
40
20
70
90
100
0.2
0.4
0.8
10
BLEU
11
BLEU 111
(a) Validation BLEU vS F-PPL.
(b) Validation BLEU vs R-PPL.
(c) Validation BLEU vs proportion
retained vectors v.
Figure Reconstruction versus generation trade-off and regularisation for Wikitext-2 validation.
Test results The best seed models from Table are selected for the baselines and NVAE and then evalu- ated the test and shown in Table and subset of this a = (0.75, 0.4, 0.3, 0.2) is plotted in Figure
Reconstruction
Generation
Model
BLEU
PPL (L)
F-PPL (L)
R-PPL (L)
VT
1.00
99.63
1.00
1.37
1.11
VTS
0.5
0.50
99.61
1.00
1.03
1.05
VTS
0.0.75
0.25
89.18
3.72
1.00
1.07
VTS
0
0.20
79.51
15.51
1.00
1.06
VTS
0
0.10
67.04
51.47
1.01
1.08
VTP
P mar
0.05*
48.940
1386.34
1.01
1.17
NVAE
a =1
0.44
99.27
1.04
1.04
1.12
NVAE
A 00.5
0.34
99.15
1.04
1.11
1.04
NVAE
A =05
0.28
96.35
1.33
1.02
1.05
NVAE
A 0044
0.28
95.96
1.41
1.01
1.03
NVAE
a 003
0.19
75.83
17.18
1.01
1.05
NVAE
Q 0.2
0.16
70.60
23.46
1.01
1.08
NVAE
a -.1
0.14
54.35
267.59
1.01
1.11
Table 4: Results for regularisation and generation on Wikitext-2. The average proportion of latent vectors retained during evaluation is reported by v. *The VTP models only use a single vector.
B. No REGULARISATION ABLATION
In this ablation consider the behaviour of the NVAE without any regularisation. T5 that the model is able to ignore the noise and revert to standard Transformer model. We can see this by the model retaining all vectors y= and being able perfectly reconstruct the input.
Reconstruction
Model
BLEU (1)
PPL (L)
T
99.63 =0.00
1.00 10.00
VT
1
99.63 =0.00
1.00 +0.00
NVAE
99.63 =0.00
1.00 +0.00
Table 5: Reconstruction results Wikitext-2 validation data without regularisation.
B.2 CONDITIONAL PRIOR EXPERIMENT
In this experiment we consider the effect of the conditional prior A We select a subset of hyperparameters that managed achieve good reconstruction and generation performance: Xa {le- le -e- (1, 10), ش = {11, 2, 5) each across random seed samples.
16

---

Published as a conference paper at ICLR 2023
100
70-
1.50
0.75
1100
1075
1050
1.81
1025
(a) Validation BLEU proportion of (b) Validation F-PPL, proportion (c) Validation R-PPL vs proportion retained vectors v.
of retained vectors
retained vectors
Figure 5: Conditional prior a controlling the proportion of retained vectors t/.
Figure displays the validation reconstruction and validation generation metrics versus the average proportion of vectors retained. We that the conditional prior hyperparameter a roughly corresponds to the proportion retained vectors v, across all the selected hyperparameter subsets. Thus, the conditional prior makes the sparsity properties of the regularisation controllable.
B.3 NUMBER OF SAMPLES EXPERIMENT
= In this experiment consider the effect of the number of samples = (1,2,5) component. We select
subset of hyperparameters X6= {le -1e-3, le-4), Xo = [1, 10), a (0.1, 0.2, 0.3, 0.4, 0.5, 0.75, 1} each across 5 random seed samples.
1250
1.02
1.05
1330
1325
1.03
1.03
1075-
1030
40
50
60
70
60
60
100
70
300
BLEU
BLEU 11
(a) Validation BLEU vs F-PPL.
(b) Validation BLEU vs R-PPL.
Figure 6: Number of samples x: effect reconstruction and generation trade off.
Figure 6 shows validation generation metrics F-PPL and R-PPL against reconstruction BLEU for different numbers of samples X per component. Increasing the number of samples does not result in significant improvement in the generation or reconstruction ability of the models. Hence, we use a single sample per component because it is more efficient.
B.4 LARGER SCALE MODELS
We conduct a larger scale experiment to test whether NVAE is able to generalise to models with more param- eters and larger datasets such Wikitext- 103. Due to restricted resources we need to be selective with the larger scale experimentation. 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 models of approx- imately 76 million trainable parameters. We use a constant learning rate of le and trained for 11 epochs (x= 150K steps). Each model experiment takes approximatcly 24hrs to which was the largest compute within budget. Otherwise the same training details as Appendix A are used.
single NVIDIA Tesla v100,
We initially trained the NVAE model with the best parameters (based lowest R-PPL score) from Table Thereafter we trained the VTS model with the closest proportion vectors retained, for comparability.
In T we see that larger model capacity allows the VTP model to be more competitive in reconstruction. but it still struggles in generation. The NVAE model able to dynamically regularise the number of vectors and is competitive with all baselines scale with respects to both reconstruction and generation.
17

---

Published as a conference paper at ICLR 2023
Reconstruction
Generation
Model
V
PPL (f)
F-PPL (t)
R-PPL (1)
VT
1.00
99.56
1.00
1.06
1.08
VTP
P P= mar
0.05*
97.42
1.54
2.41
1.01
VTS
5
0.20
99.52
1.00
1.04
1.01
NVAE
a =04
0.15
99.51
1.01
1.01
1.01
Table 6: Large scale model results for regularisation and generation on test Wikitext- 103.
Samples (Wikitext-I)
VI
each lought. cconsidered in per resulting cups.
No =0=1c=
video gumes's 's reprinted le nes 2010 allowing track video gamee browns passengers for the third race. [UNK] he, and trypical confroented were were level prime move criminal discussed color topical liberty camp dated confronted so topical series camph created the located better move topical replacement from confronted disnupted destiny newmarket thrust the nine confrented confreused destiny camp topical topical contrelled future great future camp near rumway game cupacity list to him atbitudes forwards pugeani grand grand, flush bugs forwards at during made
winds. australian forwardsed to strongth on the wall choose him capacity theory game.
VTP
down in their episodes and mccel orginally sequel the simpsons [UNK] the second episode of these episodes [UNK]
vocals responded
=1=1-=2
and is for the only synoeym of children and 22 m.
maz
there pogular other 28 june 29 its radios, protomy rad and field leave his site.
sgnificamlly are the boundaries of music association [UNK], J. love hy hertam voized the w. ed abilities.
VIS
this feature innis light or lands and diamond, is hit by campus meant their buckhist standing and causes and remained
from remained the fun.
majer living stage its chicago and rugby one scottish discovery, germany, of 50 the number athletes of aba best
resoet of the place the its social to tom unglesey the number similar at niss 00, analysts their the 50 new warrioss
8-0.2 0.5
and camlina and his confession
ione liner adapted derby, a old, all peninsula of ion body guitar with the conflicts with groups of two.
NVAE
an growth were substituted reservors in bewau lidden below south rugby to bascball flesh as front 105bery a [UNK]
level take such quest.
X/2 1g==1"1
the dua informs law of the then' called on [UNK] yuko, the minor alert firecast in nixon and compurable kicking
meyer without asin strong im tag, on any diet.
X =1
[UNK], prosecution destroyers believed motes other, collective carrier all dark newell of scientology and
further cutting for
a 0.44.4
stone as the plans the other distinct celehrities forms ever developed of the non combinations 2010 the likeness
temple.
A
Table 7:
The first samples drawn from the best models (lowest generation scores F-PPL and R-PPL) trained on Wikitext-2 dataset.
Samples (Wikitext- 103))
the [UNK], saleseam action mighteno by
Na =1===""
was found on night commission
at al process, an was the castwosdd, henan anniversary, and of, to to years at
example, incomectly, to served
VIP
place charles had largely the that they were any sections the temms of the work during the lease were the staging
because the consiructoon of the staging meant bi, lissed ahead.
Xa 1e 2
verpino sa ya ram, specimen, thatakous, that november, draft the hack and sentenced. april. [UNK], defeated
builey. which retureed on king, jr., had promouned..
maz
undertook the new milne coteutlam by observer that serbia quateemaster wasuka had anticipated the western infantry and the baltic offensive also been lirtered to the french indochina to make the interoepted sastier.
her deposit at es was one on the world, on sola, il lap 12 il dj monitor ph on lop solo zola on the top lap and fifth number.
formula on lap video
filth op score gates source for pre framework is chamber topales documeent countered bencath austrianles document countored
VIS
bencath in configuration source for cast pre justicedley ayeton 2008.
N -1= -2
the the lines secret when exception cleared better when contrary surplus cleared
0=0
lucy even described valid, land ofle bombed along attack aggressive.
fearing webber had 2011 shows that delly the shows that webber had 2011 shows from 186ua racet from 1866m ra
bollywood
NVAE
or the courtncy burns left wayne for manimaze dumaged for nor the line, the triple life the acts in blood a Ssisus 20 even frames [UNK].
N/=
rock tank dovetion now and differing actor and dump turbinestully present pop penctrated fantasy of the
dumnner reuched buil in camera while altomey. detja after
=1=1
denmark pring theodoreka, reveals curope carries and allooddyy final culture and havinginer shared forept and free
extends ground that all lost in whose nurse between state named
=0.4
newport horse said various connor tatiana, founder's is or experienced swanting artist and proposed, where fear alternative.
R
Table 8: The first 4 samples drawn from the large scale models trained on Wikitext- Vikitext-1033 dataset.
18

---

Published as a conference paper at ICLR 2023
C GENERATED SAMPLE EXAMPLES
Tables 7 and [ give examples of text generated by sampling from the prior.
D ALIGNMENT ANALYSIS
This experiment highlights the problem of latent alignment non-NVIB models, and evaluates alignment based on position, which use the interpolation experiments. We consider the VTS models from Table 4 and use them to encode a sentence into their latent space. For each latent component retained by the VTS model, we perturb it with Gaussian noise and consider the resulting autoregressively decoded output. We plot the percentage of the time given position in the output is changed by perturbing given latent component, discarding sentences where the length is changed (only 2% over 100 samples).
0.14
0.12
Denzro
5.3
5.00
0.04
0.02
Core
from
Todans
(a) VTS 0.5
(b) VTS 0.5
Figure 7: Latent vector alignment with generated output
5.05
535
0.68
610
5.01
505
6.00
5.00
core
Tokens
Tokans
(a) VTS S= 0.75
(b) VTS S 0.75
Figure 8: Latent vector alignment with generated output.
Figure 7 shows that the latent space for the VTS S 0.5 model are approximately aligned by location. However Figure 8 for more condensed representations where 0.75, there an unclear alignment of latent vectors to their position. As discussed in Section 5.3 the use of mixture distributions instead of sets vectors allows NVIB representations to avoid this problem of alignment.
INTERPOLATION EXAMPLES
Tables 9 through 12 give examples of the text generated by interpolations the latent space.
19

---

Published as a conference paper at ICLR 2023
S1
0
they were keen to instend move on with the next film, casino royale.
VT
0.25
they were koon to instcad move on with the next film, casino royale.
they were appointed the move over the maritime evacustion himself, landing coloration.
0.75
squadiom was engaged in convoy esour and mariime recornaisancee duties ofl south castern australia.
VIP
0.25
they were koen to instced move on with the next film, casino royale.
0.5
they were keen result instead on dd with the guardian over post 1904.
0.75
squadron was eriageed in convoy cscort and maritime recomsissance dutics off south castern australia. they were keen instead move on with next film, casino myale.
VTS
0.25
0=088
0.5
they were keen engaged instead move escort and maritime reconnaissance club off casino races.
0.75
squadio was engaged in convoy escort and maritame reconnsissance duties off south eastern australia.
NVAE
0.25
they were keen plans deployed with amnual cumnons posstion the cummissiomed
=0.4
0.5
they were koen convyy escert and the caribican offierr off south and operation australia,
0.75
they were keen run in convoy escort and maritime recoendisaance duties off south castern australia.
52
squadron was engaged in convoy escort and maritime recomnalssance duties off south castern australia,
Table 9: Interpolation results by varying 7 using sentences with different number of input tokens.
S1
0
the king furious at the demann and kept the [UNK] envoys wailing for weeks. the king was furious the demand and keps the [UNK] envoys waiting for weekss. the time, working the shot led social [UNK] toes waiting close.
VT
0.25
0,75
this time, the [UNK] king received the imperial eavoys but still refused to submit.
VTP
0.25
the king was furiious the demand and kept the [UNK] envoys waiting for weeks.
max
0.5
the time, the [UNK] saw the source and the [UNK] envoys still refused for colebration.
0.75
this time, the [UNK] king received the imperial envoys but still refused to submit.
VIS
0.25
the king was furous the denunn and keps the |UNK] envoys waiting for weeks.
8 S=0. 0.8
0.5
the time was furious king dermand the during emary [UNK] ability still calling for going.
0.75
this time, the [UNK] king received the imperial enrays but still refused to submit.
NVAE
0.25
the murriage was furious [UNK] king received imperial enveys still refived weeks.
=0.4
time, announce [LIK]] king received the imperial envoys but still refused weeks.
0.75
this time, the [UNK] king received the imperial envays but still refused to twice.
52
this time, the [LNK] king received the imperial envoys but still refused to submit..
Table 10: Interpolation results by varying T using sentences with the same number of input tokens.
S1
0
known damage was caused by the Rood.
VT
0.25
no known damage was coused by the flood.
known species was by the twa.
0.75
the two specios can be distinguished by number of characeristics
VTP
0.25
no known damage was caused by the flood.
max
0.5
the damage was coused by the flood of characteristies
0.75
the two species can be distinguissbed by number of characteristics
VTS
0.25
no known damage was caused by the flood.
0=0.8
0.5
no two damage would distinguished the flood number of characteristies
0.75
the two species can be distinguished by number of chanctoristics
NVAE
0.25
no known these damago be distinguished by number of characteristes
0.44
no the species can be distinguisted by number characteristics twi species can be distinguished by a number of chanateristiic
0.75
82
the two species can he distinguished by a number of characteristics
Table 11: Interpolation results by varying T using sentences with a different number of input tokens.
S1
the palace has audient graffiti and possesses windows,
VT
0.25
the palace has oncient graffiri and possesses low windows.
0.5
the palace modem class and enthusiasticaly and cinema.
0.75
smoke signals a history of native americans cinema.
VTP
0.25
the palace has ancient graffiti and possesses low windows.
mas
soul room ancient and repahlicansium downtown.
0.75
smoke signals a history of native americans in cinema.
VTS
0.25
the palace has ancient graffiti and possesses low windows.
S S=00 0.8
0.5
smoke miners has the graffni native villagers low schools.
0.75
smoke signals a history of uative americans in cinema.
NVAE
0.25
the paiace has ancient economics and larger war butterfly.
=0.4
smokeide historical gills german language in snikkepers is judicial topics of german language cinema. smmoke signals history of native americans in cinema.
0.75
S2
Table 12: Interpolation results by varying using sentences with the same number of input tokens.
F DENOISING ATTENTION FOR THE MEAN OF THE POSTERIOR
As introduced in Section 2.3 during test time evaluation VAEs use the mean of the sampling distribution instead of random sample. The mean of our BFDP posterior distribution over mixture distributions
20

---

Published as a conference paper at ICLR 2023
is its base distribution Gg- However, the base distribution is not discrete distribution, whereas at training time all samples are discrete distributions. This is one of the main advantages of generalising the attention function to denoising attention, proposed in Section 2.1 Denoising attention equally well be applied to the mixture of Gaussians of the base distribution as to the mixture of impulse distributions from sampling,
depicted in Figure Tc
To understand why the continuous base distribution looks to the model like a typical sample of discrete distribution, we need to consider how this representation is being interpreted by the attention function. 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 the mean of the vectors returned for the same query at training time. This is what the base distribution is giving us. In contrast, with standard attention there no finite set of vectors we can use 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.
More precisely, we can consider the latent representation parameterisation of an attention function from query vectors to result vectors. 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 's base distribution Gg and total pseudo-count op. the base distribution G. specifies the mean of a distribution over such query-result functions, and the pseudo-count og 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. 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.
To efficiently compute denoising attention applied the mixture of Gaussians Gi. we take advantage of the fact that the multiplication of two Gaussian distributions is a Gaussian distribution, giving us:
DAttn(u; 6
Σ La ( ((T))) V(vuu
J..
du So E. E, 1. ((o)))))) s(v,u,Vr)) de
gleu(
Va u+
lo
))
Vi 10
Vd الاي)
v du
WL g(0)(
Vd
Tap
))
dv
Vil
Vd
Iq L (V
Va t 1
(11)
va
where all algebraic calculations over vectors of are done componentwise. See Appendix K of this formula which is useful for implementation.
for a version
G EQUIVALENCE OF VECTOR ATTENTION AND DENOISING ATTENTION
In this section we show the equivalence of standard attention, expressed as a sum over vectors z; in Z, and denoising attention, expressed as an integral over a distribution Fz which is only nonzero at the 2,. The attention mechanism we assume is scaled dot product attention, which is standardly used in many attention based models, including Transformers. For simplicity, we consider cross attention, where a single query vector is mapped to single result vector. In this work we view cross attention 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. The decoder then apply this function to as many queries as it likes. 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 E xxxP into keys (ZW*) (wrK)eR" =ndd and values (ZWV), E Rn via weight matrices W WV e Rpxd, respectively, and maps the query u' e Ix into key space (u' (u'Q)) E RIxd via the weight matrix Wº Roxd The key space's 's dimensionality d
21

---

Published as a conference paper at ICLR 2023
is used for scaling. Scaled dot product attention is then defined as: Attention(u', N.W W
(u'
(ZW*) k7
softmax
ZWV
Vd
uz
softmax
ZWV
where u
vd
Z)
where u E: KP is the input query u' projected into Z space. In the last line we rewrite scaled dot product attention in terms of core dot product attention function Attn Z) where all operations are done in the space of Z:
Attn(u.Z)
softmax 1 Juz" vi
expl vi
u1z))
i=
exp(
Vi
uz 7)
Zi
(12)
where 2; is the hth row of Z.
We interpret Z as specifying probability distribution over a vector space, and we interpret the function Attn(u, Z) Bayesian denoising of using this distribution, depicted in Figure 1b The vector is interpreted as an observation of some true vector v I which has been corrupted by Gaussian noise. The true vector v was generated from a prior probability distribution specified by Z. The result of Attn (u is the expected value of this true vector after seeing the observation u, which can be considered the best guess of the true vector given the noisy observation, and thus a form of Bayesian denoising.
To derive this interpretation of attention as Bayesian query denoising, we interpret the set Z of vectors 2; as specifying mixture distribution Fz over vectors v which consists of one impulse distribution 82, at each vector weighted by the softmax over their scaled I.] norms:
exp( 2 | 1
Fz
62.
(13)
exp 2/a |1)
Then we can derive this interpretation by replacing attention's sum over i with an integration over Attn(u, =
exp(
u.z.")) va
n
exp(
á uz ?)
2
exp(
| (2) exp( 2 2: (2) exp( uzz T)
2v
2im1
exp
zz (2)) exp 2vd || 2) exp( Va u.z"" "²
2vd
exp( 2v4
zi {²) 82
exp( 2v8 | (2) exp( 1/a wo0 do
1 exp(
|
exp( 2vα 12) exp( va un? du
2vd
expl
2 11
(v)
Valigeepi exp(
2vd km (un -iw=-
DP
Di=1
Venva
|
En cxpp 2 (
do
expl 2/2 ||2,|12)
Jo
En mn exp( 2 12:113) 82(0)
Verva 2/a E (uk (1k- -vk))) dv
exp(-
fz(v) g(u: /d
dv
(14)
(v) glu; v,vdr du
where fz(-) is the probability density function for distribution Fz. and g(u; vdi) = exp( 2vd (u )²) is the multivariate Gaussian function with diagonal vari- V
ance of vd. The first step just adds terms which don't effect the value. The second step changes some instances of z; into an integral over which is only nonzero when æ=z; (i.e. 6z,). Thereafter, the terms are rearranged such that the formula reduces to an expected value over a with weights proportional to the probability of generating with the distribution F' and generating the query with Gaussian noise (o(oy added to Scaling the variance of the multi-dimensional Gaussian noise by vd reduces the impact of the dimensionality d on the similarity g(u; u, v.Vdd.) between and e.
22

---

Published as a conference paper at ICLR 2023
The above derivation is inspired by the interpretation of softmax Bayesian classification with normally distributed classes (Bishop 1995). and it is similar to the interpretation attention keys as latent mixture distributions provided by guyen al. 2022). However, here the Gaussian represents uncertainty about the observation or query vector instead of uncertainty about the class or key vectors, exploiting the fact incorporate the value
that Gaussian function is symmetric in its argument and mean. This allows part of the attention function into Bayesian denoising interpretation, so we have a Bayesian interpretation of the entire attention function. To the best of knowledge, this interpretation attention is novel.
The function from equation 14 is a special case of the definition of denoising attention DAtta(u;; F) given above in equation 3 where Fz= The construction above indicates that any scaled dot product attention function is an example of the DAttn(u: u; F) function e However, while attention is only defined defined over any probability distribution F over a vector
over sets-of -eeetors Z, denoising- space, not just finite sets of impulse distributions.
H DERIVING THE FACTORISED DIRICHLET PROCESS
In this section derive alternative factorisation of DP which helps with the sampling method in Section (.2 For notational convenience, in this section we use as the number of components for the base distribution instead of n++1, 1. This is still intended include both the output of the encoder and the prior component.
Here we provide the proof that
FDP = DP(G)
where G² (?",,.... G.) a² = defined
(.) G -- -i=l ag G. aj = -i and F=FDP(("),")) is
F = P.F.
~ Dir(af F. DP(G),af for for.=1.....
We start with the definition of a DP infinite symmetric Dirichlet distribution. A Dirichlet process F DG), can be defined as the limit of a sequence of finite Dirichlet distributions (see Teh (2010):
F
TROZZ
k=
010
lim Dir(
Ro
No0
Kg
KD
2/
Go for 1,...,00
Note that the weights π and the vectors 2y are independent of each other, can treat these two issues separately.
For the vectors, we know that after generating an infinite number of 2 from Gp. a proportion of exactly of of them will be generated from Gy. For a finite number of vectors Ki, let K; be the number of z og
generated from G.,, for each So we have
No 000
15g
Given the exchangeability of Dirichlet distributonn,, we can renumber the Ky categories Dir( Ko No, do so that π = (TI1 Ties S.,, Ten.e and the TT...., are all weights for vectors zij generated from component
G for i= 1......; 1,,..,
Zij
'There may be other constructions which could equally well be used to implement Attn(u, ,, Z) in terms of DAttn(u; F), but all that is important here is the existence of one. Equally, we do not intend to claim that all Dtttn(u: F) functions can be implemented in terms of Attn(u, Z).. Indeed, the greater generality of DAttn(u;; F) crucial in this work.
23

---

Published as a conference paper at ICLR 2023
For the weights, again consider the case of finite Rg before taking the limit goes to infinity, using the above indexing where categories ?j are partitioned according to their vector's base distribution component i. We define (e.) to be the vector of total weights p; -ド= Tij for each of these partitions i. By the rule for merging categories in a Dirichlet distribution, we know that these total weights are themselves distributed according to a Dirichlet distribution.
(pr)
Dir(af,...a9)
Now we take advantage of the neutrality property of Dirichlet distributions. states that this vector (Pe) of partition weights and all of the vectors ,a,, Fins) of normalised weights inside each partition are independent. In essence, this means that the only way that the weights inside each partition constrain each other through normalisation, so when normalisation is factored out they become independent. This independence allows to compute the distribution over (T Mi, miss) by simply marginalising out all the other categories. We first merge all the categories outside partition i into single category. whose weight thus 1- Pt- This gives us the marginalised Dirichlet distribution M (1-9:)) Dir( e8 8i og ND 1" KD
og(1 (st ))- Let d(;;; ox) be the probability density function for the distribution Dir(a:): xo
00
15,, 00
Ki
Mim (1 (1- pz);
00
60
K6
((og)
12 80
(1 Pi)
T(ag(1
1 j=
No
j=1
r(ag)
K4
21 HD
p:) 00
Tay
(1-)(1- o
Ps
Ng
Now we can marginalise out the weight of the outside category by integrating over pi.
[(ag))
(1 (1-p)) -20 -1 (p:) 2
Túj
1)
dp:
I
Pt
I(ag)
(1- p:)0000 -
(8
-
Ri
Tej
28 )
dp
T(ag(1
Kg
5=1
d( Til Ny
Tine
Pi Pi
K0
K0
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
B( 08
0d This gives us
Tires
Dir(
P
Ps
Ko
Ko
Now that we have all the individual distributions, we can put them together to get the factorised distribution for the case of finite Ko-
Ttj = pittys for ior. -
p ~ Dir(a?,c,,
2 π Dir
ag
for fori=1.....
Kgg
Noting that lim +00 No* the weights for the factorised Dirichlet distribution.
we can then take the limit as No goes to infinity to get our definition of
Kij
PIT for in......;jjl........
Dir(aq.e..of)
lim Dir
for lori=1,mp 1
+20
Thus the weights of DP can be rewritten as the weights of an equivalent FDP using the above construction.
24

---

Published as a conference paper at ICLR 2023
Putting the vectors and weights together, get the distribution F; over the weighted vectors in each partition
Zij G for i- (iri=
0g Dir( No
o fori=la :uri=1... L....
n lim
for
Ki
Kj
and thus
Fi DP(G, iori=
This concludes our proof that, if DNDD( then:
F. = Ee.
i=
~ Dir(a),,s,a))
F; DP(G), .g) for fori=1,...p
and thus FDP(G",a))= DDPP((G)))
I DERIVING THE KL DIVERGENCE
In this section derive the KL divergence between the prior and the posterior. We also argue that this function is approximately linear in the number of sampled vectors for each component.
To directly compare the posterior with the prior, first reformulate the prior factorised DP with the same form as the posterior. We can do this without changing the distribution specified by the prior, simply by making
copies of the base distribution G and weighting those copies
proportionately to the weights dj of the components of the posterior base distribution. This gives the ag
prior BFD ool ,xA) where GP (G%: nit1 (b) and a² =0
" ai ag
of = (og
nt1
ag
The formulation of both the prior and posterior as bounded factorised DPs of the same form simplifies the computation the KL divergence, because the KL divergence for each respective pair of factors can be computed separately, and then combined.
First consider the factors for the Dirichlet distributions over the partitions for the different components i. There is closed-fome solution to the KL divergence between two Dirichlet distributions.
DKL (Dir(a), nt1 o.ni)
Dir(ag
aj n=1 1::,a oh44
aB
T(a
n+1
T(a
o6 (W(aa))v(aa)
= log r(a
log l'a o/ og
+ a
where the gamma function and the digamma function.
For the bounded DPs for each individual component i, there are two factors, symmetric Dirichlet distribution over the weights and a Gaussian distribution over each vector. For the symmetric Dirichlet distribution, in the case where = then the KL for this term is zero, since there is no choice make for this weight. In the case where 6; 1, the KI. divergence between the posterior and prior versions of these weight distributions again has closed-form solution.
DKL. Dir( por of || Dir(ar of of Ki,ah agee
I(a I(ag
T(2)) af oyب.
+a (
at
log
Rglog
I'aa
This term then needs to be summed across components 15i5n+1.
For the factors for generating vectors from each individual component of the base distribution, because the different components have been factorised, there is also closed-fomm solution computing these KL divergences. Each Gaussian component the posterior's 's base distribution is compared independently the Gaussian of the prior's base distribution. The KI. divergence between two Gaussians (with diagonal
25

---

Published as a conference paper at ICLR 2023
covariance with values σr) is:
d
(0 -1 -1-log( )) (0 (o%))
DKL(G;
Gg)
+
(o))2
(0%)2
hml
(2 + -1-logg(((0),)
where the last step assumes =0, 2 1. This term then needs be multiplied by the number Ka of vectors for this component, and summed across components 15i5++1.
Given these exact dlosed-form solutions for each pair of factors,
compute the full KL divergence.
We start by combining the formulas for the weight factors, where some terms cancel:
DKL (Dir(a], mt! .02+1) || Dir(ag
of nt ag"
nttl.o o0
o og
n+1
DKL Dir
al) ||
og Dir(ag
M agk.)
age,
I(aa)) T(ag)
n+1
log
Σ Kolog
(((())
af OogNs
Now we can put all these pieces together.
DRI ,x,,))) | BFDP(GP,
ag
= [. al of
n+1
n+1 Ki
k. of
I (za) ==1-=
(FD os K.)
log BFDP(G",,OO..r dp d'' de
= DKL =Dki (Dir(aff ntl o Dir(ag of af "g ogg (ag)))
ntl of
a+1
DKL Dir ot p
)) Dir(ay o' of mg agn!. )) +r, tRL 6%)
age.
1=1
lagT(ag) logT(a)")
Σ
Kg logT( og ag antis
) logT log -aag-ab 1888-o8'11 (- >> 1/
((2))
ki
1=1
(p (o
th 12
(0 )2 log
(15)
(o
Equation| 15 gives the KL portion of the loss when we are given the full set x of numbers of vectors N2 generated for each component i.
If we assume that the K; are chosen stochastically, then we can take advantage of the fact that equation 15 is
approximately linear in K,, when the variation in K; fairly small relative to the values of Kg The Gaussian
tenn is exactly linear in Kg, and the terms < o? and K; (logT( af k;
(a logg of ks ))
both approximately
linear in Ki. This allows us to approximate the expectation over K; of this loss as this loss of the expectation over Ni, as discussed in Section 3.1 In this case, this approximation of the full KL divergence is:
(BFDP(G a². .R) BFDP a .x))
logT(ag)) Logg) + of (1(((()
of (55 - tNo logI
logI
8i) No
(16)
No
n+1
(ng -1
(0
log (0 )2
1 +1 Kip
(o
+
(0
-1-log
(o
i=1
h=1
26

---

Published as a conference paper at ICLR 2023
J REPARAMETERISATION TRICK AND SAMPLING
In this section we consider the reparameterisation trick to allow backpropagation through the sampling step. We consider the component Gaussians of the base distribution and the weights generated by the DP separately.
J.1 SAMPLING VECTORS FROM A COMPONENT OF THE BASE DISTRIBUTION
Each vector ZІz is sampled independently from some specific component G of the base distribution Go Since we assume that all these components are distributed according to a Gaussian C [ ((o PP)
we can sample from this distribution using location-scale shifting (Kingma & Welling 2014)
Zi p1 +o
(17)
6K (0,1))
Since the random sampling comes from an unparameterised unit Gaussian, there is no need to backpropagate enror into this sampling step, but the error can be backpropagated into pet and of given a specific sample. This is the reparameterisation trick for Gaussian distributions.
J.2 SAMPLING FROM A DIRICHLET DISTRIBUTION
A Dirichlet distribution over category weights be sampled by sampling from Gamma dis- tribution for each category and then normalising. A sum-normalised set of x random variables Th... (ir.on........) follows Dirichlet distribution if the unnormalised random variables 73 each follow a Gamma distribution.
7
(18)
EN
7/
T(a,,111)
where ((c,,,==1) is the Gamma distribution with 8=11, whose PDF f(z) expp(aaa
r). There no closed-form explicit reparameterisation trick for the Gamma distribution, but there are for approximations. We propose to use combination of two approximations for the Gamma distribution which have a reparameterisstioo trick, one for small values of ag and one for larger values of ax Fl
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. Hence, (Knowles) (015)) suggests sampling using approximation to inverse CDF the Gamma distribution the tollowing form:
%
9 (ua,aaaaaa
(19)
U(0.1).
This approximation allows the inverse CDF of the Gamma distribution to be a function the parameters and independent noise from a uniform distribution U(0, 1). However, this approximation is only recommended when the value of ag < 1 and 3= In our case 3 3= but we sometimes have large ay
Gaussian approximation to Gamma distributions Knowles (015)) further mentions that the Gaussian distribution can used to approximate a Gamma distribution for larger a. |Bahlleyan et al. (2018) uses a similar approach to approximate variational attention weights. The Gaussian distribution is symmetric distribution which can be sampled by location-scale shifting, as discussed above. The Gamma distribution
I' I(a,3=1) can be approximated with Gaussian of the form >~aa(a (Knowles) 2015). which gives us:
200
(0,1))
The Gaussian distribution is symmetric and can take negative values. Hence, this approximation is inappropriate for the Gamma distribution unless the a; parameter is sufficiently large, otherwise the sample will need to be truncated to a value greater than zero.
'We leave the investigation of implicit reparameterisation gradients FFguurnvv et al. 2018 to future work. This approach is alternative to explicit reparameterisation for cases like Gamma distributions, but first we investigate the approach which more standard in the VAE literature.
27

---

Published as a conference paper at ICLR 2023
The combined reparameterisation of Gamma distributions To visualise the for these two approximations, their average L1 distance from the true Gamma inverse CDF is plotted in Figure 9 for different values of a. The plot shows that the approximation emror is equal when a = 0.6363. To take advan- tage of the strengths of both these approximations, we propose to reparameterise the Gamma distribution blend of these two approximations with switch = -..363 and truncate negative samples to
PORTE
220
Figure 9: The average absolute difference between a Gamma inverse CDF function and our approximations:
inverse CDF approximation and Gaussian inverse CDF are plotted for values of
K PRACTICAL IMPLEMENTATION OF DENOISING ATTENTION
this section we provide the equations used to allow denoising attention to implemented in a deep learning framework training time and test time.
Denoising attention at training time During training, the set of vectors E R=xp and their log-probability values log(T)) e 1xx are both sampled and output by the NVIB layer, thereby specifying the sampled mixture distribution F. For each use of denoising attention, the query u' 1XPP is projected (1=
by the grouped matrices W W Rexd for scaling. Denoising attention can then be computed as:
The keys' dimensionality is used
DAttn(u;F) = softmax (224 Lzz +log(T) 2v7 1212)
We define this for u uER !xP, but this can easily be extended to multiple queries.
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 Gj- The NVIB layer takes its input from the encoder and maps it to the parameters (nL" ,r a of this base distribution Gj = Σ a" (("))) Nu, For convenience let DAttn(u;;G)
Test-time denoising attention can then be computed as:
(
T
a
4
T 1 (log (1)((()))))
(00g 2 (1
Vd (orje o0²
= softmax u
log
(or)²
where 1, is a row vector of p ones.
A caveat of this derivation is that it applies only for single-head attention and is not trivial to extend for multi-head attention. We leave this future work.
28


Reviews:
[{"review_id": "h6i3LoXYVhr", "reviewer": "Reviewer_w4Qq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a probabilistic mixture model representation of the attention mechanism in transformers. The probabilistic representation  allows the authors to define priors and posteriors over the mixture distributions and ultimately redefine the transformer as a nonparametric variational autoencoder. They show experimental results for the reconstruction of input sentences, generation by sampling from the prior, regularize the size of the latent space, and interpolate.", "summary_of_the_review": "The paper is promising although it could be reworked to be stronger with a more clear focus and motivation. ", "strength_and_weaknesses": "Strengths:\n- The paper seems interesting and a rather novel concept. Their approach is sound and the results seem ok. There are no glaring flaws in the paper. \n\nWeaknesses:\n- The point / takeaways / motivations are not very clear. Is the main benefit the regularization ability and the computational savings of having a lower dimensional latent space?\n- While transfomer based models e.g. BERT / GPT-2 have been shown to be very useful, the paper is showing results on a a two layer Transformer encoder and decoder with a single attention-head. While this is the original proposed structure for using attention, it does not mimic either the encoder stack models e.g. BERT or the decoder stack models e.g. GPT-2 that are now the SOTA.\n- Furthermore, it is not clear whether learning a NVAE for each attention mechanism (BERT would have 12*12=144 head) is computationally feasible and whether this training process is significantly more expensive than the standard learning of the query, key, and value weights. \n- In the experimental results, the NVAE model is compared against VTP and VT (which don't seem to have any hyperparameters as only 1 best model was chosen?) and a hand-coded solution VTS. The performance of NVAE and VTS seems relatively comparable, but is VTS supposed to be an \"oracle\" or is there any further description on what \"hand coded\" means?\n- In tables 7 and 8, all generated samples seem equally senseless? Similarly the interpolation examples (tables 9-12) seem to show all models are comparable? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "gjmllzZVrIq", "reviewer": "Reviewer_tsdp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this manuscript authors embedding space of Transformer encoders as mixture distributions and as such are able to formulate transformer encoder - decoder as a VAE model. Number of mixture distributions is variable to authors decide to use Bayesian nonparametrics to solve the modeling issue, namely they use bounded Dirichlet prior. ", "summary_of_the_review": "All in all, I find the paper quite interesting and potentially very useful. I think the way how transformer is reformulated as a VAE model is quite neat. More large-scale experiments obviously would make the paper much better as authors also state in the Conclusions. ", "strength_and_weaknesses": "Strengths: \n- Novel probabilistic formulation of attention model in transformers. \n- Full derivation is shown. \n- Model has potential usefulness beyond the experiments shown in the paper. \n\nWeaknesses: \n- Implicit reparametrization, that provides exact reparametrization for Dirichlet is not used ( https://arxiv.org/abs/1805.08498) \n- Manuscript seems to be typeset in a hurry, for example p(F|x) and p(F) are both named as prior. AFAIK, prior should not have data-term conditioning. Please clarify this point. \n- Even though mathematical development is convincing, authors have not motivated their work. They should clearly explain in Introduction that for what reason this development was done. It cannot be that just this formulation has never been done before!", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RSuylogLQx", "reviewer": "Reviewer_fw7U", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a VAE framework for that formulates attentions in transformers as a mixture of gaussians, where variational inference can be used to learn the framework. The proposed VAE with transformers as the encoder and decoder is used in text reconstruction and generation tasks. ", "summary_of_the_review": "The idea of the paper is interesting but I feel the technical depth is not enougth.", "strength_and_weaknesses": "Strength:\n\n1. The (re) formulation of attention to a gaussian mixture and the application of variational inference are interesting and intuitive.\n\n2. The paper provides comprehensive experiments on parameter sensitivity and ablation study.\n\nFeedbacks (not all about weakness)\n\n1. Although the formulation of a gaussian mixture of attention is interesting and contributes to the significance of the paper, I think the technical depth is not large enough for an accept yet. Using variational inference or VAE with gaussian mixtures is a well studied technique, which is directly applied given the formulation.\n\n2. I don't see it is necessary to highlight Bayesian nonparametric or Dirichlet processes. First of all, DPs are usually applied into latent variables so that the dimensions of them can grow with data. In the case of this paper, the dimension of Z is the batch size of data, which does not usually grow in training or testing. Secondly, although discussing Bayesian nonparametric or Dirichlet processes, the paper ends up with a vanilla Dirichlet gaussian mixture model, which does not have strong connections with Dirichlet processes. It is not even a truncated approximation (e.g. stick breaking process) of DP. The paper names this unbounded DP and claims it as an approximation of DP, which requires theoretical justification.\n\n3. In the experiments, the paper mainly compares with the baselines of different configurations of VAE with transformers. I think it is expected to put the proposed method in the category of text generation methods and compare it with other kinds of methods in the same category. This is to show how the method is positioned in the task.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "L2S46RkpbUC", "reviewer": "Reviewer_Lu5K", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper mainly proposed the nonparametric variational autoencoder (NVAE) by 'combining' VAEs and the transformer.\nFurther, the transformer encoder and decoder are used in the VAEs framework.\nDifferent from VAEs whose prior is the standard Gaussian distribution, the proposed NVAE applied the nonparametric variational information bottleneck (NVIB) regulariser for the latent embedding.\nFinally, the experiment shows NVAE could do reconstruction, generation, and regularisation tasks. The interpolation further explored the meaningfulness of the learned representation.", "summary_of_the_review": "As above, my main concern is about the experiment, and the second is the motivation.", "strength_and_weaknesses": "My biggest concern is in the experiment part. \n\nFirst, maybe I am wrong, but in my opinion, this work may be compared with SOTA works to support their claim. I am not convinced at least at the current reviewing stage. I am not convinced why we should choose NVAE rather than the original VAE. Because, the VAE could also do the reconstruction, generation, interpolation, etc. It is clear to me the VAE is more complex compared with VAE, with more inductive bias, but it is not clear to me NVAE whether NVAE could reach a significant improvement in performance or could do something VAE couldn't.\n\nSecond, the prior is a mixture distribution, could different distributions learn different aspects of the input object? e.g. a disentangled latent representation? Some ablation studies or visualizations may be needed. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster