altslate/JugnuLM-110M-R2plus

JugnuLM-110M-R2+ 🪰✨

A sub-150M language model pretrained from scratch by AltSlate Labs, for the Tiny-ML Leaderboard. The flagship of the Jugnu family: the kept R2 recipe (Qwen3 arch + value residuals + Muon) scaled to 25.2B tokens under a WSD schedule with modest decay-phase educational upweighting.

Requirements

pip install "transformers>=4.51" torch safetensors

transformers>=4.51 is required (the model builds on the Qwen3 architecture). It's a standard AutoModelForCausalLM otherwise — no extra packages.

⚠️ Load with trust_remote_code=True

This model uses value residuals (a custom attention pathway: v_i = v_proj_i(x) + λ_i·v0). Stock from_pretrained would silently drop that pathway and degrade the model (~6 pts ARC-Easy, ~0.18 byte-ppl). It ships custom modeling code with auto_map, so load it VR-aware (trust_remote_code=True):

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("altslate/JugnuLM-110M-R2plus")
model = AutoModelForCausalLM.from_pretrained(
    "altslate/JugnuLM-110M-R2plus",
    trust_remote_code=True,      # required — rebuilds the value-residual pathway
).eval()
# loads in fp32 by default; pass torch_dtype=torch.bfloat16 (transformers ≥5: dtype=...) to halve memory

ids = tok("The router will not connect to wifi, so I", return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=40, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))

Sanity check that the value-residual pathway loaded (22 vr_lambda params, mean ≈ 0.48):

lam = [p.item() for n, p in model.named_parameters() if n.endswith("vr_lambda")]
assert len(lam) == 22, "value-residual pathway not loaded — did you pass trust_remote_code=True?"

Results

metric	JugnuLM-110M-R2+
Params	109.7M
BLiMP (acc)	82.52
ARC-Easy (acc)	55.13
WikiText-2 (byte-ppl)	1.8735

Beats the JugnuLM-110M (R0) baseline on all three leaderboard metrics (BLiMP +1.3, ARC-Easy +2.65, byte-ppl 1.8735 vs 1.95), and posts the family's best BLiMP and byte-ppl. On the leaderboard's efficiency score it ranks #1 (EFF ≈ 80.21) — a narrow, within-noise lead over GPT-X2-125M (80.06) and Haidass-143M (79.83), winning on the size bonus as the smallest of the three. Numbers are from a VR-aware eval (BLiMP / ARC-Easy / WikiText via lm-eval-harness, acc; wikitext byte_perplexity).

Architecture

• Qwen3 architecture (Llama + built-in QK-Norm), deep-thin 23 layers × 576 hidden, GQA, tied embeddings.
• Value residuals (ResFormer): each layer's value gets a learned-gated residual from layer 0's value; 22 learned vr_lambda scalars (mean ≈ 0.48 in this checkpoint).
• SmolLM2 tokenizer (49,152 vocab). z-loss for logit stability.

Training

• 25.2B tokens (48,000 steps × 524,288 tok/step) on 2× NVIDIA RTX PRO 4500 Blackwell GPUs.
• Muon optimizer on 2D hidden matrices (attn + MLP); AdamW for embeddings / head / norms / vr_lambda.
• WSD schedule (stable → decay over the last ~21% of steps), with decay-phase upweighting of educational data (FineWeb-Edu). Final checkpoint (step 48000) is the best; val perplexity bottomed at end of decay.

License

Apache-2.0. Training recipe and code: https://github.com/AltSlate-Labs/jugnu