In causal attention Rfa attends only to the prefix.It is also sometimes called “decoder self-attention” or “autoregressive attention.”
This allows for a recurrent computation.
Tuple \(({\mathbf {S}}_t\in \mathbb {R}^{2D\times d},{\mathbf {z}}_t\in \mathbb {R}^{2D})\) 
is used as the “hidden state” at time step \(t\)  to keep track of the history, similar to those in RNNs.
Then
\(\textsc {Rfa} ({\mathbf {q}}_t,\lbrace {\mathbf {k}}_i\rbrace _{i\le t},\lbrace {\mathbf {v}}_i\rbrace _{i\le t}) = \phi ({\mathbf {q}}_t)^\top {\mathbf {S}}_t / (\phi ({\mathbf {q}}_t)\cdot {\mathbf {z}}_t)\) ,
where
\({\mathbf {S}}_t = {\mathbf {S}}_{t-1} + \phi \left({\mathbf {k}}_t\right)\otimes {\mathbf {v}}_t, \quad {\mathbf {z}}_t = {\mathbf {z}}_{t-1} + \phi \left({\mathbf {k}}_t\right).\) 


\(2D\)  denotes the size of \(\phi (\cdot ).\) 
Appendix REF  summarizes the computation procedure of Rfa,
and Figure REF  compares it
against the softmax attention.
Appendix REF  derives causal Rfa in detail.
Analogously to the softmax attention, Rfa has its multiheaded variant [1]}.
In our experiments we use causal Rfa in a transformer language model (§REF ), and both cross and causal Rfa in the decoder of a sequence-to-sequence machine translation model.
<FIGURE>Rfa-Gate: Learning with Recency Bias
The canonical softmax attention does not have
any explicit modeling of distance or locality.
In learning problems where such inductive bias
is crucial [2]}, [3]}, [4]}, [5]},
transformers heavily rely on positional encodings.
Answering to this, many approaches have been proposed,
e.g.,
learning the attention spans [6]}, [7]},
and enhancing the attention computation with recurrent [8]}, [9]}
or convolutional [10]}, [11]} components.
Rfa faces the same issue, but its causal attention variant (Eq. REF )
offers a straightforward way of learning with recency bias.
We draw inspiration from its connections to RNNs,
and augment Rfa with a learned gating mechanism [12]}, [13]}, [14]}:
\(\begin{split}g_t &= \operatorname{sigmoid}( {\mathbf {w}}_g\cdot {\mathbf {x}}_t+ b_g), \\{\mathbf {S}}_t &= g_t\, {\mathbf {S}}_{t-1} + (1-g_t)\,\phi \left({\mathbf {k}}_t\right)\otimes {\mathbf {v}}_t, \\{\mathbf {z}}_t &= g_t\, {\mathbf {z}}_{t-1} + (1-g_t)\,\phi \left({\mathbf {k}}_t\right).\end{split}\) 

\({\mathbf {w}}_g\)  and \(b_g\)  are learned parameters,
and \({\mathbf {x}}_t\)  is the input representation at timestep \(t\) .In multihead attention [1]},
\({\mathbf {k}}_t\)  and \({\mathbf {v}}_t\)  are calculated from \({\mathbf {x}}_t\) 
using learned affine transformations.
By multiplying the learned scalar gates \(0< g_t<1\)  against the
hidden state \(({\mathbf {S}}_t, {\mathbf {z}}_t)\) ,
history is exponentially decayed, favoring more recent context.
The gating mechanism shows another benefit of Rfa:
it would be otherwise more difficult to
build similar techniques into the softmax attention,
where there is no clear sense of “recurrence” (Appendix REF ).
It proves useful in our language modeling experiments (§REF ).

