論文解説: SAS — 疎注意の「どのブロックを読むか」を言語モデル損失で端から端まで学習する

長い文脈のどのブロックを読むかを決める「セレクター」は、従来は元モデルの注意分布を真似る蒸留で訓練されていた。SASはセレクターの点数を対数ゲートとしてsoftmaxの中に注入し、Top-Kで止まっていた勾配を通して、次トークン予測の損失だけで選択を学習させる。1024トークン予算でGPQA-Diamondが10.6〜15.5点改善した論文を、4つの設計判断のアブレーションと勾配の式まで追って読む。

SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking

論文の主張を一段落で

今回読むのは Tencent HY LLM Frontier と香港科技大学のチームによる "SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking"（arXiv:2609.13141、2026年9月11日公開）です。abstract の主張はこうです。事前学習済みTransformerの注意計算を減らす「事後の疎化」では、クエリごとに少数の文脈ブロックを選ぶ軽量なセレクターを訓練するが、ハードなTop-K選択が言語モデル損失からの勾配を遮断するため、既存手法は元モデルの層ごとの注意分布を蒸留して代用してきた。しかしその順位付けは「限られた注意予算のもとで予測にどれだけ効くか」とは一致せず、予算を無駄にしうる。SASはセレクターの連続スコアを訓練中の注意ロジットに注入し、標準の逆伝播でセレクターを更新する。うまく動かすには、ゲートを対数の形でsoftmaxの内側に置くこと、softmaxで正規化したゲートで履歴ブロックを常時保持される現在ブロックに対して較正すること、連続スコアを保って相対的な優先順位を学ばせることが重要だった。FlashAttention型のTritonカーネルも実装し、推論・長文理解・エージェント課題で訓練型の疎注意ベースラインを一貫して上回り、特に予算が小さいときの伸びが大きい。

以下では、なぜ勾配が止まるのか、対数ゲートがどうやってそれを通すのか、そして4つの設計判断が実験でどう裏づけられたのかを順に見ていきます。

比喩: 「誰の話を聞くか」を決める係の育て方

Attention機構を1から理解するでは、自己注意を会議室のメモ取りに例えました。全員の名札（Key）と自分の質問（Query）を見比べて、関連の深い人の発言（Value）を長く聞く仕組みです。

ところが会議が長引いて参加者が数万人になると、全員の名札を毎回見比べるのは無理です。そこで「この質問なら、あのあたりの席の人たちだけ聞けばいい」と絞り込む係を置く。これがブロック疎注意のセレクターで、席をまとめた「ブロック」ごとに点数を付け、上位K個だけを注意の対象にします。

問題は、この係をどう育てるかです。従来の代表的なやり方は「先生（元の密な注意）がどこを見ていたか」を真似させることでした。先生の視線を層ごとに記録して、係の点数がそれと合うように訓練する。これが注意分布の蒸留です。

論文が指摘するのは、この育て方のズレです。先生の視線が向いた場所は、「予算K個しか聞けない状況で、最終的な答えを正しくするために本当に必要な場所」とは限りません。理想は、係の選び方が悪ければ答えが間違う、その責任を係に直接負わせることです。SASはまさにそれをやります。係の点数を注意計算の中に組み込み、最終的な次トークン予測の損失から勾配を係まで流します。

前提: ブロック疎注意と、Top-Kが勾配を止める理由

まず論文 §3 の記法を押さえます。現在トークンのクエリを $\mathbf{q}$、過去 $n$ トークンのKeyとValueを $\mathbf{K},\mathbf{V}$ とすると、通常の注意は次のとおりです（スケール係数は論文にならって省略）。

「質問と全部の名札の内積を取り、割合に直し、その割合で全部の発言を混ぜる」という式です。生成時は毎トークンこれをやるので、長さ $n$ の系列全体では累積コストが $n^2$ に比例します。

