Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: This paper contributes to the literature on efficient [transformer][]
architectures. Specifically, papers that focus on reducing the quadratic
computational complexity of the self-attention mechanism between tokens.

The problem this paper focuses on is the incompatibility of [relative
positional encodings (RPE)][rpe] and sub-quadratic efficient attention
mechanisms. It adapts the kernelized attention used in the [Performer][]
architecture to allow the use of RPE. In the process, the authors find that
this adaptation allows use of the FFT algorithm to perform the approximate
attention mechanism, resulting in an $O(N \log N)$ algorithm, where $N$ is
the number of tokens.

Benchmarks are provided on the GLUE dataset according to various metrics
and on the ImageNet dataset using Image Transformer variants. Performance
to competing methods chosen is better on most of the metrics presented.

[transformer]: https://papers.nips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf
[rpe]: https://arxiv.org/abs/1803.02155
[performer]: https://arxiv.org/abs/2009.14794
main_review: ## Originality

> Are the tasks or methods new? Is the work a novel
> combination of well-known techniques? (This can be valuable!) Is it clear
> how this work differs from previous contributions? Is related work
> adequately cited?

This work is a novel combination of two existing methods: [relative
positional encodings][rpe] and [kernelized efficient
self-attention][performer]. In the process they make an independent
contribution, discovering that combining the two methods enables use of the
FFT algorithm to compute the efficient self-attention mechanism. This fills
its own niche among the many existing efficient self-attention mechanisms.

Related work cited is a thorough review of efficient self-attention
methods, with each correctly sorted into their respective types. The review
makes it clear where this method fits into this set.

In addition the authors make various observations that will be useful to
machine learning practitioners:

1. Ablation comparisons of [Performer][] and [Random Feature
Attention][rfa]
2. Analysis of the function class of transformers using [RPE][]
3. Analysis of the approximation error of kernelized attention
    1. Fourth paragraph of Section 3.3 makes a useful observation about the
training dynamics of transformers using kernelized attention
    2. This is further put into context in the literature by the footnotes
in the following paragraph on "Normalized Kernel Attention with RPE"
4. Address the limitations of existing kernelized attention mechanisms,
such as [Performer][], that are unstable when trained from scratch by using
a normalization step

[rpe]: https://arxiv.org/abs/1803.02155
[performer]: https://arxiv.org/abs/2009.14794
[rfa]: https://openreview.net/forum?id=QtTKTdVrFBB

## Quality

> Is the submission technically sound? Are claims well supported
> (e.g., by theoretical analysis or experimental results)? Are the methods
> used appropriate? Is this a complete piece of work or work in progress? Are
> the authors careful and honest about evaluating both the strengths and
> weaknesses of their work?

All claims in the paper are supported by theoretical inquiry and
experiment. The experimental results presented demonstrate the claims
made in the text. On the benchmarks chosen (GLUE, machine translation,
ImageNet) the network performs well in comparison to the methods chosen for
comparison.

The technical analysis is in depth and considers a limitation of the work.
In a footnote of Section 3.2, the authors mention that the method cannot be
used to improve the efficiency of the inference step. In practice the
memory explosion of the kernel matrix during training is a larger problem
when training transformers than the inference step memory usage. But, it
may affect this method's ability to scale to large sequence lengths.

Lemma 2 in particular is an important observation about the variance of
kernelized attention and is backed up in the paper by theoretical analysis
and precise experimental exploration illustrated in Figure 1. Both of these
together are very effective in convincing the reader that this Lemma is
correct and worthwhile.

[performer]: https://arxiv.org/abs/2009.14794

## Clarity

> Is the submission clearly written? Is it well organized? (If not, please
> make constructive suggestions for improving its clarity.) Does it
> adequately inform the reader? (Note that a superbly written paper provides
> enough information for an expert reader to reproduce its results.)

The submission is clearly written and well organised. Reading this paper in
sequence each topic is introduced in a way that is easy to understand and
each idea follows from the previous. I am confident that any researcher
familiar with transformers will be able to understand this paper.

Section 2 deserves praise for providing a concise and useful description of
kernelized attention. It helps to make the contribution of the paper clear
to the reader by drawing together various prior work in kernelized
attention under a common framework. I would use it as a reference if I were
writing about kernelized attention in future.

Some minor suggestions I would make:

