SGK86/mor-mini-fineweb-edu

MoR mini -- scaled-down Mixture-of-Recursions reproduction

From-scratch reproduction of Mixture-of-Recursions (arXiv 2507.10524), Recipe 1: expert-choice router, sigmoid weighted gating alpha=0.1, aux BCE coeff 0.001, Middle-Cycle layer sharing. Reference implementation: github.com/raymin0223/mixture_of_recursions (router fidelity fixes taken from router.py in that repo).

Config

• vocab_size=49152, d_model=512, n_heads=8, intermediate=1408
• n_middle=6, n_recursions=3, max_len=2048, rope_theta=10000.0
• ~50.9M parameters (25.7M non-embedding)

Data

HuggingFaceFW/fineweb-edu sample-10BT, tokenized with the SmolLM2-135M tokenizer, uint16 token cache in SGK86/mor-fineweb-tokcache (655M train tokens, built from parquet shards directly, no streaming).

Training

Total batch 128 x 2048 = 262k tokens/step, lr 1e-3 warmup-stable-decay (warmup 125 steps, linear decay over final 20%), aux BCE 0.001, z-reg 1e-5, AdamW beta (0.9, 0.95), weight decay 0.1, bf16, grad clip 1.0. 2500 steps (~655M tokens, one pass over the cache), a10g-small, ~22 min.

Results

Final: val NLL 5.90 (pure CE on 20M held-out FineWeb-Edu tokens, exp(5.90) ~= 365 ppl). Train CE ~5.86-5.97, stable through the entire decay phase.

What the four attempts showed

• v1/v2: no streaming in-container -- data cache build stalled. Fixed by building the token cache in a separate CPU job and mounting it as a read-only dataset volume.
• v4 (lr 3e-3, z-reg off): diverged. CE fell to ~7.7 then snapped back to ~9.6 with router logit explosion (aux spikes ~600-1000, router logit norm e+01..e+16). Root cause identified by comparing to the reference config: the reference trains at effective batch 1024 x 2048 = 2.1M tokens/step at lr 3e-3 -- 8x our batch. lr 3e-3 at 262k tokens/step is simply too hot.
• v5 (lr 1e-3, z-reg 1e-5, pure-CE eval): stable, converged, no snap-back. The v4 eval numbers were also contaminated by the aux term (reported val NLL 9.09 was overestimated); v5 evals CE only.

Routing findings (honest caveats)

• depth_mean pinned at exactly 2.000 throughout training; final depth histogram [0.0, 0.333, 0.333, 0.333] -- the expert-choice router stayed essentially uniform and never learned token-level depth specialization. The capacity fractions (1.0, 1/2, 1/3 per depth 0/1/2) combined with the target depth distribution make depth 0 receive zero tokens.
• Router logit norm grew from ~67 to ~69 over the tail -- large but stable, no explosion. Worth watching, not a failure.
• At 51M parameters / 655M tokens this is far below the paper's regime: the paper itself shows MoR only matches the Recursive baseline at 135M and wins at >=360M, and it pre-filters weights from a SmolLM-360M checkpoint (not reproduced here). So this run demonstrates the mechanics (stable expert-choice recursion training at small scale), not the paper's efficiency claims.

Files

• model.safetensors + config.json: final checkpoint (step 2500)
• train_v5.py: the exact training script used
• Trackio dashboard: SGK86/mor-fineweb-mini-trackio