Depth-Staggered Fibonacci Spacing for Sparse Attention: Static Schedules Beat Learned Dilation and Extrapolate Where Dense Attention Fails

問題提起

本論文は、スパース・アテンション機構を通じて、言語モデルにおける自己注意（self-attention）の二次的な計算コストを削減するという課題に取り組んでいる。既存の多くの手法は、クエリを特定のサブセット（局所的なウィンドウ、ストライド・パターン、あるいはヘッドごとのバリエーションなど）に制限しているが、制御された設定において、単一のスペーシング・パラメータをネットワークの深さに沿って最適に変化させる方法についての理解は限られている。具体的には、フィボナッチ数列に基づいたオフセットをスケーリングする層ごとのスカラー $\alpha$ を、受容野のカバー範囲とモデル性能のバランスを取るために、固定すべきか、動的に学習させるべきか、あるいは静的なスケジュールによって設定すべきかを調査している。

手法

著者らは、同一のレシピ（パラメータ数60M、隠れ次元512、レイヤー数16、トークン数426M）で訓練された21の言語モデルを用いて、制御された比較を行った。

アーキテクチャおよびアテンション機構:

• アテンデッド・セット（注目対象集合）: 各位置 $i$ におけるクエリに対し、モデルはサイズ $W$ の高密度なローカル・ウィンドウと、フィボナッチ数列（$fib = [1, 2, 3, 5, \dots, 987]$）に基づく一連のスパースなオフセットに注目する。
• スケーリング: 各オフセット $f_k$ は、ターゲットとなる距離 $d_k = \alpha f_k$ となるように、レイヤーごとのスカラー $\alpha$ によってスケーリングされる。スカラー $\alpha$ は $[0.5, 1.5]$ の範囲に制約される。
• 実装: オフセットが非整数であっても勾配が $\alpha$ に流れるようにするため、モデルはキー/値（key/value）の行間の線形補間を使用する。回転式位置エンコーディング（RoPE）は、集約（gathering）の前に行われる。

実験変数:
研究では、16個のレイヤーに対してスカラー $\alpha$ を設定する4つの手法を比較した：

1. Fixed（固定）: すべてのレイヤーで $\alpha = 1$。
2. Learned（学習型）: レイヤーごとの自由パラメータ $\theta$ を持ち、初期値 $\alpha=1$ として勾配降下法により最適化される。
3. Staggered (Linear)（スタッガード・リニア）: $\alpha$ がレイヤー0で0.5からレイヤー15で1.5へと線形に増加する固定されたランプ（傾斜）状のスケジュール。
4. Staggered (Coprime/HDC)（スタッガード・コプライム/HDC）: リニア・ランプと同じ $\alpha$ のマルチセットを用いるが、隣接するレイヤー間の間隔の差異を最大化するために、互いに素なストライド（$\ell \to (7\ell \mod L)$）を用いてレイヤーに再割り当てしたもの。

コントロール（対照群）:

• スペーシングの密度による影響を分離するため、ベースとなる「リーチ・マッチド（reach-matched）」な2の累乗（ログ・スパース）を用いた。
• 同一のアーキテクチャと訓練レシピを持つデンス（高密度）なベースラインを、パフォーマンスの上限として使用した。

主な貢献と結果

1. 静的なスタッガリングが支配的なレバーである
すべてのウィンドウサイズ（$W \in \{6, 8, 10, 12\}$）およびデータセットにおいて、静的なレイヤーごとのスタッガード・スケジュールは、固定または学習型のスケジュールよりも優れた性能を示した。

• パフォーマンスの順序: Staggered $\approx$ Coprime $>$ Learned $>$ Fixed。
• 最良のモデル: $W=12$ を用いたリニア・スタッガーが最も低いパープレキシティを達成した（FineWeb-Edu: 42.02、Wikipedia: 37.59）。
• ベースへの非依存性: スタッガリングの利点は、ベースとなる数列の種類に依存しない。2の累乗ベースにリニア・スタッガーを適用すると大幅に改善され、固定されたフィボナッチ・ベースを上回り、学習型のフィボナッチ・ベースと同等の性能に達した。これは、ベースとなる数列の具体的な選択よりも、レイヤーごとのスペーシングの多様性の方が重要であることを示している。