1. The related work is discussed in the Introduction, which may not draw
the eye for those searching for it.  However, putting it in it's own
section would ruin the flow and probably push the paper over the page
limit. I would suggest putting a `\paragraph` bold title indicating that
paragraph 2 of the Introduction is the Related Work section but if the
author's find it worse then I would not argue.
2. Section 3 mentions the authors will "develop new techniques" but does
not indicate what those techniques are. It would be better to be precise
about what the new techniques are.
3. Sentence 3 of Section 3.3 has a grammar mistake that impedes
readability, "those kernelized attention" should probably be "kernelized
attention"

## Significance

> Are the results important? Are others (researchers or practitioners) likely
> to use the ideas or build on them? Does the submission address a difficult
> task in a better way than previous work?  Does it advance the state of the
> art in a demonstrable way? Does it provide unique data, unique conclusions
> about existing data, or a unique theoretical or experimental approach?

The specific niche this work fits into in the literature on efficient
self-attention mechanisms is useful. This seems like a natural question
when considering [RPE][] and [kernelized attention][performer]. 

In addition, the demonstration of how the FFT algorithm may be used to
perform a self-attention mechanism appears valuable. It seems like there
may be ways to build upon it, such as incorporating it with other relative
positional encodings, such as [rotary encodings][rotary] and the result may
continue to improve performance. There are likely other ways this work
could be adopted and improved into Transformer design.

[rpe]: https://arxiv.org/abs/1803.02155
[performer]: https://arxiv.org/abs/2009.14794
[rfa]: https://openreview.net/forum?id=QtTKTdVrFBB
[rotary]: https://arxiv.org/abs/2104.09864
limitations_and_societal_impact: Efficient self-attention methods with sub-quadratic computational
complexity are typically applied to problems that have a large number of
tokens. An example of this can be seen in the experiment in Section 4.5 of
the [Performer][] paper, on the ImageNet64 task that cannot be addressed by
traditional transformers. This paper does not include experiments with
numbers of tokens that would be infeasible for a normal transformer.

