メインコンテンツまでスキップ

03. アーキテクチャ設計とスケーリング

02で見た3つの Hyena 演算子を、どう積み重ねるか（block layout）、どう効率化するか（filter grouping）、どこまで文脈を伸ばせるか（context extension）、そして 実際どれだけ速いか（throughput） を見ていきます。

1. ブロックレイアウト：演算子をどう積むか​

マルチハイブリッドは、SE・MR・LI・MHA（multi-head attention）を 縞状（striped）に積み重ねて 全体を構成します。論文は、7B のマルチハイブリッドを OpenGenome2 の 400B トークンで学習し、レイアウトを変えて事前学習の品質（perplexity, PPL）を比較しました。

ブロックレイアウトPPL@400B（低いほど良い）
MHA-MHA-MHA（純 Transformer 相当）3.09
LI-LI-LI（純・長畳み込み）2.87
SE-SE-LI2.88
SE-MR-LI2.83 ✅

ここから2つの重要な知見が得られます。

  • SE-MR-LI が最良。短・中・長の異なる距離スケールを組み合わせるのが効く。
  • 純・長畳み込み（LI-LI-LI）は SE-SE-LI で置き換えても品質ほぼ同じ で、スループットは大きく向上する。つまり「全部を長い暗黙フィルタにする」必要はない。

実際の StripedHyena 2 は、この SE-MR-LI を基本ブロックとして深さ分だけ繰り返し（7B なら深さ 32）、さらに 5 つの MHA を畳み込みブロックに interleave します。

SEMRLISEMRLIMHASEMRLI…MHASE-MR-LI を深さ分（7B なら 32 層）繰り返し、合計 5 つの MHA を散りばめるSEMRLIMHA
SE-MR-LI を基本単位に積み、要所に MHA を挟む縞状レイアウト（論文 Fig. 2.1・Table 2.1 に対応）
新しいドメインでは ablation を

SE-MR-LI は安定したベースラインですが、論文は「フィルタ長・減衰の強さ・初期化を変えるなら、新しいタスク/ドメインでは block layout を ablation（比較検証）すべき」と注意しています。万能の固定解ではなく、設計の出発点という位置づけです。

2. フィルタのグループ化：GEMV から GEMM へ​

StripedHyena 2 は、入力依存畳み込みに グループ化（grouping） という設計を採り入れます。フィルタを チャネルのグループ全体で共有 するのです。グループ G\mathcal{G}（サイズ dgd_g）に属するチャネル α\alpha は、同じフィルタ hGh^{\mathcal{G}} で畳み込まれます。

∀α∈G:ytα=∑j=0tht−jG xjα\forall \alpha \in \mathcal{G}: \quad y^\alpha_t = \sum_{j=0}^{t} h^{\mathcal{G}}_{t-j}\, x^\alpha_j

なぜこれが効くのか。グループ化により、離散畳み込みを GEMV（行列×ベクトル）の集まりではなく GEMM（行列×行列）の系列として表現 できるようになります。GEMM は GPU の テンソルコア が最も得意とする演算で、ここが高速化の土台です（詳細は 04）。

グループ化なし各チャネルが別フィルタGEMV（テンソルコア活用しにくい）グループ化ありグループ1（共有）グループ2（共有）グループ内で同じフィルタを共有GEMM（テンソルコアをフル活用）
フィルタのグループ化。共有することで畳み込みを GEMM にまとめ、テンソルコアで高速化できる

しかも 品質への影響は最小 です。グループサイズ 1（共有なし）と 16 で収束に差はなく、64 を超えると小さな劣化が出る程度。ほぼ無料で効率を得られる 設計です（なお、これは従来の grouped CNN とは別物で、グループ内でチャネルを混ぜるわけではありません）。

3. 文脈拡張：8k から 1M へ​

StripedHyena 2 は、まず短い文脈で base モデルを学習し（7B なら 2T トークン・8192 文脈＝ Evo 2 7B）、その後 midtraining で 100 万トークンまで文脈を拡張 します。拡張には rotary attention 向けの技術——position interpolation（PI） と adjusted base frequency（ABF）、およびその組み合わせ——を流用します。

