Augmenting Self-attention with Persistent Memory
Authors: Sainbayar Sukhbaatar
Edouard Grave
Guillaume Lample
Herve Jegou
Armand Joulin
Venue:  (2020)
Keywords: transformer
language modeling
self-attention

Abstract: Transformer networks have lead to important progress in language modeling and machine translation. These models include two consecutive modules, a feed-forward layer and a self-attention layer. The latter allows the network to capture long term dependencies and are often regarded as the key ingredient in the success of Transformers. Building upon this intuition, we propose a new model that solely consists of attention layers. More precisely, we augment the self-attention layers with persistent memory vectors that play a similar role as the feed-forward layer. Thanks to these vectors, we can remove the feed-forward layer without degrading the performance of a transformer. Our evaluation shows the benefits brought by our model on standard character and word level language modeling benchmarks.

Under review as a conference paper at ICLR 2020
AUGMENTING SELF-ATTENTION WITH PERSISTENT MEMORY
Anonymous authors
Paper under double-blind review
ABSTRACT
Transformer networks have lead to important progress in language modeling and machine translation. These models include two consecutive modules, a feed- forward layer and a self-attention layer. The latter allows the network to capture long term dependencies and are often regarded as the key ingredient in the success of Trannsfomers. Building upon this intuition, we propose new model that solely consists of attention layers. More precisely, we augment the self-attention layers with persistent memory vectors that play a similar role as the feed-forwart layer. Thanks to these vectors, we can remove the feed-forward layer without degrading the performance of a transformer. Our evaluation shows the benefits brought by our model on standard character and word level language modeling benchmarks.
INTRODUCTION
Transformer networks Vaswani et al. 2017) are sequence models that rely on the attention mecha- nism (Bahdanau et al. 2015) to capture long term dependencies. Since their introduction in the context of machine translation, they have been applied many natural language processing tasks, such as language modeling (AI-Rfou et al. 2019) or sentence representation (Devlin et al. 2019). On most of them, they are now surpassing the former wate-ooftheean art models based on recurrent (Hochreiter Schmidhuber 1997) or convolutional networks (Dauphin et al.) 2017). At their core, transformers use self-attention layer that forms a representation of the current input by gathering the most relevant information from its context. This layer is repeated along the network depth, allowing for information to flow for long distances and to form rich sequence representations. The self-attention mechanism is often considered as the key component of their success and many have worked on improving
transformers by increasing the size the context captured by those layers (Wu et al. 2019 Dai et al.
2019 Sukhbaatar et al., 2019).
However, self-attention layers are not the only component of transformer networks and they do not explain the effectiveness of transformers by themselves. Each of these layers is followed by a feedforward layer. These feedforward layers contain most of the parameters of the model. This suggests that their role is probably as important as the self-attention mechanism. In fact, the transformer layer, i.e., the sequence of self-attention and feedforward sublayers, should be regarded as a single mechanism that gathers information from the context and transforms it into rich representation. Having such two different layer types of at the core makes Transformer models harder to analyse and understand. particular, there are not many works exploring the properties of feedforward layers.
In this work, we simplify the transformer architecture by revisiting its mechanism, while keeping its properties. We introduce new layer that merges the self-attention and feedforward sublayers into single unified attention layer, as illustrated in Figure D As opposed to the two-step mechanism of the transformer layer, it directly builds its representation from the context and persistent memory block without going through feedforward transformation. The additional persistent memory block stores, in the form of key-value vectors, information that does not depend on the context. In terms of parameters, these persistent key-value vectors replace the feedforward sublayer. This modification dramatically simplifies the structure of the network with no loss of performance.
We evaluate the resulting architecture on standard word level and character level language modeling benchmarks and report performances that are competitive with transformers.

---