Discussion
On query and key norms, and learned random feature variance.
Eq. REF  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 \(\sigma ^2\) .
Empirically, this normalization step
scales down the logits [1]}
and enforces that \(-1\le {\mathbf {q}}^\top {\mathbf {k}}\le 1\) .
In consequence, the softmax outputs would be “flattened”
if not for \(\sigma \) ,
which can be set a priori as a hyperparameter [17]}, [18]}, [19]}.
Here we instead learn it from data with the reparameterization trick [20]}:
\(\widetilde{{\mathbf {w}}}_i\sim {\mathcal {N}}({\mathbf {0}}, {\mathbf {I}}_d), \quad {\mathbf {w}}_i = \sigma \circ \widetilde{{\mathbf {w}}}_i.\) 

\({\mathbf {I}}_d\)  is the \(d\times d\)  identity matrix,
and \(\circ \)  denotes elementwise product between vectors.
\(d\) -dimensional vector \(\sigma \)  is learned,
but random vectors \(\widetilde{{\mathbf {w}}}_i\)  are not.This departs from Eq. REF 
by lifting the isotropic assumption imposed on the Gaussian distribution:
note the difference between the vector \(\sigma \) 
in Eq. REF  and the scalar \(\sigma \)  in Eq. REF .
We find this improves the performance in practice (§),
even though the same result in Theorem REF  may not directly apply.
This norm-1 constraint is never mandatory.
Rather, we employ it for notation clarity and easier implementation.
In preliminary experiments we find
it has little impact on the performance
when \(\sigma \)  is set properly
or learned from data.
Eq. REF  in Appendix 
presents Rfa 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 [21]}
can be approximated with feature map:
\(\phi _{\arccos }({\mathbf {x}}) = \sqrt{1/D}[\operatorname{ReLU}({\mathbf {w}}_1\cdot {\mathbf {x}}),\dots ,\operatorname{ReLU}({\mathbf {w}}_D\cdot {\mathbf {x}})]^\top \)  [22]}.Apart from replacing the sinusoid functions with \(\operatorname{ReLU}\) ,
it constructs \({\mathbf {w}}_i\)  in the same way as Eq. REF .
In our experiments, the Gaussian and arc-cosine variants achieve similar performance.
This supplements the exploration of alternatives to softmax in attention [23]}, [24]}.
Relations to prior work.
[25]} inspire
the causal attention variant of Rfa.
They use a feature map based on the exponential linear unit activation [26]}:
\(\operatorname{elu}(\cdot ) + 1\) .
It significantly underperforms both the baseline and Rfa 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 [27]},
with applications in masked language modeling for proteins.
They propose positive random features to approximate softmax,
aiming for a lower variance in critical regions.
Rfa instead normalizes the queries and keys before random projection to reduce variance.

Going beyond both, Rfa establishes the benefits
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 Rfa and fast weights [28]}, [29]}, [2]}, [4]}.
Emphasizing recent patterns,
they learn a temporal memory to store history
similarly to Eqs. REF .
The main difference is that
Rfa additionally
normalizes the output using \(\phi ({\mathbf {q}}_t)\cdot {\mathbf {z}}\)  as in Eq. REF ,
a by-product of approximating softmax's partition function.
It is intriguing to study the role of this
normalization term, which we leave to future work.

Complexity Analysis
Time.
Scaling linearly in the sequence lengths,
Rfa needs less computation (in terms of number of operations) for long sequences.
This implies speedup wherever the quadratic-time softmax attention
cannot be fully-parallelized across time steps.
More specifically:
Significant speedup can be expected in autoregressive decoding,
both conditional (e.g., machine translation)
and unconditional (e.g., sampling from a language model).
For example, 1.9\(\times \)  speedup is achieved
in our machine translation experiments (§REF );
and more for longer sequences (e.g.,
12\(\times \)  for 2,048-length ones; §).

Some applications (e.g., language modeling, text classification)
reveal inputs to the model in full.A causal masking is usually used to prevent the model from accessing future tokens in language models.
When there are enough threads to parallelize softmax attention across time steps,
hardly any speedup from Rfa can be achieved;
when there are not, typically for very long sequences (\(>\) 1,000), substantial speed gain is possible.
For example, Rfa does not achieve any speedup
when working with 512-length context (§REF ),
but achieves a \(5.3\times \)  speedup with 4,000-length context (§REF ).