文脈長検証 perplexity（PI + ABF, 7B）
32k2.782
65k2.763
131k2.748
262k2.707
524k2.663
1M2.597

文脈を伸ばしても perplexity の悪化はなく（むしろ低下）、すべてのモデルが目標の最大文脈長で in-context recall を達成 しました（needle-in-a-haystack）。

4. スケールでのスループット​

02の FIR 演算子（SE・MR）のおかげで、StripedHyena 2 は前世代ハイブリッドや Transformer に対して一貫した高速化を示します。

  • 7B・40B のいずれでも、H100 クラスタで最適化 Transformer 比 1.2〜2.9 倍 高速（dense 層・正規化層は FP8 精度）。
  • 短い系列長でも Transformer と前世代 StripedHyena の両方より高速。
  • 40B でピーク MFU（Model FLOPs Utilization）は 16K 文脈で約 34%。
長文脈で MFU が下がるのは「悪いこと」ではない

長い系列長では、ハイブリッドの MFU は下がります。しかしこれは、subquadratic（準二次）なスケーリングによってモデル全体の FLOPS 自体が減る ためです。Transformer は O(n2)O(n^2) で「無駄に」FLOPS を積み上げて MFU を稼いでいるとも言えます。実際の壁時計時間（iteration time）では StripedHyena 2 が一貫して速く、分散設定をさらに調整すれば長文脈での高速化はもっと伸びる余地があります。

5. 学習の詳細（再現のために）​

論文は再現性のため、学習の実務的な設定も詳しく開示しています。

インフラと精度

  • 学習はオープンソースの Savanna（設定ファイルは凍結コミットで公開）。
  • 混合精度学習で、dense 層（SwiGLU・射影）と正規化層は FP8。
  • 7B のバッチサイズは critical batch size estimation（McCandlish et al.）で決定。

並列化の構成（コンテキスト並列を含む）

StripedHyena 2 は、Tensor Parallel・Sequence Parallel・Context Parallel を、系列長に応じて組み合わせます。

設定7B（256 H100）40B（2,048 H100）
Tensor Parallel2 → 328 → 64
Context Parallel1 → 21 → 2
Global Batch Size4M トークン8M トークン

短い文脈では Tensor Parallel 中心、長い文脈になるほど Context Parallel を効かせる、という切り替えです（数値は文脈長 16K → 1M で増加）。

グループ化の ablation

前述のフィルタ・グループ化について、論文は品質への影響を検証しています。

  • group size 1（共有なし）と 16 で収束に差なし。
  • 64 を超えると小さな劣化（小さいモデルで顕著）。
  • グループ数が少ないと、Hyena-MR の指数減衰正則化（チャネル横断で初期化）の 粒度が粗くなる ため。

実用上は、品質を保ちつつテンソルコア効率を上げられる範囲で group size を選びます。「ほぼ無料で速くなる」工夫が、実験でも裏づけられている わけです。

早期実験：FFN を畳み込みで置き換える

論文は初期設計で、各 Hyena/MHA 演算子の後の FFN（MLP・SwiGLU）を Hyena-SE で置き換えた 変種も試しています。スループットがわずかに下がる代わりに収束が改善する傾向があり、Hyena-SE が MHA や状態空間モデルより高速なことと合わせ、「将来のマルチハイブリッドは MLP と Hyena-SE の比率も最適化の対象になる」と展望を述べています。

6. まとめ​

  • SE-MR-LI ブロックを基本に MHA を散りばめるレイアウトが、品質・速度のバランスで最良（PPL 2.83）。
  • フィルタのグループ化 で畳み込みを GEMM にまとめ、品質をほぼ落とさずテンソルコアで高速化。
  • 8k で事前学習 → 1M まで文脈拡張（PI・ABF）。perplexity は悪化せず recall も維持。
  • 最適化 Transformer 比 1.2〜2.9 倍、短文脈でも高速。

「グループ化で畳み込みを GEMM にする」と述べましたが、その GEMM を GPU でどう実装するか が次の核心です。04. ハードウェア対応の畳み込みカーネルへ進みましょう。