Under review as a conference paper at ICLR 2020
Self-attention
Feedforward
All-attention
value
value
context
attention
ReLU
context
attention
key
key
query
query
Add-norm
Add-norm
Add-norm
Figure 1: On the left panel, the standard transformer layer is composed of self-attention sublayer followed by a feedforward sublayer. On the right panel, our all-attention layer merges the weights of the feedforward sublayer with the self-attention sublayer. We represent both models in the case of single head, but in the general case, both the self-attention sublayer and our all-attention layers have multiple heads.
RELATED WORK
Neural language modeling. Different network architectures have been proposed for language modeling, such as feed-forward networks (Bengio et al. (2003a), recurrent networks (Mikolov et al.) 2010), gated convolutional networks (Dauphin et al.] (017) and transformer networks (Vaswani et al. 2017). Of particular interest, Al-Rfou et al. (2019) apply deep transformers to character level language modeling. Dai et al. (2019) introduces a caching mechanism, relying on the relative position embeddings from Shaw et al. (011)) which makes inference in these models much more efficient for unbounded sequences. More recently, Sukhbaatar et al. (209) add a learnable self-attention span to extend the size of the context.
Word level language models deal with large vocabularies and computing the most probable word computationally demanding. Solutions are to either replace the softmax loss with approxima- tion (Goodman] [200]] Morin & Bengio 2005), to sample from the vocabulary during training (Bengio) al. [2003b] Jozefowiez et 2016) or include subword units (Sennrich al., 2016). simple yet effective solution is to replace the loss by hierarchical softmax designed to better take advantage of the GPU specificities (Grave al. 2017a).
Finally, many works focus on the regularization of large language models. In particular, Zaremba et al. (2014) show that dropout (Srivastava et al., (014)) is effective for recurrent networks. More recently, Press & Wolf (2017) show that tying the embedding and classifier weights significantly improves generalization. [Bevskki & Auli (2019) further show that combining this regularization technique with the adaptive softmax of (Grave et al. 2017a) reduces the memory footprint of transformer while improving its performance.
Attention based models. The attention mechanism was first introduced in the context of mixture of experts by Jordan & Jacobs (1994). It is only recently that Bahdanau et al. (2013) have shown their potential when used in neural networks in the context of machine translation. Since then, this mechanism is commonly incorporated within many models, with applications in natural language processing and computer vision, besides transformers. (Sukhbattar et al. (2015) apply the attention mechanism on the same sequence, i.e., the so-called self-attention, in an auto-regressive model called end-to-end memory network. They show their potential in the context of language modeling. Graves) al. (2014) use the attention mechanism for reading from and writing to internal memory for solving algorithmic tasks. |Vinyals et al.] (2015) combine this self-attention mechanism with a recurrent network to solve simple algorithmic problems. Later, Merity et al.) (2017) show that these networks can be used as language models if combined with a cache mechanism (Grave et al., 20176). The attention mechanism has been also applied to question answering (Miller et al. 2016) and image
2

---

