The quadratic complexity of attention in transformer architectures remains a
big bottleneck in scaling up large foundation models for long context. In fact,
recent theoretical results show the hardness of approximating the output of
softmax attention mechanism in sub-quadratic time assuming Strong Exponential
Time Hypothesis. In this paper, we show how to break this theoretical barrier
by replacing softmax with a polynomial function and polynomial sketching. In
particular we show that sketches for Polynomial Kernel from the randomized
numerical linear algebra literature can be used to approximate the polynomial
attention which leads to a significantly faster attention mechanism without
assuming any sparse structure for the attention matrix that has been done in
many previous works.
In addition, we propose an efficient block-based algorithm that lets us apply
the causal mask to the attention matrix without explicitly realizing the $n
\times n$ attention matrix and compute the output of the polynomial attention
mechanism in time linear in the context length. The block-based algorithm gives
significant speedups over the \emph{cumulative sum} algorithm used by Performer
to apply the causal mask to the attention matrix. These observations help us
design \emph{PolySketchFormer}, a practical linear-time transformer
architecture for language modeling with provable guarantees.
We validate our design empirically by training language models with long
context lengths. We first show that the eval perplexities of our models are
comparable to that of models trained with softmax attention. We then show that
for large context lengths our training times are significantly faster than
FlashAttention.

该研究论文通过使用多项式函数和多项式取样法，替代基于 softmax 的注意力机制，以及提出基于区块的算法，实现了一个能够在线性时间内计算多项式注意力机制的实用线性变换器架构，从而显著提高了大规模上下文下的注意力机制的速度。