The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry

📄 arXiv: 2402.04347v1 📥 PDF

作者: Michael Zhang, Kush Bhatia, Hermann Kumbong, Christopher Ré

分类: cs.LG, cs.CL

发布日期: 2024-02-06

备注: 30 pages, 20 figures, 15 tables, ICLR 2024

💡 一句话要点

提出Hedgehog以解决线性注意力性能不足问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 线性注意力 Transformer 深度学习 自然语言处理 模型压缩 性能优化 可学习模型

📋 核心要点

  1. 现有线性注意力方法在性能上往往低于标准softmax注意力，导致在实际应用中效果不佳。
  2. Hedgehog通过简单的可训练MLP生成模仿softmax注意力的注意力权重，解决了线性注意力的性能不足问题。
  3. 实验结果显示，Hedgehog在多个基准上显著提升性能，尤其是在WikiText-103和GLUE任务中表现优异。

📝 摘要（中文）

线性注意力在提高Transformer效率方面展现出潜力，将注意力的二次复杂度降低至线性。然而，线性注意力在质量上往往不及标准的softmax注意力。为缩小这一性能差距，本文提出Hedgehog，一种可学习的线性注意力，保留了softmax注意力的低熵和点积单调性特性，同时保持线性复杂度。实验表明，Hedgehog在从头训练和微调转换设置中恢复了超过99%的标准Transformer质量，并在WikiText-103和GLUE基准上显著超越了先前的线性注意力方法。

🔬 方法详解

问题定义：本文旨在解决线性注意力在性能上低于标准softmax注意力的问题。现有的线性注意力方法缺乏softmax注意力的低熵和单调性特性，导致其在任务性能上表现不佳。

核心思路：Hedgehog的核心思路是通过可学习的线性注意力机制，保留softmax注意力的关键特性，同时保持线性复杂度。通过简单的多层感知机（MLP）生成注意力权重，使其能够模仿softmax的行为。

技术框架：Hedgehog的整体架构包括输入特征的处理、可训练MLP生成注意力权重、以及基于这些权重进行的线性注意力计算。该框架确保了计算的高效性和性能的提升。

关键创新：Hedgehog的主要创新在于其能够在保持线性复杂度的同时，保留softmax注意力的低熵和单调性特性。这一设计使得Hedgehog在性能上显著优于以往的线性注意力方法。

关键设计：Hedgehog的关键设计包括使用简单的可训练MLP来生成注意力权重，并在训练过程中优化这些权重以模仿softmax注意力的特性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，Hedgehog在WikiText-103上恢复了超过99%的标准Transformer质量，较之前的线性注意力方法在困惑度上提升了6个点，在GLUE任务上提升了8.7分。此外，将预训练的GPT-2转换为线性注意力变体后，达到了16.7的困惑度，创造了新的性能记录。

🎯 应用场景

Hedgehog的研究成果在自然语言处理、计算机视觉等多个领域具有广泛的应用潜力。通过提高Transformer模型的效率和性能，Hedgehog可以在大规模数据处理和实时应用中发挥重要作用，推动智能系统的进一步发展。

📄 摘要（原文）

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.