Loop-Index Embedding は深さ外挿と浅ループ解を Trade-off する — OpenMythos における Inference-time Scaling の 2 つのレジーム

要旨（Abstract）

Recurrent-Depth Transformer（RDT）の重要な性質に depth extrapolation（深さ外挿） がある：訓練時の loop 回数 T で学習したモデルを、推論時に T 回以外の loop 数で走らせても性能を維持する能力。これは RDT の「推論時計算量を動的にスケールできる」という魅力の理論的根拠となっている（Saunshi et al., 2025）。しかし、depth extrapolation が どの条件で成立し、どの条件で破綻するか は実装依存性が強く、系統的には特徴づけられていなかった。

本研究は前 2 報（Paper 1: A1 幅×ループ sweep / Paper 2: LIE ablation）で得られた 8 つの checkpoint を eval_loops ∈ {1, 2, 3, 4, 6, 8, 12, 16, 24, 32} で再評価し、depth extrapolation の挙動を系統的に観察した。結果、十分な容量（dim = 128）では LIE 有無にかかわらず全 loop 数で 100 % が維持されるが、タイト容量（dim = 64）では LIE on / off の選択が外挿の可否を決定的に分ける ことが判明した。

特に dim = 64, loops = 2, LIE on の組み合わせは、訓練時と完全一致する eval_loops = 2 でのみ 100 %、それ以外では 2–14 % に崩壊 する「ピンポイント特化」を示す。一方 LIE off は eval_loops ≥ train_loops の任意値で 100 % を維持する fixed-point 収束 型の汎化を学習している。

主要な結果

• dim = 128 では extrapolation は non-issue：LIE on/off、train_loops 2/4 のいずれも eval_loops ∈ [1, 32] の全域で test_acc 100 %。
• dim = 64 で 4 つの異なるメカニズムが分離観察された：LIE × train_loops の組み合わせにより、ピンポイント特化 / fixed-point 収束 / 浅ループ一発解 / 最小反復要求 の 4 レジームが明瞭に分離。
• LIE on は train_loops に過剰適応し外挿を破壊する：d64L2 + LIE on は eval_loops = 3 で 13.6 %、eval_loops ≥ 4 で 2 %（クラス数 97 におけるランダム未満）まで崩壊。
• LIE on は eval_loops = 1 で驚くべき性能：d64L4 + LIE on は 4 loop で訓練されたにもかかわらず 1 loop で 93 % を達成。LIE が初期 loop に実質的な計算を強制した結果と解釈。
• LIE off は fixed-point 学習：d64L2 + LIE off は eval_loops ∈ [2, 32] 全てで 100 %。同じ operator を任意回数適用して収束する構造を内部的に学んでいる。

1. 背景と問題意識

1.1 研究の背景

Depth extrapolation は RDT の設計上の中核機能で、Saunshi et al. (2025) は looped transformer を用いた合成タスクで、train_loops より多い eval_loops でも性能が落ちず、むしろ向上するケースすらあることを示した。これは通常の固定深さ transformer にはない性質で、inference-time compute scaling を可能にする。

Inference-time scaling は、近年の o1-style large reasoning models が示した「テスト時により長く考えると精度が上がる」パラダイムの核心的な仕組み。RDT 版 o1 を実装するためには、depth extrapolation が 確実に 成立することが必須条件となる。

一方、OpenMythos のような実装では、depth extrapolation を 支援する機構 として以下が搭載されている：

1. LTI-stable injection：ρ(A) < 1 を構造的に保証することで、任意の loop 数で発散しないことを保証（Parcae 2026）。
2. per-loop LoRA adapter：max_loop_iters まで loop ごとの scale vector を学習、超える場合は最終 loop の scale を使い回し（clamp）。
3. Loop-Index Embedding (LIE)：loop 番号を正弦波 encoding として hidden に加算することで、各 loop で異なる計算を学習する余地を与える。

ここで興味深い構造的緊張がある：

• LIE は loop ごとの差別化を促す → 外挿時に「見たことのない t」が来ると問題になる可能性。
• LTI injection は 任意 loop 数での安定性を保証 → 外挿を支援。

これらが競合しているのか協調しているのかは実証されていなかった。

1.2 研究目的

本研究の目的は以下の点に集約される：

