Skip to content

Back to the catalog

Towards Looped Models Done Right, Part II: Rethinking at Fixed Points
arxiv.orgpaper

Towards Looped Models Done Right, Part II: Rethinking at Fixed Points

This paper studies looped language models — architectures that reuse a shared block of layers recurrently — and argues that as recurrent hidden states approach fixed points, the specific path taken to reach them matters less, enabling several efficiency tricks: truncated backpropagation through time during training, terminal key-value (KV) cache sharing across loop iterations for faster decoding, a distilled student model that prefills up to 1.79x faster, and RL policy gradients computed from saved rollout states that run 2x faster than backpropagating through the full trajectory. To make these fixed points reliable, the authors replace fixed-depth training with a learned depth prior (regularized with an entropy term to stay broad, improving on Huginn's fixed broad prior) and introduce orthogonal input injection, which removes the component of the injected input that the hidden state could otherwise amplify or cancel. Across model scales from 100M to 1.6B parameters, both changes lower perplexity, and at 1.6B the learned prior with a 3x smaller KV cache matches the downstream performance of full-cache fixed-depth training.

TL;DR A follow-up paper on looped language models shows that training recurrent states toward fixed points enables KV-cache sharing, faster distillation-based prefill, and cheaper RL gradients, via a learned depth prior and orthogonal input injection.
Open original ↗
Document

What it claims

The efficiency tricks only work if the looped model's hidden states actually converge to fixed points. Huginn's states converge, and terminal KV sharing barely changes its zero-shot GSM8K accuracy. Ouro's states do not converge, and the same sharing collapses its accuracy at four loops. The authors' own fixed-depth model trained with full backpropagation through time (BPTT) at 1.6B shows the same split: sharing drops its GSM8K accuracy from 50.6 to 21.2. The shortcuts therefore depend on how the model was trained, and you cannot bolt them onto an arbitrary looped checkpoint.

Measured savings:

  • Up to 2.6× less pretraining activation memory.
  • 3× smaller decoding KV cache.
  • 1.79× faster prefill.
  • 2× faster RL scoring plus backward pass.

How it works

Experiments use the Huginn layout:

  • a non-shared "prelude" that embeds the input,
  • a shared recurrent core that starts from a random state and re-injects the embedding on every loop,
  • a non-shared "coda" that produces logits.

A token counts as converged from depth r when its relative hidden-state change stays below 2% for every later step. Tokens converge at different depths, not in causal order. Convergence depth correlates only weakly with position (Spearman 0.14).

  • Truncated BPTT: gradients flow only through the last b recurrences. Unlike implicit (equilibrium) gradients, it needs no warm-up phase. Endpoint gradients applied from initialization destabilize training. Truncation by itself is enough to produce fixed points.
  • Terminal KV sharing: each prefix token keeps only its final-loop KV, and every recurrence of later tokens attends to that. Prefill and decoding can run at different recurrence budgets. Prefill gains level off at moderate depth, while decoding gains continue to grow with depth.
  • Distilled prefill:
  • A student is initialized from the teacher's recurrent core and predicts the endpoint state directly from the prelude output.
  • Prefill then applies one teacher recurrence plus the coda.
  • The loss is normalized hidden-state regression plus next-token KL against depth-R teacher targets computed online.
  • Decoding is unchanged: the teacher runs at full depth with terminal KV sharing.
  • RL rollout-state reuse: saved rollout endpoints z_R are fed through the recurrent block once. The backward pass is then looped b−1 times over that single recorded graph, which is a Neumann-series approximation of the full-trajectory gradient. Update cost and memory no longer depend on depth R.
  • Learned depth prior: the distribution over training depths is learned from detached cross-entropy, using an EMA-centered advantage. It adds no extra forward pass, and an entropy term keeps it broad.
  • Orthogonal injection (OrthoInj): the part of the carried-over state that lies along the input is projected out before the input is added. Every recurrence therefore receives the input at the same strength.

Evidence

Truncation

  • With fixed depth R=5 and full BPTT, extrapolating to R=64 raises perplexity by more than 25%, and KV sharing raises it by 21%.
  • With truncated BPTT at b=2, the same two changes raise perplexity by at most 2% and 3%.
  • At 400M with sampled depths, b≥4 matches full BPTT within 1% perplexity.

KV sharing

  • At R=5, sharing keeps 4 of 12 banks and adds at most 0.85% perplexity.
  • Learned-prior models stay within 0.4 points of their full-cache downstream average.

Distilled prefill

  • The student is trained on a quarter of the teacher's pretraining tokens.
  • It is up to 1.79× faster on 8K prompts.
  • It trails the teacher's downstream average by 1.0–4.5 points.
  • At matched latency, it beats the teacher prefilling with two recurrences by 0.5–0.9 points.

RL reuse: GSM8K pass@1 stays within 1.6 points of full BPTT replay.

OrthoInj vs Parcae's injection (100M–1.6B): lower training loss, 0.3–1.4% lower validation perplexity, and a higher downstream average.

How to use it

Code: https://github.com/ifm-ai/xllm-loop. The page states no license, install steps, or checkpoint names.

Limits

  • Results cover 100M–1.6B parameters only.
  • Distilled prefill costs 1.0–4.5 downstream points, and its parity claim holds only at matched latency.
  • The 2× RL speedup covers scoring and backward only. Rollout generation still takes most of the RL training time.
  • The fetched text cuts off partway into Section 4. Full details of the depth-prior algorithm and the result tables were not available here.

Technologies named here

Techniques named here

  • Truncated backpropagation
  • Terminal key-value sharing
  • Distilled student for prefill
  • RL updates from saved rollout states
  • Learned depth prior with entropy term
  • Orthogonal injection

Related in research