README.md

---
license: apache-2.0
library_name: transformers
pipeline_tag: text-generation
language:
- en
tags:
- jugnu
- tiny-lm
- value-residual
- muon
- pretrained-from-scratch
datasets:
- HuggingFaceFW/fineweb-edu
base_model: altslate/JugnuLM-110M
---

# JugnuLM-110M-R2+ 🪰✨

A **sub-150M** language model pretrained **from scratch** by [AltSlate Labs](https://github.com/AltSlate-Labs),
for the [Tiny-ML Leaderboard](https://huggingface.co/spaces/Glint-Research/Tiny-ML-Leaderboard). The flagship of the
[Jugnu](https://github.com/AltSlate-Labs/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

```bash
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`):

```python
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):

```python
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](https://arxiv.org/abs/2410.17897)): 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