1. Depth extrapolation の成立条件の特徴づけ：どの (dim, train_loops, LIE, eval_loops) 組み合わせで成立するか。
2. A1 / A2 知見の統合：前 2 報の所見（容量閾値・LIE の害）は extrapolation 挙動にどう反映されるか。
3. Inference-time scaling のための実装指針：実用的に使える RDT を構築するためのベストプラクティス。

1.3 関連研究との位置づけ

• Saunshi et al. (2025)：depth extrapolation 現象の実証。本研究はその 成立条件の境界 を初めて明示化する試み。
• Parcae et al. (2026)：LTI-stable 注入で任意 loop 数の安定性を保証。本研究は 安定 ≠ extrapolation を示し、安定性は必要条件だが十分条件ではないことを実証する。
• Bae et al. (2024)：depth-wise LoRA の外挿効果。本研究はそれと LIE の相互作用を ablation で分離する。
• o1 / Deep Thinking：inference-time scaling paradigm。RDT 実装で同等の機能を得るための設計的論点を提供。

2. 研究手法

2.1 使用する Checkpoint

前 2 報（Paper 1 / Paper 2）で訓練された 8 モデル：

全て 5,000 step 訓練済み、test_acc 100 % 到達済み（2 checkpoint は初期段階で grokking 済みの final state）。

2.2 Checkpoint の保存と再読み込み

A1 / A2 sweep のコードを改修し、訓練終了時に model.state_dict() と MythosConfig を model.pt として各 run ディレクトリに保存するようにした。これにより inference-time の loop sweep を容易に実施できる。

2.3 評価プロトコル

各 checkpoint に対し、eval_loops ∈ {1, 2, 3, 4, 6, 8, 12, 16, 24, 32} の 10 点で以下を評価：

1. モデルを .eval() モード、torch.no_grad() 下で読み込み。
2. MythosConfig.act_threshold = 1.01 に上書き（ACT halting を無効化）。これにより、要求した n_loops が必ず全て実行される。
3. Test 全 6,586 examples に対し model(x, n_loops=L) を呼び、logits[:, -1, :] と正解 c の accuracy / loss を測定。

ACT 無効化は重要：LIE on モデルは ACT を使って早期停止する可能性 があり、その場合 eval_loops の意味が曖昧になる。固定深さでの evaluation を行うため、halting を抑制する。

2.4 評価指標

• test_acc(eval_loops)：6,586 test examples に対する top-1 accuracy。
• extrapolation range：100 % を維持する eval_loops の連続区間。
• degradation threshold：test_acc < 50 % に落ちる最小 eval_loops（下方崩壊の指標）。

2.5 実験環境

• Hardware：NVIDIA GPU (CUDA, 8 GB VRAM)
• Software：Python 3.10.9 + PyTorch 2.11.0+cu128
• Wall-clock：8 checkpoints × 10 eval points × 6586 examples ≒ 7 秒（全 inference）

3. 結果

3.1 全貌：Extrapolation 挙動の 2 層構造

表1: 8 checkpoint × 10 eval_loops の完全 matrix（test_acc, %）

3.2 dim = 128：Extrapolation は解決済み問題

dim = 128 の 4 config は LIE on/off、train_loops = 2/4 にかかわらず、eval_loops = 1 から eval_loops = 32 まで全てで 100 %。train_loops = 2 で訓練して eval_loops = 32 で 16 倍の iteration をしても性能は一切劣化しない。

eval_loops = 1（下方外挿）についても同様に 100 %。これは 1 loop でも十分な処理ができることを示し、recurrent block が実効的に identity に近い fixed-point operator を学習している可能性を示唆。

この結果は 十分な容量があれば depth extrapolation は自然に成立する ことを意味する。実装上の工夫（LIE、LoRA、LTI injection）は、この regime では無用ないし等価。

3.3 dim = 64：4 つの異なる外挿メカニズム

dim = 64 では LIE と train_loops の組み合わせにより 全く異なる 4 つの挙動 が分離される。

パターン A: ピンポイント特化（d64L2 + LIE on）

eval_loops:    1     2     3     4     6     8    12    16    24    32
test_acc:    1.8  100.0  13.6   3.8   2.1   2.0   2.0   2.0   2.0   2.0

訓練時と完全一致する eval_loops = 2 でのみ 100 %、それ以外では 97 クラスのランダム（1 %）並か以下に崩壊。eval_loops = 3 での 13.6 % は 2 loop 目までで偶然解ける examples の残響と思われる。長い eval_loops では 完全に崩壊して 2 % に張り付く。

