Loop-Index Embedding は homogeneous タスクの Grokking を遅延させる — OpenMythos RecurrentBlock の Ablation 実験

要旨（Abstract）

OpenMythos の RecurrentBlock には、loop ごとの差別化のために正弦波ベースの loop-index embedding（以下 LIE） が搭載されている。この機構は、同一の weight-tied ブロックを T 回適用する RDT において「どの loop にいるか」を hidden state に注入することで、ブロックが iteration-specific な処理を学習することを狙ったものである。しかし OpenMythos README は「この機構の効果検証は未解決」と明記しており、実際に LIE が grokking に与える効果は不明だった。

本研究は、modular-arithmetic grokking タスク上で use_loop_index_embedding フラグを on/off した 4 config の ablation を実施し、4 ケース全てで LIE をオフにした方が汎化到達が 1.5–3.3 倍速い ことを示す。最大効果は最低容量設定（dim = 64, loops = 2）で観察された 3.3 倍高速化で、これは LIE が dim / 8 チャネル分の表現容量を消費している、という channel-budget 仮説 と整合的である。最終 test_acc は両設定とも 100 % のため、LIE は「到達点」ではなく「速度」のみに効く。

主要な結果

• 全 4 config で LIE off が高速：d64L2 で 3.3×、d64L4 で 1.5×、d128L2 で 1.9×、d128L4 で 1.7× の grokking 高速化。
• 容量が小さいほど LIE の害が大きい：dim = 64 で最大効果、dim = 128 で縮小。Channel-budget 仮説と一致。
• 最終精度は不変：LIE on / off どちらも最終 test_acc 100 %。LIE は asymptote ではなく transition のタイミングのみを変える。
• LoRAAdapter との機能重複の可能性：LoRAAdapter が既に per-loop scale vector を学習しており、LIE の信号は冗長で optimizer の negative noise になっている疑い。

1. 背景と問題意識

1.1 研究の背景

OpenMythos の RDT 実装では、同一の TransformerBlock を RecurrentBlock 内で n_loops 回繰り返し適用する。このままでは全ての loop で文字通り同じ計算が行われるため、loop-specific な挙動 を学習させるための差別化機構が設計上必要になる。

OpenMythos は2つの差別化機構を併用している：

1. LoRAAdapter：per-loop scale vector を nn.Embedding(max_loops, rank) で学習。low-rank の重み微分を各 loop 毎に微調整する。
2. Loop-Index Embedding (LIE)：hidden state h の先頭 dim / 8 チャネルに、loop 番号 t を sinusoidal encoding した信号を 加算 する。RoPE の深さ版という位置づけ。

LoRAAdapter は明確な per-loop パラメータを持つのに対し、LIE は パラメータを持たない固定変換 で、h の一部チャネルを「loop index 専用」に割り当てる形で機能する。

OpenMythos README (lines 297-303) は「LIE の効果検証は未完了」と明記しており、本機構が実際に grokking や extrapolation にプラスに働くのかは未実証だった。

1.2 研究目的

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

1. LIE の単独効果の分離：ほか全ての条件（dim、loops、optimizer、data、seed）を固定して use_loop_index_embedding フラグのみ切り替えた ablation を行う。
2. 効果の方向性と大きさ：LIE は grokking を速めるのか遅めるのか、効果量は構成によってどう変わるか。
3. メカニズムの同定：もし効果があるなら、それはどの回路（attention / MoE router / residual stream のどこ）で起きているか。

1.3 関連研究との位置づけ

• Rotary Position Embedding (Su et al., 2021)：token position を正弦波複素数で hidden に注入する設計。LIE はこの「深さ版」として設計されている。
• Universal Transformer (Dehghani et al., 2018)：同じブロックを繰り返し適用する古典 RDT。loop index の扱いはモデルごとに異なる。
• LoRA (Hu et al., 2021)：low-rank adaptation による効率的な fine-tuning。OpenMythos の LoRAAdapter は per-loop の scale vector のみを学習する軽量版。
• Bae et al. (2024)：depth-wise LoRA の depth extrapolation 効果の実証。

本研究は、同じ目的（loop 差別化）を持つ2つの機構（LoRA と LIE）が共存していること を指摘し、それらの分離実験として位置づけられる。

2. 研究手法

2.1 Ablation 設計

前報 A1 sweep のうち、loops の効果が実際に観察された 4 config（dim = 64, loops = 2, dim = 64, loops = 4, dim = 128, loops = 2, dim = 128, loops = 4）を LIE on / off の 2 バージョンで再訓練する。loops = 1 は定義上 LIE が t = 0 のみで動作するため除外。dim = 256 も LIE 有無の差異が最も出にくいと予想されるため省略。

合計 4 configs × 2 (LIE on / off) = 8 runs。

2.2 LIE の実装と ablation

loop_index_embedding(h, t, loop_dim) は以下の計算を行う：

