JugnuLM-110M-R2 (value residuals + Muon)
What "R2" means. This is rung 2 of the JugnuLM phase-2 ablation ladder. We take rung R1 — JugnuLM-110M-R1 (value residuals) — and add one change: we swap the optimizer to Muon for the 2D hidden matrices, holding everything else fixed (same value-residual architecture, data, tokens, schedule shape, batch, seed). The ladder keeps a lever only if it beats the previous rung. R2 does, so it's kept.
What changed: the Muon optimizer
R0 and R1 trained with AdamW on every parameter. R2 uses Muon — Newton-Schulz
orthogonalized momentum updates — on the 2D hidden weight matrices (attention + MLP,
81.4M params), and keeps AdamW for the rest (tied embedding/head, RMSNorm gains, and
the 22 value-residual vr_lambda scalars, 28.3M params). Muon's advantage is known to be
largest at exactly this scale (sub-150M) and to shrink as models grow.
Both optimizers ride one shared cosine LR multiplier: Muon peak 2e-2, AdamW peak 1.5e-3,
each annealed to 10% of peak.
Result vs the R1 and R0 rungs
| 110M rung | BLiMP ↑ | ARC-Easy ↑ | WikiText-2 byte-ppl ↓ |
|---|---|---|---|
| R0 — baseline | 81.25 | 52.48 | 1.95 |
| R1 — + value residuals | 81.10 | 54.67 | 1.940 |
| R2 — + Muon | 80.78 | 56.10 | 1.932 |
Muon lifts ARC-Easy by +1.43 over R1 (54.67 → 56.10) and lowers perplexity, at a small
BLiMP dip (−0.32, ~2× SE). Across the two ladder rungs ARC-Easy has climbed +3.6
(52.48 → 56.10) while BLiMP stayed ~flat — steadily closing the reasoning gap to the
125–143M leaderboard models. Evaluated with EleutherAI lm-evaluation-harness.
Training-efficiency note. Muon's convergence lead was largest early — val-perplexity was −26% vs R1 at step 1000 — and compressed to ~−1% by the end of the fixed 8.4B-token budget (the classic "Muon wins the sprint, converges toward AdamW" pattern). Notably, the downstream ARC gain persisted even as the final perplexity gap narrowed.
⚠️ Loading — requires trust_remote_code
Value residuals are a custom attention pathway, so this model must be loaded with the included modeling file. Loading it as a stock Qwen3 silently drops the residual and understates quality.
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("altslate/JugnuLM-110M-R2")
model = AutoModelForCausalLM.from_pretrained("altslate/JugnuLM-110M-R2", trust_remote_code=True)
ids = tok("The theory of relativity was developed by", return_tensors="pt").input_ids
print(tok.decode(model.generate(ids, max_new_tokens=30, repetition_penalty=1.3)[0], skip_special_tokens=True))
Architecture & training
Identical architecture to R1: Qwen3-arch + value-residual pathway, 23 layers × 576, GQA 9/3, RoPE, SwiGLU, RMSNorm, QK-Norm, tied embeddings, SmolLM2 tokenizer (49,152), z-loss 1e-4. Trained from scratch on ~8.4B tokens of FineWeb-Edu (sample-10BT), shared cosine schedule (Muon 2e-2, AdamW 1.5e-3, → 10% floor), ~0.5M-token global batch, bf16, DDP on 4× RTX PRO 4500 Blackwell. Only the optimizer differs from R1.
Intended use & limitations
Research base model (not instruction-tuned). Fluent English continuation and reasonable recall for its size; limited multi-step reasoning, occasional repetition, English-only. Not for production.
Links
- Previous rung: altslate/JugnuLM-110M-R1
- Baseline: altslate/JugnuLM-110M
- Training code & ablation ladder: https://github.com/AltSlate-Labs/jugnu