この挙動は モデルが「step-0 でこれをする、step-1 でそれをする、終了」 という特化したプログラムを学習してしまったことを示す。LIE の sinusoidal signal が各 step を意味的にラベル付けし、t ≥ 2 ではラベル外のため未知の挙動に落ち込む。

パターン B: 完全な fixed-point 上方外挿（d64L2 + LIE off）

eval_loops:    1     2     3     4     6     8    12    16    24    32
test_acc:    2.9  100.0 100.0 100.0 100.0 100.0 100.0 100.0 100.0 100.0

eval_loops = 2 から 32 まで全て 100 %。同じ recurrent operator を任意回数適用して answer に収束する、真の fixed-point 学習。これこそが RDT の理想形で、inference-time scaling の基盤となる挙動。

eval_loops = 1 で 2.9 % なのは、prelude + 1 recurrent iteration + coda だけでは収束が不十分であることを示す（次項のパターン D と同系）。

パターン C: 浅ループ一発解 + 無限上方外挿（d64L4 + LIE on）

eval_loops:    1     2     3     4     6     8    12    16    24    32
test_acc:   93.0  100.0 100.0 100.0 100.0 100.0 100.0 100.0 100.0 100.0

訓練時 loops = 4 にもかかわらず、eval_loops = 1 で既に 93 %、全体として極めて robust。LIE on なのに外挿が効く理由は、以下の仮説で説明できる：

• 4 loop で訓練するとき LIE の t=0,1,2,3 の signal が全 step で一貫して解けるようモデルが学習
• その結果 loop 1 だけでも task が解ける degenerate な解 に収束（LIE off ではこれが起きない）
• かつ 4 loops 以上でも各 t の signal が extrapolation の "最後の learned t" のまま反復されるため崩れない

つまりこれは luck の要素が強く、毎 seed で再現するかは要確認。しかし「LIE on でも外挿が効くケースがある」ことの確かな例。

パターン D: 最小反復要求 + 無限上方外挿（d64L4 + LIE off）

eval_loops:    1     2     3     4     6     8    12    16    24    32
test_acc:    1.1   1.1  74.5 100.0 100.0 100.0 100.0 100.0 100.0 100.0

eval_loops = 1, 2 では 1.1 %（ランダム）、eval_loops = 3 で 74.5 %、eval_loops ≥ 4 で 100 %。最低 3–4 loops 必要だが、それ以上はいくらでも。

これは「真の fixed-point 学習」だが、4 loop で訓練されたせいで 収束に最低 3 loops 必要な attractor を学んでしまったケース。d64L2 + LIE off が 2 loops から収束したのと対照的。Train_loops が収束に必要な step 数の下限を決める ことを示唆する。

3.4 結果のまとめ

1. dim = 128 では extrapolation は非問題 — LIE や train_loops の選択によらず universal に成立。
2. dim = 64 では 4 つの異なる外挿レジームが分離される — LIE × train_loops の組み合わせごとに異なる学習機序。
3. LIE off は fixed-point 収束を学習する傾向がある — eval_loops ≥ 収束最小 step で universal に 100 %。
4. LIE on はピンポイント特化のリスクを持つが、浅ループ一発解 enable の御利益もある。

4. 考察

4.1 2 つの学習レジームの理論的解釈

観察された挙動は、以下の 2 つの学習レジームに分類できる：

レジーム 1: Fixed-Point 学習

Recurrent block を operator $F$ とみなしたとき、$F$ が contractive になるよう学習されるケース。つまり、$F(h^*) = h^*$ を満たす不動点 $h^*$ に向かって任意の初期値から収束する。

• LIE off モデルはこのレジームに落ち込みやすい
• eval_loops に依存しない、または下限以上ならばどこでも同じ answer を出す
• inference-time scaling の理想形

レジーム 2: Finite-Step プログラム学習

各 loop で異なる step-specific な処理を学習し、合計 $T$ step で解く有限プログラムを書くケース。

• LIE on モデルはこのレジームに落ち込みやすい（LIE が step index を露出するため）
• T = train_loops でのみ動作する 可能性が高い
• 上方外挿 (T > train_loops) が崩壊する

この 2 レジームは tradeoff 関係にある：

• レジーム 1 は外挿に強いが浅ループ（T < train_loops）で未収束
• レジーム 2 は浅ループで解けるかもしれないが（d64L4 + LIE on のケース）、外挿が脆弱

