JugnuLM-110M-R4b (rebalanced logit distillation)

What "R4b" means. The rebalanced counterpart of JugnuLM-110M-R4a. Same offline top-16 logit distillation from SmolLM2-1.7B, but with the KD term dialed down to stop it dominating:

loss = α·CE + (1−α)·τ²·KL(student ‖ teacher_top16) + z-loss R4a: α=0.5, τ=2 (KD ≈4× CE) → R4b: α=0.7, τ=1 (KD ≈0.4× CE)

R4b was run to test whether a gentler KD keeps R4a's ARC-Easy gain without its perplexity collapse. It does not — an honest negative.

Result — the rebalance fixed perplexity but lost the ARC gain

110M rung	BLiMP ↑	ARC-Easy ↑	WikiText-2 byte-ppl ↓
R2 — value residuals + Muon (base)	80.78	56.10	1.932
R4a — + KD, heavy (α=0.5, τ=2)	80.39	56.99	2.178
R4b — + KD, light (α=0.7, τ=1)	79.30	55.47	1.9165

R4b recovers R2-level perplexity (1.9165, best of the three), but ARC-Easy falls to 55.47 — below the no-KD R2 baseline (56.10) — and BLiMP drops to 79.30. So R4b is worse than R2 on BLiMP, ARC, and the blended efficiency score; it is not kept. Together R4a/R4b bracket a narrow, unfriendly KD operating point: heavy KD buys ARC but wrecks perplexity; light KD fixes perplexity but loses the ARC benefit. Offline logit KD from a 1.7B teacher (≈15× the student) did not yield a keepable win at this 8.4B-token scale.

Caveat / confound: R4a and R4b trained on the 6B-token KD region (the span with precomputed teacher logits), whereas R2 sampled the full ~10B FineWeb-Edu pool — so the KD runs saw less unique data, which likely contributed to R4b's BLiMP drop independent of the KD change. A fair KD comparison would precompute teacher logits over the full corpus. Evaluated with EleutherAI lm-evaluation-harness.

⚠️ Loading — requires trust_remote_code

from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("altslate/JugnuLM-110M-R4b")
model = AutoModelForCausalLM.from_pretrained("altslate/JugnuLM-110M-R4b", 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 to R4a except the KD weights (α=0.7, τ=1): Qwen3-arch + value residuals, 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 over the 6B-token KD region, 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.

Intended use & limitations

Research base model, published as an ablation record. For actual use prefer JugnuLM-110M-R2 (better BLiMP + ARC) or R4a (best ARC-Easy). English-only. Not for production.

Links

• Base rung / recommended: altslate/JugnuLM-110M-R2
• Heavy-KD counterpart: altslate/JugnuLM-110M-R4a
• Training code & ablation ladder: https://github.com/AltSlate-Labs/jugnu