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

Towards Looped Models Done Right, Part II. Rethinking at Fixed Points. By Benhao Huang and six others, from the Institute of Foundation Models, the University of Southern California, and one other institution. Published on the fifth of October. Its code repository is worked out: about three hundred and fifty source files, with a readme and a licence.

Every recurrence of a looped language model adds cost, and it adds it everywhere: in training, in decoding, in prefill, and in reinforcement learning. Each extra recurrence is a bill paid four times over, once at each of those stages, so any saving that touches the recurrence itself is multiplied.

The organising observation is about fixed points. The closer the recurrent states get to them, the less the path to them matters. Once the destination is what counts, the route can be cut short, shared or skipped, and that pays off four ways. Training can use truncated backpropagation. Decoding can share terminal key-value states with almost no loss in accuracy. A distilled student prefills up to one point seven nine times faster. And RL updates can compute gradients from saved rollout states, twice as fast as backpropagating through the replayed trajectory.

All four lean on states actually landing near their fixed points. So the work turns to the two components of training that shape those fixed points: the depth prior, and input injection.

First, the depth prior. Fixed-depth training breaks key-value sharing. Huginn's broad depth prior supports sharing, but dilutes supervision at the target depth more than sharing requires. The authors instead learn the prior from prediction feedback, with an entropy term that keeps it broad. The aim is a prior broad enough for sharing, without spreading supervision thinner than needed.

Second, injection. In existing schemes, the component of the state that lies along the input can amplify the injection, or cancel it. Orthogonal injection removes that component, so what is injected is no longer at the mercy of the state it is added to.

From a hundred million to one point six billion parameters, the learned prior and orthogonal injection each lower perplexity at every scale, relative to Huginn's prior and to existing injection schemes respectively. At one point six billion parameters, the learned prior with a key-value cache three times smaller matches the downstream average of fixed-depth training with the full cache.

The argument runs in one direction. Convergence to fixed points is what makes cheap training, decoding, prefill and RL possible, and the two changes are aimed at the parts of training that decide how well that convergence happens. The cache result shows the payoff in a single comparison: a third of the memory, the same downstream average.