SAS、Transformerの長文Attentionを予算内で最適化

要約

SASは、言語モデリング損失でcontext rankingをエンドツーエンド最適化する疎Attention機構を提案した。従来手法はhard Top-K選択で勾配が遮断され、元モデルのdense attention重みに近い順位を学習していた。SASはselectorの連続スコアをattention softmax内のlog形式のゲートとして注入し、予測への影響に基づく順位を直接学習する。正規化softmaxゲートで過去のcontextと常時保持するcurrent blockを調整し、Triton kernelでFlashAttention型の長系列処理にも対応した。推論、長文理解、エージェントタスクで既存の学習型疎Attentionを上回り、特に厳しいattention予算で差が大きくなった。