ブロック疎注意は、$n$ 個の文脈を $b$ トークンずつの $C=n/b$ 個のブロックに分け、セレクター $\mathcal{R}_\theta$ がブロックごとの点数 $\mathbf{s}\in\mathbb{R}^C$ を出し、上位 $K$ 個のブロックの和集合 $\mathcal{S}$ だけを見ます。

「点数の高いK個のブロックを選び、そこに入るトークンだけで注意を計算する」という意味です。1クエリあたりのコストは $n$ から $|\mathcal{S}|$ に落ち、系列全体では $n|\mathcal{S}|$ に比例するようになります（§3）。

ここで肝心なのが、Top-Kは点数に対して区分的に定数だという事実です。点数を少し動かしても、上位K個の顔ぶれが変わらない限り出力は1ミリも動きません。微分がゼロ、つまり「係の点数をどう変えれば損失が下がるか」という情報が損失側から一切届かない。論文はこれを Figure 1(a) の「Gradient Blockage in Discrete Selection」として図示しています。既存手法が蒸留や発見的な代理信号に頼るのは、この壁を迂回するためです（§1, §2）。

蒸留のどこがズレるのか

論文 §1 は、注意分布の蒸留の限界を2つ挙げています。

1つ目は層をまたいだ補完性を見ていないことです。ある層で見逃したブロックが別の層で拾われる、という分担があってもいいのに、層ごとに先生を真似させると各層が独立に「先生と同じ場所」を目指してしまいます。

2つ目はValueを見ていないことです。注意の重みが一致していても、その先で混ぜられる $\mathbf{V}$ の中身が最終予測にどう効くかは蒸留の目的関数に入っていません。図で言えば、Figure 1 の $\mathbf{V}$ 行列がまるごと無視されているわけです。

SASのアイデアは、この2つを一気に解消する最も素朴な方法、つまり最終的な言語モデル損失そのもので係を訓練することです。層をまたいだ分担も、Valueの中身も、最終損失には自然に織り込まれています。あとは勾配の通り道さえ作ればいい。

仕組み: 点数を対数ゲートとしてsoftmaxの中に入れる

論文 §4.1 は、まず疎化を「選択」ではなく順位付け（context ranking）の問題として捉え直します。Top-Kは点数の順序で決まるので、係が学ぶべきは「どのブロックがどれだけ大事か」という連続的な順序です。訓練中はその順序を連続量のまま注意計算に参加させ、推論時に初めてK個に切り落とします。

具体的には、常に保持する現在ブロックを $B_0$、セレクターが順位付けする履歴ブロックを $\mathcal{H}=\{B_1,\ldots,B_C\}$ とし、履歴ブロックの点数 $\mathbf{s}$ を微分可能な活性化 $\phi$ で正のゲートに変換します。

「履歴ブロックには点数から作った正の重みを付け、現在ブロックの重みは1に固定して偏りを与えない」という式です。各ゲート $g_m$ はブロック内の全トークンに同じ値で配られます。

そしてこのゲートを、$\log$ を取って注意ロジットに足し込みます。§4.3 の最終形（式(9)）はこうです。

$\mathbf{g}_{\mathcal{S}}$ は選ばれたブロックのゲートをトークン単位に展開したもの、$\mathcal{S}=B_0\cup\bigcup_{m\in\mathcal{I}}B_m$ は現在ブロックと選ばれた履歴ブロックの和集合です。読み方は「内積の点数表に、係が付けた重みの対数を足してからsoftmaxする」。指数関数の中で足し算をすると外では掛け算になるので、$\exp(z_i+\log g_m)=g_m\exp(z_i)$、つまりsoftmax後の注意確率に係の重みが掛かったのと同じことになります。これで係の点数 $\mathbf{s}$ は注意確率の中に連続的に入り込み、損失からの勾配が $\mathbf{g}$ を経て $\theta$ まで届きます。

推論時は Algorithm 2 のとおり、点数 $\mathbf{s}$ でTop-Kを取り、ゲートは使わず通常の疎注意を計算します。つまり推論時の計算はSeerAttention-Rと完全に同じで、違うのは係の訓練方法だけです（§5.1）。