freqs = 1 / theta ** (range(0, loop_dim, 2) / loop_dim)
angles = t * freqs                                # shape: (loop_dim/2,)
emb = cat(sin(angles), cos(angles))[:loop_dim]    # shape: (loop_dim,)
h[..., :loop_dim] += emb                          # 加算（非破壊的な拡張ではない）

loop_dim = dim / 8（dim = 64 なら 8 チャネル、dim = 128 なら 16 チャネル）。theta = 10000 は RoPE と同様の base frequency。

Ablation は MythosConfig.use_loop_index_embedding: bool フラグを追加し、RecurrentBlock.forward 内で以下のように gating：

for t in range(n_loops):
    if self.cfg.use_loop_index_embedding:
        h_loop = loop_index_embedding(h, t, self.loop_dim)
    else:
        h_loop = h                                 # 何もしない
    ...

2.3 訓練条件

前報 A1 と完全同一。(a + b) mod 97、train_frac = 0.3、AdamW WD = 1.0、LR 1e-3 → 1e-4、5,000 steps、fp32、SEED = 42、full-batch。唯一の違いは use_loop_index_embedding フラグのみ。

2.4 評価指標

• step_grok：test_acc ≥ 99 % を初めて達成した step。
• speedup = step_grok (LIE on) / step_grok (LIE off)。> 1 なら LIE off が速い。
• 最終 test_acc：5,000 step 時点での test accuracy。asymptote 差を検出。

2.5 実験環境

• Hardware：NVIDIA GPU (CUDA, 8 GB VRAM)
• Software：Python 3.10.9 + PyTorch 2.11.0+cu128
• Wall-clock：LIE off 側 sweep（4 runs × 5000 steps）を 24.4 分で完了。LIE on 側は前報 A1 の結果を流用。

3. 結果

3.1 定量比較：全 4 config で LIE off が勝利

表1: step_grok 比較（4 configs × LIE on/off）

4/4 すべてで LIE off が勝利。効果は単調でないが、最大値は低容量設定（dim = 64, loops = 2）で観察された 3.3× 高速化。

3.2 最終精度：両設定とも 100 %

LIE on / off いずれも 5,000 step 終了時に test_acc = 100 % に到達する。つまり LIE は汎化の最終到達点には影響せず、そこに至る速度のみを変える。

この性質は実用的に重要：長く訓練すれば LIE on でも問題ないが、wall-clock を短縮したいなら LIE off の方が常に好ましい。

3.3 効果量の dim 依存：容量が小さいほど LIE の害が大きい

dim 軸でまとめると：

また loops 軸でまとめると：

最大効果は dim = 64, loops = 2 = 2つの最低設定の交差点。これは容量が最もカツカツで、かつ LIE の信号が loop 全体に占める割合が大きい config。「容量が足りないほど LIE の channel 消費が効く」 という仮説に整合的。

3.4 Memorization への影響は軽微

step_mem（train_acc ≥ 99 %）も比較：

dim=64 でのみ LIE off が 25 % 速く memorization する（400→300）が、dim=128 では差がない。LIE の負の効果は主に memorization → generalization の transition 部分にかかっていることを示唆。

3.5 結果のまとめ

1. Modular arithmetic タスクで LIE は 一貫して grokking を遅延させる（4/4）。
2. 効果量は 1.5× 〜 3.3× で、容量が小さいほど悪化する。
3. 最終精度は変わらず、LIE の影響は transition timing のみ。
4. Memorization への影響は軽微で、害は generalization phase に集中する。

4. 考察

4.1 Channel-Budget 仮説

LIE は hidden state の先頭 dim / 8 チャネルに sinusoidal 信号を 加算 する。これは、対応するチャネルが「loop index 情報の運搬」に恒久的に占有されることを意味する。

dim = 64 の場合：loop_dim = 8。total 64 チャネルのうち 8 チャネル（12.5 %）が loop 情報専用。残り 56 チャネルでアルゴリズム構造（加算演算子、mod 演算、embedding lookup, …）を学ばなければならない。dim = 128 では loop_dim = 16、占有率は同じ 12.5 % だが、利用可能チャネル数が 112 → 許容度が高い。

この見方で観察を整理すると：

• dim = 64, loops = 2：56 チャネルで task を解くのに加え、8 チャネルに注入された sinusoidal ノイズが attention pattern を乱す → 最大の害（3.3×）
• dim = 128, loops = 2/4：112 チャネルは余裕があるが、attention weights が LIE 信号を denoise する必要は残る → 中程度の害（1.7–1.9×）

4.2 LoRAAdapter との機能重複仮説

OpenMythos の LoRAAdapter は以下を学習する：

self.scale = nn.Embedding(max_loops, rank)    # per-loop scale vector
delta(x, t) = (down(x) * scale[t]) @ B

