looped-nanochat-d24
A two‑pass, shared‑weight routing‑loop Transformer built on
karpathy/nanochat @ 92d63d4,
trained at the d24 (~1.38 B param) scale. This is the looped model from a
controlled comparison against the stock nanochat baseline
(kz919/baseline-nanochat-d24).
Code: https://github.com/kyleliang919/looped_nanochat
What is the routing loop?
The Transformer trunk is run twice per forward:
Pass 1 runs the ordinary causal Transformer and captures the normalized, RoPE‑applied Q/K features from selected deeper layers.
Pass 2 re‑runs the trunk from the token embeddings. Each pass‑2 layer adds a causal attention‑logit prior from one pass‑1 layer:
attn_logits = q2·k2ᵀ/√D + g · q1·k1ᵀ/√Dwith a per‑layer learned gate
g = softplus(gate_logit). The additive bias is computed via augmented‑width Q/K + ordinary attention (no N×N matrix materialized), and the training path runs it through a fused, compiled FlexAttention kernel. Everything is causal and differentiable.
Extra parameters over baseline: routing_gate_logits (per layer) and
pass_embeddings (2×n_embd).
Results (vs. baseline, identical training budget)
| Metric | Baseline | Looped |
|---|---|---|
| Base CORE | 0.2552 | 0.2783 |
| ChatCORE (full sets) | 0.2253 | 0.2521 |
| ARC‑Easy | 62.04% | 68.01% |
| ARC‑Challenge | 50.00% | 52.22% |
| MMLU | 37.06% | 37.96% |
| GSM8K | 2.88% | 4.17% |
| HumanEval | 10.98% | 10.98% |
The looped model wins or ties every metric, at ~2.4× the training compute (inherent to the two passes). See the full report.
Model details
| Layers | 24 |
n_embd |
1536 |
| Heads / KV heads | 12 / 12 |
| Vocab | 32768 |
| Context | 2048 |
| Window pattern | SSSL |
| Routing pattern | progressive |
routing_gate_init |
0.05 |
| Params | ~1.38 B |
| Precision | bf16 (fp8 matmuls in training) |
Training: pretrained for 5,568 steps on ~5.84 B tokens (--target-param-data-ratio=8,
1,048,576‑token batches), then chat‑SFT on SmolTalk + MMLU(×3) + GSM8K(×4).
This checkpoint is the SFT model (step 467).
Files
model_000467.pt— model weights (state_dict)meta_000467.json—model_config(incl. routing fields) + training configtokenizer.pkl,token_bytes.pt— nanochat tokenizer
Usage
This uses the nanochat runtime (not 🤗 Transformers). Clone the repo, drop the
files into ~/.cache/nanochat/chatsft_checkpoints/looped-d24/ (weights/meta) and
~/.cache/nanochat/tokenizer/, then:
python -m scripts.chat_cli -g looped-d24 -i sft -p "Why is the sky blue?"
The routing config is stored in meta_000467.json, so the runtime rebuilds the
two‑pass model automatically.
Limitations
A ~$100‑tier micro‑model: it hallucinates, is weak at math/code, and is meant for research on architecture, not production use.
License
MIT (inherited from nanochat).