The authors discuss how normalizing the kernelized attention reduces the
capacity of the model to express sharp attention distributions in the final
paragraph of Section 3.3. As the goal of this work is not to lose the extra
capacity offered by [RPE][] when using kernelized attention it is a
limitation that the authors do not present an experiment investigating how
the model performs without normalization and with RPE (ie PRF w/ RPE in the
notation of the paper's experiments). It may be necessary to address the
stability using the same tricks employed by [Performer][] to run this
experiment.

The FFT algorithm proposed to perform the Toeplitz matrix multiplication
receives little discussion, save that it is possible. I think this is
appropriate but there are some concerns that could be addressed:

1. FFT algorithms are not well supported in all deep learning frameworks
and FFT Toeplitz matrix multiplications are not supported directly
2. FFT algorithms, unless carefully implemented, require the inputs to be a
power of 2 number of dimensions
3. Implementing this using FFT operations in an autograd system will have a
relatively large number of intermediate variables that will be stored for
the backward pass and could increase memory usage. This could be addressed
with custom operations that define the appropriate backward pass directly.

[performer]: https://arxiv.org/abs/2009.14794
[rpe]: https://arxiv.org/abs/1803.02155
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The authors propose a method to compute kernelized attention with relative position encoding (RPE) efficiently in $O(n\log n)$. They show that the RPE family of attention is richer than attention without RPE, and the usual $O(n)$ implementation of kernelized attention is no longer possible.  The derivation gives rise to a convolution, which the authors propose to compute in the Fourier domain, providing $O(n\log n)$ complexity, in contrast to $O(n^2)$  for a naive implementation.

The paper provides a comprehensive derivation of the algorithm, mathematically and in pseudo code. Interestingly, it is found that with the help of RPE training stability of kernelized attention can be improved, thus allowing for the first time to train certain models end-to-end.

Experiments with random toy data where performed to assess the computational load of the algorithm, and experiments on datasets coming from language pre-training, machine translation, and image classification are used to show the effectiveness and the improved training stability of the algorithm on real world problems



 

main_review: ## Originality
The paper brings three well known techniques together: kernelized attention, RPE, and the Convolution Theorem, to enable kernelized RPE attention. It's a nice idea that leads to an elegant algorithm. To my knowledge it is the first time that this has been proposed. All prior work has been cited appropriately.

## Quality 
The overall quality of the paper is very good. The derivations are sound and the experiments support the claims made by the authos.


There are two things missing that I think should be discussed: 

* How does the model behave with causal attention. I think it can be implemented due to the RPE by setting $c_i=0$ for $i\ge 0$ in (11) and (12). Has this been done in the experiments and found to work? If so, then I think footnote 1 is not a limitation unique to this algorithm. Sequence generation at inference time always gives another factor $n$, and so does it here. The convolution does not improve on that, bat at least it does not make things worth (at least compared to non-kernelized attention) 

* For the *normalized kernelized attention with RPE* ( l 203 - 209) it should be noted that this is not the same as *unnormalized kernelized attention*. You can produce sharp distributions but only at specific distances for which for which $b_{i-j}$ happens to be large. So NPRF w/ RPE has lower expressivity than PRF, still showing better performance (Tab. 3). It would add to the strength of the paper if the authors could elaborate on that.

## Clarity
The paper is very well written. The mathematical derivation is comprehensive an easy to follow. All important information is in the main paper

## Significance
RPE have been shown to have a big impact on the performance of attention based models on various problems. On the other hand, kernelized attention can save a big deal of computational cost and make transformer models scalable to real world large scale problems. To link both approaches in an efficient way clearly advances the field and has a big impact for practitioners working with transformer models.  

**I vote for accepting the paper.**

Minor issues
* l 79: indicate the axis of the softmax operation
* Sec. 2.1 ff:  all $x_i$ are row vectors which is unusual and can lead to confusion, in particular when writing $x=(x_1,\ldots,x_n)$. You should either revise the notation or state explicitly that $x_i\in\mathbb R^d$ is a row vector 
* l 88 - 94: wherever you have running variables $j'$ in summations, make sure to also use $j'$ in the summation terms.
* Fig. 1a would gain from one more sequence length step ($n=48k$ or $n=64k$) showing the (almost) linear relation better
 


limitations_and_societal_impact: The authors don't explicitly address limitations of their model, other than additional speedup for sequence generation is not provided.
Another potential limitation, that should be mentioned, could come from the fact that NPRF even with RPE cannot generate sparse attention to the same extend as PRF or plain dot-then-exponentiate attention  (see above)

Although not explicitly addressed in the paper,  I don't see any negative societal impact of the work.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper introduces a novel way to add relative positional encoding to efficient transformers with kernelized attention. The authors make use of the efficient matrix multiplication of Toeplitz matrices using the fast Fourier transform.  The authors perform several experiments on masked language modeling, machine translation and even vision transformers.
main_review: Strengths
-------------

- Using the fast multiplication of Toeplitz matrices based on the Fourier transform to implement relative positional encodings is brilliant and very interesting
- The experimental evaluation shows that kernelized transformers with the proposed RPE can actually train and often achieve better results than the same model without RPE
- The theoretical analysis regarding the norm of the queries and keys as well as the ablations and analyses in section 4.4 are great and of independent interest. In particular, I find Figure 2 very informative that concisely showcases both the fact that reducing the norm of the queries and keys helps the random Fourier feature approximations as well as that RPE accounts for a large part of the attention distribution (when used) and helps even more when approximating.
- The authors provide clean code with their submission

Weaknesses
-----------------

### Asymptotic complexity vs real world speed

The main weakness of the paper stems from the misrepresentation of the complexity of the attention computation using the proposed algorithm. The algorithm in its provided form, requires storing in memory a matrix of size $O(ND^2)$ (namely $\phi(K_i) V_i^T$ for all $i$) per sample per head. This is the same issue as showcased in [15] for autoregressive kernelized attention. The result is a theoretically fast algorithm but a slow one in practice.

This can be seen also from Figure 1a. Even at 32k sequence lengths and using only the forward pass, the proposed method needs to down-project the queries and keys to half the dimensions in order to be faster than the **full softmax** attention. In comparison, the simplest kernelized attention by [15] is **10x times** faster in the same experiment. Down-projecting using PRF could make it even faster.

Finally, evaluation of the computational cost is not provided even for a single one of the real world experiments. For instance what is the computational cost of an epoch in GPU wall clock time for the models in tables 1, 2 and 3. Since the proposed model is an approximation of the softmax transformer it should at least be verified that it is faster than that.

### Minor comments

1. The wording and notation in section 3.2 is unnecessarily complicated. The $vec(\cdot)$ notation can be omitted and if needed intermediate matrices can be properly defined together with their dimensions.
limitations_and_societal_impact: The authors do not discuss the societal impact of their work.

The limitations are not adequately discussed as mentioned in the weaknesses section of the main review.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: There are two major contributions of this paper. First, it proposes a method to incorporate relative positional embeddings (RPE) into kernelized attention, which is based on Fast Fourier Transformation (FFT) and Toeplitz matrix multiplication and leads to a $O(n \log n)$ time complexity compared to $O(n^2)$ using the conventional self-attention with RPE. Second, it proposes to normalize the queries and keys by their `$\ell_2$ norms, which makes the optimization more stable.
main_review: First, I summarize the pros and cons of this paper.

Pros:
1. The proposed method is simple and novel. 
2. Adding the proposed RPE improves the performance over the linear kernelized attentions.
3. The authors conduct experiments on a variety of tasks: language pre-training, image classification, and machine translation and show that the proposed method is more accurate than the kernelized attention baselines.
4. The authors demonstrate that the proposed normalized kernelized attention is not more more accurate and more stable.


Cons:
1. No inference time provided on real datasets (See Limitations And Societal Impact for details)
2. Not showing the performance-efficiency tradeoff (See Limitations And Societal Impact for details)
3. The proposed method operates at only $O(n \log n)$ time, but the authors do not compare with simple linear baselines which has relative positional information like just using convolutions [1] or a combination of convolutions and kernelized attention. This makes me wonder why we should prefer the proposed $O(n \log n)$ method rather than a $O(n)$ baseline.
4. Figure 3 (b) is very hard to read. It would be better to show y-axis between 30 and 35 or 33 and 35.

Overall, the paper is written in high quality and I recommend accept it; however, there are a couple of critical concerns which hinder me from rating this paper higher.

Reference
[1] Tay, Yi, et al. "Are Pre-trained Convolutions Better than Pre-trained Transformers?." ACL 2021.

limitations_and_societal_impact: While I believe reducing the time complexity is impactful, the impact of the paper is sort of limited by the experimental results it presents for the following two reasons:
1. The authors only show the forward speed in Fig. 1 (a). In all tables in Sec. 5, the authors only show the time complexity instead of the real inference time on the real datasets. This worries me that this paper may be yet another paper with great time complexity but relatively slow inference time on real use cases.
2. While the paper is proposing an efficient alternative to self-attention, The authors do not show the performance-efficiency trade-off. I wish I could see that the method is at the Pareto frontier. For example, showing with a similar inference speed, this method outperforms linear attention and softmax attention. Linear attentions have $O(n)$ time complexity, so it fairer to compare with a larger model with linear attention and the same inference time as the proposed PRF + RLE method. Similarly, comparing the proposed method to a smaller softmax attention model with the same inference time is also expected.

I am willing to raise my score to 8 if both of these concerns are addressed.
review_title: null

Meta-review: This is an interesting paper extending linear attention mechanisms proposed in Performers to work with an *arbitrary* relative position encoding (RPE) mask. Previous results showed that it is possible for special instantiations of RPE. This paper shows that as long as a mask has Toeplitz structure (this can be actually further extended to a low displacement rank structure), efficient computation of the attention module can be achieved by combining random feature map tricks with Fast Fourier Transformer mechanism (every RPE is by definition a Toeplitz matrix). The authors furthermore claim, that by incorporating RPE into linear attention modules, they can substantially improve their overall accuracy. According to them, this is possible since by adding RPEs, one can constrain queries and keys to be L2-normalized (resulting in the substantial reduction of the variance of the random feature map estimator that grows exponentially in the lengths of keys/queries). Normally such a normalization would affect the expressiveness of the model, but according to the authors, the RPE (which in principle is not bounded) provides extra expressiveness that would be lost otherwise. 
This claim is controversial and additional detailed experiments showing that RPE can indeed carry out the otherwise-lost expressiveness would further strengthen the paper (in particular an extensive ablation studies over L2-normalized variants with and without RPE would provide extra clarity).  The paper would also benefit from providing more theoretical understanding of this phenomenon.
The other important point that should be discussed in more detail is an efficient implementation on TPUs of Fast Fourier Transform. Unless such an implementation is provided, practical usage of this algorithm will be limited. 
Nevertheless, it is an important result in the field and an elegant approach to incorporating structured masks into kernelizable attention modules. This algorithmic observation is the main contribution of the paper.
Decision: Accept (Poster)