これは per-loop の identity-adjacent 変換を学習する仕組み。scale[0]、scale[1]、... が loop ごとに異なる値を取ることで、各 loop で微妙に違う計算が可能。

LoRAAdapter は： - 学習可能な per-loop 差別化 - low-rank なので優雅な degradation (rank=4 なので 128×4+4×128 = 1,024 params) - inference 時にも scale[t] を lookup

対して LIE は： - 非学習な per-loop 差別化（固定 sinusoidal） - 全 loop に対して 同じ信号形状（theta, freqs は固定） - hidden state に加算 するので、task 表現と干渉

つまり LoRAAdapter が既に 「loop ごとの処理差別化」 という役割をより適応的に果たしており、LIE の信号は redundant で denoise コストだけがかかる 可能性が高い。

LoRAAdapter の learned scale は task-relevant な differentiation を学べるが、LIE の sinusoidal は task と無関係な「座標」を固定的に与える。後者の方が最適化に不利、という解釈。

4.3 Task Homogeneity の観点

(a + b) mod p は 各 loop で同じ計算 を繰り返すべきタスク：

• Loop 1: 部分和の計算（？）
• Loop 2: 残りの整形（？）
• … どの loop も似たような線形演算の反復

このように 各 loop で似た処理 が自然なタスクでは、loop index を知ること自体に情報価値がない。LIE の「どの loop にいるかのシグナル」は、モデルにとって 使い道がなく、denoise すべきノイズ に過ぎない。

反対に loop ごとに異なる処理が望ましい heterogeneous なタスク（例：k-hop reasoning — loop 1 で最初の hop、loop 2 で 2 番目 …）では LIE が有用な可能性がある。これは本研究の重要な 未検証領域。

4.4 限界と制約

1. 単一タスク：(a+b) mod p は極めて homogeneous なタスク。LIE の効果は heterogeneous な iterated computation で逆転する可能性があり、結論を他 task に一般化してはならない。
2. 固定 seed：全 runs SEED = 42 単一。grokking の stochasticity を考慮すると、3–5 seeds での追試が望ましい。
3. LoRAAdapter の分離実験未実施：「LIE は LoRA と重複」という仮説は、LoRA も off にした 追加 ablation で直接検証できる。現在は間接的な argumentation のみ。
4. Depth extrapolation 未検証：LIE の設計意図は「train_loops ≠ eval_loops 条件下での性能維持」。そのケースで LIE が真価を発揮する可能性があり、本研究は train_loops = eval_loops の条件下のみを扱う。→ 本シリーズ第3報で検証。

4.5 実用的な示唆

• OpenMythos 実装：homogeneous な grokking タスクを扱う場合は MythosConfig.use_loop_index_embedding = False を設定するのが合理的。Flag はデフォルト True のまま温存してよい（heterogeneous task や extrapolation 実験では on が必要かもしれない）。
• 新 RDT 設計者への教訓：同じ目的（loop 差別化）を持つ機構が複数ある場合、それらの相互作用を ablation で確かめる必要がある。2 つ足して 1 より悪い可能性は十分にある。
• 診断ツール：「LIE を外すと速くなる」こと自体が、タスクが homogeneous であることの診断に使える。

5. 結論

5.1 知見の要約

1. Modular arithmetic grokking において、Loop-Index Embedding は 1.5–3.3 倍 grokking を遅延させる（4/4 config）。
2. 効果量は容量が小さいほど大きく、dim = 64, loops = 2 で最大の 3.3× 高速化が観測された。
3. LIE は asymptote（最終精度）には影響せず、transition のタイミングのみを変える。
4. 仮説：LIE の dim / 8 チャネル占有が表現容量を圧迫、かつ LoRAAdapter が既に per-loop 差別化を担うため LIE は redundant な denoise コスト。

5.2 今後の課題

1. Heterogeneous task での追試：k-hop reasoning、program execution など loop ごとの処理が異なるタスクで LIE の符号が逆転するか検証。
2. LoRA もオフにした 3-way ablation：{LIE + LoRA, LIE only, LoRA only, neither} の 4 条件で LoRA との重複性仮説を直接検証。
3. Depth extrapolation 実験：LIE の設計意図は depth-extrapolation 時の loop differentiation。train_loops ≠ eval_loops 条件で効果が発現するか。→ 本シリーズ第3報 Depth Extrapolation で検証
4. Mechanistic interpretability：LIE 信号が注入されるチャネルと attention head の相互作用を linear probe や causal intervention で分析。

付録

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

B. 補足データ

各 config の詳細 metrics（抜粋）

LIE off は memorization phase からの移行が滑らか（test_acc 35 % 段階を経由）なのに対し、LIE on は test_acc 10 % → 99 % への jump が急で、到達まで長い。

C. 用語集

LIE on 側は A1 sweep の結果を流用、LIE off 側は追加で 25 分。grokking は seed 依存の stochastic な現象のため、正確な再現には SEED = 42 の保持が必要。