Under review as a conference paper at ICLR 2020
captioning (Xu et al. 2015). Finally, Shazeer et al. (2017) uses the attention mechanism a mixture of experts in a recurrent network.
TRANSFORMER LAYER
A transformer model is made of stack of identical layers, called transformer layers. Each layer is composed of a multi-head self-attention sublayer followed by a feedforward sublayer. Each sublayer is also followed by an add-norm operation, i.e., skip-connection (He et al., (0166), and layer normalization (Lei Ba et al. (0066). In this section, we review the structure of the transformer layer and refer the reader to Vaswani al. (2017) for additional details of the overall model.
Multi-head self-attention sublayer. A core mechanism a transformer network is the multi-head self-attention layer, which consists of multiple attention heads applied in parallel. Each attention head applies the attention mechanism of Bahdanau et al. (2013) on an input sequence of vectors. More formally, given sequence X1, ..... XT of d-dimensional input vectors, each head applies two linear transformations to these vectors to form the key and value vectors:
k; V2 = W.
WEXt;
(2) where We and W. x the "key" and "value" matrices of size da x d, where d, d/H is the dimension a head and H is the number of heads. The key vectors are then used to compute a similarity score between an element of the input sequence and all the elements of its context C,. The context can be, for instance, the elements of the sequence that precede in the case of language modeling, or the whole sequence in the encoder for machine translation, The similarity score between t and an element of its context C₁ is defined as
Stc x W 9 (k. ct),
where W, q E Rddxxd is the "query" matrix, and p(t, e) is a position encoding function. There are several ways to encode positions: fixed absolute (Vaswani et al. 2017). learned absolute (Al-Rfou
et al. 2019), and learned relative (Sukhbaatar et al. 2015 Shaw et al. 2018). The relative position encoding function improves the efficiency for unbounded sequences, making them useful for language modeling (Dai et al., 2019). In this paper, we thus use the relative position encoding defined as c(l, M,-es where n1 position embeddings learned during training. The head then outputs a vector y, by taking the average of the context representations weighted by attention weights Oze obtained by applying a softmax function the similarity scores:
exp (ste/ Vdn)
y: Σ Cveclve+ + p(t, e)) and atc Σ exp (Sts/ Vdh) cEC.
IEC;
Note that one can use different position encoding functions for the key and value sides. Finally, the outputs from the different heads concatenated for each timestep and multiplied by the d d "output" matrix W.. The final output of this sublayer is thus a sequence of I'' vectors of dimension d.
Feedforward sublayer. The second element of a transformer layer is fully connected feedforward layer. This sublayer is applied to each position / in the input sequence independently. and consists of two affine transformations with pointwise non-linear function in between:
(5) where o (I) max(0, æ) is the ReLU activation function; V and U are matrices of dimension d x dj and dy x d respectively; b and e are the bias terms. Typically, d is set to be 4 times larger than d.
FF(x;) = U a (Vx, + bb + c,
Add-norm. Both the multi-head self-attention and the feed-forward layer are followed by add- norm operation. This transformation is simply a residual connection (He et al. 2016) followed by layer normalization (Lei Ba et al. 2016). The layer normalization computes the average and standard deviation of the output activations of a given sublayer and normalizes them accordingly. This guarantees that the input y, of the following sublayer is well conditioned, i.e., that y 0 and y? yTy. Vd. More precisely, the AddNo rm operation is defined as:
AddNorm(X xx) LayerNorm(X, + Sublayer (x:)), where Subl Sublayer ayer is either multi-head self-attention feedforward sublayer.
(6)
3

---

Under review as a conference paper at ICLR 2020
Transformer layer. The overall transformer layer has the following set of equations:
= AddNorm(Multinn Heeett())))))
(7)
Y: AddNorm(FF((z))
(8)
where Mult Head is the multi-head self-attention sublayer. This is shown on the left panel of Fig.
OUR APPROACH
In this section, we first show that feedforward sublayer can be viewed attention layer. Then, take advantage of this interpretation of feedforward model to concatenate with the self-attention layer, forming novel layer that relies solely multi head attention layer without the need for feedforward sublayer.
4.1 FEEDFORWARD SUBLAYER AS AN ATTENTION LAYER
We transform the feedforward sublayer into an attention layer by replacing the ReLU non-linear function in Eq. by So Softmax function and removing the biases:
Yt USoftmax (otetmxxxxxxx))) (Vx)
Σ as U.,--
Here we use notations U.,; and Vi,, to denote column and row vectors respectively. The activation au is thus the attention weight computed with Vi,* and X,. The vectors Xt, V,,- and U-,, are equivalent the query, key and value vectors respectively. The Eq. 9 is also equivalent to the self-attention sublayer of Eq. with the context vectors k,, set to zero and the vectors V.. and U.,; are used key and value side position embeddings respectively. This allows for a similar implementation for the feedforward and the self-attention sublayers, and opens the possibility of merging them into a single layer.
4.2 PERSISTENT MEMORY AUGMENTED SELF-ATTENTION LAYER
Here we propose a single attention layer that can replace both self-attention and feedforward layers in Transformers, which we call all-attention layer. Our layer applies the attention mechanism simultaneously on the sequence of input vectors, in the standard self-attention layer, and on set of vectors not conditioned on the input. These vectors are added to capture information that does not depend on the immediate context, like general knowledge about the task. They are shared across the data and, in some sense, forms persistent memory similar to the feedforward layer. Therefore we call them persistent vectors. More precisely,
persistent vectors are a set N pairs of key-value vectors, respectively stacked in two dn N dimensional matrices M; and M.. As discussed in Section M2 and M, can be interpreted as and of a feedforward sublayer.
These persistent vectors are simply added to the pool of key and value vectors conditioned on the input:
Concat ([Wxx),.. WAXT], MR), [V1. VT+Nl Concat ([W, XI,.... W.xy M,).
k kT+N]
(10)
(11)
Let us denote by C++ the concatenation of the context C and the indices corresponding to the N persistent vectors. The similarity score between an element of the input sequence and an element c of its extended context Of+ is computed the same way as in Eq. (3) i.e.:
Stc x W (k. p(t,c),,
(12)
where the position encoding corresponding to persistent vector is equal to zero. The all-attention then outputs a vector yt with the same attention function as in Eq. i.c..
exp (Stc/ Vdh)
y: Σ avelve + p((,c)) and aect
(13)
Σ exp (Sts/ (ds/vdd))
Eect
4

---

Under review as a conference paper at ICLR 2020
As with self-attention sublayer, all-attention layer can have multiple heads, where outputs from the different heads are concatenated for each timestep t and multiplied Wo. Note that persistent vectors not shared between heads. Our overall layer is then simply this new MultiHeedAAlAttt sublayer followed by the AddNorm operation as defined in Eq. (6) i.e.
Yt AddNorm wwa.ttHHeaAAAAAte
(14)
The right panel of Fig. D summariie the all-attention layer the case of single head: we remove the feedforward sublayer and add unconditioned persistent vectors to the self-attention sublayer. While the persistent vectors are directly comparable to a feedforward sublayer in the case of a single head, a multi-head version is more comparable to multiple small feedforward layers working in parallel. If there are as many persistent vectors as the ReLU units, an all- al-tteeetiio layer has the same number of parameters the standard transformer layer regardless of the number of heads (ignoring bias terms).
Note that using attention mechanism to address unconditioned persistent vectors has been previously proposed in the context of question answering with knowledge bases (Miller et al. 2016
4.3 LANGUAGE MODELING
Language modeling is the problem of assigning a probability to sequence of tokens (wi,...; Wr):
P(w1...... WT) II P(w, wi-1,...,,, w1).
In this paper, we focus on tokens that are either words or characters. Language modeling has been dominated by neural networks with models either based on feedforward networks (Bengio et al. 2003a) or recurrent networks (Mikolov et al. al.]|20110 (010)) Recently auto-regressive versions transformers have been achieving the best performance on standard benchmarks (AI-Rfou et al. 2019: Dai et al., 2019] Baevski & Auli 2019). In this section, we describe several specificities of these models that we borrow to make our model work on language modeling, especially with a large vocabulary and
long context.
Relative position embeddings and caching. The relative position embeddings are learnable vec- tors u; that are encoding the relative positions in the sequence by setting p(t, c) = u,-c in Eq. 3 They replace the fixed absolute position embeddings of the original transformer to allow these models to work on unbounded sequences. When the input sequence is processed in small blocks for efficiency, caching mechanism (Dai et al. 2019) is necessary to ensure that every token t has the same context length regardless of its position in the block.
Adaptive context size. In adaptive attention span (Sukhbaatar al. 2019), each attention head separately learns its context size from data. This allows few heads to have a very long attention span, while others focus only on recent past. As result, it becomes possible to extend the maximum attention span without increasing memory footprint and computation time significantly. The method works by multiplying the attention weights in Eq. 4 by soft-masking function m, (t r) that maps values to [0, 1]. The real parameter: E [0, T' controls how much of the attention stays the same, and it is learned together with the rest of the model. Since our attention weights in Eq. [3 contain additional values corresponding to the persistent vectors, we simply pad the masking function with 1 on the locations corresponding to those persistent vectors. This ensures that we only adapt the context size, while the persistent vectors are always included in the attention.
Adaptive input and output. In word level language modeling, the size of the vocabulary is very large, making the use of a softmax loss function prohibitive both in terms of running time and memory footprint. A standard solution to circumvent this issue is to replace the full softmax function by the adaptive softmax Grave et al. (2017a). The idea of the adaptive softmax is to split the vocabulary into disjoint clusters and compare words only within the same cluster. The clusters Vi,..., Vx are formed by partitioning the vocabulary V by following word frequency. The most frequent words are in the first cluster V while the least frequent ones are in the last cluster. The size of each cluster is picked to minimize the overall running time, leading to small clusters of frequent words and large clusters of infrequent words. Finally, they further reduce the rumning time and the memory footprint by adapting the capacity of the classifiers according to their cluster assignment: The words in the k-th cluster have classifier that is 4 smaller than the one in the first cluster. The underlying motivation
5

---

Under review as a conference paper at ICLR 2020
that infreeueent words are hard to predict and there is thus need to many parameters for them. The memory footprint of the model is further reduced by tying up the embedding weights with the classifier weights (Inan et al. [2017 Press & Wolf 2017). In the case of the adaptive softmax, this leads to a special form of embeddings called adaptive input (Baevski & Auli 2019).
is
5
EXPERIMENTS
5.1 EXPERIMENTAL SETUP
In this section, we describe hyperparameters choices, our optimization scheme as well as the details of the datasets we consider.
Implementation details. We initialize token and position embeddings from (0, 1), and the matrices W Wg...ooo from u(-vd, vd). The position embeddings are shared accross all the heads. Persistent vectors are reparameterized by k; v vdikk, and v. = vNv'. where the parameters kj and v are initialized from (0, 1/d,) and (0, 1/N) respectively. This way the persistent vectors have the same unit variance as the context vectors initially, while the underlying parameters k/ and v are initialized similar the weights of feedforward sublayer.
For character level language modeling, we set the model dimension to d = 512, and the number of heads to 8. Our small (large) models have (36) all-attention layers, = 1022 (2048) persistent vectors and dropout rate of 0.3 (0.4) applied to attention weights. The adaptive span has the same hyperparameters as Sukhbaatar et al. (2019) with a maximum span of 8192, except the loss coefficient is set to 10-?. We use Adagrad (Duchi et al. (011)) 1 with a learning rate of 0.07. We clip individual gradients with a norm larger than 0.03 Pascaan et al. 2013). We warmup the learning rate linearly for 32k timesteps (Vaswani et al., 2017). A training batch consists of 64 samples, each with 512 consecutive tokens. When the loss on validation stops decreasing, we divide the learning rate by 10 for additional 20-30k steps. Training large models takes about day 64 V100 GPUs.
For word level language modeling, we use a model with = 512 and 36 layers, each with 8 heads and 2048 persistent vectors. We use Adam with learning rate of 0.00025 and 8k warmup steps. The whole gradient norm is clipped at 1. A batch consists of 64 samples, each with 256 tokens. We use an adaptive span of 2048 with a loss of 10-7 The dropout rate set to 0.3 for attention weights, and 0.1 for input embeddings and the final representation.
Datasets and metries. For character level language modeling, we consider the enwik8 and t datasets from Mahoney (20TT). Both datasets have training set of 100M tokens and vocabulary of 28 and 205 unique characters respectively (ineluding the end-of-sentence token). Both datasets are made of Wikipedia articles split at the character level. The text 8 dataset is preprocessed by lowering casing and retaining only whitespaces and the letters that are in the ISO basic Latin alphabet. We report bit per character (bpc) dev and test sets.
For word level language modeling, we consider the Wi k wix.eet-il 0.3 dataset introduced by Merity (et al. (0177). The training set of Wik. Text- 103 contains around 100M tokens and a vocabulary of about 260k words. Each word in the vocabulary appears at least 3 times in the training data. The dataset is made of Wikipedia articles. We report perplexity (ppl) on the dev and test sets.
Dataset specific implementation details. Following Baevski & Auli (2019) we use tied adaptive softmax and adaptive input with 3 clusters of size 20k, 40k and 200k. The dimensions of the classifiers in each cluster are consecutively divided by 4, leading to the following dimensions d, d/4 and d/16.
1ikiiext-103.
5.2
MAIN RESULTS
We compare our approach to the state of the art on several standard benchmarks on both word level and character level language modeling.
Character level language modeling. In Table I we report the results on enwik8. Our small model outperforms all other models of similar sizes. Our large model matches the otate-of-the-ar performance with significantly fewer parameters. On text 8, our small model also matches the best performing model from (Sukhbattar et al. (2019) shown in Table 2 Our large model is 0.01 bpc

---

Under review as a conference paper at ICLR 2020
Table 1: Comparison with the state of the character level language modeling enwi k8. We report bpc for the test set as well as the number of parameters.
Model
#Params
test bpc
Small models
Ha al. (017)
LN HyperNetworks
27M
1.34
Chung et al. (2017)
LN HM-LSTM
35M
1.32
Zilly
2017
Recurrent highway networks
46M
1.27
Mujika ct al.
2011
Large FS-LSTM-4
47M
1.25
Krause et
(2017
Large mLSTM
46M
1.24
AI-Rtou et al.
(2019)
T12
44M
1.11
Dai al. (2019)
Transformer-Xi
41M
1.06
Sukhbaatar et al.
(019))
Transformer + adaptive span
39M
1.02
All-attention network adaptive span
39M
1.01
Large models
AI-Rfou al.
(201
T64
235M
1.06
Dai et al.
(2019)
Transformer-X 181
88M
1.03
Dai
(2019)
Transformer-XL 241
277M
0.99
Child el al. (201
Sparse Transformer (fixed)
95M
0.99
Sukhbaatar et al.
2019 Transformer + adaptive span
209M
0.98
All-attention network + adaptive span
114M
0.98
Table 2: Comparison with the state of the art character level language modeling text 8. We
report bpc for the dev and test sets as well as the number parameters.
Model
#Params
dev bpc
test bpc
Small models
Chung et al. (2017)
LN HM-LSTM
35M
1.29
Zilly et al. (2017)
Recurrent highway networks
45M
1.27
Krause et (2017
Large mLSTM
45M
1.27
AI-Riou et al. 2019
T12
44M
1.18
Sukhbaatar al.
2019
Transformer adaptive span
38M
1.05
1.11
All-attention network
adaptive span
38M
1.05
1.11
Large models
Al-Rfou al.
2019 T64
235M
1.06
1.13
Dai et al. (2019)
Transformer-XX XL
277M
1.08
Sukhbaatar et al.
(2019) Transformer + adaptive span
209M
1.01
1.07
ranstorrmr adaptie span
116M
1.02
1.08
All-attention network + adaptive span
114M
1.02
1.08
below the state-of-the-art but it matches the performance of "Transformer adaptive span" baseline I of similar size.
Word level language modeling. In Table 3 we compare the all-attention network with the state of the art among small models the Wiki Wixt 103 dataset. Our network is 3.4 ppl better than the previous best, which was Transformer-2 XL of a comparable size For completeness, we also report the state of the obtained with larger models, that is about 2 perplexity points better than us. In Appendix A we show sample attention maps from our model.
5.3 ABLATION STUDY
In this section, we compare different variations of our large model on character level language modeling on 8. 8. First. we vary the number of persistent vectors N in each layer as shown in
'This baseline has 22 layers, 512 hidden units and 4096 ReLU units. "ee Appendix B for Transformer baselines trained using the code and settings model.
7

---

Under review as a conference paper at ICLR 2020
Table 3: Comparison with the state of the art word level language modeling Wiki Text-103
We report perplexity (ppl) for the dev and test sets as well as the number of parameters.
Model
#Params dev ppl
test ppl
Small models
Grave et al. (2017b) LSTM
48.7
Bai et al. (2018) TCN
45.2
Dauphin al. (2017) GCNN-8
44.9
Grave ct al.
Merity (2017b)
- LSLMSSMM-- + Neural cache
40,8
(2018)
4-layer QRNN
151M
32.0
33.0
Rae et al.
(2018)
LSTM Hebbian Cache
29.7
29.9
Dai et al.
2019
Transformer-X XL Standard
151M
23.1
24.0
All-attention network + adaptive span
133M
19.7
20.6
Best published result with large model (Dai et al.
2019
257M
17.7
18.3
1.100
1.04
Dev. (bpc)
1.075
Dev. (bpe)
1.03
1.050
1.02
1.025
1.01
1024
2048
3072
FF
single
head
attn
all
The number of persistent vectors (N)
attn
head
split
split
The performance our large model on Text we vary (left) the number of persistent vectors, or (right) the way how persistent vectors integrate with self-attention.
Figure 2:
Figure 2 left). The result shows that persistent vectors are crucial for performance, already reaching good performance at N = 1024. A model without persistent vectors (i.e. N 0) is equivalent to transformer model without feedforward sublayers, and it performs poorly. This also demonstrates the importance of feedforward layers in transformer models. However, it maintains decent performances because still has lot of parameters (38M) in the W W.,..n, matrices.
We also compare several different ways of integrating persistent vectors into self-attention:
All-attn:
this is our default model presented in Section 4 where persistent vectors are simply concatenated to context vectors.
Attn-split: this is the same as "all-attn" except the attention over context and persistent vectors
are computed separately. In other words, replace the softmax in Eq. 13 with two separate softmax functions: one for context vectors only and one for persistent vectors only.
this is the same as "all-attn" except we constrain half of the heads to attend only to context vectors, and the other half to attend only to persistent vectors.
Head-split:
Single-head: this the same "attn-split",
now persistent vectors are not split into multiple heads. Instead, each layer has single set of persistent key-value vectors of dimension d.
FF-attn: Transformer model where the ReLU of feedforward sublayers is replaced with So coftmax function as discussed in Section 4.1 This is the same as "single-head" above except persistent vectors are kept separate sublayer that comes after self-attention sublayer. Since this will double the depth of model, we decrease the number of layers to 24 and increase the feedforward size 3072 maintain the number of parameters same.
Note that all those versions have the same number of parameters except "head-split" which has fewer parameters because half its persistent vectors are not used. The result is shown in Figure rright. There few things to notice: (i) "all-attn" outperforms "attn-split", which indicates that there is
8

---

Under review as a conference paper at ICLR 2020
benefit in computing attention jointly persistent and context vectors; (ii) "single-head" is worse than "attn-split", which means persistent vectors with more heads are better; and (iii) dividing the heads into context-only and persistent-only groups does not work well; and (iv) "FF-attm" does not work as good as "all-attn" which means the switch from ReLU to Soft fmmax alone is not sufficient.
CONCLUSION
In this paper, we propose a novel attention layer that presents a unified mechanism to aggregate general and contextual information. It extends the self-attention layer of a transformer with set of persistent vectors that are capable of storing information that is complementary to the short term information in contexts. We also show that these persistent vectors can replace the feedforward layers in transformer network with no loss of performance. We think that this simplified layer can help better understand how information is processed and stored transformer- sequence models.
REFERENCES
Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-level language modeling with deeper self-attention. In Proceedings of the 33rd AAAI Conference on Artificial Intelligence, 2019.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In ICLR, 2019.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In ICLR, 2015.
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Jauvin. A neural probabilistic language model. Journal of machine learning research, 137-1115, 2003a.
Yoshua Bengio, dean-Sébastien Senécal, et al. Quick training of probabilistic neural nets by impor- tance sampling. In AISTATS, pp. 1-9, 2003b.
Rewon Child, Scott Gray, Alec Radford, and Nya Sutskever. Generating long sequences with sparse transformers. arXiv preprini arXiv: 1904. 10509, 2019.
Junyoung Chung. Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural networks. In ICLR, 2017.
Zihang Dai, Zhilin Yang, Yiming Yang, William W Cohen, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint 1901.0090000000 2019.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In ICML, 2017.
Jacob Devlin, Ming-Wee Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT (1), 2019.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(ul):2111-2159, 2011.
Joshua T Goodman. bit of progress in language modeling. Computer Speech Language, 15(4): 403-434, 2001.
Edouard Grave, Armand Joulin, Moustapha Cissé, and Hervé Jégou. Efficient softmax approximation for gpus. In ICML, 2017a.
Edouard Grave, Armand Joulin, and Nicolas Usunier. Improving neural language models with continuous cache. In ICLR, 2017b.
arXiv preprint Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. vrXiv:1410.5444 2014.

---

Under review as a conference paper at ICLR 2020
David Ha, Andrew M. Dai, and Quoe V. Le. Hypernetworks. In ICLR, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition,
pp. 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735 1780, 1997.
Hakan Inan, Khashayar Khosravi, and Richard Socher. Tying word vectors and word classifiers: loss framework for language modeling. In ICLR, 2017.
Michael I Jordan and Robert A Jacobs. Hicrarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-2144 1994.
Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint arXiv:1602.02410, 2016.
Ben Krause, lain Murray, Steve Renals, and Liang Lu. Multiplicatiie LSTM for sequence modelling. In ICLR (Workshop), 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.064450, 2016.
Matt Mahoney. Large text compression benchmark. URL: http://www. mattmahoney. net/text/text. html, 2011.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In ICLR, 2017.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. An analysis neural language modeling multiple scales. arXiv preprint arXiv: 1803.08240 2018.
Tomášs Mikolov, Martin Karafiát, Luká's Burget, Jan Černocký, and Sanjeev Khudanpur. Recurrent neural network based language model. Eleventh annual conference of the international speech communication association, 2010.
Alexander H. Miller, Adam Fisch, Jesse Dodge, Weston. Key-value memory networks for directly reading documents. In EMNLP, 2016.
Karimi, Antoine Bordes, and Jason
Frederic Morin and Yoshua Bengio. Hicrarchical probabilistic neural network language model. In AISTATS. 2005.
Asier Mujika, Florian Meier, and Angelika Steger. Fast-slow recurrent neural networks. NIPS, pp. 5915-5924, 2017.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In ICML, 2013.
Ofir Press and Lior Wolf. Using the output embedding to improve language models. In EACL (2), 2017.
Jack W. Rac, Chris Dyer, Peter Dayan, and Timothy P. Lillicrap. Fast parametric learning with activation memorization. In ICML, 2018.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In ACL (/), 2016.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position repressntatiions.. In NAACI-HIT (2), 2018.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey E. Hinton, and Jeff Dean. Outragcously large neural networks: The sparsely- gated mixture-of -f-extssts layer. In ICLR, 2017.
10

---

Under review as a conference paper at ICLR 2020
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1):1929-15)) 1958, 2014.
Sainbayar Sukhbaatar, Arthur Szlam, Jason Weston, and Rob Fergus. End-to-end memory networks. In NIPS, 2015.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span transformers. In ACL, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. In NIPS, 2015.
Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. In ICLR, 2019.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron C. Courville, Ruslan Salakhutdinov, Richard S. Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In ICML, 2015.
Wojeiech Zaremba, Ilya Sutskever, and Oriol Vinyals. Recurrent neural network regularization. arXiv preprint arXiv:1409.255: 23299, 2014.
Julian Georg Zilly, Rupesh Kumar Srivastava, Jan Koutník, and Jürgen Schmidhuber. Recurrent highway networks. ICMI, 2017.

---

Under review as a conference paper at ICLR 2020
A
ATTENTION MAPS
context
pe menory
02
30
180
1508
2006
3ie
3089
500 1900 150B
2006 300 3500
Figure 3: Sample attention maps from our model that trained on the WikiText-103 dataset. The plots correspond to 4 different attention heads in the model. The -aiis is different samples from short sequence, and the X-axis shows all the vectors the attention. The first 2048 vectors come from the context, and the remaining 2048 persistent vectors. In the top heads, few persistent vectors are dominating the attention, although the 2nd head has some attention weights in the context part well. The 3rd head has to very recent context. The last head is mostly attending to about last 500 tokens in the context, but there are some activations in the persistent vectors.
diverse activations the persistent vectors, while also attending
12

---

Under review as a conference paper at ICLR 2020
B BASELINE TRAINING
We trained baseline Transformer models on the Wiki Text-103 dataset using the same code and settings as our model to check if some of the training details (e.g. weight initialization, position embeddings, adaptive span, etc.) affected our result in Table 3 We considered two baselines with roughtly the same number of parameters as our model: 1) a 22-layer model with hidden size of 512 and feedforward layer size of 4096; 2) a 36-layer model with hidden size of 512 and feedforward layer size of 2048. The 22-layer model is more comparable to our model because it has similar number of nonlinear layers as our model, while 36-layer has twice as much non-linear layers as our model. The training of those models are plotted in Figure against our model (excluding the finetuning part). As which is not surprising as such deep Transformer models are known to be unstable during training. The same thing happenned to the 22-layer model but near the end of its training. However, is almost certain that its final performance would have been worse than our model even if its training did not diverge.
see, the training of the 36-layer model has diverged early in the training,
allattentoon
80
transformer 122
tronsformer L36
55
45
20
109k
200k
300k
400k
50ck
600k
7004
B00k
training steps
Figure 4: Training of baseline Transformer models on wikiText-103 dataset.
13


Reviews:
[{"review_id": "S1lXnkxpKB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper considers an architecture change to the transformer in which they swap the feedforward subcomponent of the standard transformer with an \"attention only\" variant that includes persistent \"memory\" vectors. The model is evaluated against a suite of baselines on the tasks of character- and word-level language modeling. Combining this \"all attention\" approach with adaptive span yields results about equivalent to the SOTA, in some cases with fewer parameters than existing models. The authors do a nice job of presenting ablation results. A key finding here, for example, is that the a model stripped of both persistent vectors and the feedforward sublayer performs poorly. \n\nOverall, I'm on the fence regarding this submission. This is solid work, and the idea of exploiting persistent representations in the transformer seems promising. But the architecture change here is relatively minor, and the gains seem somewhat minor (the exception may be in Table 2 which shows equivalent performance with half the parameters to previous SOTA, but then no other model is in the range of O(100m) parameters, so hard to know what's going on here). \n\nOne thing I would have liked is more motivation. If equivalence with fewer parameters is the main aim, then the model seems to fair reasonably well but the results are not really compelling. If, on the other hand, the authors are primarily interested in exploiting persistence, then I think this could have been investigated a bit more exhaustively, and perhaps the focus need not be on necessarily replacing the feedforward subcomponent (although that is one reasonable strategy). I do not see a huge inherent advantage to removing the feedforward layer, and it seems like there are alternative strategies --- at least equally as good --- to reduce parameters. \n\nA question for the authors: did you consider a vanilla Transformer with persistent memory vectors added? This would be something like adding a constant dummy input (independent of the example) that would be passed forward and arbitrarily transformed. I guess the meta-point here is that it doesn't seem to me that the persistent representations and the feedforward sublayer are necessarily mutually exclusive.\n\nAs a minor comment, I think Eq 13 is redundant since it literally repeats Eq. 4 save for swapping in $C_t^+$ for $C_t$.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "ByxAgnZl9B", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n\nThis paper mainly proposes a modification of well-studied Transformer architecture that is widely used for the text generation tasks, i.e., language modeling and machine translation.\nThe main idea is to consist of Transformer architecture only with self-attention layers. In other words, the proposed method discards the feed-forward layers and augment the self-attention layers with persistent memory vectors.\nThey conduct experiments on character and word-based language modeling and show the on-par or slightly better performance comparing with the standard Transformer language model and other similar Transformer modifications.\n \nThe proposed method is just a modification of the existing neural network architecture.\nMoreover, their method did not significantly improve the performance of language modeling in the experiments.\nFrom these perspectives, the proposed method is not innovative.\n \nHowever, the Transformer architecture is currently applied to a wide variety of tasks in text generation. Therefore, the proposed method can be largely influential to the community. \nActually, I like an idea discarding the feed-forward (sub-)layers in the Transformer, whose intuitive (and theoretical) role is not much discussed in the literature.\n\n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "BkehdZb_uB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a simple modification to the ubiquitous Transformer model. Noticing that the feed-forward layer of a Transformer layer looks a bit like an attention over \"persistent\" memory vectors, the authors propose to explicitly incorporate this notion directly into the self-attention layer. This involves concatenating the contextual representations with global, learned memory vectors, which are attended over. \n\nThe model is tested on widedly-utilized character/word-level language modeling benchmarks, where it is found to outperform, or be on par with, existing models while using fewer number of parameters. \n\nInsofar as architectural advancements can translate to general improvements across multiple NLP tasks, this paper could be seen as important. However, I am not sure that in 2019, demonstrating arguably-marginal perplexity improvements on standard datasets is enough. I would love to see if this type of layer can result in better conditional generation models (e.g. translation, summarization), or can train GPT2/BERT/XLNet-style models whose representations better transfer to other tasks.\n\nI had some further questions/comments:\n\n- I found the motivation of the persistent memory vector as replacing the FF-layer somewhat tenuous. Eq(5) is definitely different from Eq(9)! In my opinion this work can be better motivated/presented as just a standalone modification to the Transformer layer.\n\n- While it is impressive that the proposed approach performs better (or on par with, in the case of character-level language modeling) than the previous state-of-the-art models which are larger, to me it is not immediately clear if the benefit is coming from the proposed modifications, or something else (e.g. some of the things mentioned in 4.3). I understand most of this is taking from prior work, but as we all know, in deep models various architectural/hyperparameter modifications can interact in unexpected ways. Therefore, at a minimum, I would like to see the performance of a comparable model with the same exact setting, except for the persistent attention layer (i.e. N = 0 but using the feedforward sublayers). (If I understand the work correctly, this baseline should have comparable number of parameters?)\n\n- The ablation studies are good but it would be good to see them on the word-level task as well.\n\n- What is the performance of the FF-attn baseline with the same depth? (I understand the number of parameters would be  larger, but does this model perform as well as the Dai et al. 2019 work?)\n\n- What if you combine this with the FF sublayer as well?\n\n- Have you tried qualitatively analyzing the attention distributions? What are some examples in which the persistent memory vectors are attended to most? What are some examples in which the attention distribution for this model differs considerly versus regular self-attention layers?\n\n[EDIT after author rebuttal]\nThank you very much for the rebuttal. I have updated my score to reflect the latest iteration of the paper.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a modification to the Transformer architecture in which the self-attention and feed-forward layer are merged into a self-attention layer with "persistent" memory vectors. This involves concatenating the contextual representations with global, learned memory vectors, which are attended over. Experiments show slight gains in character and word-level language modeling benchmarks. 

While the proposed architectural changes are interesting, they are also rather minor and had a small impact in performance and in number of model parameters. The motivation of the persistent memory vector as replacing the FF-layer is a bit tenuous since Eqs 5 and 9 are substantially different. Overall the contribution seems a bit thin for a ICLR paper. I suggest more analysis and possibly experimentation in other tasks in a future iteration of this paper.
Decision: Reject