要旨（Abstract）
Recurrent-Depth Transformer（以下 RDT）の設計原理の中核には「ループ回数は層パラメータの代替になる」という主張がある（Parcae 2026）。しかしこの交換可能性がどの条件下で成立するかは実験的に十分に検証されてこなかった。本研究は、OpenMythos 実装上のモデルを (a + b) mod 97 という古典的な grokking タスクで訓練し、dim ∈ {64, 128, 256} × loops ∈ {1, 2, 4} の 3×3 グリッドを同一条件で sweep することで、ループ数と幅の相補関係を実証的に特徴づけた。
主要な結果
- 容量閾値（capacity floor）の存在：
dim = 64, loops = 1は 5000 step の訓練でも test_acc 30 % に留まり汎化失敗。loopsを追加することでのみ grokking が可能になる。 - 閾値超過後の逆転：
dim ≥ 128ではloops = 1が最速で grokking に到達（step 800）。追加の loops はむしろ grokking を 2.4 倍遅延させる。 - 同 FLOPs でも「幅 vs 深さ」は非対称：FLOPs 1.6 倍差に対し step_grok は 2.1–2.4 倍差 — 一単位の幅追加は一単位の loop 追加より常に優位。
- Memorization は幅に比例してクリーンにスケール：step_mem は 600 → 200 → 100 と dim を倍にするたび約半減。loops 依存はほぼゼロ。
1. 背景と問題意識
1.1 研究の背景
Recurrent-Depth Transformer は、同一の Transformer ブロック重みをテスト時に任意回数適用できる設計 (Saunshi et al., 2025; Parcae et al., 2026) により、推論時の計算量を動的にスケールすることを狙った近年のアーキテクチャ群である。OpenMythos はこの考え方を MLA / GQA / MoE / ACT / LTI injection といった現代的な部品で再構成した実装で、本研究の実験基盤となる。
RDT 系の理論的主張で特に魅力的なのは、「ループ回数 T は層パラメータ数 N の代替として機能する」 という Parcae 2026 の主張である。仮にこれが成立するなら、小さいモデルでも多ループ回すことで大きなモデルと同等の表現力が得られ、メモリに厳しい個人研究環境においても実用的な RDT 研究が可能になる。
1.2 研究目的
本研究の目的は以下の点に集約される：
- 交換可能性の実証的検証：Parcae の主張「loops は params の代替」は本当に任意の条件で成立するのか。
- 容量閾値の特定：もし交換可能性に条件があるなら、それはどのようなメトリックで規定されるのか。
- FLOPs 単位の比較：計算量を揃えた条件下で、幅と深さはどちらが効率的か。
1.3 関連研究との位置づけ
- Saunshi et al. (2025)：深さ外挿 (depth extrapolation) 現象の基礎研究
- Parcae et al. (2026)：LTI-stable 再帰で任意 loop 数の安定学習を示し、loops ↔ params の交換を主張
- Power et al. (2022), Nanda et al. (2023)：grokking（記憶 → 汎化の相転移）の実証と機序解析
- Coconut (Hao et al., 2024)：CoT を潜在空間で展開するエージェントにおける RDT 活用
本研究は、Parcae の理論的主張を grokking の clean な phase transition 上で FLOPs-matched に sweep する という形で、「交換可能性の境界条件」 を初めて明示的に特徴づけるものである。
2. 研究手法
2.1 モデル
OpenMythos v0.4.0 の標準 RDT 構造：
Input → Prelude(1 layer) → RecurrentBlock(loops=T) → Coda(1 layer) → Output
↓↑ h_{t+1} = A·h_t + B·e + f(h_t, e)
RecurrentBlock は LTI-stable injection（ρ(A) < 1 保証）・ACT halting・MoE FFN・per-loop LoRA adapter を内包する。
| 項目 | 値 | 備考 |
|---|---|---|
| Attention type | GQA | n_heads = dim/32, n_kv_heads = n_heads/2 |
| MoE experts | 4 routed + 1 shared | top-2 routing, expert_dim = dim |
| Prelude / Coda | 1 / 1 layers | dense SwiGLU FFN |
| LoRA rank | 4 | per-loop depth adaptation |
| Loop-index embedding | ON | dim/8 チャネルに正弦波注入（本研究では固定） |
| ACT threshold | 0.99 | 学習時 halting 閾値 |
Sweep 軸は dim ∈ {64, 128, 256}、loops (= max_loop_iters) ∈ {1, 2, 4} の 3×3 グリッド。
2.2 タスク：(a + b) mod p
古典的 grokking タスク。p = 97（素数）とし、全ての (a, b) ∈ [0, p)^2 の 9,409 ペアを列挙。ランダム分割で train fraction = 0.3（2,823 ペア）、残 6,586 を test に割り当てる。
各 example の入力は 4 トークン [a, +, b, =]、正解は 1 トークン c = (a + b) mod p。Vocab は数値 0..96 + + + = + PAD = 100 トークン。
因果 LM として扱い、logits[:, -1, :]（= 位置の出力）と c の cross-entropy を損失とする。
2.3 最適化
| 項目 | 値 | 備考 |
|---|---|---|
| Optimizer | AdamW | betas = (0.9, 0.98) |
| Weight decay | 1.0 | grokking 成立に必須の高 WD |
| Learning rate | 1e-3 → 1e-4 | cosine decay, warmup 200 steps |
| Batch | Full-batch | 2,823 sequences every step |
| Grad clip | 1.0 | gradient norm |
| Total steps | 5,000 | |
| Precision | fp32 | 小モデル・短系列で bf16 の旨味なし |
| Seed | 42 | 固定 |
2.4 評価指標
step_mem：train_acc が初めて 99 % を超える step。記憶完了時期。step_grok：test_acc が初めて 99 % を超える step。汎化完了時期。未達なら-(DNG = did not grok)。FLOPs/step：Kaplan 2020 の6 · N_active · T近似による概算。MoE は top-k 活性化パラメータ数のみを計上。
2.5 実験環境
- Hardware：NVIDIA GPU (CUDA, 8 GB VRAM)
- Software：Python 3.10.9 + PyTorch 2.11.0+cu128 + transformers 5.5.4
- Wall-clock：sweep 全 9 runs を 46.8 分で完了
3. 結果
3.1 Sweep の全体像
表1: 9 configs の step_grok / 最終 test_acc / FLOPs/step
| dim | loops | params | FLOPs/step | step_mem |
step_grok |
test_acc |
|---|---|---|---|---|---|---|
| 64 | 1 | 151,110 | 8.9×10⁹ | 600 | — (DNG) | 30.2 % |
| 64 | 2 | 151,114 | 1.3×10¹⁰ | 400 | 2,300 | 100.0 % |
| 64 | 4 | 151,122 | 2.2×10¹⁰ | 400 | 1,700 | 100.0 % |
| 128 | 1 | 588,678 | 3.6×10¹⁰ | 200 | 800 | 100.0 % |
| 128 | 2 | 588,682 | 5.3×10¹⁰ | 200 | 1,900 | 100.0 % |
| 128 | 4 | 588,690 | 8.9×10¹⁰ | 200 | 1,900 | 100.0 % |
| 256 | 1 | 2,324,742 | 1.4×10¹¹ | 100 | 800 | 100.0 % |
| 256 | 2 | 2,324,746 | 2.1×10¹¹ | 100 | 1,200 | 100.0 % |
| 256 | 4 | 2,324,754 | 3.6×10¹¹ | 100 | 1,200 | 100.0 % |
3.2 容量閾値：dim = 64, loops = 1 の単一失敗点
9 runs 中唯一、dim = 64, loops = 1 のみが 5000 step を費やしても test_acc 30.2 % に留まり grokking に到達しなかった。train_acc は step 600 で 99 % に達しているので、記憶には成功しているが汎化ができない 状態。同じ dim = 64 でも loops = 2 では step 2300 で、loops = 4 では step 1700 で grokking する。
これは次のように要約できる：
- 単体の幅 = 64：汎化不能
- 単体の幅 = 64 + 1 追加 loop：汎化可能（step 2300）
- 単体の幅 = 64 + 3 追加 loops：汎化可能（step 1700）
→ ループは「不足する幅」を代替する capacity 供給源として機能する。これは Parcae 主張の positive evidence。
3.3 閾値超過後の逆転：幅十分なら loops = 1 が最速
dim = 128 以上では状況が反転する：
| config | step_grok |
|---|---|
dim=128, loops=1 |
800 ← 最速 |
dim=128, loops=2 |
1,900 |
dim=128, loops=4 |
1,900 |
dim=256, loops=1 |
800 ← 同速 |
dim=256, loops=2 |
1,200 |
dim=256, loops=4 |
1,200 |
同一幅で loops を増やすと grokking が 2.4 倍遅くなる。dim = 128 での loops = 2 と loops = 4 は完全に同じ step_grok = 1900 で差がなく、loops を増やしてもメリットがない。
dim = 128 → 256 の幅倍増でも loops = 1 の step_grok は変わらず 800 のまま。幅にも閾値超過後の diminishing return がある。
3.4 FLOPs-matched 比較：幅と深さの非対称性
「同じ FLOPs 予算なら、どっちに使う方が得か」を確かめる。3.56×10¹⁰ FLOPs 近傍と 8.9×10¹⁰ FLOPs 近傍で比較できるペアが存在：
| 比較ペア | FLOPs/step | step_grok | 勝敗 |
|---|---|---|---|
d64L4 |
2.2×10¹⁰ | 1,700 | |
d128L1 |
3.6×10¹⁰ | 800 | ← 2.1× 速い |
| FLOPs 差 | 1.6× | — | |
d128L4 |
8.9×10¹⁰ | 1,900 | |
d256L1 |
1.4×10¹¹ | 800 | ← 2.4× 速い |
| FLOPs 差 | 1.6× | — |
両ケースとも、FLOPs 1.6 倍差に対し幅側の step_grok は 2.1–2.4 倍速い。言い換えれば、同 FLOPs における「1 単位の幅追加」は「1 単位の loop 追加」より常に優位 である。
3.5 Memorization のクリーンなスケーリング
step_grok と違い、step_mem は幅に対して綺麗に単調：
| dim | step_mem (平均) |
|---|---|
| 64 | 400–600 |
| 128 | 200 |
| 256 | 100 |
loops 依存はほぼゼロ。記憶速度は純粋に幅で決まる。これは記憶が Prelude + Coda の表現力に依存し、recurrent block の iteration 回数とは独立であることの傍証。
3.6 結果のまとめ
- ループ数は、幅が足りないときに限り、参照的に幅の代替として機能する。
- 幅が十分 (
dim ≥ 128) な場合、追加の loop は grokking を遅延させる。 - 同 FLOPs における幅追加と loop 追加は 非対称：幅が常に優位。
- Memorization 速度は幅のみに依存し、loops には非依存。
4. 考察
4.1 容量閾値仮説（Capacity Floor Hypothesis）
観察された単一の汎化失敗点 d64L1 と、loop 追加による救済から、以下の仮説を提示する：
モジュラー算術の汎化に必要な表現容量には閾値 $C^*$ が存在する。 実効容量 $C(\text{dim}, \text{loops}) \geq C^*$ であれば汎化可能。 $C(\text{dim}, \text{loops})$ は dim に対して強く、loops に対して弱く単調増加する関数で近似できる。
d64L1 が $C < C^*$、d64L2 以降は $C \geq C^*$ と解釈できる。なお $C^*$ はタスク依存で、より複雑なタスク（k-hop reasoning, SCAN, COGS など）では $C^*$ が高くなり、loops の役割が大きくなると予想される。この検証は heterogeneous なタスク群での本研究の追試 が次のステップになる。
4.2 Loop 追加で grokking が遅くなる機序
dim = 128, loops = 1 → 2 で step_grok が 800 → 1900 と 2.4 倍悪化する現象は、以下のいずれか（もしくは複合）で説明可能：
- 最適化結合の増大：weight-tied な recurrent block を T 回繰り返す構造では、1 step の勾配に T 回分の誤差伝搬が畳み込まれる。T が大きいほど optimizer がアラインさせるべき制約が増え、grokking 相転移が遅延する。
- タスクの homogeneity：
(a+b) mod pは loop ごとに異なる処理が不要。1 回で解ける。追加 loop はむしろ「何もしないこと」を学ぶ必要があり、Hilbert 空間内の局所 minimum を増やす。 - MoE router の二重最適化：routed experts は per-loop に再 routing されるため、loops が多いと router が不安定化する可能性（今回の sweep では直接未検証）。
なお最終 test_acc は loops によらず 100 % に到達するため、loop は asymptote ではなく transition の速度 にのみ影響する。
4.3 幅 vs 深さの非対称性の物理的解釈
FLOPs が同等なのに幅が勝つ理由として考えられるのは：
- 幅を増やした 1 層：新たな表現方向を獲得（並列に情報を保持できるチャネル数の増加）
- 同じ層を T 回適用：既存の表現方向を反復変換する（時間軸の拡張）
前者は「情報容量」を増やし、後者は「計算時間」を増やす。本タスクは 1-step で解ける計算なので、時間軸拡張の恩恵がない。タスクが本質的に iterative な場合（例：k-hop reasoning chains）には逆転する可能性 があり、これは本研究の direct follow-up。
4.4 限界と制約
- 単一タスク：
(a + b) mod pの結果を一般化するには他の synthetic task（parity、k-hop、SCAN、COGS）での追試が必要。 - 固定乱数 seed：SEED = 42 単一の結果。grokking は stochastic な側面があるため、step_grok の 100 単位の差は noise の可能性もある。3–5 seeds での平均化が望ましい。
- LIE が on：本研究は
use_loop_index_embedding = True固定。この設定は loops の効果と混線している可能性（次報告 A2 で独立検証）。 - FLOPs 近似の粗さ：Kaplan の
6·N·Tは dense transformer 向けで、MoE の sparse activation や recurrent の重み共有を正確には表さない。より精密な profiling は今後の課題。
4.5 実用的な示唆
- 小型モデル設計：幅不足が疑われる場合に限り loops を 2–4 に増やすのが合理的。幅で解ける範囲なら loops = 1 で十分。
- FLOPs 予算配分：compute-bound な状況では 幅優先 が step efficiency の観点で有利。
- タスク複雑度推定：「loops を増やしても速くならない」こと自体がタスクの homogeneity を示唆する診断に使える。
5. 結論
5.1 知見の要約
- ループは容量閾値以下の補償装置：
dim = 64, loops = 1のみが汎化不能で、loops を追加すると救済される。 - 閾値超過後の loops は grokking を遅延させる：
dim ≥ 128ではloops = 1が常に最速（step_grok = 800）。 - FLOPs が等しくても幅が常に優位：幅 1 単位 > loop 1 単位の非対称性が 1.6× FLOPs 差でも 2.1–2.4× step_grok 差として出る。
- Memorization は純粋に幅依存：step_mem は幅を倍にするたび約半減、loops には非依存。
5.2 今後の課題
- Heterogeneous task での追試：k-hop reasoning、SCAN、COGS など iterative な計算を要するタスクで逆転が起きるか検証。
- Multi-seed 化：3–5 seeds で step_grok の confidence interval を確立。
- Loop-index embedding の ablation：loops が効く / 効かない分水嶺に LIE が関わっているか。→ 本シリーズ次報告 A2 で独立検証
- 深さ外挿との連関：本研究は train_loops = eval_loops の条件下のみ。inference-time に loops を変えた場合の挙動は別論文で扱う。→ 本シリーズ第3報で検証
付録
A. 実行環境
- GPU：NVIDIA GPU (CUDA, 8 GB VRAM)
- 言語：Python 3.10.9
- 主要ライブラリ：PyTorch 2.11.0+cu128, transformers 5.5.4, datasets 4.8.4, loguru 0.7.3
B. 補足データ
Step-by-step 経過の抜粋（baseline run: dim=128, loops=2, LIE on）
| step | train_loss | test_loss | train_acc | test_acc | ρ(A) mean |
|---|---|---|---|---|---|
| 1 | 4.63 | 4.63 | 0.9 % | 1.1 % | 0.368 |
| 200 | ~1e-3 | ~3.6 | ≥99 % | 30 % | 0.375 |
| 800 | ~1e-4 | 0.16 | 100 % | 96 % | 0.366 |
| 1100 | ~1e-4 | 0.03 | 100 % | 99 % | 0.363 |
| 5000 | 0 | 0 | 100 % | 100 % | 0.362 |
| 30000 | 0 | 0 | 100 % | 100 % | 0.400 |
C. 用語集
| 用語 | 説明 |
|---|---|
| Grokking | 長期訓練の途中で test_acc が急激に立ち上がる相転移現象。train → test のギャップが暴力的に埋まる。 |
| RDT | Recurrent-Depth Transformer。同一の transformer block 重みを T 回繰り返し適用する設計。 |
| LTI injection | Linear Time-Invariant 入力注入。h_{t+1} = A·h_t + B·e + f(h_t, e) で ρ(A) < 1 を構造的に保証し、任意 T で安定。 |
step_grok |
test_acc ≥ 99 % に初めて到達した step。 |
step_mem |
train_acc ≥ 99 % に初めて到達した step。 |
| MoE | Mixture of Experts。複数の小型 FFN を router がルーティングする sparse 層。 |
seed 固定・全データ決定的・full-batch 訓練のため、同一環境であれば数値まで一致する。8 GB VRAM クラス GPU で sweep 全体が 47 分で完了する軽量さも特徴。