R ANDOM F EATURE A TTENTION
Anonymous authors
Paper under double-blind review
A BSTRACT
Transformers are state-of-the-art models for a variety of sequence modeling tasks.
At their core is an attention function which models pairwise interactions between the inputs at every timestep.
While attention is powerful, it does not scale efﬁciently to long sequences due to its quadratic time and space complexity in the sequence length.
We propose R FA , a linear time and space a ttention that uses r andom f eature methods to approximate the softmax function, and explore its applications in transformers.
R FA offers a straightforward way of learning with recency bias through an optional gating mechanism and can be used as a drop-in replacement for conventional softmax attention.
Experiments on language modeling and machine translation demonstrate that R FA achieves similar or better performance compared to strong transformer baselines.
In the machine translation experiment, R FA decodes twice as fast as a vanilla transformer.
Our analysis shows that R FA ’s efﬁciency gains are especially notable on long sequences, suggesting that R FA will be particularly useful in tasks that require working with large inputs, fast decoding speed, or low memory footprints.
I NTRODUCTION
Transformer architectures (Vaswani et al., 2017) have achieved tremendous success on a variety of sequence modeling tasks (Ott et al., 2018;
Radford et al., 2018; Parmar et al., 2018;
Devlin et al., 2019; Parisotto et al., 2019, inter alia ).
Under the hood, the key component is attention (Bahdanau et al., 2015).
Attention models pairwise interactions of the inputs, regardless of their distance.
This comes with quadratic time and memory costs, making the transformers computationally expensive, especially for long sequences.
A large body of research has been devoted to improving their time and memory efﬁciency (Tay et al., 2020d).
Some are able to achieve better asymptotic complexity (Lee et al., 2019; Child et al., 2019; Sukhbaatar et al., 2019; Beltagy et al., 2020, inter alia ), while it is more challenging to improve on shorter sequences: the additional computation steps required by some approaches can overshadow the time and memory they save (Kitaev et al., 2020; Wang et al., 2020; Roy et al., 2020, inter alia ).
This work proposes r andom f eature a ttention (R FA ).
It scales linearly in sequence length in terms of both time and space.
R FA builds on a kernel perspective of softmax (Rawat et al., 2019).
Using the well-established random feature maps (Rahimi & Recht, 2008; Avron et al., 2016; § 2),
R FA approximates the dot-then-exponentiate function with a kernel trick (Hofmann et al., 2008): exp( x · y )
≈
φ ( x ) · φ ( y ) .
Inspired by its connections to gated recurrent neural networks (Hochreiter & Schmidhuber, 1997; Cho et al., 2014) and differentiable plasticity (Ba et al., 2016;
Miconi et al., 2018), we further augment R FA with an optional gating mechanism, offering a straightforward way of learning with recency bias when locality is desired.
R FA and its gated variant ( § 3) can be used as a drop-in substitute for the canonical softmax attention, and increase the number of parameters by less than 0.1%.
We explore its applications in transformers on language modeling and machine translation ( § 4).
Our experiments show that R FA achieves comparable performance to vanilla transformer baselines in both tasks, while outperforming a recent related approach (Katharopoulos et al., 2020).
The gating mechanism proves particularly useful in language modeling: the gated variant of R FA outperforms the transformer baseline on WikiText-103.
R FA shines in decoding, even for shorter sequences.
In our head-to-head comparison on machine translation benchmarks, R FA decodes around 2 × faster than a transformer baseline, without accuracy loss.
Our analysis shows that more signiﬁcant time and memory efﬁciency improvements can be achieved for longer sequences: 12 × decoding speedup with less than 10% of the memory for 2,048-length outputs.
B ACKGROUND 2. A TTENTION IN S EQUENCE M ODELING
The attention mechanism (Bahdanau et al., 2015) has been widely used in many sequence modeling tasks.
Its dot-product variant is the key building block for the state-of-the-art transformer architectures (Vaswani et al., 2017).
Let { q t } Nt
=
denote a sequence of N query vectors, that attend to sequences of M key and value vectors.
At each timestep, the attention linearly combines the values weighted by the outputs of a softmax:
[Equation]
τ is the temperature hyperparameter determining how “ﬂat” the softmax is (Hinton et al., 2015).
Calculating attention for a single query takes O ( M ) time and space.
For the full sequence of N queries the space amounts to O ( MN ) .
When the computation cannot be parallelized across the queries, e.g., in autoregressive decoding, the time complexity is quadratic in the sequence length.
2. R ANDOM F EATURE M ETHODS
The theoretical backbone of this work is the unbiased estimation of the Gaussian kernel by Rahimi & Recht (2008).
Based on Bochner’s theorem (Bochner, 1955), Rahimi & Recht (2008) proposed random Fourier features to approximate a desired shift-invariant kernel.
The method nonlinearly transforms a pair of vectors x and y using a random feature map φ ; the inner product between φ ( x ) and φ ( y ) approximates the kernel evaluation on x and y .
More precisely:
Theorem 1 (Rahimi & Recht, 2008) .
Let φ : R d → R 2 D be a nonlinear transformation:
[Equation]
When d -dimensional random vectors w i are independently sampled from N ( 0 , σ 2 I d ) ,
[Equation]
Variance of the estimation is inversely proportional to D (Yu et al., 2016).
Random feature methods proved successful in speeding up kernel methods (Oliva et al., 2015; Avron et al., 2017; Sun, 2019, inter alia ), and more recently are used to efﬁciently approximate softmax (Rawat et al., 2019).
In § 3.1, we use it to derive an unbiased estimate to exp( (cid:104) · , · (cid:105) ) and then an efﬁcient approximation to softmax attention.
M ODEL
This section presents R FA ( § 3.1) and its gated variant ( § 3.2).
In § 3.3 we lay out several design choices and relate R FA to prior works.
We close by practically analyzing R FA ’s complexity ( § 3.4).
3. R ANDOM F EATURE A TTENTION
R FA builds on an unbiased estimate to exp( (cid:104) · , · (cid:105) ) from Theorem 1, which we begin with:
[Equation] [Equation]
[Table]
/Qj9CFVQBU2hKTlAqIaS4yaHfQ29ht5Ab6CfoJ+gI+gIegu9Hesdh+/B7/9tDlvQFnQVugpdg65B16Hr0A3oBnQTugndgm5Bt6Hb0B3oDnQXugvdg+5B96H70APoAfQQeghtQ9vQI+gR9Bh6DD2BnkBPoafQDrQD7UK70B60B+1D+9Az6Bn0HHqO8Mj78HAWkPTIFmHER64SRn7kGmEESK4TRoLkBmFESG4SRobkFmGESG4TRorkDmHESO4SRo7kHmEESe4TRpLkAWFESR4SRpZkmzDCJI8II03ymDDiJE8II0/ylDACJTuEkSjZJYxIyR5hZEr2CSNU8owwUiXPCU9iZUfihsdhyCI3t6NYhcWFNcjtQHjaDrpC6QXLVr4caluNfxXlS5P18BVpdtFbWrQai5Z11Fh415q+Pz0xXho/Gq8My3hjvDO2jLbRMbhxa/xi/Gr8Vvlc+b3yZ+Wvu9ZHc9N7nhv/uSp//wPhaM5q</latexit> (
[Figure]
not have
[Equation]
⊗ denotes the outer product between vectors, and σ 2 corresponds to the temperature term τ in Eq
R FA can be used as a drop-in-replacement for softmax-attention.
(a) The input is revealed in full to cross attention and encoder self attention . Here R FA calculates attention using Eq. (b) In causal attention R FA attends only to the preﬁx. This allows for a recurrent computation.
Tuple ( S t ∈ R 2 D × d , z t ∈ R 2 D ) is used as the “hidden state” at time step t to keep track of the history, similar to those in RNNs.
Then R FA ( q t , { k i } i ≤ t , { v i } i ≤ t ) = φ ( q t ) (cid:62) S t / φ ( q t ) · z t , where
[Equation]
2 D denotes the size of φ ( · ) .
Algos.
1 and 2 (Appendix A) summarize the computation procedure of R FA , and Figure 1 compares it against the softmax attention.
Analogously to the softmax attention, R FA has its multiheaded variant (Vaswani et al., 2017).
Note that in the sequence-to-sequence case, the input to R FA cross-attention is of ﬁxed-length size as opposed to variable-length size in the standard transformer.
In the experiments we use causal R FA in a transformer language model ( § 4.1), and both cross and causal R FA in the decoder of a sequenceto-sequence machine translation model.
3. R FA -G ATE : L EARNING WITH R ECENCY B IAS
The canonical softmax attention does not have any explicit modeling of distance or locality.
In learning problems where such inductive bias is crucial (Ba et al., 2016; Parmar et al., 2018; Miconi et al., 2018; Li et al., 2019, inter alia ), transformers heavily rely on positional encodings.
Answering to this, many approaches have been proposed, e.g., learning the attention spans (Sukhbaatar et al., 2019; Wu et al., 2020), and enhancing the attention computation with recurrent (Hao et al., 2019;
Chen et al., 2019) or convolutional (Wu et al., 2019; Mohamed et al., 2019) components.
R FA faces the same issue, but its causal attention variant (Eq. 6) offers a straightforward way of learning with recency bias.
We draw inspiration from its connections to RNNs, and augment R FA with a learned gating mechanism (Hochreiter & Schmidhuber, 1997; Cho et al., 2014):
[Equation]
w g and b g are learned parameters, and x t is the input representation at timestep t .
By multiplying the learned scalar gates 0 < g t < 1 against the hidden state ( S t , z t ) , history is exponentially decayed, favoring more recent context.
The gating mechanism shows another beneﬁt of R FA : it would be otherwise more difﬁcult to build similar techniques into the softmax attention, where there is no clear sense of “recurrence.”
It proves useful in our language modeling experiments ( § 4.1).
3. D ISCUSSION
On query and key norms, and learned random feature variance.
Eq.
5 assumes both the query and keys are of norm-1.
It therefore approximates a softmax attention that normalizes the queries and keys before multiplying them, and then scales the logits by dividing them by σ 2 .
Empirically, this normalization step scales down the logits (Vaswani et al., 2017) and enforces that − 1 ≤ q (cid:62) k ≤ 1 .
A consequence of this is that the softmax outputs would be “ﬂattened” if not for σ , which can be set a priori as a hyperparameter (Yu et al., 2016; Avron et al., 2017; Sun, 2019, inter alia ).
Here we instead learn it from data with the reparameterization trick (Kingma & Welling, 2014):
[Equation] [Equation]
I d is the d × d identity matrix, and ◦ denotes elementwise product between vectors.
d -dimensional vector σ is learned.
This norm-1 constraint is never mandatory.
Rather, we employ it for notation clarity and easier implementation.
In preliminary experiments we ﬁnd it has little impact on the performance when σ is set properly or learned from data.
Eq. 9 in Appendix A presents R FA without imposing it.
Going beyond the Gaussian kernel.
More broadly, random feature methods can be applied to a family of shift-invariant kernels, with the Gaussian kernel being one of them.
In the same family, the order-1 arc-cosine kernel (Cho & Saul, 2009) can be approximated with feature map: φ arccos ( x ) =
(cid:112) 1 /D
[ReLU( w 1 · x ) , . .
, ReLU( w D · x )] (cid:62) (Alber et al., 2017).
In our experiments, the Gaussian and arc-cosine variants achieve similar performance.
This supplements the exploration of alternatives to softmax in attention (Tsai et al., 2019; Gao et al., 2019).
Relations to prior works.
Katharopoulos et al.
(2020) inspires the causal attention variant of R FA .
They use a feature map based on the the exponential linear units activation (Clevert et al., 2016): elu( · ) + 1 .
It signiﬁcantly underperforms both the baseline and R FA in our controlled experiments, showing the importance of a properly-chosen feature map.
Random feature approximation of attention is also explored by a concurrent work (Choromanski et al., 2020), with applications in masked language modeling for proteins.
They point out that with further assumptions (that this work does not enforce), such approximation uniformly converges to the softmax attention.
Going beyond both,
R FA establishes the beneﬁts of random feature methods as a more universal substitute for softmax across all attention variants, facilitating its applications in, e.g., sequence-to-sequence learning.
There are interesting connections between gated R FA and fast weights (Ba et al., 2016) and differentiable plasticity (Miconi et al., 2018).
Emphasizing recent patterns, they learn a temporal memory to store history similarly to Eqs.
The main difference is that R FA additionally normalizes the output using φ ( q t ) · z as in Eq.
6, an output from approximating softmax’s partition function.
It is intriguing to study the role of this normalization term, which we leave to future work.
3. C OMPLEXITY A NALYSIS
Time.
Scaling linearly in the sequence lengths, R FA implies speedup wherever the complexity is quadratic for softmax attention.
More speciﬁcally:
• Signiﬁcant speedup can be expected in autoregressive decoding , both conditional (e
g., machine translation) and unconditional (e g., sampling from a language model).
A 1.9 × decoding speedup is achieved with the relatively short outputs in our machine translation experiments ( § 4.2); and more for longer sequences (e
g., 12 × for 2,048-length ones; § 5).
• For some applications, the softmax attention can be parallelized across the time steps, and thus hardly any speed improvement can be achieved.
These include language modeling, text classiﬁcation, and teacher forcing training of sequence-to-sequence models.
Memory.
Asymptotically, R FA achieves a better memory efﬁciency than its softmax counterpart (linear vs. quadratic).
To reach a more practical conclusion, we include in our analysis the cost of the feature maps.
φ ’s memory overhead largely depends on its size D .
For example, let’s consider the cross attention of a decoder.
R FA uses O (2 D + Dd ) space to store φ ( q t ) , z , and S (Eqs. ?? ; line 12 of Algo. 2).
In contrast, softmax cross attention stores the encoder outputs with O ( Md ) memory, with M being the source length.
In this case R FA has a lower memory overhead when
D (cid:28) M .
Typically D should be no less than d in order for reasonable approximation (Yu et al., 2016);
In a transformer model, d is the size of an attention head, which is usually around 64 or 128 (Vaswani et al., 2017; Ott et al., 2018).
This suggests that R FA can achieve signiﬁcant memory saving with longer sequences, which is supported by our empirical analysis in § 5.
Further, using moderate sized feature maps is also desirable, so that its overhead does not overshadow the time and memory R FA saves.
We experiment with D at d and 2 d ; the beneﬁt of using D > 2 d is marginal.
Table 3 in Appendix A discusses the time and space complexity in more detail.
E XPERIMENTS
We evaluate R FA on language modeling and machine translation.
4. L ANGUAGE M ODELING
We experiment with WikiText-103 (Merity et al., 2016).
It is based on English Wikipedia.
Tablein Appendix B summarizes some of its statistics.
Setting.
We compare the following models:
• B ASE is our implementation of the strong transformer-based language model by Baevski & Auli (2019).
• R FA builds on B ASE , but replaces the softmax attention with random feature attention.
We experiment with both Gaussian and arc-cosine kernel variants.
• R FA -G ATE additionally learns a sigmoid gate on top of R FA ( § 3.2).
It also has a Gaussian kernel variant and a arc-cosine kernel one.• φ elu is a baseline to R FA .
Instead of the random feature methods it uses the elu( · ) +feature map, as in Katharopoulos et al.
To ensure fair comparisons, we use comparable implementations, tuning, and training procedure.
All models use a 512 block size during both training and evaluation, i.e., they read as input a segment of 512 consecutive tokens, without access to the context from previous mini-batches.
R FA variants use 64-dimensional random feature maps.
We experiment with two model size settings, small (around 38M parameters) and big (around 242M parameters); they are described in Appendix B. along with other implementation details.
Results.
Table 1 compares the models’ performance in perplexity on WikiText-103 development and test data.
Both kernel variants of R FA , without gating, outperform φ elu by more than 2.4 and 2.
test perplexity for the small and big model respectively, conﬁrming the beneﬁts from using random is less intuitive to apply it in B ASE
[Table]
feature approximation.
Yet both underperform B ASE , with R FA -Gaussian having a smaller gap.
Comparing R FA against its gated variants, a more than 1.8 perplexity improvement can be attributed to the gating mechanism; and the gap is larger for small models.
Notably, R FA -G ATE -Gaussian outperforms B ASE under both size settings by at least 1.2 perplexity.
In general, R FA models with
Gaussian feature maps outperform their arc-cosine counterparts.
11 From the analysis in § 3.4 we would not expect speedup by R FA models, nor do we see any in the experiments.
Closing this section, we explore a “stateful” variant of R FA -G ATE -Gaussian.
It passes the last hidden state ( S t , z t ) to the next mini-batch during both training and evaluation, a technique commonly used in RNN language models (Merity et al., 2018).
This is a consequence of R FA ’s RNN-style computation, and is less straightforward to be applicable in the vanilla transformer models.
13 From the last row of Table 1 we see that this brings a more than
1.5 test perplexity improvement.
4. M ACHINE T RANSLATION
Datasets.
We
experiment
with three standard machine translation datasets
• WMT14 EN-DE and EN-FR (Bojar et al., 2014).
Our data split and preprocessing follow those of Vaswani et al.
We share the source and target vocabularies within each language pair, with 32,768 byte pair encoding types (BPE; Sennrich et al., 2016).
• IWSLT14 DE-EN (Cettolo et al., 2014) is based on TED talks.
The preprocessing follows Edunov et al.
Separate vocabularies of 9K/7K BPE types are used for the source and target.
Table 4 in Appendix B summarizes some statistics of the datasets.
Setting.
We compare the R FA variants described in § 4.1.
Here the B ASE model they build on is our implementation of the base-sized transformer by Vaswani et al.
All R FA models apply random feature attention in decoder cross and causal attention, but use softmax attention in encoders.
By § 3.4, this setting yields the greatest decoding time and memory savings.
The cross attention feature maps are 128-dimensional, and the causal ones are 64-dimensional.
The gated variants learn sigmoid gates in the decoder causal attention.
The φ elu baseline uses the same setting and applies feature map in both decoder cross and causal attention, but not in the encoders.
We
−
222222222 11 Sequence Length
D ec o d
i n g S p ee d : T o k e n s / s
Softmax RFA- arccos
RFA-Gaussian (a) Speed vs. lengths.
222222222 11 Sequence Length
M e m o r y :
G B
Softmax RFA- arccos
[Figure]
evaluate the models using SacreBLEU (Post, 2018) beam search with beam size 4 and
6 is used.
Further details are described in
[Table]
Results.
Table 2 compares the models’ test set BLEU on three machine translation datasets.
Overall both Gaussian and arc-cosine variants of R FA achieve similar performance to B ASE on all three datasets, signiﬁcantly outperforming Katharopoulos et al.
Differently from the trends in the language modeling experiments, here the gating mechanism does not lead to substantial gains.
Notably, all R FA variants decode more than 1 . 8 × faster than B ASE .
A NALYSIS
Decoding time and memory varying sequence lengths.
§ 3.4 shows that R FA can potentially achieve more signiﬁcant speedup and memory saving for longer sequences, which we now explore.
We use a simulation experiment to compare R FA
’s sequence-to-sequence decoding speed and memory overhead against the baseline’s.
Here we assume the input and output sequences are of the same length.
The compared models are of the same size as those described in § 4.2.
All models are tested using greedy decoding with the same batch size of 16.
From Figures 2 (a) and (b) we observe clear trends.
Varying the lengths, both R FA variants achieve consistent decoding speed with nearly-constant memory overhead.
In contrast, the baseline decodes slower for longer sequences taking an increasing amount of memory.
Notably, for 2,048-length sequences, R FA decodes around 12 × faster than the baseline while using less than 10% of the memory.
These results suggest that R FA can be particularly useful in sequence-to-sequence tasks with longer sequences, e.g., document-level machine translation (Miculicich et al., 2018).
Figure 3 in Appendix C.1 compares the speed and memory consumption in unconditional decoding (e.g., sampling from a language model).
The overall trends are similar to those in Figure 2.
Notes on decoding speed.
With a lower memory overhead, R FA can use a larger batch size than the baseline.
As noted by Katharopoulos et al.
(2020) and Kasai et al.
(2020), if we had used minibatches as large as the hardware allows, R FA would have achieved a more signiﬁcant speed gain.
Nonetheless, we control for batch size even though it is not the most favorable setting for R FA , since the conclusion translates better to common applications where one generates a single sequence at a time (e.g., instantaneous machine translation).
For the softmax attention baseline, we follow Ott et al.
and cache previously computed query/key/value representations, which signiﬁcantly improves its decoding speed (over not caching).
Further analysis results.
R FA achieves comparable performance to softmax attention.
Appendix C.2 empirically shows that this cannot be attributed to that R FA learns a good approximation to softmax: we train with one attention but evaluate with the other, the performance is hardly better than randomly-initialized untrained models.
Yet, a R FA model initialized from a pretrained softmax transformer achieves decent training loss after a moderate amount of ﬁnetuning steps (Appendix C.3).
This suggests some potential applications, e.g., to transfer knowledge from a softmax transformer pretrained on large-scale data (e.g., GPT-3; Brown et al., 2020) to a R FA model that is efﬁcient to sample from.
R ELATED W ORK
One common motivation across the following studies, that is shared by this work and the research we have already discussed, is to scale transformers to very long text, audio, or image sequences.
Note that there are plenty orthogonal choices for improving efﬁciency such as weight sharing (Dehghani et al., 2019), quantization (Shen et al., 2020), knowledge distillation (Sanh et al., 2020), and adapters (Houlsby et al., 2019).
For a detailed overview we refer the reader to Tay et al.
(2020d).
Sparse attention patterns.
The idea behind is to limit the reception ﬁeld of attention computation to reduce memory overhead.
It motivates one of the earliest attempts in improving attention’s efﬁciency, and still receives lots of interest.
The sparse patterns can be set a priori (Liu et al., 2018;
Qiu et al., 2019; Child et al., 2019; Beltagy et al., 2020; Ho et al., 2020; You et al., 2020, inter alia ) or learned from data (Sukhbaatar et al., 2019; Roy et al., 2020; Tay et al., 2020b; Kitaev et al., 2020, inter alia ).
For most of these approaches, it is yet to be empirically veriﬁed that they are suitable for large-scale sequence-to-sequence learning; few of them have recorded decoding speed beneﬁts.
Compressed context.
Wang et al.
(2020) propose to compress the context along the timesteps so that the effective sequence length for attention computation is reduced.
Another line of work aims to store past context into a memory module with limited size (Lee et al., 2019; Ainslie et al., 2020;
Rae et al., 2020, inter alia ), so that accessing longer history only moderately increase the overhead.
Related to these approaches, Dai et al.
(2019) use a segment-level recurrence mechanism to attend to a ﬁxed size of past memories, a technique that proves successful in transformer language models.
Reminiscent of RNN language models, R FA attends beyond a ﬁxed context window through a stateful computation, without increasing time or memory overhead.
To the best of our knowledge, R FA is the ﬁrst to perform on par with standard transformer on language modeling and machine translation while providing signiﬁcant decoding time and memory efﬁciency improvements in practical settings.
C ONCLUSION
We presented random feature attention (R FA ).
It views the softmax attention through the lens of kernel methods, and approximates it with random feature methods.
With an optional gating mechanism, R FA provides a straightforward way of learning with recency bias.
R FA ’s time and space complexity is linear in the sequence length.
We use R FA as a drop-in substitute for softmax attention in transformer models.
On machine translation and language modeling benchmarks, R FA achieves comparable or better performance than strong baselines.
In the machine translation experiment, R FA decodes twice faster.
Further time and memory efﬁciency improvements can be achieved for longer sequences.
Appendices A R ANDOM F EATURE A TTENTION IN M ORE D ETAILS
Algos.
1 and 2 describe causal and cross random feature attention’s computation procedures.
1: procedure RFA-C AUSAL ( { q i } Ni =1 , { k i } Ni =1 , { v i } Ni =1 ) 2: (cid:46) S is a D × d matrix 3: (cid:46) z is a D -dimensional vector 4:
S , z ← 0 , 0 5: for i = 1 to N do 6: (cid:101) q i , (cid:101) k
i ← φ ( q i ) , φ ( k i ) (cid:46)
Random feature maps 7:
S ← S + (cid:101) k
i ⊗ v i 8:
z ← z + (cid:101) k i 9: h (cid:62) i ← (cid:101) q (cid:62) i S / (cid:101) q i · z 10: end for 11: return { h i } Ni =1 12: end procedure
1: procedure RFA-C ROSS ( { q i } Ni =1 , { k i } Mi =1 , { v i } Mi =1 ) 2: (cid:46) S is a D × d matrix 3: (cid:46) z is a D -dimensional vector 4:
S , z ← 0 , 0 5: for i = 1 to M do 6: (cid:101) k
i ← φ ( k i ) (cid:46)
Random feature map 7:
S ← S + (cid:101) k
i ⊗ v (cid:62) i 8:
z ← z + (cid:101) k
i 9: end for 10: for i = 1 to N do 11: (cid:101) q i ← φ ( q i ) (cid:46) Random feature map 12: h (cid:62) i ← (cid:101) q (cid:62) i S / (cid:101) q i · z 13: end for 14: return { h i } Ni =1 15: end procedure
RFA without norm-1 constraints.
§ 3.1 assumes that the queries and keys are unit vectors.
This norm-1 constraint is not a must.
Here we present a R FA without imposing this constraint.
Let
C ( x ) = exp( (cid:107) x (cid:107) 2 / 2 σ 2 ) .
From Eq.
4 we have attn ( q t , { k i } , { v i } ) = [Equation] i
[Equation]
The speciﬁc attention computation is similar to those in § 3.1.
In sum, lifting the norm-1 constraint brings an additional scalar term C ( · ) .
Detailed complexity analysis.
Table 3 considers a sequence-to-sequence model, and breaks down the comparisons to training (with teacher forcing; Williams & Zipser, 1989) and autoregressive decoding.
R FA has a lower space complexity, since it never explicitly populates the attention matrices.
As for time, R FA trains in linear time, and so does the softmax attention: in teacher-forcing training a standard transformer decoder parallelizes the attention computation across time steps.
The trend of the time comparison differs during decoding: when only one output token is produced at a time,
R FA decodes linearly in the output length, while softmax attention quadratically.
[Table]
w
[Table] [Table]
B E XPERIMENTAL D ETAILS
Table 4 summarizes some statistics of the datasets used in our experiments.
Our implementation is based on JAX.
B. L ANGUAGE M ODELING
We compare the models using two model size settings, summarized in Table 5.
We use the ﬁxed sinusoidal position embeddings by Vaswani et al.
All models are trained for up to 150K gradient steps using the Adam optimizer (Kingma & Ba, 2015).
No (cid:96) 2 -regularization is used.
We apply early stopping based on development set perplexity.
All models are trained using 16 TPU v accelerators, and tested using a single TPU v2 accelerator.
B. M ACHINE T RANSLATION
WMT14.
We use the ﬁxed sinusoidal position embeddings by Vaswani et al.
For both ENDE and EN-FR experiments, we train the models using the Adam (with β 1 = 0 .
1 , β 2 = 0 .
98 , and (cid:15) =
10 − 9 ) optimizer for up to 350K gradient steps.
We use batch size of 1,024 instances for ENDE, while 4,096 for the much larger EN-FR dataset.
We apply early stopping based on development set BLEU.
Following standard practice, we average 10 most recent checkpoints.
The learning rate follows that by Vaswani et al.
No (cid:96) 2 regularization or gradient clipping is used.
All models are trained using 16 TPU v3 accelerators, and tested using a single TPU v2 accelerator.
Other hyperparameters are summarized in Table 6.
[Table]
[Table]
C M ORE A NALYSIS R ESULTS C. M ORE R ESULTS ON D ECODING S PEED AND M EMORY O VERHEAD
Figures 3 compare the R FA ’s unconditional decoding speed and memory against the softmax attention.
The setting is the same as that in § 5 except that here the models do not have an encoder.
This experiment aims to simulate the applications such as sampling from a language model
[Table]
[Figure]
d i n
[Figure] [Figure]
o k e n s / s
[Figure]
[Table]
[Figure] [Figure]
C. T RAIN AND E VALUATE WITH D IFFERENT A TTENTION F UNCTIONS
R FA achieves comparable performance to its softmax counterpart.
Does this imply that it learns a good approximation to the softmax attention?
To answer this question, we consider:
(i) an R FA -Gaussian model initialized from a pretrained softmax-transformer; (ii) a softmax-transformer initialized from a pretrained an R FA -Gaussian model.
If R FA ’s good performance can be attributed to that it learns a good approximation to softmax, both, without ﬁnetunining, should perform similarly to the pretrained models.
However, this is not the case on IWSLT14 DE-EN.
Both pretrained models achieve more than 35.2 development set
BLEU.
In contrast, (i) and (ii) respectively get 2.3 and 1.1 BLEU without ﬁnetuning, hardly beating a randomly-initialized untrained model.
This result aligns with the observation by Choromanski et al.
(2020), and suggests that it is not the case that R FA performs well because it learns to imitate softmax attention’s outputs.
C. K NOWLEDGE T RANSFER FROM S OFTMAX A TTENTION TO RFA
We ﬁrst supplement the observation in Appendix C.2 by ﬁnetuning (i) on the same pretraining data.
Figure 4 plots the learning curves.
It takes R FA roughly 1,500 steps to reach similar training loss to the pretrained model.
As a baseline, “R FA Reset” resets the multihead attention parameters (i.e., those for query, key, value, and output projections) to randomly initialized ones.
Its learning curve is similar to that of (i), suggesting that the pretrained multihead attention parameters are no more useful to R FA than randomly initialized ones.
To further conﬁrm this observation, “softmax Reset” resets the multihead attention parameters without changing the attention functions.
It converges to the pretraining loss in less than 200 steps.
Takeaway.
By the above results on IWSLT14, pretrained knowledge in a softmax transformer cannot be directly transferred to an R FA model.
However, from Figure 4 and a much larger-scale experiment by Choromanski et al.
(2020), we do observe that R FA can recover the pretraining loss, and the computation cost of ﬁnetuning is much less than training a model from scratch.
This suggests some potential applications.
For example, one might be able to initialize an R FA language model from a softmax transformer pretrained on large-scale data (e.g., GPT-3; Brown et al., 2020), and ﬁnetune it at a low cost.
The outcome would be an R FA model retaining most of the pretraining knowledge, but is much faster and more memory-friendly to sample from.
We leave such exploration to future work.
C. R ESULTS ON C HARACTER - LEVEL T EXT C LASSIFICATION
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]