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