looped-qwen3-9.4M-T32-fineweb

looped-qwen3-9.4M · 32 useful loops

A 9.44M-parameter looped transformer in which all 32 loops earn their keep. A 2-layer Qwen3-style core is applied 32 times with shared weights — 64 effective layers from 2 layers of parameters — trained from scratch on 100M FineWeb tokens.

Naive weight-tied looping does not merely saturate: past T=8 extra depth actively hurts, and T=64 is worse than T=2 at 24× the compute. Two parameter-free changes remove that wall:

1. depth-RoPE — the state is conjugated by a step-dependent channel rotation, so the block sees a different coordinate frame at every loop and the effective weight matrix differs per step. Zero learned parameters.
2. Constant-share input injection — the token embedding keeps a fixed share of the block input at every loop, instead of decaying like 1/t as the state norm grows. One learned scalar.

With the wall removed, trained-and-evaluated-at-T quality is monotone through T=64 (4.118 → 4.081 → 4.059 nats at T = 16/32/64 on 25M tokens), where at 25M tokens it beats an unshared 16-layer model with 4.9× the parameters (4.078). The learned versions of the same ideas — per-step embeddings, per-step adaLN — consistently hurt: step symmetry has to be broken structurally, not bought with capacity.

Results

Robustness: 3 seeds per side at 25M tokens give 4.093 ± 0.011 (recipe, T=32) against 4.220 ± 0.016 (naive optimum, T=8) — an 11σ gap.

Perplexity here is not comparable to other models. The tokenizer is a byte-level BPE with an 8,192-entry vocabulary trained on this dataset, because the 10M-parameter budget counts embeddings — Qwen3's 151,936-entry vocabulary would consume the entire budget before the transformer got a single layer. Compare bits per byte instead, which is tokenizer-independent. On this validation text the tokenizer yields 3.78 bytes/token (GPT-2: 4.48).

Why it works

Trajectory diagnostics of the naive loop show the failure is not DEQ-style convergence to a fixed point — update magnitude never decays. Instead successive updates become collinear (cosine → 0.999 by loop 16) while the state norm grows exactly linearly: the loop spends its compute making one vector longer. In the coordinates every consumer actually sees (pre-norm), that is a fixed point — a ray in projective space. Factorial-shuffle probes show the block reads the state increasingly (state-driven share of its output: 0.25 → 0.94 across loops) yet maps it to a nearly constant direction.

A step-dependent rotation makes repeating the same update geometrically unproductive. In the trained champion the successive-update cosine stays at 0.93 at loop 32 (naive: 0.999), the read-out improves at every loop (3.780 → 3.722 over the last six), and inference-time depth transfers: the same weights run at any T, degrading gently beyond the trained depth.

Quick start

The architecture is custom, so load it with the code from the GitHub repository rather than AutoModel:

import json, torch
from safetensors.torch import load_file
from loopedlm.config import ModelConfig
from loopedlm.model import LoopedQwen3
from tokenizers import Tokenizer

cfg = ModelConfig.from_dict(json.load(open("model_config.json")))
model = LoopedQwen3(cfg).eval()
sd = load_file("model.safetensors")
sd["lm_head.weight"] = sd["embed_tokens.weight"]        # tied embeddings
model.load_state_dict(sd)
tok = Tokenizer.from_file("tokenizer.json")

ids = torch.tensor([tok.encode("The looped transformer").ids])
out = model(ids, n_loops=32)          # depth is a runtime choice
logits = model.head(out["hidden"])

n_loops is an inference-time argument — that is the point of the architecture. Run the same checkpoint shallower for cheap drafts or deeper for quality.

Architecture

Qwen3 recipe — pre-RMSNorm, per-head QK-norm, GQA, SwiGLU, RoPE, no biases, tied embeddings — with the middle of the network looped:

e   = EmbedNorm(Embed(idx))
h_0 = e
h_t = h_{t-1} + R_t^{-1} · Core( R_t · (h_{t-1}/rms + α·e) )     t = 1..32
logits = Head(RMSNorm(h_32))

Training

• Data: FineWeb sample/10BT, shard 000 for training, shard 014 for validation — disjoint documents; the tokenizer never saw the validation shard.
• Budget enforcement: non-overlapping windows in a seeded permutation — every training token is seen exactly once, asserted against the measured count; the trainer refuses to start if the model exceeds 10M parameters.
• Optimisation: AdamW, lr 2e-3 (the symmetry-broken loop prefers a hotter LR than the naive one — the optimum moves up, not down), cosine schedule, wd 0.1, clip 1.0, bf16 autocast, batch 32×512, gradient checkpointing per loop.
• Hardware: one A100 80GB, ~2.5 h wall-clock, torch.compile on the shared block (2.5× measured).

Files

model.safetensors (weights, fp32), model_config.json / train_config.json (exact configs), tokenizer.json (HF tokenizers file), summary.json (final metrics + full test-time depth sweep), log.jsonl (complete training log).

Limitations

At 9.44M parameters and 100M tokens this is a research artefact for studying depth recurrence, not an assistant: it produces locally plausible English and nothing more. No instruction tuning, no safety alignment, English only, 512-token context.

Code, 94-run experiment log and the full report: github.com/Embim/Looped-lm