Memory.
Asymptotically, Rfa has a better memory efficiency
than its softmax counterpart (linear vs. quadratic).
To reach a more practical conclusion,
we include in our analysis the cost of the feature maps.
\(\phi \) 's memory overhead largely depends on its size \(D\) .
For example, let's consider the cross attention of a decoder.
Rfa uses \({\mathcal {O}}(4D + 2Dd)\)  space to store
\(\phi ({\mathbf {q}}_t)\) ,
\(\sum _i \phi ({\mathbf {k}}_i)\otimes {\mathbf {v}}_i\) ,
and \(\sum _i\phi ({\mathbf {k}}_i)\) 
(Eq. REF ; line REF  of Algo. REF ).Rfa never constructs the \(M\times 2D\times d\)  tensor
\([\phi ({\mathbf {k}}_i)\otimes {\mathbf {v}}_i]_i\) ,
but sequentially processes the sequence.
In contrast,
softmax cross attention
stores the encoder outputs with
\({\mathcal {O}}(Md)\)  memory, with \(M\)  being
the source length.
In this case Rfa has a lower memory overhead
when \(2D \ll M\) .
Typically \(D\)  should be no less than \(d\)  in order for reasonable approximation [17]};
In a transformer model, \(d\)  is the size of an attention head,
which is usually around 64 or 128 [1]}, [34]}.
This suggests that
Rfa can achieve significant memory saving with longer sequences,
which is supported by our empirical analysis in §.
Further, using moderate sized feature maps is also desirable,
so that its overhead does not overshadow the time and memory Rfa saves.
We experiment with \(D\)  at \(d\)  and \(2d\) ;
the benefit of using \(D>2d\)  is marginal.
Appendix REF 
discusses the time and space complexity in more detail,
and Appendix REF  studies
the effect of random feature size on performance.
Experiments
We evaluate Rfa on language modeling, machine translation, and long text classification.
Language Modeling
Setting.
We experiment with WikiText-103 [35]}.
It is based on English Wikipedia.
Table REF  in Appendix  summarizes some of its statistics.
We compare the following models:
Base is our implementation of the strong transformer-based language model by [36]}.

Rfa builds on Base, but replaces the softmax attention
with random feature attention.
We experiment with both Gaussian and arc-cosine kernel variants.

Rfa-Gate additionally learns a sigmoid gate on top of Rfa (§REF ). It also has a Gaussian kernel variant and a arc-cosine kernel one.This gating technique is specific to Rfa variants,
in the sense that it is less intuitive to apply it in Base.

