Empirical test of the depth-extrapolation claim: U-shape with fixed-loop training, flat plateau with random-loop training

[tonyzdev]

TL;DR

I reproducibly trained three variants of a 117.8 M-parameter OpenMythos (MLA + MoE, dim=768, n_experts=16, 2 prelude + 2 coda layers, GPT-2 tokenizer) on ~491 M tokens of HuggingFaceFW/fineweb-edu on a single H100 SXM, then swept n_loops at inference. The result contradicts the depth-extrapolation claim on this codebase: neither training style produces a monotonically decreasing PPL curve, which is the defining empirical signature of depth extrapolation.

All three runs share the exact same backbone, data, tokenizer, optimizer, and step count. The only difference is how n_loops is chosen during training:

Run	Training n_loops	Final train loss	Best eval PPL	PPL(n_loops=16)
baseline_1	fixed at 1	4.19	70.0 @ 1	—
looped_8	fixed at 8	3.99	57.9 @ 8	170.7 (!)
looped_random	uniform random ∈ {4, 6, 8, 12, 16}	4.07	65.2 @ any of {4,6,8,12,16}	65.2

ρ(A) stays at ~0.357 across all three runs (Parcae LTI stability behaves as advertised, see rho_A.png).

The figure

Three qualitatively different shapes:

• looped_8 (blue): a sharp V minimum exactly at the training loop count. Moving inference depth by one step in either direction (7 or 9 would interpolate between the points, 6 and 12 already tripled PPL) degrades drastically. PPL(1)=281, PPL(8)=58, PPL(16)=171.

• looped_random (green): catastrophic below the minimum training depth (PPL(1)=1217, PPL(2)=401); a knee at n_loops=4 (the minimum training value); then a completely flat plateau for any loop count in [4, 16]. PPLs at 4/6/8/12/16 agree to ≤ 0.1%:

  n_loops  PPL(looped_random)
        4   65.347
        6   65.237
        8   65.237
       12   65.240
       16   65.236
  

• baseline_1 (orange): a single point at n_loops=1, PPL 70.0.

Why the sharp V in looped_8

Two mechanisms currently make each loop step functionally distinct:

1. loop_index_embedding injects sin(t · ω) / cos(t · ω) into the first dim/8 channels of h at each step t (main.py loop_index_embedding)
2. LoRAAdapter uses self.scale = nn.Embedding(max_loops, rank), i.e. a per-loop-depth learnable scale vector that multiplies the LoRA intermediate representation (main.py LoRAAdapter.forward)

