looped-fineweb-10m

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

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