\(\phi _{\operatorname{elu}}\)  is a baseline to Rfa.
Instead of the random feature methods it uses the
\(\operatorname{elu}(\cdot ) + 1\)  feature map,
as in [25]}.

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.
Rfa 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 REF 
along with other implementation details.
<TABLE>Results.
Table REF  compares the models'
performance in perplexity
on WikiText-103 development and test data.
Both kernel variants of Rfa, without gating,
outperform \(\phi _{\operatorname{elu}}\) 
by more than 2.4 and 2.1 test perplexity for the small and big model respectively,
confirming the benefits from using random feature approximation.All models are trained for 150K steps;
this could be part of the reason behind
the suboptimal performance of \(\phi _{\operatorname{elu}}\) :
it may need 3 times more gradient updates
to reach similar performance to the softmax attention baseline [25]}.
Yet both underperform Base, with
Rfa-Gaussian having a smaller gap.
Comparing Rfa 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, Rfa-Gate-Gaussian outperforms Base under both size settings by at least 1.2 perplexity.
In general, Rfa models with Gaussian feature maps outperform
their arc-cosine counterparts.We observe that Rfa Gaussian variants
are more stable and easier to train than the arc-cosine ones as well as \(\phi _{\operatorname{elu}}\) .
We conjecture that this is because the outputs of the Gaussian feature maps have an \(\ell _2\) -norm of 1,
which can help stabilize training.
To see why, \(\sin ^2(x) + \cos ^2(x)=\cos (x-x)=1\) .
From the analysis in §REF  we would not expect
speedup by Rfa models, nor do we see any in the experiments.In fact, Rfa trains around 15% slower than Base due to the additional overhead from the feature maps.
Closing this section, we explore a “stateful” variant of Rfa-Gate-Gaussian.
It passes the last hidden state \(({\mathbf {S}}_t, {\mathbf {z}}_t)\) 
to the next mini-batch during both training and evaluation,
a technique commonly used in RNN language models [39]}.
This is a consequence of Rfa's RNN-style computation,
and is less straightforward to be applicable in the vanilla transformer models.Some transformer models use a
text segment from the previous mini-batch as a prefix
[36]}, [41]}. Unlike Rfa, this gives the model access to only a limited amount of context, and significantly increases the memory overhead.
From the last row of Table REF 
we see that this brings a
more than 1.5 test perplexity improvement.
Machine Translation
Datasets. We experiment with three standard machine translation datasets.
[nosep,leftmargin=1em,labelwidth=*,align=left]
WMT14 EN-DE and EN-FR [42]}.
Our data split and preprocessing follow those of [1]}.
We share the source and target vocabularies within each language pair,
with 32,768 byte pair encoding types (BPE; [44]}).

IWSLT14 DE-EN [45]} is based on TED talks.
The preprocessing follows [46]}.
Separate vocabularies of 9K/7K BPE types are used for the source and target.

Table REF  in Appendix  summarizes some statistics of the datasets.
Setting.
We compare the Rfa variants described in §REF .
They build on a Base model that is our implementation of the base-sized
transformer [1]}.
All Rfa models apply random feature attention
in decoder cross and causal attention, but use softmax attention in encoders.
This setting yields the greatest decoding time and memory savings (§REF ).
We use 128/64 for \(D\)  in cross/causal attention.
Rfa-Gate learns sigmoid gates in the decoder causal attention.
The \(\phi _{\operatorname{elu}}\)  baseline
uses the same setting and applies
feature map in both decoder cross and causal attention,
but not in the encoders.
Further details are described in Appendix REF .
<TABLE>Results.
Table REF  compares the models' test set BLEU
on three machine translation datasets.
Overall both Gaussian and arc-cosine variants of Rfa achieve similar performance to Base on all three datasets,
significantly outperforming [25]}.
Differently from the trends in the language modeling experiments,
here the gating mechanism does not lead to substantial gains.
Notably, all Rfa variants decode more than \(1.8\times \) 
faster than Base.

Long Text Classification
We further evaluate Rfa's accuracy and efficiency when used as text encoders on three NLP tasks from the recently proposed
Long Range Arena benchmark [49]},
designed to evaluate efficient Transformer variants on tasks that require processing long sequences.https://github.com/google-research/long-range-arena
Experimental setting and datasets.
We compare Rfa against baselines on the following datasets:
ListOps (LO; [50]}) aims to diagnose the capability of modelling hierarchically structured data.
Given a sequence of operations on single-digit integers, the model predicts the solution, also a single-digit integer.
It is formulated as a 10-way classification.
We follow [49]} and consider sequences with 500–2,000 symbols.

Character-level text classification with the IMDb movie review dataset [52]}.
This is a binary sentiment classification task.

Character-level document retrieval with the ACL Anthology Network (AAN; [53]}) dataset.
The model classifies whether there is a citation between a pair of papers.

