ズイウェイ・リー (Zhiwei Li)氏らは9月11日(現地時間)、事前学習済みTransformerの累積Attention（アテンション）コストを削減する新たなスパースAttentionメカニズム「SAS (Simple Attention Sparsification)」を提案した。これは、同日付けでarXiv cs.CLに掲載された論文で報告されたもの。既存手法が抱えていたコンテキストランキングと予測への影響の不整合を解消し、言語モデリングロスとのエンドツーエンド最適化を実現することで、計算効率と性能の向上を図る。
Transformerは、そのAttentionメカニズムによりシーケンス長に対し二次関数的に計算コストが増大するという課題を持つ。この計算負荷を軽減するため、ポストトレーニングAttention Sparsification（スパース化）という手法が注目されている。これは、Transformerモデルにおいてクエリごとに少数の関連性の高いコンテキスト単位を選択することで、累積するAttentionコストを効果的に削減することを目的としていると、ズイウェイ・リー氏らの論文は指摘している。
既存の学習可能なスパースAttention手法は、軽量なセレクターがコンテキスト単位を評価し、その後のハードなTop-K選択で最も高いスコアを持つK個の単位を選別する構造を持つ。しかし、このハードな選択は言語モデリングロスからの勾配（学習信号）を遮断し、セレクターの最適化を妨げていた。結果として、コンテキスト単位のランキングが、固定されたAttention予算（クエリごとにAttention対象となるコンテキスト単位の数）における予測への直接的な影響と整合しないという深刻な課題が生じていた。これは、限られた計算予算が、モデルの予測にほとんど寄与しない有用性の低いコンテキスト単位に不必要に浪費される可能性を意味した。
この根本的な不整合に対処するため、ズイウェイ・リー氏らは、言語モデリングロスとエンドツーエンドでコンテキストランキングを直接最適化するゲート付きスパースAttentionメカニズムであるSASを考案した。SASの主要な概念は、トレーニング中にセレクターの連続的なスコアをAttentionロジットに直接注入し、標準的なバックプロパゲーション（誤差逆伝播）を通じて言語モデリングロスがセレクターの重みを更新できるようにすることにある。これにより、セレクターがモデルの最終的な予測性能に与える影響を直接的に学習し、最適なコンテキスト選択を可能にする。
このシンプルな設計を実用的に機能させるためには、いくつかの重要な技術的選択肢が特定された。具体的には、セレクターから供給されるゲートをAttentionソフトマックスの対数形式の内部に戦略的に配置すること、正規化されたソフトマックスゲートを用いることで、履歴コンテキストを常に保持される現在のブロックと効果的にキャリブレーションすること、そして連続的なセレクタースコアを保持し、モデルがハードな選択だけでなく、コンテキスト間の相対的な優先順位も学習できるようにすることが挙げられる。これらの要素が統合されることで、SASは既存手法の限界を乗り越える。
長シーケンス学習を効率的にサポートするため、SASはFlashAttention（フラッシュアテンション）スタイルの計算アーキテクチャに統合された、メモリ効率の高いTriton（トリトン）カーネルとして実装された。この実装により、大規模モデルや長文入力に対してもスケーラブルなパフォーマンスを実現する。推論段階での効率性、長文コンテキスト理解タスク、およびエージェントタスクなど、様々な評価シナリオにおいて、SASは既存の学習可能なスパースAttentionベースラインをAttention予算全体で一貫して上回る性能を示したという。特に厳しい予算条件下では、SASは大幅な性能向上を達成し、ダウンストリームタスクにおいてより効果的なコンテキストランキングが実証された形となった。
参考: arXiv cs.CL (アーカイブ) — 2026年9月12日 02:58 (JST)
原文ハイライト"Simple Attention Sparsification via End-to-End Optimization of Context Ranking"