ahiok/looped-fineweb-10m

A looped decoder-only language model trained under a hard budget of 9,441,152 parameters (6,295,424 non-embedding) and 100,000,000 training tokens of FineWeb.

The block that gets looped is Qwen3-style (RMSNorm pre-norm, GQA, QK-norm, SwiGLU, RoPE). The same 2 layers are applied R times; R is chosen at inference, so the same weights can be run cheap or deep.

Results

eval loops R	val loss	perplexity	bits/byte
1	7.4054	1644.79	2.8394
2	6.4121	609.16	2.4586
4	5.2127	183.59	1.9987
8	4.1340	62.43	1.5851
16	3.7898	44.25	1.4531
24	3.8698	47.93	1.4838
32	3.9913	54.13	1.5304
48	4.1892	65.97	1.6063
64	4.3275	75.76	1.6593
96	4.5078	90.72	1.7284
128	4.6195	101.44	1.7712

Validation is a held-out document split of the same FineWeb shard, 0 tokens, tokenised with the 8192-entry byte-level BPE included in this repo. Bits-per-byte is reported alongside perplexity because perplexity alone is not comparable across tokenizers.

The recurrence is s <- Block(s + e): the embedded input is added back into the state at the start of every iteration. That one tensor add is the entire difference from an unlooped model of identical parameter count, and it is worth 0.10 nats here. Run it at R=16, the depth it was trained at: this variant buys quality rather than depth robustness and degrades sharply on either side (3.79 at R=16, 4.13 at R=8, 3.99 at R=32).

For reference, an unlooped 4-layer model of the same size trained on the same 100M tokens reaches 3.8965 / 49.23 / 1.4940, and a plain looped model with no update rule reaches 3.8315 / 46.13 / 1.4691.

This is a research artefact for studying test-time depth scaling under a hard budget, not a usable text generator. At 9.4M parameters and 100M tokens it produces the statistics of English, not sentences you would want to read.

Usage

import importlib.util, json, sys, torch
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer

repo = "ahiok/looped-fineweb-10m"
src = hf_hub_download(repo, "model.py")
spec = importlib.util.spec_from_file_location("loopllm_model", src)
mod = importlib.util.module_from_spec(spec)
sys.modules["loopllm_model"] = mod       # required: @dataclass resolves via sys.modules
spec.loader.exec_module(mod)

cfg = mod.ModelConfig(**json.load(open(hf_hub_download(repo, "model_config.json"))))
model = mod.LoopedLM(cfg)
model.load_state_dict(torch.load(hf_hub_download(repo, "pytorch_model.bin"), map_location="cpu"))
model.eval()

tok = Tokenizer.from_file(hf_hub_download(repo, "tokenizer.json"))
ids = torch.tensor([tok.encode("The capital of France is").ids])
out = model(ids, n_loops=32)          # spend more or less compute here
print(tok.decode([int(out["logits"][0, -1].argmax())]))

Training

Code, full ablations and the report: https://github.com/ahiokk/looped-models