To ensure fair comparisons,
we implement Rfa on top of the transformer baseline by [49]},
and closely follow their preprocessing, data split, model size, and training procedure.
Speed and memory are evaluated on the IMDb dataset.
For our Rfa model, we use \(D=64\)  for the IMDb dataset, and \(D=128\)  for others.
We refer the readers to [49]} for further details.
Results.
From Table REF  we can see that
Rfa outperforms the transformer
baseline on two out of the three datasets,
achieving the best performance on IMDb with 66% accuracy.
Averaging across three datasets,
Rfa outperforms the transformer by 0.3% accuracy,
second only to [56]} with a 0.1% accuracy gap.
In terms of time and memory efficiency,
Rfa is among the strongest.
Rfa speeds up over the transformer by \(1.1\) –\(5.3\times \) , varying by sequence length.
Importantly, compared to the only two baselines that perform comparably to the baseline transformer model [57]}, [56]}, Rfa has a clear advantage in both speed and memory efficiency, and is the only model that is competitive in both accuracy and efficiency.
<TABLE>Analysis
<FIGURE>Decoding time and memory varying by sequence length.
§REF  shows that Rfa can potentially
achieve more significant speedup and memory saving for longer sequences, which we now explore.
We use a simulation conditional generation experiment on to compare
Rfa'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 §REF ,
with 6-layer encoders and decoders.
Other hyperparameters are summarized in Appendix REF .
All models are tested using greedy decoding with the same batch size of 16, on a TPU v2 accelerator.
From Figures REF  (a) and (b) we observe clear trends.
Varying the lengths,
both Rfa 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,
Rfa decodes around 12\(\times \)  faster than the baseline while using less than 10%
of the memory.
Rfa-\(\arccos \)  slightly outperforms Rfa-Gaussian in terms of speed and memory efficiency.
This is because when using the same \(D\)  (as we do here),
the \(\phi _{\arccos }\)  is half the size of \(\phi _{\text{Gaussian}}\) .
These results suggest that Rfa can be particularly useful
in sequence-to-sequence tasks with longer sequences,
e.g., document-level machine translation [59]}.
Figure REF  in Appendix REF 
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 REF .
Notes on decoding speed.
With a lower memory overhead, Rfa can use
a larger batch size than the baseline.
As noted by [25]} and [61]},
if we had used mini-batches as large as the hardware allows,
Rfa could have achieved a more significant speed gain.
Nonetheless, we control for batch size
even though it is not the most favorable setting for Rfa,
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 [34]} and cache previously computed query/key/value representations,
which significantly improves its decoding speed (over not caching).
Further analysis results.
Rfa achieves comparable performance to softmax attention.
Appendix REF  empirically shows that
this cannot be attributed to
Rfa learning a good approximation to softmax:
when we train with one attention but evaluate with the other,
the performance is hardly better than randomly-initialized untrained models.
Yet, an Rfa model
initialized from a pretrained softmax transformer
achieves decent training loss after a
moderate amount of finetuning steps (Appendix REF ).
This suggests some potential applications,
e.g., transferring knowledge
from a pretrained transformer
(e.g., GPT-3; [63]})
to an Rfa model that is more efficient to sample from.

Related Work
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 long sequences.
Note that there are plenty orthogonal choices for improving efficiency such as weight sharing [64]}, quantization [65]}, knowledge distillation [66]}, and adapters [67]}.
For a detailed overview we refer the reader to [68]}.
Sparse attention patterns.
The idea behind these methods
is to limit
the reception field of attention computation.
It motivates earlier attempts in
improving attention's efficiency,
and still receives lots of interest.
The sparse patterns can be set a priori [69]}, [70]}, [71]}, [72]}
or learned from data [6]}, [74]}.
For most of these approaches,
it is yet to be empirically verified that
they are suitable for large-scale sequence-to-sequence learning;
few of them have recorded decoding speed benefits.
Compressed context.
[75]} 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 [76]}, [77]}, [78]},
so that accessing longer history only moderately increases the overhead.
Reminiscent of RNN language models,
Rfa attends beyond a fixed context window
through a stateful computation, without increasing
time or memory overhead.

