fix: stabilize deep recurrence (T>=32) and correct scaling strategy

[skyhighbg22-jpg]

Summary

Fixes five compounding training instabilities triggered by the non-linear recurrent iteration scaling (1B:16 -> 1T:64) and corrects the scaling strategy for trillion-parameter regimes.

Problem

The original 64-loop design at dim=16384 triggers:

• Residual drift -- 64 unnormalized transformer_out terms compound, pushing the hidden state out of its useful encoding region
• Norm explosion -- residual stream grows as O(sqrt(T)), reaching ~8x initial magnitude at T=64
• Gradient circulation -- 64 unrolled steps through shared weights amplifies the dominant Jacobian eigenvector exponentially
• Representation collapse -- sinusoidal loop-index embedding only covers dim//8 channels (12.5%)
• Compute synchronization -- ACT halted positions still require full loop execution with KV cache

Solution -- Five Stabilization Mechanisms

1. Recurrent state normalization -- RMSNorm after every LTI injection step bounds residual magnitude
2. Stochastic depth recurrence -- linearly increasing dropout with depth breaks gradient circulation (training-only, dtype-matched for mixed-precision)
3. ACT ponder cost regularization -- penalizes total steps per position; model._last_ponder_cost added to LM loss
4. Learned iteration embeddings -- nn.Embedding table covers all dim channels per iteration, preventing representation collapse; clamps for depth extrapolation
5. muP initialization -- width-dependent scaling (1/width for head, 1/sqrt(width) for hidden layers) enables hyperparameter transfer across scales

Corrected Scaling

Variant	Old loops	New loops	Stabilization
1B-3B	16	16	learned embed, ponder
10B	24	24	+ stochastic depth (0.05), muP
50B-100B	32	32	+ stochastic depth (0.10)
500B	48	40	+ stochastic depth (0.15)
1T	64	48	+ stochastic depth (0.15)

Depth extrapolation to 64+ loops at inference preserved via embedding clamping.

Additional Fixes

• Initialize _last_ponder_cost in __init__ to prevent AttributeError
• muP: only head gets 1/width; attention .wo and FFN down correctly get 1/sqrt(width) as hidden layers
• Tests: slice freqs_cis[:T] in standalone attention/block tests (pre-existing RoPE dimension mismatch)
• LTI test: use <= 1.0 + 1e-6 for spectral radius after large gradient step (float precision boundary)

Testing

82/82 tests passing. All new components verified:

• TestLearnedIterationEmbedding (4 tests)
• TestStochasticDepthRecurrence (3 tests)
• TestMupInit (3 tests)
• TestRecurrentBlock extended (5 new tests)
• All pre-existing tests updated and passing

Files Changed

• open_mythos/main.py -- +218 lines (new classes + RecurrentBlock refactor + muP)
• open_mythos/variants.py -- +94 lines (corrected scaling + stabilization params)
• open_mythos/__init__.py -- +2 lines (export LearnedIterationEmbedding)
• tests/test_main.py -- +174 lines (20 new tests + pre-existing test fixes)