SWAA: 効率的かつ品質を維持するための長文処理に向けたスライディングウィンドウ注意の適応（Sliding Window Attention Adaptation）

Abstract

TransformerベースのLLMにおける自己注意の二次計算量は、長いコンテキスト推論を非常に高コストなものにしています。最も単純な疎注意パターンであるSliding Window Attention（SWA）は線形計算量の代替手段を提供しますが、それでも壊滅的な長コンテキスト性能の崩壊に悩まされます。この崩壊は、2つの根本的な要因に起因します。すなわち、(1) Full Attention（FA）で事前学習されたモデルに対して何の工夫もなくSWAを適用する際の学習・推論不一致、そして、(2) SWAを常時すべてのモジュールに適用することによる、遠距離情報へのアクセスに関する構造的な必然的な制約です。これら2つの課題に対処するために、私たちは、コストのかかる事前学習なしでFAモデルをSWAへ適応させる、プラグ・アンド・プレイ型のレシピ集であるSliding Window Attention Adaptation（SWAA）を提案します。SWAAは、これらの異なる問題に対処するために4つの中核戦略を体系的に組み合わせます：(1) Full Attention（FA）デコードと(2) FA層とSWA層のインタリーブ（交互挿入）であり、選択的に遠距離情報へのアクセスを許可することで構造的な欠陥を緩和します。さらに、(3) ``sink''トークンの保持と(4) 軽量な微調整により、学習・推論不一致を緩和します。実験の結果、個別の戦略だけでは不十分ですが、特定の相乗的な組み合わせによって長コンテキスト性能が効果的に回復することがわかりました。計算オーバーヘッドは変動するものの、性能効率のトレードオフ分析により、多様なシナリオに対する最適なSWAA構成を特定でき、品質の許容可能な維持を伴いながら長コンテキスト推論で30%〜100%の高速化を達成します。コード、データ、モデル重みは https://github.com/yuyijiong/sliding-window-attention-adaptation で利用可能です