Conclusion
We presented random feature attention (Rfa).
It views the softmax attention through
the lens of kernel methods,
and approximates it with random feature methods.
With an optional gating mechanism,
Rfa provides a straightforward way
of learning with recency bias.
Rfa's time and space complexity
is linear in the sequence length.
We use Rfa as a drop-in substitute for softmax attention
in transformer models.
On language modeling, machine translation,
and long text classification benchmarks,
Rfa achieves comparable or better performance
than strong baselines.
In the machine translation experiment,
Rfa decodes twice as fast.
Further time and memory efficiency
improvements can be achieved for longer sequences.

Acknowledgments
We would like to thank
Phil Blunsom,
Chris Dyer,
Nando de Freitas,
Jungo Kasai,
Adhiguna Kuncoro,
Dianqi Li,
Ofir Press,
Lianhui Qin,
Swabha Swayamdipta,
Sam Thomson,
the language team at DeepMind
and the ARK group at the University of Washington
for their helpful feedback.
We also thank Tay Yi
for helping run the Long Range Arena experiments,
Richard Tanburn
for the advice on implementations,
and the anonymous reviewers
for their thoughtful comments.
This work was supported in part by NSF grant 1562364 and a Google Fellowship. Nikolaos Pappas was supported by the Swiss National Science Foundation under grant number P400P2_183911 “UNISON.”

Random Feature Attention in More Detail
Detailed Computation Procedure
Algorithms REF  and REF 
describe causal and cross random feature
attention's computation procedures.
<FIGURE><FIGURE>
Variance of Random Fourier Features
The following result is due to [17]}.
Using the same notation as in §REF :
\(\operatorname{Var}(\phi \left({\mathbf {x}}\right)\cdot \phi \left({\mathbf {y}}\right))=\frac{1}{2D}\left(1 - e^{-z^2}\right)^2,\) 

where \(z=\left\Vert {\mathbf {x}}-{\mathbf {y}}\right\Vert  / \sigma \) .

Derivation of Causal Rfa
This section presents a detailed derivation of causal Rfa as in §REF .
Following Eq. REF  but changing the attended keys and values to the prefix:
\(\textsc {Rfa} ({\mathbf {q}}_t,\lbrace {\mathbf {k}}_i\rbrace _{i\le t},\lbrace {\mathbf {v}}_i\rbrace _{i\le t})&= \frac{\phi \left({\mathbf {q}}_t\right)^\top \sum _{i\le t} \phi \left({\mathbf {k}}_i\right)\otimes {\mathbf {v}}_i}{\phi \left({\mathbf {q}}_t\right)\cdot \sum _{j\le t} \phi \left({\mathbf {k}}_j\right)}\) 

Let \({\mathbf {S}}_t \triangleq \sum _{i\le t} \phi ({\mathbf {k}}_i)\otimes {\mathbf {v}}_i\) ,
and \({\mathbf {z}}_t \triangleq \sum _{i\le t} \phi ({\mathbf {k}}_i)\) ;
both can be calculated recurrently.
Assuming \({\mathbf {S}}_0 = \mathbf {0}\)  and \({\mathbf {z}}_0 = \mathbf {0}\) :
\({\mathbf {S}}_t = {\mathbf {S}}_{t-1} + \phi \left({\mathbf {k}}_t\right)\otimes {\mathbf {v}}_t, \quad {\mathbf {z}}_t = {\mathbf {z}}_{t-1} + \phi \left({\mathbf {k}}_t\right), \quad t\ge 1.\) 

This completes the derivation of causal Rfa as in §REF .

