SpiralFormer loops coarse-to-fine, beating Pythia-1.4B 5-shot by 2.4 points on fewer FLOPs

SpiralFormer: Looped Transformers Can Learn Hierarchical Dependencies via Multi-Resolution Recursion

Chengting Yu, Xiaobo Shu, Yadao Wang, Yizhen Zhang, Haoyi Wu, You Wu, Rujiao Long, Ziheng Chen, Yuchi Xu, Wenbo Su, Bo Zheng

cs.LG

2026-02-12

SpiralFormer runs a shared looped core on a coarse-to-fine resolution schedule. At 1.4B, the L variant uses ~7% fewer FLOPs than Pythia and gains +2.44 5-shot average accuracy.

What problem this solves

Mechanisms like MeSH already let looped Transformers match or beat a vanilla stack at equal compute. Every iteration still runs attention at full token length. If early loops only need coarse global structure and later loops do local refinement, paying quadratic cost at every step is waste.

SpiralFormer treats sequence resolution as a first-class axis inside the loop. One shared core runs on a schedule of compressed lengths, coarse to fine, so hierarchical dependencies can form across iterations.

Method

The backbone is Middle-cycle: pre-loop, shared loop, post-loop. Each iteration has four moves: downsample token states into chunk-level latents at the current resolution, run the shared core on that short sequence, upsample back to a token-level update, then apply a causal right-shift and a topology update (Anchor or MeSH).

The default schedule is coarse-to-fine, doubling each time, e.g. {1/8, 1/4, 1/2, 1} or starting at 1/16. Chunk size is g = floor(1/r). Downsampling uses a per-iteration linear scorer and softmax weights; upsampling uses an output-dependent allocation with gain √g so update scale stays comparable across chunk sizes.

Causality is the hard part. Chunk aggregation sees "future" tokens inside a block, so the upsampled update is shifted right by st. Default st = g−1 is the smallest shift that keeps strict autoregressive causality and leaves a one-token overlap between the producing chunk and the receiving chunk. Chunk boundaries use a half-chunk offset by default, so decode-time compute does not clump on a rigid grid.

Two product points. SpiralFormer-B keeps LoopedFormer's layer allocation and only changes the resolution schedule: almost the same parameters, fewer FLOPs. SpiralFormer-L matches a non-recursive Pythia's parameter count and replaces the middle full-resolution compute with coarse-to-fine loops.

Results

Same recipe as the companion paper: from-scratch pretraining on 250B deduplicated Pile tokens, sequence length 4096. Downstream is 0-shot and 5-shot averages over 9 tasks.

1.4B modelNon-emb paramsPrefill FLOPs0-shot5-shotPile PPL
Pythia 24-layer1208.6M14.08e1249.5051.937.44
LoopedFormer+MeSH805.8M14.08e1250.5652.797.39
SpiralFormer-B+MeSH805.9M12.92e1251.4853.227.30
SpiralFormer-L+MeSH1208.8M13.13e1251.7554.377.14

Relative to full-resolution looping, B cuts FLOPs by about 7-11% (410M: 4.59→4.11; 1B: 9.67→8.95). At matched parameters, L's 1.4B 5-shot is +2.44 over Pythia and +1.58 over LoopedFormer+MeSH.

On 410M, flipping the schedule to fine-to-coarse raises Pile PPL from 9.00 to 9.24 and drops 0-shot from 44.31 to 43.61. MeSH topology beats Anchor. A no-overlap parallel shift (st=g) loses quality; the appendix argues it lets low-resolution work leave the per-token critical path. Learned down/up-scaling beats mean pooling plus uniform broadcast. Recurrence-ratio vs validation loss is U-shaped: best around 30-40% of layers in the shared core; zero looping and aggressive sharing both do worse.

Attention probes on 410M (500 Pile validation sequences): as resolution rises, key-marginal entropy falls and Local Attention Mass rises. Coarse loops are diffuse; fine loops are local. The same probes on full-resolution LoopedFormer show weaker, less ordered cross-loop shifts.

Why it matters

Prior looped work mostly tweaked how state is passed. This paper adds a third axis: besides parameter depth and compute depth, sequence length can shrink per iteration. For pretraining under a parameter or prefill-FLOP budget, B beats a larger Pythia with one-third fewer non-embedding weights; L beats the matched-parameter vanilla on fewer FLOPs. It is the same group's follow-up to MeSH. Stacking the two is the best setting; Anchor alone still runs, just weaker.

This is not a free speedup on existing weights. Everything is trained from scratch. Decode needs chunk-triggered multi-resolution updates and caches, more moving parts than a plain loop.

Limitations

Main results stop at 1.4B. There is no 2.8B/6.9B comparison, so large-scale claims are unearned. The parallel no-overlap regime clearly regresses, and the appendix only shows that pipelining is possible, not how to recover quality. Attention probes are correlational: statistics move with resolution, which does not by itself prove hierarchical reasoning. There is no head-to-head with Huginn or Ouro, and no instruction-tuning or long-context eval. Chunk size is locked by the schedule; it does not segment by content.

Terms

Source

What people are saying

Related papers

All paper explainers