ハリネズミとヤマアラシ：Softmax模倣による表現力豊かな線形アテンション

要旨

線形アテンションはTransformerの効率化において潜在能力を示しており、アテンションの二次的な計算量をシーケンス長に対して線形に削減します。これは、(1)線形Transformerをゼロから訓練する、(2)タスク固有のTransformerを線形バージョンに「ファインチューニング変換」し、タスク性能を回復する、(3)大規模言語モデルなどのTransformerを線形バージョンに「事前学習変換」し、下流タスクでファインチューニング可能にする、といったエキサイティングな可能性を秘めています。しかし、線形アテンションは品質面で標準的なソフトマックスアテンションに劣ることが多いです。この性能差を埋めるため、従来の線形アテンションには、良好な性能に関連するソフトマックスアテンションの重要な特性、すなわち低エントロピー（または「スパイキー」）な重みと内積単調性が欠けていることがわかりました。さらに、これらの特性を保持しつつソフトマックス性能に匹敵する驚くほど単純な特徴マップを観察しましたが、線形アテンションでは計算効率が悪いことが判明しました。そこで、ソフトマックスアテンションのスパイキー性と単調性を保持しつつ線形計算量を維持する学習可能な線形アテンション「Hedgehog」を提案します。Hedgehogは、ソフトマックスアテンションを模倣するアテンション重みを生成するために単純な学習可能なMLPを使用します。実験では、Hedgehogがゼロからの訓練およびファインチューニング変換の設定において標準Transformerの品質の99%以上を回復し、因果的GPTモデルではWikiText-103において従来の線形アテンションを最大6パープレキシティポイント、双方向BERTモデルではGLUEスコアで最大8.7ポイント上回りました。Hedgehogは事前学習変換も可能にします。事前学習済みGPT-2を線形アテンションバージョンに変換することで、125MパラメータのサブクアドラティックデコーダモデルにおいてWikiText-103で最先端の16.7パープレキシティを達成しました。最後に、事前学習済みLlama-2 7Bを実用的な線形アテンションLlamaに変換しました。低ランク適応を用いることで、Hedgehog-Llama2 7Bはベースの標準アテンションモデルに対して28.1ポイント高いROUGE-1スコアを達成し、従来の線形アテンションでは16.5ポイントの低下が生じていました。

English

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.

ハリネズミとヤマアラシ：Softmax模倣による表現力豊かな線形アテンション

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

要旨

Support