2. 学習されたダイレーション（膨張）は不活性である
レイヤーが最適な「バネのような」スペーシングを「学習」するという仮説に反して、学習されたスカラー $\alpha$ は初期値付近（$\alpha \approx 1$）でほぼ静止したままとなった。

• パフォーマンス: 学習型はFixedよりはわずかに改善したが、静的スタッガーには及ばなかった。
• 効率性: 学習型は、非整数オフセットのために融合されていない（unfused）補間集約が必要であったため、推論レイテンシが約5倍高くなった。一方、静的バリアントは高速な融合済みブールマスク・パスを使用できた。
• メカニズムの調査: 制御された検索タスクにより、勾配降下が特定のスペーシングを報酬とする場合でも、スカラーを大きく移動させないことが確認された。この恩恵は適応的なものではなく、構造的なもの（多様なスペーシング）である。

3. 長さの外挿：スパースの堅牢性 vs デンスの崩壊
最も重要な発見は、訓練長（1024トークン）の4倍の長さ（4096トークン）への外挿に関するものである。

• デンス・ベースライン: 壊滅的に崩壊し、パープレキシティが（FineWeb-Eduにおいて33.37から100.36へと）201%増加した。
• スパース・バリアント: 堅牢性を維持し、パープレキシティは横ばい、あるいはわずかに改善した（例：リニア・スタッガーは1.3%低下）。
• メカニズム: 固定オフセットのスパース・アテンションは、訓練中に見られた相対位置のみをクエリする。これらのオフセットに対するRoPEの角度は、任意の長さにおいても分布内（in-distribution）にあるため、モデルは自然に外挿できる。一方、デンス・アテンションはすべての相対位置に注目するため、訓練中に一度も現れなかったRoটি角度を含む位置にも注目してしまい、失敗を引き起こす。
• ベースへの感度: 2の累乗ベースは、フィボナッチ（+0% ～ +2.8%）よりも劣化が大きかった（+11.3%）。これは、そのスパースな段差が、長い評価長において未訓練の大きなオフセット（例：1024, 2048）を導入してしまうためである。

4. 正直なネガティブな結果と限界
本論文は、以下の2つのネガティブな発見によって自らの主張の範囲を明確に限定している。

• 分布内品質のギャップ: 訓練長（1024）において、最良のスパース・モデルはデンス・ベースラインよりもパープレキシティが約26%高い。スパース・アテンションは、長さの外挿性のために品質を犠牲にするものである。
• 一様な利得: スタッガリングによる改善は、すべてのコンテキスト位置において一様であり、長距離に集中しているわけではない。この利得は、特定の「長距離到達」効果ではなく、一般的なレイヤーごとの表現の多様性から生じている。

意義と主張

本論文は、スパース自己注意において、静的で深さごとに変化するスケジュールは、固定または学習型のスケジュールよりも優れていると主張している。著者らは、「学習されたバネ」メカニズムは不要かつ非効率であり、代わりに手動でスタッガードなスペーシングの範囲を課すことで、必要な受容野の多様性を提供できると論じている。

主な実用的な意味合いはトレードオフである。構造化されたスパース・アテンションは、訓練長における分布内の品質を犠牲にするが、デンス・アテンションが完全に失敗するような長さへのほぼゼロコストの外挿を実現する。これは、デンス・アテンションの計算負荷や、学習された位置適応の複雑さを回避しつつ、長いコンテキスト処理を必要とするアプリケーションにおいて、有効な戦略となる。

著者らは、これらの結果がテストされたスケール（パラメータ数60M）およびスケーリングなしのバニラRoPEに特有のものであること、そして最適なスタッガード・スケジュール（リニアかコプライムか）は依然として大規模モデルにおける未解決の課題であることを強調している。