JugnuLM-110M-R2: value residuals + Muon (ablation rung R2)

README.md

@@ -0,0 +1,88 @@
+---
+license: apache-2.0
+language:
+  - en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+  - tiny-lm
+  - qwen3
+  - value-residual
+  - resformer
+  - muon
+  - ablation
+datasets:
+  - HuggingFaceFW/fineweb-edu
+---
+
+# JugnuLM-110M-R2 (value residuals + Muon)
+
+**What "R2" means.** This is **rung 2 of the JugnuLM phase-2 ablation ladder.** We take
+rung **R1** — [JugnuLM-110M-R1](https://huggingface.co/altslate/JugnuLM-110M-R1) (value
+residuals) — and add **one** change: we swap the optimizer to **Muon** for the 2D hidden
+matrices, holding *everything else fixed* (same value-residual architecture, data, tokens,
+schedule shape, batch, seed). The ladder keeps a lever only if it beats the previous rung.
+**R2 does, so it's kept.**
+
+## What changed: the Muon optimizer
+
+R0 and R1 trained with AdamW on every parameter. R2 uses **Muon** — Newton-Schulz
+**orthogonalized momentum** updates — on the 2D hidden weight matrices (attention + MLP,
+81.4M params), and **keeps AdamW** for the rest (tied embedding/head, RMSNorm gains, and
+the 22 value-residual `vr_lambda` scalars, 28.3M params). Muon's advantage is known to be
+**largest at exactly this scale** (sub-150M) and to shrink as models grow.
+
+Both optimizers ride one shared cosine LR multiplier: Muon peak `2e-2`, AdamW peak `1.5e-3`,
+each annealed to 10% of peak.
+
+## Result vs the R1 and R0 rungs
+
+| 110M rung | BLiMP ↑ | ARC-Easy ↑ | WikiText-2 byte-ppl ↓ |
+|---|---|---|---|
+| R0 — baseline | 81.25 | 52.48 | 1.95 |
+| R1 — + value residuals | 81.10 | 54.67 | 1.940 |
+| **R2 — + Muon** | 80.78 | **56.10** | **1.932** |
+
+**Muon lifts ARC-Easy by +1.43 over R1 (54.67 → 56.10) and lowers perplexity, at a small
+BLiMP dip (−0.32, ~2× SE).** Across the two ladder rungs ARC-Easy has climbed **+3.6**
+(52.48 → 56.10) while BLiMP stayed ~flat — steadily closing the reasoning gap to the
+125–143M leaderboard models. Evaluated with EleutherAI `lm-evaluation-harness`.
+
+**Training-efficiency note.** Muon's convergence lead was largest early — val-perplexity
+was **−26% vs R1 at step 1000** — and compressed to ~−1% by the end of the fixed 8.4B-token
+budget (the classic "Muon wins the sprint, converges toward AdamW" pattern). Notably, the
+**downstream ARC gain persisted** even as the final perplexity gap narrowed.
+
+## ⚠️ Loading — requires `trust_remote_code`
+
+Value residuals are a **custom attention pathway**, so this model must be loaded with the
+included modeling file. Loading it as a stock Qwen3 **silently drops the residual** and
+understates quality.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+tok = AutoTokenizer.from_pretrained("altslate/JugnuLM-110M-R2")
+model = AutoModelForCausalLM.from_pretrained("altslate/JugnuLM-110M-R2", 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 to R1: Qwen3-arch + value-residual pathway, **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 (sample-10BT), shared
+cosine schedule (Muon 2e-2, AdamW 1.5e-3, → 10% floor), ~0.5M-token global batch, bf16,
+DDP on 4× RTX PRO 4500 Blackwell. Only the optimizer differs from R1.
+
+## Intended use & limitations
+
+Research **base** model (not instruction-tuned). Fluent English continuation and
+reasonable recall for its size; limited multi-step reasoning, occasional repetition,
+English-only. Not for production.
+
+## Links
+
+- Previous rung: [altslate/JugnuLM-110M-R1](https://huggingface.co/altslate/JugnuLM-110M-R1)
+- Baseline: [altslate/JugnuLM-110M](https://huggingface.co/altslate/JugnuLM-110M)
+- Training code & ablation ladder: https://github.com/AltSlate-Labs/jugnu

config.json

@@ -0,0 +1,62 @@
+{
+  "architectures": [
+    "JugnuVRForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "modeling_jugnu_vr.JugnuVRConfig",
+    "AutoModelForCausalLM": "modeling_jugnu_vr.JugnuVRForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "bfloat16",
+  "eos_token_id": 0,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 2048,
+  "max_window_layers": 28,
+  "model_type": "jugnu_vr",
+  "num_attention_heads": 9,
+  "num_hidden_layers": 23,
+  "num_key_value_heads": 3,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.2.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 49152
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 0,
+  "transformers_version": "5.2.0",
+  "use_cache": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba6f6d2cf935d4bcba4f50c12e0a1af027c2be8bf737dd6a870c8eca8b0c66ac
+size 219505412

modeling_jugnu_vr.py

@@ -0,0 +1,54 @@
+"""JugnuLM-110M-R1 — Qwen3 decoder + value residuals (ResFormer).
+
+Each layer's value projection gains a learned-gated residual from the FIRST layer's
+value:  v_i = v_proj_i(x) + lambda_i * v_0  (v_0 = layer-0 value; lambda_i learned).
+Implemented by replacing each layer's v_proj with a residual Linear that shares a
+per-forward context; layer 0 writes v_0, later layers add lambda_i * v_0.
+
+Loadable via:  AutoModelForCausalLM.from_pretrained("altslate/JugnuLM-110M-R1", trust_remote_code=True)
+(Loading as a stock Qwen3 silently drops the residual pathway.)
+"""
+import torch
+import torch.nn as nn
+from transformers import Qwen3Config, Qwen3ForCausalLM
+
+
+class JugnuVRConfig(Qwen3Config):
+    model_type = "jugnu_vr"
+
+
+class _VResidualLinear(nn.Linear):
+    def __init__(self, in_features, out_features, is_first, bias=False):
+        super().__init__(in_features, out_features, bias=bias)
+        self._vr_is_first = is_first
+        self._vr_ctx = None
+        if not is_first:
+            self.vr_lambda = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        v = super().forward(x)
+        ctx = self._vr_ctx
+        if ctx is None:
+            return v
+        if self._vr_is_first:
+            ctx["v0"] = v
+        else:
+            v0 = ctx.get("v0")
+            if v0 is not None:
+                v = v + self.vr_lambda * v0
+        return v
+
+
+class JugnuVRForCausalLM(Qwen3ForCausalLM):
+    config_class = JugnuVRConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        ctx = {}
+        for i, layer in enumerate(self.model.layers):
+            old = layer.self_attn.v_proj
+            new = _VResidualLinear(old.in_features, old.out_features,
+                                   is_first=(i == 0), bias=(old.bias is not None))
+            new._vr_ctx = ctx
+            layer.self_attn.v_proj = new
+        self.post_init()

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 8192,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}