SAS、予測結果に直結するコンテキスト順位付けで疎な注意機構を改善

長いコンテキストを扱う Transformer では、注意機構の累積計算量が系列長に対して二次的に増加します。そのため、各クエリが参照するトークンやブロックを少数に絞る疎な注意機構が研究されています。しかし、計算対象を減らすだけでは十分ではありません。固定された注意予算の中で、どのコンテキストを残すかを正しく順位付けする必要があります。

既存手法が抱えるずれ

一般的な学習型手法では、軽量な選択器がコンテキスト単位にスコアを付け、その後にハードな Top-K 選択を行います。これは予算を明確に制御できますが、離散的な選択によって言語モデリング損失の勾配が選択器に届きにくくなります。その結果、選択器は元の密なモデルが各層で示す注意分布を模倣するように学習されることが多くなります。

密な注意重みの再現は有用な近似ですが、限られた予算の下で最終予測に最も貢献するコンテキストを直接最適化しているわけではありません。密なモデルで高い重みを持つ要素が、他の要素を削除した後にも最も価値が高いとは限らないからです。保持できる単位が少ないほど、この順位のずれは大きな問題になります。

SAS の仕組み

Simple Attention Sparsification、略して SAS は、このずれを比較的単純な構成で解消しようとします。学習時、選択器は連続的なスコアを生成し、その値をゲートとして注意 logits に注入します。特に、ゲートを対数形式で注意 Softmax の内部に配置する点が重要です。これにより、言語モデリング損失が通常の逆伝播を通じて選択器を直接更新できます。

論文では、実用上重要な設計として次の点を挙げています。

• Softmax 内の対数ゲート: 選択処理を別の不可微分な枝刈り段階にせず、注意の正規化と一体化します。
• 正規化された Softmax ゲート: 過去のコンテキストは選択対象ですが、現在のブロックは常に保持されます。両者のスケールを調整し、どちらか一方が不当に支配することを抑えます。
• 連続スコアの維持: 単純な保持・破棄のラベルではなく、候補間の相対的な優先度を学習します。予算が変わる場合にも順位情報を利用しやすくなります。

長系列学習への対応として、著者らは SAS を FlashAttention 風の計算に統合するメモリ効率の高い Triton カーネルも実装しました。これにより、選択器の設計だけでなく、実際の学習時の実行コストも考慮されています。

意義と今後の見方

論文によれば、SAS は推論、長コンテキスト理解、エージェント系タスクにおいて、学習型の疎注意ベースラインを上回り、特に注意予算が厳しい条件で大きな改善を示しました。重要なのは、疎注意の価値が削減率だけでなく、最終予測に対するコンテキストの有用性を順位付けできるかどうかで決まるという点です。

この考え方は、長コンテキスト推論、コンテキスト圧縮、高効率な Transformer 推論に示唆を与えます。一方、今回の素材にはモデル規模、データセット、圧縮率、具体的なスコアは含まれていません。そのため、改善幅や適用範囲については、完全な論文の実験条件と合わせて判断する必要があります。

出典：Hugging Face Daily Papers