Linear attentions have shown potential for improving Transformer efficiency,
reducing attention's quadratic complexity to linear in sequence length. This
holds exciting promise for (1) training linear Transformers from scratch, (2)
"finetuned-conversion" of task-specific Transformers into linear versions that
recover task performance, and (3) "pretrained-conversion" of Transformers such
as large language models into linear versions finetunable on downstream tasks.
However, linear attentions often underperform standard softmax attention in
quality. To close this performance gap, we find prior linear attentions lack
key properties of softmax attention tied to good performance: low-entropy (or
"spiky") weights and dot-product monotonicity. We further observe surprisingly
simple feature maps that retain these properties and match softmax performance,
but are inefficient to compute in linear attention. We thus propose Hedgehog, a
learnable linear attention that retains the spiky and monotonic properties of
softmax attention while maintaining linear complexity. Hedgehog uses simple
trainable MLPs to produce attention weights mimicking softmax attention.
Experiments show Hedgehog recovers over 99% of standard Transformer quality in
train-from-scratch and finetuned-conversion settings, outperforming prior
linear attentions up to 6 perplexity points on WikiText-103 with causal GPTs,
and up to 8.7 GLUE score points on finetuned bidirectional BERTs. Hedgehog also
enables pretrained-conversion. Converting a pretrained GPT-2 into a linear
attention variant achieves state-of-the-art 16.7 perplexity on WikiText-103 for
125M subquadratic decoder models. We finally turn a pretrained Llama-2 7B into
a viable linear attention Llama. With low-rank adaptation, Hedgehog-Llama2 7B
achieves 28.1 higher ROUGE-1 points over the base standard attention model,
where prior linear attentions lead to 16.5 point drops.

线性注意力可以提高 Transformer 的效率，将注意力的二次复杂性降低到与序列长度线性相关。线性注意力在训练线性 Transformer、将任务特定的 Transformer 转换为线性版本以恢复任务性能以及将大型语言模型转换为可在下游任务上进行微调的线性版本等方面具有潜力。然而，线性注意力在质量上经常不如标准的 softmax 注意力。为了弥合这一性能差距，我们发现之前的线性注意力缺少与优秀性能相关的低熵（或 “尖峰”）权重和点积单调性等关键特性。我们进一步观察到保留这些特性并且与 softmax 性能相匹配的非常简单的特征映射，但在线性注意力中计算效率低下。因此，我们提出了 Hedgehog，它是一个可学习的线性注意力，保留 softmax 注意力的尖峰和单调性，同时保持线性复杂性。Hedgehog 使用简单的可训练的多层感知机（MLP）生成模仿 softmax 注意力的注意力权重。实验证明，Hedgehog 在从头训练和微调转换设置中恢复了超过 99% 的标准 Transformer 质量，在 WikiText-103 上通过因果 GPT 模型的困惑度高达 6 点，在微调的双向 BERT 上 GLUE 分数高达 8.7 点。Hedgehog 还实现了预训练转换。将预训练的 GPT-2 转换为线性注意力变体，在 125M 的次二次解码模型上，在 WikiText-103 中实现了 16.7 困惑度的最佳结果。最后，我们将预训练的 Llama-2 7B 转换为可行的线性注意力 Llama。通过低秩适应，Hedgehog-Llama2 7B 相比于基本标准注意力模型，ROUGE-1 分数提高了 28.1 个百分点，而之前的线性注意力则导致了 16.5 个百分点的降低。