技術的概要：確率的アテンション

問題定義

効率的なアテンション機構は、トランスフォーマーモデルのスケーリングに不可欠であるが、既存のアプローチは計算コストと受容野の網羅性の間でトレードオフに直面している。**スライディングウィンドウアテンション（SWA）**は、トークンをサイズ $w$ の局所ウィンドウに制限することで、層あたりの複雑さを $O(nw)$に削減する。しかし、この決定論的な局所性は、有効な受容野を線形成長（$\ell$層後で$\ell w$）に制限し、$w \ll n$の場合、シーケンスの大部分に到達できないままとなる。グローバルトークン、手作りのスパースパターン、またはブロックレベルのルーティングなどの対策が存在するが、これらはしばしばアーキテクチャの複雑さを導入するか、証明可能な網羅性の保証を提供できない。一方、フルアテンションはグローバルな接続性を提供するが、許容しがたい$O(n^2)$ のコストを伴う。

本論文は、約 13 万個のニューロンを持ち、接続確率がわずか 0.02% である**ショウジョウバエ（Drosophila melanogaster）の結合体（connectome）**から着想を得ている。このスパースさにもかかわらず、ネットワークは平均経路長が約 4.4 ホップという短い値を達成している。この「スモールワールド」トポロジーは、密な局所クラスタリングと、確率的ショートカットとして機能する広範に分布した長距離接続の相互作用によって生じている。著者らは、効率的なグローバル情報フローは、密な全対全接続を必要とするのではなく、構造化された局所計算に分散型・確率的な長距離接続を付加することで創発し得ると仮定している。

手法：確率的アテンション（SA）

著者らは、結合体のルーティング戦略を模倣するパラメータフリーの拡張である**確率的アテンション（SA）**を提案する。これは SWA に対する拡張である。

中核メカニズム

SA は、標準的な SWA 計算の前後に適用される 3 段階のプロセスを通じて、固定された局所ウィンドウを確率的なグローバルウィンドウに変換する。

1. 順列（Permute）: 各層ごとに独立して、ランダムな順列 $\sigma \sim \text{Uniform}(S_n)$ を引き当てる。トークンシーケンス表現（$Q, K, V$）を $\sigma$ に従って並べ替える。
2. ウィンドウアテンション: 順列された空間で標準的な SWA を適用する。順列がランダムであるため、サイズ $w$ の固定ウィンドウは、元のシーケンスのランダムな部分集合にまたがる。
3. 順列の解除（Undo Permute）: 出力に $\sigma^{-1}$ を適用して、元の順序を復元する。

元のトークン空間において、位置 $i$ は、元の距離 $|i-j|$ に関係なく、シーケンス全体に均一に分布し、期待サイズが $w$ のランダムな近傍 $\tilde{N}^\sigma_w(i)$ にアテンションを向ける。これは結合体における分散型の長距離ショートカットに類似し、任意のトークン対に対して約 $w/n$ のペアワイズ接続確率を生み出す。

理論的性質

• 受容野の成長: SWA の受容野は線形に成長する（完全網羅に $O(n/w)$ 層が必要）のに対し、SA は層ごとに独立してサンプリングされた順列により、指数関数的に成長する受容野を実現する。本論文は、完全なシーケンス網羅が $O(\log_w n)$ 層で達成されることを証明している。
• 複雑性: 順列およびその逆順列の操作は、$O(n)$ のインデックスの並べ替えである。アテンション計算は $O(nw)$ のままである。したがって、SA は標準的な SWA と同じ $O(nw)$ の層あたりの予算を維持する。
• 因果的制約: 自己回帰設定では、因果マスクを順列された空間に適用し、トークンが因果的にアクセス可能な先行トークンのみにアテンションを向けつつ、確率的な長距離接続を可能にする。

ゲート付き SA と SWA の組み合わせ