Dim = 128 では両レジームが成立する冗長性があり、LIE の選択によらず外挿が効く。Dim = 64 では冗長性がなく、LIE の選択で学習レジームが固定される。

4.2 なぜ d64L4 + LIE on は外挿できたのか（パターン C の分析）

このケースは最も理解困難。d64L4 + LIE on は LIE on にもかかわらず外挿が完璧で、さらに eval_loops = 1 でも 93 % を達成する。

仮説 A: Degenerate solution

4 loop の全てで似たような処理を学習した結果、モデルが各 t において「収束済みの h を identity pass する」挙動になった可能性。LIE の signal は ε オーダーで効いているが、既に収束した状態では影響がない。これは一種のラッキーな perfect generalization。

仮説 B: Prelude/Coda が主役

eval_loops = 1 で 93 % ということは、prelude + 1 recurrent iteration + coda だけで 93 % 解けている。つまり recurrent block はせいぜい 7 % の refinement しか提供していない。この場合、LIE の sinusoidal も meaningful な役割を果たしておらず、結果として全 eval_loops で似たような答えになる。

どちらの仮説も、LIE の設計意図（step-specific な学習）が発現していないから外挿できている という皮肉な解釈。この挙動が seed robust かどうかは要検証。

4.3 A1 / A2 知見との統合

本研究を前 2 報と統合すると、OpenMythos RDT の設計論として以下の統一的な図が浮かび上がる：

        │ Dim が十分 (≥128)       │ Dim がタイト (=64)
────────┼────────────────────────┼─────────────────────────
Grokking│ loops=1 が最速          │ loops≥2 が必須
 (A1)   │ 追加 loops は 2.4× 遅延 │ loops が capacity を補償
────────┼────────────────────────┼─────────────────────────
LIE off │ grokking 1.5-1.9× 高速 │ grokking 1.5-3.3× 高速
 (A2)   │ 最終精度は不変         │ 最終精度は不変
────────┼────────────────────────┼─────────────────────────
Depth   │ どの設定でも 100 %      │ LIE off → fixed-point 外挿
Extrap  │ 外挿が効く             │ LIE on → ピンポイント特化で崩壊
 (本報) │                        │ 例外あり（D64L4+LIE on の lucky case）

統合的な実装指針：

1. 容量に余裕があるなら LIE off が常に合理的：grokking も速く、外挿も効く、何も失わない。
2. タイト容量なら状況依存：extrapolation が要件なら LIE off 必須、train_loops と eval_loops が一致することが保証されるなら LIE on でも可。
3. 未知タスクへの適用時は LIE off でスタート：homogeneous / fixed-point 型の学習を優先する方が、汎用性が高い。

4.4 Inference-time Scaling のための実装要件

o1-style の「推論時にもっと考えると精度向上」を OpenMythos で実現するには、本研究から以下の設計要件が導ける：

1. LIE off が必須条件：fixed-point 収束を学習させないと、eval_loops を増やすメリットがない。
2. LTI injection は前提：ρ(A) < 1 保証がなければ、eval_loops を増やすほど h が発散する。本研究のチェックポイントは全てこれを満たす（事後観測で ρ(A) < 0.5 程度）。
3. Train_loops ≥ 3–4 が推奨：収束する operator を学習させるためには、複数回の適用で同じ答えになることを訓練中に要求する必要がある。
4. Dim は task の capacity_floor 以上：本研究の modular arithmetic では dim = 128 が安全側。より複雑なタスクでは閾値も高くなる。

4.5 限界と制約

1. 単一タスク：modular arithmetic は homogeneous task の代表例。heterogeneous な iterated computation（k-hop reasoning、program execution、tree-structured tasks）では結論が変わる可能性。
2. 単一 seed：すべて SEED = 42。特に d64L4 + LIE on の lucky extrapolation が seed robust かは未検証。
3. ACT を無効化した評価：実用では ACT を活用する方が実効 compute が抑えられる。ACT on での extrapolation 挙動は別検証が必要。
4. 内部状態の観測を行っていない：Fixed-point 学習と finite-step 学習を仮説しているが、内部表現の直接観測（linear probing、PCA trajectory、Hopfield energy の時系列）は未実施。機構的解釈（mechanistic interpretability）は future work。

4.6 実用的な示唆