Rfa without Norm-1 Constraints
§REF  assumes that the queries and keys are unit vectors.
This norm-1 constraint is not a must.
Here we present a Rfa without imposing this constraint.
Let \(C({\mathbf {x}}) = \exp (\left\Vert {\mathbf {x}}\right\Vert ^2/2\sigma ^2)\) .
From Eq. REF  we have
\(\operatorname{attn}\left({\mathbf {q}}_t,\lbrace {\mathbf {k}}_i\rbrace ,\lbrace {\mathbf {v}}_i\rbrace \right)=\)
\(\begin{split}\sum _i\frac{\exp \left({\mathbf {q}}_t\cdot {\mathbf {k}}_i /\sigma ^2\right)}{\sum _j\exp \left({\mathbf {q}}_t\cdot {\mathbf {k}}_j /\sigma ^2\right)}{\mathbf {v}}_i^\top &\approx \sum _i\frac{C({\mathbf {q}}_t)\,C({\mathbf {k}}_i)\,\phi \left({\mathbf {q}}_t\right)^\top \phi \left({\mathbf {k}}_i\right){\mathbf {v}}_i^\top }{\sum _j C({\mathbf {q}}_t)\,C({\mathbf {k}}_j)\,\phi \left({\mathbf {q}}_t\right) \cdot \phi \left({\mathbf {k}}_j\right)}\\&= \frac{\phi \left({\mathbf {q}}_t\right)^\top \sum _i C({\mathbf {k}}_i)\,\phi \left({\mathbf {k}}_i\right)\otimes {\mathbf {v}}_i}{\phi \left({\mathbf {q}}_t\right)\cdot \sum _jC({\mathbf {k}}_j)\,\phi \left({\mathbf {k}}_j\right)}.\end{split}\) 

The specific attention computation is similar to those in §REF .
In sum, lifting the norm-1 constraint
brings an additional scalar term \(C(\cdot )\) .

Relating Rfa-Gate to Softmax Attention
Drawing inspiration from gated RNNs,
§REF  introduces a gated variant of Rfa.
Now we study its “softmax counterpart.”
\(\begin{split}\widetilde{\mathbf {k}}_i &= \mathbf {k}_i \ (1-g_i) \prod _{j=i+1}^{t}g_j,\quad \widetilde{\mathbf {v}}_i = \mathbf {v}_i \ (1-g_i) \prod _{j=i+1}^{t}g_j, \quad i=1,\dots ,t \\\mathbf {h}_t&=\operatorname{attn}(\mathbf {q}_t, \lbrace \widetilde{\mathbf {k}}_i\rbrace _{i\le t}, \lbrace \widetilde{\mathbf {v}}_i\rbrace _{i\le t}).\end{split}\) 

\(\mathbf {h}_t\)  is the output at timestep \(t\)  and is used for onward computation.
At each step, all prefix keys and values are decayed by a gate value before calculating the attention.
This implies that the attention computation for \(\mathbf {q}_{t+1}\)  cannot start until that of \(\mathbf {q}_t\)  is finished.
Combined with the linear complexity of softmax normalization, this amounts to quadratic time in sequence length,
even for language modeling training.
The above model is less intuitive and more expensive in practice, without the Rfa perspective. This shows that Rfa brings some benefits in developing new attention models.

