スライディングウィンドウ注意機構は線形注意機構を上回るのか

導入

大規模言語モデルの長文脈推論では、計算量だけでなくメモリ使用量も大きな課題になります。通常の自己注意機構は、文脈が長くなるほど多くのトークン間の関係を扱う必要があります。さらに生成時には、過去のキーとバリューを保持し続けなければなりません。入力が長くなるほど、推論のコストとエネルギー消費は増えていきます。

この問題への候補として、線形注意機構が注目されてきました。注意計算を組み替え、系列長に対する増加をより緩やかにする考え方です。しかし、理論上の計算量が小さくても、既存のモデルへ導入した際に能力を維持できるとは限りません。標準的な注意機構を前提に学習されたモデルを改造するには、ポストトレーニングが必要になる場合があります。

今回の研究は、こうした実用面の比較を行いました。対象は、ポストトレーニングによって線形注意機構へ改造したモデルと、アテンションシンクを組み合わせたスライディングウィンドウ注意機構です。

主なポイント

• 単純な手法を実際の基準として検証した。 新しい仕組みの理論的な美しさではなく、既存モデルの推論コストを下げる方法としてどちらが有効かを比較しています。
• 保持する範囲を限定できる。 スライディングウィンドウでは、直近の一部のトークンに注目するため、全履歴のキーとバリューを無期限に保存する必要がありません。アテンションシンクは、初期部分にある少量の情報を残し、単純な切り捨てによる不安定化を抑えます。
• 追加のポストトレーニングが不要。 報告された結果では、複数の言語モデルと下流タスクで、スライディングウィンドウ方式はポストトレーニング済み線形注意機構と同等、またはそれ以上の性能を示しました。
• 長文脈推論で差が拡大した。 Needle-in-a-HaystackとBABILongでは、スライディングウィンドウ方式の性能が、評価された線形注意モデルの約2～10倍に達しました。
• 導入コストが小さい。 再学習工程を必要とせず、速度が速く、メモリ使用量も少ない点が実装上の利点です。

意義と限界

線形注意機構はスケーリングの改善を期待できますが、情報の集約方法そのものを変更するため、既存モデルの能力を保つには十分な適応が必要です。今回の結果は、効率的な注意機構を評価する際、漸近計算量だけでなく、改造費用、性能の移行、長文脈での信頼性も見るべきだと示しています。

モデル提供者にとって、スライディングウィンドウ方式はより直接的な選択肢になり得ます。モデルアーキテクチャを全面的に置き換えず、追加のポストトレーニングも行わずに、履歴キャッシュの負担を抑えられるからです。長文書の質問応答や文書分析、長時間の生成では特に有用でしょう。

ただし、線形注意機構全体が否定されたわけではありません。今回比較されたのはポストトレーニング済みの方式であり、最初から線形注意を前提に学習したモデルまで含むものではありません。現時点での実務的な結論は、既存モデルの推論メモリを低コストで削減するなら、スライディングウィンドウ注意機構をまず試す価値が高いということです。

出典：Hugging Face Daily Papers