ICLRFeb, 2024
刺猬与箭猪：具有 Softmax 模拟的表达性线性注意力
The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry
Michael Zhang, Kush Bhatia, Hermann Kumbong, Christopher Ré
TL;DR线性注意力可以提高 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 个百分点的降低。