Detailed Complexity Analysis
Table REF  considers
a sequence-to-sequence model,
and breaks down the comparisons
to training (with teacher forcing; [80]})
and autoregressive decoding.
Here we assume enough threads to fully parallelize
softmax attention across timesteps when the inputs are revealed to the model in full.
Rfa has a lower space complexity,
since it never explicitly populates the attention matrices.
As for time,
Rfa 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,
Rfa decodes linearly in the output length,
while softmax attention decodes quadratically.
<TABLE><TABLE>
Experimental Details
Table REF  summarizes
some statistics of the datasets used in our experiments.
Our implementation is based on JAX.https://github.com/google/jax.
<TABLE>During training, we sample a different random projection matrix for each attention head.
Preliminary experiments suggest this performs better than using the same random projection throughout training (Table REF ).
Our conjecture is that this helps keep the attention heads from “over committing” to any particular random projection [81]}.
To avoid the overhead of sampling from Gaussian during training,
we do this in an offline manner.
I.e., before training we construct a pool of random matrices (typically 200),
at each training step we draw from the pool.
At test time each attention head uses the same random projection, since no accuracy benefit is observed by using different ones for different test instances.
Language Modeling
We compare the models using two model size settings, summarized in
Table REF .
We use the fixed sinusoidal position embeddings by [1]}.
All models are trained for up to 150K gradient steps
using the Adam optimizer [83]}.
No \(\ell _2\) -regularization is used.
We apply early stopping based on development set perplexity.
All models are trained using 16 TPU v3 accelerators,
and tested using a single TPU v2 accelerator.
<TABLE>
Machine Translation
WMT14.
We use the fixed sinusoidal position embeddings by [1]}.
For both EN-DE and EN-FR experiments,
we train the models using the Adam (with \(\beta _1=0.1\) , \(\beta _2=0.98\) , and \(\epsilon =10^{-9}\) ) optimizer
for up to 350K gradient steps.
We use a batch size of 1,024 instances for EN-DE,
while 4,096 for the much larger EN-FR dataset.
The learning rate follows that by [1]}.
Early stopping is applied based on development set BLEU.
No \(\ell _2\)  regularization or gradient clipping is used.
All models are trained using 16 TPU v3 accelerators,
and tested using a single TPU v2 accelerator.
Following standard practice, we average 10 most recent checkpoints at test time.
We evaluate the models using SacreBLEU [86]}.https://github.com/mjpost/sacrebleu
A beam search with beam size 4 and length penalty 0.6 is used.
Other hyperparameters are summarized in Table REF .
<TABLE>
More Analysis Results
More Results on Decoding Speed and Memory Overhead
<FIGURE>Figure REF 
compares the Rfa's unconditional decoding speed and memory
against the softmax attention.
The setting is the same as that in §
except that here the models do not have an encoder.
This experiment aims to simulate the applications
such as sampling from a language model.

Effect of Random Feature Size
This section studies how the size of \(\phi (\cdot )\)  affects the performance.
Table REF  summarize Rfa-Gaussian's performance on WMT14 EN-DE development set.
The model and training are the same as that used in §REF  except random feature size.
Recall from §REF  that the size of \(\phi (\cdot )\)  is \(2D\)  for Rfa-Gaussian.
When the size of \(\phi (\cdot )\)  is too small (32 or 64 for cross attention, 32 for causal attention),
training does not converge.
We observe accuracy improvements by using random features sufficiently large (256 for cross attention and 128 for causal attention);
going beyond that, the benefit is marginal.
<TABLE>
Train and Evaluate with Different Attention Functions
Rfa 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 Rfa-Gaussian model initialized
from a pretrained softmax-transformer;
(ii)
a softmax-transformer initialized
from a pretrained an Rfa-Gaussian model.
If Rfa's good performance can be attributed
to learning a good approximation to softmax,
both, without finetunining,
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 finetuning,
hardly beating a randomly-initialized untrained model.
This result aligns with the observation by [27]},
and suggests that it is not the case
that Rfa performs well because it learns
to imitate softmax attention's outputs.
<FIGURE>Knowledge Transfer from Softmax Attention to RFA
We first supplement the observation in Appendix REF 
by finetuning (i) on the same pretraining data.
Figure REF  plots the learning curves.
It takes Rfa roughly 1,500 steps to reach similar training
loss to the pretrained model.
As a baseline, “Rfa 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 Rfa than randomly initialized ones.
To further confirm 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.
From the above results on IWSLT14,
pretrained knowledge in a softmax transformer
cannot be directly transferred to an Rfa model.
However, from Figure REF  and
a much larger-scale experiment by [27]},
we do observe that
Rfa can recover the pretraining loss,
and the computation cost of finetuning
is much less than training a model from scratch.
This suggests some potential applications.
For example, one might be able to initialize an Rfa language model
from a softmax transformer pretrained on large-scale data
(e.g., GPT-3; [63]}),
and finetune it at a low cost.
The outcome would be an Rfa 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.
