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-lossR4a: α=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