DoSA：選択結果を保ちながら疎なAttentionの参照先探索を高速化

長い文書やコードを入力したとき、疎なAttentionは、使う参照先を少数に絞る。その選別を担うのがIndexerだ。ここでは、いま処理している側の表現をクエリ、参照候補の表現をキーと呼ぶ。DoSAが対象にするDSA型のIndexerでは、少数のキーを選ぶために、全キーを複数のヘッドで採点する計算が残る。ヘッドは、それぞれ異なる方向から候補を評価する単位と考えればよい。

たとえば、長いコードを入力してレビューさせる場面を考える。モデルが使う参照先は絞られていても、その候補を探す段階で入力全体を詳しく採点していれば、選別自体が重くなる。この例は計算の流れを理解するためのもので、DoSAがレビューの正確さを高めるという意味ではない。論文が検証する中心は、元の選択を保ちながら選別を速くできるかだ。

比較対象のSGLangは、全ヘッドによる採点後に、上位K個を選ぶTop-Kを行う。LiteTopKは採点、候補のフィルタ、選択をまとめ、中間結果の書き出しを減らす。それでも全キーを全ヘッドで評価する仕事は残る。

少数のヘッドだけで候補を絞り、残った候補を詳しく採点する近似も考えられる。しかし、最初の段階で本来の上位候補を落とすと、その後に完全採点しても取り戻せない。DoSAが変えるのは除外の根拠である。少数ヘッドで低く見えたという理由ではなく、残りを考慮しても上位に届かないと証明できた場合だけ、完全採点を省く。

対象となる元のスコアは、次の形を持つ。

ここで、$s_i$は候補iのスコア、hはヘッド、$w_h$はその重み、$q_h$はヘッドごとのクエリ、$k_i$は候補キーである。内積$q_h^{\mathsf T}k_i$は両者の対応を測り、$[x]_+$は負の値をゼロにするReLUを表す。重みには正と負があり、活性化したヘッドがスコアを加算する場合と、減算する場合がある。

この符号の違いがDoSAの設計を左右する。採点を省くなら、見ていない正の寄与を過小評価せず、負の寄与を過大に差し引かない必要がある。そうして作る安全な最大見積もりが、スコアの「上界」だ。DoSAはこの構造を利用するため、あらゆるAttentionにそのまま通用する一般的な高速化とは捉えない方がよい。

最初に、現在参照できる候補から、異なるキーを少なくともK個選び、元の採点器で完全に採点する。このシード集合の中でK番目に高いスコアを、閾値τとする。全候補にはシードも含まれるため、全体のK番目のスコアがτを下回ることはない。τは、最終的な選択境界の下界、つまり確実に届いている最低線になる。実装では走査中にこの閾値を固定する。

次に、各候補のスコア上界を求める。上界がτを厳密に下回れば、その候補はTop-Kへ入れない。コードレビューの例なら、候補を全て詳しく比べ切る前に、「最大限に見積もっても、既に確保した上位候補に届かない」と判断する形だ。実際には人がコードを読む手順ではなく、ベクトルの計算でこの判定を行う。

上界が閾値以上の候補は、元の全ヘッドで再採点する。その実スコアがτ以上の候補だけを書き出し、シードと合わせて最後のTop-Kを選ぶ。境界と等しい候補は両段階で残す。上界による完全採点の省略と、実スコアによる書き出しの削減は、別々の働きである。

DoSAは一部のヘッドを直接評価し、省略するヘッドの寄与をまとめて見積もる。正の重みを持つヘッドが予算内に収まるなら全て保持する。収まらない場合は、重みとクエリの大きさに基づいて選ぶ。負のヘッドの選択には、シード上でどれだけ活性化したかという頻度も使う。この頻度は効率を左右する情報であり、選択の正しさを保証する条件ではない。

省略する正のヘッドについては、保持したヘッドへの非負の射影と、そこでは表せない残差に分ける。射影は、既に求めた内積を使って寄与を見積もるためのものだ。残差はベクトルの大きさで上から抑える。残差同士の幾何関係も使うと上界をさらに狭められるが、その分クエリ側の計算が増える。

負のヘッドでは、実際に差し引く量の下界を作る。差し引き過ぎなければ、全体の上界は安全に保てる。DoSAはシードで半分以上活性化した負ヘッドを集約する。全てを一律にまとめると、内積の相殺で境界が緩むためだ。ここでも活性化情報の役割は、除外しやすい上界を作ることにある。

同じTop-Kを保つには、実数の式で上界が成立するだけでは足りない。論文が扱うFP8という低精度の数値形式では、丸めや演算誤差を含めても、上界が元の計算済みスコアを下回らないようにする必要がある。保証の基準は、元の採点器との整合性だ。

候補を保存する領域が満杯になったときも、残った候補を単に捨ててはいけない。正確なTop-Kへの統合、容量拡張、または全領域の再計算が必要になる。不正な数値で上界を保証できない場合は枝刈りせず、クエリ単位のエラーでは、既に除外したキーも含めて回復する。高速化は、こうした欠落防止を含む条件付きで成立する。

また、出力を生成するdecoding段階は別の経路を使う。全FP8採点を維持したまま、最近選ばれた異なるキーを現在のクエリで採点して閾値を作り、書き出しと選択を減らす。入力処理のprefillで行う、全ヘッド採点の省略と混同してはいけない。