• OpenMythos 実装者への直接的提言：
• 深さ外挿をアピールしたいなら use_loop_index_embedding = False をデフォルトに変更する検討価値あり。
• 「外挿テストで性能が落ちた」と報告する際は、LIE on 状態で評価していないか確認。

• LoRAAdapter の clamping 挙動（t > max_loops で最終 scale を使い回す）は外挿に寄与している可能性があり、これ単体の ablation も研究価値あり。

• RDT 研究コミュニティへの示唆：

• 「Recurrent transformer は自然に extrapolation できる」という主張は、容量によって成立条件が大きく変わる ことを認識すべき。

• 「LIE（sinusoidal loop encoding）は有用」という一般論は homogeneous task では偽。タスク特性に依存した設計選択。

• Mechanistic interpretability 研究との接続：

• パターン A（d64L2 + LIE on）の崩壊時の hidden state を観察すれば、「learned program が壊れる瞬間」の clean な実例となる。
• パターン B（fixed-point）と B の attractor basin を PCA で可視化すれば、RDT の「考える軌跡」が直接見える。

5. 結論

5.1 知見の要約

1. Dim = 128 で depth extrapolation は universal — LIE、train_loops によらず eval_loops ∈ [1, 32] で 100 %。
2. Dim = 64 で 4 つの学習レジームが分離観察された — LIE on は finite-step プログラム学習、LIE off は fixed-point 学習に落ち込む傾向。
3. LIE on は train_loops にピンポイント特化するリスク — d64L2 + LIE on は eval_loops = 2 以外で 2 % に崩壊。
4. LIE off は上方外挿に強く、浅ループ下方外挿に弱い — eval_loops ≥ 収束最小 step で universal。
5. A1 + A2 + 本報の統合図：容量に余裕があれば LIE off が全指標で優位。Inference-time scaling を目指すなら LIE off を推奨。

5.2 今後の課題

1. Heterogeneous task での追試：k-hop reasoning や program execution で、LIE on の finite-step 学習が真価を発揮する（外挿が効く）可能性を検証。
2. Multi-seed 化：特にパターン C（d64L4 + LIE on の lucky extrapolation）の頑健性を 5 seeds で再現性確認。
3. Mechanistic interpretability：
4. パターン A 崩壊時の hidden state を PCA / causal intervention で分析
5. パターン B の attractor を Hopfield energy 曲線で可視化
6. LIE 信号が注入されるチャネルと attention head の関係を attention pattern 可視化
7. ACT halting との相互作用：ACT on での extrapolation 挙動は本研究では未評価。ACT が「train_loops の長さを自動検出して halt する」可能性があり、その場合 inference-time scaling の実装がさらに複雑化する。
8. 他 RDT 実装での追試：Universal Transformer、Coconut、Parcae のオリジナル実装でも同じ性質が観察されるかの比較研究。

付録

A. 実行環境

• GPU：NVIDIA GPU (CUDA, 8 GB VRAM)
• 言語：Python 3.10.9
• 主要ライブラリ：PyTorch 2.11.0+cu128

B. 補足データ

Extrapolation 時の test_loss（抜粋）

d64L2 + LIE on の loss は eval_loops が進むほど増加しており、単なる accuracy の崩壊だけでなく 確率分布そのものが悪化している（言い換えれば「違う答えに自信を持つ」方向に崩壊）。

C. 用語集

シリーズ連作としての位置づけ

本研究は以下の 3 報連作の最終篇：

1. Paper 1 (A1): 「再帰深さ型 Transformer における幅とループ数の交換可能性」 → loops は容量閾値以下でのみ幅を代替、それ以上では逆効果
2. Paper 2 (A2): 「Loop-Index Embedding は homogeneous タスクの Grokking を遅延させる」 → LIE は 1.5–3.3× の grokking 遅延、LoRA と機能重複
3. Paper 3 (本報): 「Loop-Index Embedding は深さ外挿と浅ループ解を Trade-off する」 → LIE on は finite-step 学習、LIE off は fixed-point 学習、inference-time scaling には後者必須

3 報を通じて、OpenMythos RDT の設計論における 「容量 × ループ × LIE」3 軸のインタラクション を実証的に特徴づけ、実用的な設計指針を導出した。

Checkpoint の保存さえあれば extrapolation 評価は数秒で完了するため、新規 config の検証も安価。次の課題は heterogeneous task や mechanistic interpretability による仮説検証。