JugnuLM-110M-R4a (value residuals + Muon + logit distillation)
What "R4a" means. Rung 4 of the JugnuLM phase-2 ablation ladder — the first knowledge-distillation attempt. We take rung R2 (JugnuLM-110M-R2, value residuals + Muon) and add offline top-k logit distillation from a teacher, holding data/optimizer/schedule fixed. This is the "a" variant — a deliberately KD-heavy setting; a rebalanced R4b follows.
What changed: logit KD from SmolLM2-1.7B
The frozen teacher SmolLM2-1.7B (base; shares our SmolLM2 tokenizer, 49,152 vocab) was run over ~6B tokens of FineWeb-Edu, storing its top-16 next-token logits per position. The student trains to match that distribution:
loss = α·CE + (1−α)·τ²·KL(student ‖ teacher_top16) + z-loss, with α=0.5, τ=2.
Result — distillation lifts ARC, but this setting over-weights KD
| 110M rung | BLiMP ↑ | ARC-Easy ↑ | WikiText-2 byte-ppl ↓ |
|---|---|---|---|
| R2 — value residuals + Muon | 80.78 | 56.10 | 1.932 |
| R4a — + logit KD (α=0.5, τ=2) | 80.39 | 56.99 | 2.178 |
ARC-Easy reached 56.99 — the best of any JugnuLM rung, essentially matching GPT-X2-125M's
57.07 — confirming distillation is the right lever for reasoning. However, with α=0.5 and
τ=2 the KD term carries ~4× the weight of the hard-label loss, so the student over-imitates the
teacher's softened distribution at the cost of sharp next-token prediction: WikiText perplexity
rose to 2.178 (from R2's 1.932). On the leaderboard's blended efficiency score this perplexity
regression outweighs the ARC gain, so R4a is not kept as the ladder base — but it validates
the lever and motivates the rebalanced R4b (α=0.7, τ=1), which aims to keep the ARC win
without the perplexity damage. Evaluated with EleutherAI lm-evaluation-harness.
⚠️ Loading — requires trust_remote_code
Value residuals are a custom attention pathway; loading as a stock Qwen3 silently drops them.
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("altslate/JugnuLM-110M-R4a")
model = AutoModelForCausalLM.from_pretrained("altslate/JugnuLM-110M-R4a", 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/optimizer to R2 (value residuals + Muon): Qwen3-arch, 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 with the KD loss above, Muon (peak 2e-2) + AdamW (1.5e-3) on a shared cosine, ~0.5M-token global batch, bf16, DDP on 4× RTX PRO 4500 Blackwell. Teacher logits precomputed offline (top-16 over a 6B-token region).
Intended use & limitations
Research base model (not instruction-tuned). Best-in-family ARC-Easy reasoning, but higher perplexity than R1/R2 due to the KD-heavy setting. English-only, occasional repetition. Not for production. For a more balanced distilled model, see R4b (when released).
Links
- Base rung: altslate/JugnuLM-110M-R2
- Teacher: HuggingFaceTB/SmolLM2-1.7B
- Training code & ablation ladder: https://github.com/AltSlate-Labs/jugnu