著者はGLM-5.2で、クエリとキーの準備を含む完全Indexerを評価した。入力は英語、中国語、コード。文脈長は128Kから約1Mまでの5条件、prefillチャンクは4Kから32Kまでの4条件で、全60構成を比べた。約1M入力のprefixと21層の最終チャンクを再生し、単一NVIDIA GPUでウォームアップ後にCUDA eventsで時間を測っている。比較先はSGLang v0.5.18とLiteTopKだ。

モデル全体では、英語・中国語の書籍とコードの約1Mトークン入力で、最初の出力までの時間、TTFTを測った。テンソル並列度は8で、prefillチャンクはGLM-5.2が32K、DeepSeek-V3.2が16K。リクエスト間ではprefix cacheを消し、既存の入力処理結果の再利用による影響を抑えた。

さらに、機能を取り除く部品比較で、上界による候補除外と疎な書き出しの寄与を調べている。これはどの部分が速さを生むかを確かめる評価である。一方、これらの測定が直接示すのは計算時間であり、回答品質の向上ではない。

著者報告では、完全Indexerの全60構成でDoSAが最速だった。高速化は対SGLangで1.12〜2.74倍、対LiteTopKで1.17〜1.69倍である。準備を含めた比較で改善した点は、採点部分だけの短縮よりも広い効果を示している。

約1Mトークン入力のTTFTでは、GLM-5.2が対SGLangで1.86倍、対LiteTopKで1.26倍。DeepSeek-V3.2はそれぞれ1.94倍、1.50倍だった。最大値だけを取り出すと、モデルと比較先の違いが見えなくなる。

全体への効き方は、Indexerが処理時間のどれだけを占めるかにも左右される。著者のSGLangでのGPU区間計測では、GLM-5.2のIndexer割合は128Kで23.6%、約1Mで73.0%、DeepSeek-V3.2は約1Mで90.8%だった。この割合はCPUのトークン化やリクエスト転送を除く指標で、TTFTとは異なる。それでも、長文脈で選別費用が大きくなるという問題設定を裏付けている。

GLM-5.2の評価サンプルでは、上界による枝刈り率は英語・中国語・コードの領域平均で97.2〜97.5%だった。これは判定対象のクエリ・キー対から、どれだけ候補を除外したかの割合であり、速度の改善率ではない。各領域で21層中20層が95%を超えた一方、最初の層は89〜93%と低かった。

LOFTのSciFact検索例では、16Kクエリ・896Kキー、21層平均の採点と選択を比較した。両機能を使わないBaseに対して、候補除外のみでレイテンシが37.4%、疎な書き出しのみで23.4%、併用で47.3%減った。併用は1.90倍の高速化に相当する。既に書き出しを減らしていても、採点を減らす余地があるという結果だ。

負ヘッドの扱いも重要だった。同じ候補プールを使った部品比較では、負ヘッド集約を外すと残存率が2.65%から95.06%へ上がり、容量超過からの回復を含むIndexer時間は9.51倍になった。全負ヘッドを集約しても、残存率は4.06%、時間は12.0%増となった。単に多くの情報をまとめればよいのではなく、上界を狭く保つ設計が効果を支えている。

別の層別評価では、準備済みの最終32Kクエリについて採点と選択だけを測った。252セルのうち、対SGLangでは251セル、対LiteTopKでは250セルでDoSAが速かった。最大高速化はそれぞれ3.74倍、2.08倍だが、遅かった3つの比較ではレイテンシが0.7〜9.0%増えている。これは準備込みの全60構成とは異なる評価である。

この層別評価には測定上の留保もある。候補容量の修正後に全3手法を再測定したのは5セルで、残る247セルは元の測定値を使っている。全行列を再実行した結果ではない。また、約1Mキーでの容量修正により作業用メモリは9.94 GiB増えた。速さと候補の欠落防止を両立する費用として、メモリも見る必要がある。

提供資料には、具体的なGPU型番、詳細なソフトウェア環境、主評価の分散や信頼区間がない。入力も長さで選んだ書籍・コードの連続区間であり、本番トラフィック全体を代表する設計ではない。結果の幅を広く解釈するには、こうした条件を揃えた追加検証が必要だ。

DoSAでも全キーの走査は残る。短い文脈や、上界が緩く多数の候補が残る入力では、候補除外の追加費用を回収できない。適用候補は、符号付き重みとReLUを持つDSA型Indexerで、元の選択規則を保ちたい長文脈prefillである。多くのキーが選択境界から十分遠いほど、完全採点を省きやすい。

選択の保証にも範囲がある。候補IDまで同じにするには、同点を解決する共通の順序が必要だ。同点を含む有効なTop-K集合を得られても、それだけで生成文の同一性までは保証しない。また、この研究は訓練文脈長を超える入力での精度改善を示していない。DeepSeekの宣言文脈長を超える測定も、計算のスケーリングを観察したものとして読むべきだ。

残る論点は、実際の入力分布で上界がどれほど候補を絞れるか、そのときの回復処理や追加メモリを含めても、全体の待ち時間が縮むかである。選択を保つ理論と、利用条件ごとの利益を分けて確かめる必要がある。