「スモールワールド」領域（高い局所クラスタリング＋短い経路）を再現するために、著者らは SA と SWA を組み合わせた双経路アーキテクチャを提案する。
$$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$$
ここで、$Y_{sa}$ と $Y_{swa}$ はそれぞれ確率的経路とスライディングウィンドウ経路からの出力である。ゲート重み（$g_{sa}, g_{swa}$）は独立したシグモイド射影によって計算され、モデルが局所的な一貫性（SWA 由来）とグローバルな網羅性（SA 由来）のバランスを学習することを可能にする。この組み合わせは、ほぼ無視できるパラメータを追加し、$O(nw)$ の複雑性を維持する。

主要な貢献

1. 確率的アテンション（SA）: ランダムな順列を適用することで、同じ $O(nw)$の計算予算内で指数関数的な受容野成長（$O(\log_w n)$）を実現する、SWA へのパラメータフリーの「ドロップイン」拡張。
2. ゲート付き SA + SWA: 生物学的結合体で観察されるスモールワールド領域を理論的かつ実証的に再現する双経路アーキテクチャであり、局所クラスタリングと確率的ショートカットのバランスを取る。
3. 実証的検証:
   • 事前学習: 0 から訓練された 3 億 6000 万パラメータモデルにおいて、ゲート付き SA + SWA の組み合わせは、平均ゼロショット精度で最良の結果を達成し、純粋な SWA を上回り、パープレキシティではフルアテンションに匹敵しつつ、下流タスクではそれを上回った。
   • 学習フリー推論: 事前学習済みのQwen3-8BおよびQwen3-30B-A3Bモデルに適用したところ、SA は SWA を一貫して上回り、同等の計算予算において**ブロックアテンションの混合（MoBA）**と同等かそれ以上の性能を示した。ウィンドウサイズが増加するにつれ、フルアテンションの品質を著しく速く回復させた。

結果

• 事前学習: SA + SWA モデルは、平均ゼロショット精度で35.9を達成した（SWA は 35.1、フルアテンションは 34.9）。特に、広範な談話文脈を必要とするLAMBADAの性能を大幅に向上させつつ、WikiTextでの低いパープレキシティを維持した。純粋な SA は高いパープレキシティを示し、言語モデル化には局所的一貫性が不可欠であることを確認した。
• Qwen3 における推論:
  • 有効ウィンドウサイズ 128 において、Qwen3-8B 上の SA は**70.9%**の平均精度に達し（フルアテンション基準から 1 ポイント以内）、SWA は 62.2% であった。
  • Qwen3-30B-A3B において、SA はウィンドウサイズ 64 で**73.2%**の精度に達し、SWA（47.0%）および MoBA（66.3%）を大幅に上回った。
  • SA は、SWA の性能が崩壊する非常に小さなウィンドウサイズ（例：$w=32$）でも、知識集約型タスク（MMLU、BoolQ）において堅牢性を示した。
• 効率性: SA は線形スケーリング $O(nw)$ を維持する。A100 GPU において、シーケンス長 2,048 でフルアテンションに対して1.5 倍の高速化を示し、32,768 では28 倍に増加した。

意義と主張

本論文は、結合体に着想を得た確率的ルーティングが、効率的なアテンションの表現力を向上させるための実用的なプリミティブであると主張する。著者らは、グローバル情報フローは密な全対全接続を必要とするのではなく、構造化された局所計算と、深さにわたって蓄積されるスパースかつ確率的な長距離ショートカットの相互作用から創発し得ると論じている。

この研究の意義は以下の点にある。

1. 単純性: SA はアテンション機構自体に学習可能なパラメータを必要とせず、$O(n)$ のインデックスオーバーヘッドのみで済み、既存の SWA カーネル（Mistral、Gemma 2、gpt-oss など）への「ドロップイン」アップグレードとなる。
2. 理論的保証: SWA の線形網羅に比べて大幅に改善された、証明可能な $O(\log_w n)$ の網羅深さを提供する。
3. 補完性: 既存の線形およびスパースなアプローチを補完するものとして提示され、手作りのスパースパターンやグローバルトークンの複雑さなしに、局所的効率と全球的表現性の間のギャップを埋める新たな設計原理を提供する。