Both were introduced intentionally — the README argues for them on the grounds that "each loop is not a repetition — it is a distinct computational phase". But under fixed-n_loops training, they also tightly bind each loop step to a specific role: the scale for step t is only ever gradient-updated on inputs that passed through steps 0..t-1 with their corresponding sin/cos signatures. At inference, running fewer loops cuts the reasoning chain short; running more replays the final LoRA row (thanks to the clamp in #10 / the existing t_idx = loop_t if loop_t <= max_t else max_t) but against sin/cos phases the scale has never seen, so behaviour becomes OOD.

Why the plateau in looped_random

Replacing model(x, n_loops=cfg.max_loop_iters) with model(x, n_loops=random.choice([4,6,8,12,16])) makes every in-range depth a valid training signal. The V disappears — good — but rather than collapsing onto a monotonically-decreasing curve, it collapses onto a flat one: PPL is identical to four decimal places across the entire training range.

That is also consistent with the mechanism: the model has learned a function that converges within 4 loops and produces the same output thereafter, regardless of how many additional loops you give it. Extra compute at inference buys nothing. This is the opposite of the claim that "more loops at inference = deeper reasoning chains = harder problems solved" — more loops simply idle.

What this means for the central claim

Both training strategies tested leave OpenMythos without the inference-time depth extrapolation behaviour the README advertises:

• Fixed loops → single-point optimum, no scaling in either direction.
• Random loops → robust plateau, but also no scaling.

Saunshi et al. 2025 / Parcae-style depth extrapolation appears to require architectural choices that OpenMythos has chosen not to make (purely anonymous loops without per-step embeddings/adapters, or a substantially different curriculum). The current loop_index_embedding + per-loop LoRAAdapter.scale design seems to be the proximate cause of the failure — they improve peak quality at the training loop count at the direct cost of depth extrapolation.

What works as claimed (positives)

• LTI stability is real. Across all three runs ρ(A) stays pinned at ≈ 0.357 for every one of 15,000 steps. No divergence, no drift, no residual explosion. Parcae's LTI injection behaves exactly as the README and paper promise. (rho_A.png attached below.)
• Same-params advantage of looped vs non-looped is real. looped_8@8 beats baseline_1@1 by ≈ 18 % PPL (57.9 vs 70.0) at the same parameter count (the trade being ~8× inference FLOPs per forward pass).

Reproduction

The full pipeline used to generate these numbers is in PR #27 (experiments/ directory). All commands, seeds, and default hyperparameters are in experiments/README.md; a run_all.sh orchestrator reproduces looped_8 + baseline_1 end-to-end. The looped_random run uses:

python train.py --run_name looped_random \
    --max_loop_iters 16 \
    --loop_sample_mode random_set --loop_choices 4 6 8 12 16 \
    --batch_size 16 --grad_accum_steps 2 --max_steps 15000

Wall-clock on H100 SXM: ~4 h for looped_8, ~1.5 h for baseline_1, ~3.2 h for looped_random. Requires PR #26 (bf16 dtype fix) to run in bfloat16.

Appendix: training loss and ρ(A)

Questions for maintainers / community

1. Is the sharp-V behaviour under fixed-loop training already known? The README section "The Loop Index Embedding Hypothesis" discusses the design motivation but doesn't discuss this trade-off with depth extrapolation explicitly.
2. Is there a training recipe that the author believes would produce monotonic inference-time scaling with the current architecture? (Curriculum? Larger max_loop_iters? Annealing the loop-index embedding strength?) Happy to run additional ablations if there's a specific hypothesis worth testing.
3. Would the project be open to adding an optional flag that disables loop_index_embedding / per-step LoRA scale, so users can choose the "depth extrapolation" regime vs. the "peak quality at fixed depth" regime?

Happy to rerun with any suggested variations; each ablation is ~$10 and ~4 hours on an H100.

[tonyzdev]

Follow-up: Ablation of loop_index_embedding + per-loop LoRAAdapter.scale — my diagnosis was wrong

In the original post I conjectured that the V-shape (fixed training) and the flat plateau (random training) were caused by the loop_index_embedding + per-loop LoRAAdapter.scale mechanisms. I added an opt-in anonymous flag on a local branch that gates both off, ran two new 15,000-step runs with everything else held fixed, and swept the same loop grid. The V-shape and plateau persist essentially unchanged, so my mechanistic explanation does not hold.

The five-curve picture

n_loops	looped_8	anon_fixed_8	looped_random	anon_random	baseline_1
1	281.5	503.1	1217.0	749.6	70.0
2	242.9	273.5	401.3	294.0	—
4	193.3	174.3	65.3	62.5	—
6	163.1	109.3	65.2	61.6	—
8	57.9	58.0	65.2	61.5	—
12	121.8	84.8	65.2	61.5	—
16	170.7	107.1	65.2	61.5	—

Qualitatively both shapes survive the ablation:

• anon_fixed_8 keeps the sharp V with its minimum exactly at the training n_loops=8 (PPL 58.0 vs 57.9 for looped_8 — statistically identical).
• anon_random keeps the flat plateau across the training range {4, 6, 8, 12, 16} (PPLs agree to 0.1%).
• n_loops=1/2 remain catastrophic in every looped variant.

What the ablation did show (secondary effects)

1. Anonymous + fixed is gentler at extrapolation. PPL(16) drops from 170.7 → 107.1, and PPL(12) from 121.8 → 84.8. So loop_index_embedding + LoRA are responsible for some of the sharpness of the upturn beyond the training loop count, but not for its existence.

2. Anonymous + random edges ahead. Plateau PPL is 61.5 vs 65.2 — a small (6 %) but free improvement.

3. ρ(A) starts to drift for the first time. Across every previous run and both of these, ρ(A) is pinned at 0.3574. In anon_random it actively moves:

   

   ρ(A) ends at 0.3320, i.e. A_discrete is pushed lower than its stable point. The only configuration where this happened is anonymous + random loop count — exactly when the model can neither tag-route inputs by loop depth nor assume a fixed number of steps. The LTI injection parameters apparently adapt to compensate.

So what does cause the training-depth binding?

Given this ablation, neither of the obvious "per-step distinctness" mechanisms is responsible. Remaining candidates I did not get to test within this experiment's budget:

• ACT halting learning a "convergence signature" at the trained depth. Could be ruled out by freezing / disabling ACT and re-running.
• LTI injection itself — the A and B parameters may be learning an implicit "T-step program" whose fixed-point behaviour is tied to the training loop count. The ρ(A) drift observation above is at least weakly consistent with this.
• The MoE router + shared FFN learning a depth-conditioned computation. Random-depth training implicitly teaches "converge quickly and cruise", which matches the plateau.

Happy to run (a) frozen-A ablation or (b) MoE-router-bias-off ablation if anyone has strong priors on which of these is the root cause — each is another $10 / ~4 h on an H100 and cleanly disentangles a single mechanism. Open to other experimental designs.

Methodology / reproducibility

All five runs share the same 117.8 M-param MLA + MoE backbone, same 491 M FineWeb-Edu tokens, same AdamW + warmup/cosine, same evaluation slice (--skip_docs=2_000_000). The only difference between the four looped runs is:

	fixed n_loops=8	random n_loops ∈ {4,6,8,12,16}
default (loop_index_embedding=True, use_per_loop_lora=True)	looped_8	looped_random
anonymous (loop_index_embedding=False, use_per_loop_lora=False)	anon_fixed_8	anon_random

The anonymous flag is a 3-site gate in RecurrentBlock.forward (skip loop_index_embedding(...) + skip self.lora(trans_out, t)). I haven't pushed a PR for the gate because it's only useful if the project wants it as a first-class option; the current two PRs (#26, #27) cover the bug fix and the validation harness regardless. Happy to add it if useful.

Take-home:

• LTI stability still looks robust across every configuration tested, including the one where ρ(A) actively moves.
• Depth extrapolation in the Saunshi / Parcae sense does not appear in OpenMythos with any of these four training configurations. Something architecturally deeper — likely the LTI dynamics, the MoE router, or the ACT halting — is binding the model to its training-depth distribution. Open to suggestions on which to poke next.

[tonyzdev]

Second follow-up: Freezing the LTI injection does not remove the V-shape either

Continuing the investigation. Having ruled out loop_index_embedding + per-loop LoRAAdapter.scale in the previous comment, the next obvious candidate was the LTI injection itself — specifically whether the model learns a "T-step program" by tuning log_A / log_dt / B during training. (This was also suggested by the observation that ρ(A) was the only thing that budged in anon_random.)

Ablation: freeze LTI

Added --freeze_lti in my local training script; it sets requires_grad=False on LTIInjection.log_A, log_dt, and B immediately after model construction. Initial values → ρ(A) = exp(-exp(0)) = 0.3679, which is exactly where looped_8, anon_fixed_8, etc. settle anyway, so this should be close to apples-to-apples. Ran two new variants (everything else identical to the earlier runs):

• frozen_lti_fixed_8: loop_index_embedding + per-loop LoRA ON, LTI frozen, n_loops=8 at every step
• frozen_lti_random: same two features ON, LTI frozen, n_loops ∈ {4, 6, 8, 12, 16} per step

Seven-curve picture

And the numbers (PPL at n_loops at inference):

n_loops	looped_8	anon_fixed_8	frozen_lti_fixed_8	looped_random	anon_random	frozen_lti_random
1	281.5	503.1	305.0	1217.0	749.6	1072.0
2	242.9	273.5	258.0	401.3	294.0	345.0
4	193.3	174.3	198.6	65.3	62.5	65.1
6	163.1	109.3	164.1	65.2	61.6	64.9
8	57.9	58.0	58.0	65.2	61.5	64.9
12	121.8	84.8	121.9	65.2	61.5	64.9
16	170.7	107.1	174.1	65.2	61.5	64.9

Results

Freezing LTI does not move the V-shape or the plateau. In fact frozen_lti_fixed_8's curve is almost pixel-identical to looped_8's at every depth — peak at 58 at n_loops=8, 305 at 1, 174 at 16 — even though ρ(A) has been held fixed at 0.3672 for the entire 15,000-step run. Training loss landed at 3.999 (vs 3.989 for looped_8, 3.995 for anon_fixed_8). frozen_lti_random shows a clean flat plateau at 64.94 for all n_loops ∈ {6, 8, 12, 16} — indistinguishable from looped_random's plateau at 65.24.

So the V-shape is not caused by:

1. loop_index_embedding (ruled out in previous comment)
2. per-loop LoRAAdapter.scale (ruled out in previous comment)
3. learned LTI injection dynamics (ruled out here)

All three together would seem to cover most of the "per-step distinctness" story in the architecture. With all three frozen or removed, the V-shape persists unchanged. I now think the V-shape is an emergent property of training at a fixed loop depth, not a consequence of any particular per-step signaling mechanism. The TransformerBlock + MoE + the fixed e injection per step appear to jointly learn a "T-step program" that lives inside the standard-transformer pathway itself.

A corollary worth noting

If the V-shape is purely a training-strategy artifact, then the README's reliance on Saunshi et al. 2025 and Parcae for depth-extrapolation is potentially misleading: the architectural ingredients OpenMythos advertises do not, on their own, deliver depth extrapolation. The only training recipe I've found that produces robust behavior across inference depth is random-loop training — which in turn produces a flat plateau rather than monotonic improvement, i.e. more inference compute buys nothing.

Concretely, across every 2×3 = 6 looped configuration I trained (fixed vs random × {default, anon, frozen-LTI}):

• No configuration produces monotonically decreasing PPL with more inference loops.
• No configuration produces a PPL at n_loops > training_loops that is lower than PPL at the trained loop count.

Remaining candidates

I've now exhausted the "obvious per-step distinguishing mechanisms". What's left to test, ordered by how cleanly each can be ablated:

1. ACT halting (ACTHalting module + the cumulative_p / weight accumulation logic in RecurrentBlock.forward). Intervention: replace the ACT-weighted sum with just h from the last loop iteration (i.e. no early halting, no weighted aggregation). If the V-shape flattens → ACT is implicated.
2. MoE routing. Intervention: set n_experts_per_tok=0 (use only shared experts) or freeze the router + router bias. If plateau / V-shape changes → MoE specialization is the source.
3. The hidden state h recurrence itself. Intervention: break the recurrent carry by setting A=0, B=1 and only keeping the Transformer(h_t, e) term. This is close to "8 independent forward passes" and should definitely lose the V-shape if the binding is in the recurrence structure.

Happy to run any of these if the community has priors on which is most likely or most desired. Each is another ~4 hours / ~$12 on an H100.

Meta-observation: I've now run seven variants of OpenMythos on identical data, and the central inference-time scaling claim fails to emerge under every training recipe × ablation I've tested. The architecture is stable, trainable, and consistently beats a single-loop equivalent by 10-20% PPL — which is real but modest for ~8× inference FLOPs — but its distinguishing "more loops = deeper reasoning" pitch has not replicated here. Possibly scale matters (117M may be too small for the effect to emerge), possibly the training corpus matters (general web text may not exercise the compositional reasoning where looping is supposed to win), possibly there's a recipe detail I'm missing. Open to all three possibilities; open to pushing this further.

Reproduction details: PR #27 for the pipeline. The --freeze_lti flag is a four-line addition to experiments/train.py that I can spin into a PR if useful. Full eval_ckpt_*.json and per-step train.log (including n_loops actually used per step) for all seven runs are available on request.

[tonyzdev]

Third follow-up: Six more ablations — ACT halting is the only mechanism that changes the loop-scaling shape

Continuing to search for the root cause of the training-depth binding. Previous comments ruled out loop_index_embedding + per-loop LoRA, and frozen LTI injection. This round adds three more interventions, each in both fixed-n_loops=8 and random-n_loops ∈ {4,6,8,12,16} variants, for six new runs. Everything else held constant (117.8M MLA+MoE, 491M FineWeb-Edu tokens, same AdamW schedule).

The three new interventions

1. disable_act: bypass ACTHalting entirely. Instead of the ACT-weighted sum h_out = Σ_t w_t · h_t, just return the hidden state from the final loop iteration. Tests whether the learned halting probabilities encode "I'm done thinking at step N" and bind the model to its trained depth.
2. freeze_moe_router: set requires_grad=False on MoEFFN.router.weight at init. Router scores stay fixed-random for the entire run; MoE specialization can only happen through expert weights, not routing decisions. Tests whether per-step learned routing patterns contribute to the binding.
3. break_recurrence: replace the LTI injection update h = A·h + B·e + trans_out with h = trans_out directly. Kills the state-carry completely while keeping the transformer block + e-injection via norm(h + e) (i.e. the block still sees the previous h as input, just doesn't carry an LTI-weighted history forward). Tests whether the LTI carry is the source of the implicit "T-step program".

The result

At the trained loop count, all 13 runs converge to within ~15% of each other (56.79 – 65.24 PPL). But at inference depths outside the training regime the differences are dramatic. Most interestingly:

disable_act_random is the only configuration of any I've tested that does not catastrophically collapse at n_loops < 4.

n_loops	looped_random	anon_random	frozen_lti_random	break_rec_random	freeze_router_random	disable_act_random
1	1217.0	749.6	1072.0	617.0	1129.9	131.3
2	401.3	294.0	345.0	533.0	327.5	77.7
4	65.3	62.5	65.1	60.3	64.5	62.8
6	65.2	61.6	64.9	58.9	64.2	60.3
8	65.2	61.5	64.9	58.9	64.2	59.7
12	65.2	61.5	64.9	58.9	64.2	59.5
16	65.2	61.5	64.9	58.9	64.2	59.6

Everyone else: PPL(1) in the hundreds or over a thousand, PPL(4..16) on a dead-flat plateau. disable_act_random: PPL(1)=131, smooth monotonic decrease down to 59 at n_loops=12, then essentially flat. This is the first time in the 13-run matrix where "more inference loops monotonically helps until saturation" — the qualitative shape the README and Saunshi/Parcae-style depth-extrapolation literature promises.

The fixed-n_loops=8 panel tells the same story more mildly: disable_act_fixed_8 has PPL(16)=87 versus 107-174 for all other fixed-8 runs, i.e. the softest V-shape upturn beyond the training depth.

Mechanistic interpretation

With ACT active, the model learns a per-position halting policy that effectively encodes "at loop step t I should stop contributing". When trained with fixed n_loops=8, this policy gets tuned specifically for T=8; when trained with random n_loops, it gets tuned for T ∈ {4..16}. At inference-time loop counts outside the training distribution, the halting policy is out-of-distribution and the ACT-weighted sum aggregates hidden states in ways that don't match what the model was trained to produce.

Disabling ACT (returning h from the last loop iteration) removes this policy from the output entirely. The model is forced to produce a usable hidden state at every depth it's trained at (because there's no weighting to rescue it), and — because training sees n_loops=1 implicitly (it's the trivial case where h after the first loop must be sane for the random-loop training to converge to a good output when n_loops=4 is rolled early in the schedule) — the function extrapolates much more gracefully to n_loops=1.

This is a specific, testable claim: ACT halting in OpenMythos is the primary mechanism binding the model to its training loop count. Not loop_index_embedding, not per-loop LoRA, not LTI dynamics, not MoE routing, not the h recurrence. ACT.

A second, unrelated surprise: break_recurrence is the best-performing configuration

break_rec_fixed_8 hits the lowest trained-loop PPL of all 13 runs (56.79 vs 57.89 for looped_8), and break_rec_random gets the lowest random-regime plateau (58.86). Replacing h = A·h_t + B·e + trans_out with h = trans_out — i.e. killing the LTI carry that was supposed to accumulate reasoning across loops — does not hurt and may slightly help. This is consistent with the frozen-LTI result from the previous comment: the learned LTI parameters do not appear to be essential to whatever OpenMythos is actually doing well here. Whatever "reasoning across loops" means in practice, the LTI state carry is not how it's being implemented.

Summary of 13 runs (PPL at trained loop count)

#	Run	Trained at	PPL@trained	Notable
1	break_rec_fixed_8	8	56.79	🏆 best fixed
2	freeze_router_fixed_8	8	57.22
3	looped_8 (default)	8	57.89
4	anon_fixed_8	8	58.02
5	frozen_lti_fixed_8	8	58.04
6	break_rec_random	16	58.86	🏆 best random plateau
7	disable_act_fixed_8	8	59.40	softest V
8	disable_act_random	16	59.58	⭐ only monotonic curve
9	anon_random	16	61.48
10	freeze_router_random	16	64.16
11	frozen_lti_random	16	64.94
12	looped_random	16	65.24
13	baseline_1	1	69.98	dense equivalent

Bottom line

After 13 single-GPU runs (~$150 of H100 time, 2.5B trained tokens):

• The depth-extrapolation claim in the README does not hold for OpenMythos as shipped, under five training × ablation combinations.
• It does start to hold once ACT halting is disabled and training uses random loop counts — but at the cost of giving up one of the features the README lists as a core architectural strength (adaptive per-token compute).
• LTI stability (ρ(A) < 1) holds across every configuration tested, including those where ρ(A) actively drifts during training. That part of the README's story is robust.
• "Looped beats non-looped at same parameter count" holds: every looped variant beats baseline_1's PPL 69.98 by 8-19%. The magnitude is modest for what is effectively 8× inference FLOPs per forward pass.

Practical recommendation (if the project is interested)

If the goal is to actually support model.generate(n_loops=k) for variable k at inference, the minimum viable recipe from this data is:

1. Train with stochastic-depth sampling (n_loops random per step from a wide range).
2. Disable or straight-through-aggregate the ACT halting — don't use it as a training signal at all, or treat halting probabilities as a soft attention over loop outputs that's trained jointly with depth sampling.

That combination (disable_act_random) is the only one in 13 runs that delivered a monotonic depth-scaling curve. Whether it's worth it depends on how much you value the ACT-style adaptive compute feature versus the depth-extrapolation feature — in the current implementation they are not both achievable simultaneously.

Happy to run a 14th variant if the project disagrees with this diagnosis. Code for all flags (--disable_act, --freeze_moe_router, --break_recurrence, --freeze_lti, and the existing anon flags) is a small addition to experiments/train.py that I can spin into a PR if useful. Full eval JSONs and train.logs for all 13 runs available on request.