刺猬与豪猪：具有 Softmax 模仿的表达性线性注意力

线性注意力已展现出提升 Transformer 效率的潜力，将注意力机制在序列长度上的二次复杂度降低为线性复杂度。这为 (1) 从头训练线性 Transformer，(2) 将特定任务的 Transformer 进行“微调转换”以恢复任务性能的线性版本，以及 (3) 将诸如大型语言模型等 Transformer 进行“预训练转换”为可在下游任务上微调的线性版本，带来了激动人心的前景。然而，线性注意力在质量上通常不如标准的 softmax 注意力。为了缩小这一性能差距，我们发现先前的线性注意力缺乏与良好性能相关的 softmax 注意力的关键特性：低熵（或“尖锐”）权重和点积单调性。我们进一步观察到，保留这些特性并能匹配 softmax 性能的惊人简单的特征映射，在线性注意力中计算效率低下。因此，我们提出了 Hedgehog，一种可学习的线性注意力，它在保持线性复杂度的同时保留了 softmax 注意力的尖锐和单调特性。Hedgehog 使用简单的可训练 MLP 来产生模仿 softmax 注意力的注意力权重。实验表明，在从头训练和微调转换设置中，Hedgehog 恢复了超过 99% 的标准 Transformer 质量，在使用因果 GPT 的 WikiText-103 上比先前的线性注意力最多高出 6 个困惑度点，在使用微调双向 BERT 时最多高出 8.7 个 GLUE 分数点。Hedgehog 还支持预训练转换。将预训练的 GPT-2 转换为线性注意力变体，在 WikiText-103 上实现了 125M 次二次方解码器模型中最先进的 16.7 困惑度。我们最终将预训练的 Llama-2 7B 转化为可用的线性注意力 Llama。通过低秩适应，Hedgehog-Llama2 7B 比基础标准注意力模型高出 28.1 个 ROUGE-1 点，而先前的线性注意力则导致 16.5 个点的下降。