フォゲッティング・トランスフォーマー：フォゲットゲート付きソフトマックス・アテンション

Forgetting Transformer: Softmax Attention with a Forget Gate

March 3, 2025

著者: Zhixuan Lin, Evgenii Nikishin, Xu Owen He, Aaron Courville

cs.AI

要旨

現代のリカレントシーケンスモデルにおいて重要な要素の一つが忘却ゲートです。Transformerは明示的なリカレント形式を持ちませんが、本研究では、データ依存の方法で正規化されていないアテンションスコアを減衰させることで、Transformerに自然に忘却ゲートを組み込めることを示します。このアテンションメカニズムを「Forgetting Attention」、そしてその結果得られるモデルを「Forgetting Transformer（FoX）」と名付けました。FoXは、長文脈の言語モデリング、長さの外挿、短文脈の下流タスクにおいてTransformerを上回る性能を示し、長文脈の下流タスクではTransformerと同等の性能を発揮します。さらに、FlashAttentionアルゴリズムと互換性があり、位置埋め込みを必要としません。針を干し草の中から探すテストを含むいくつかの分析により、FoXはMamba-2、HGRN2、DeltaNetなどのリカレントシーケンスモデルに対するTransformerの優れた長文脈能力も保持していることが示されています。また、リカレントシーケンスモデルで一般的に使用されるいくつかのアーキテクチャコンポーネントを組み込んだ「Pro」ブロック設計を導入し、これがFoXとTransformerの両方の性能を大幅に向上させることを見出しました。私たちのコードはhttps://github.com/zhixuan-lin/forgetting-transformerで公開されています。

English

An essential component of modern recurrent sequence models is the forget gate. While Transformers do not have an explicit recurrent form, we show that a forget gate can be naturally incorporated into Transformers by down-weighting the unnormalized attention scores in a data-dependent way. We name this attention mechanism the Forgetting Attention and the resulting model the Forgetting Transformer (FoX). We show that FoX outperforms the Transformer on long-context language modeling, length extrapolation, and short-context downstream tasks, while performing on par with the Transformer on long-context downstream tasks. Moreover, it is compatible with the FlashAttention algorithm and does not require any positional embeddings. Several analyses, including the needle-in-the-haystack test, show that FoX also retains the Transformer's superior long-context capabilities over recurrent sequence models such as Mamba-2, HGRN2, and DeltaNet. We also introduce a "Pro" block design that incorporates some common architectural components in recurrent sequence models and find it significantly improves the performance of both FoX and the Transformer. Our code is available at https://github.com/zhixuan-lin/forgetting-transformer.