Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation

Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation. By Xuyuan Xiong and seven others, from the University of Hong Kong and Yale University. Published on the twentieth of October.

Looped transformers reuse one shared block of layers several times over. That buys more computational depth without adding trainable parameters. They are trained with backpropagation through time, and to save training compute that backpropagation is often truncated to the last few loops. Truncation has a cost, though: it gives biased gradients, which may weaken the very benefit looping was meant to provide. The finding here is that, when forward and backward loop counts are chosen together, the savings from truncation are not worth what they lose. Full backpropagation through time comes out ahead.

The study is systematic. The authors jointly vary model size, training tokens, and the forward and backward loop counts, across nearly six hundred training runs. From these they fit a single scaling law that captures the effects of both looping and truncation through a quantity they call effective parameters. Despite the compute truncation saves, the fitted law favours full backpropagation when the two loop counts are jointly optimized. Matched-compute runs agree: full backpropagation through fewer loops consistently beats truncated settings that use the same compute.

The law says more. At a fixed model size, the compute-optimal loop count and the training-token budget both grow as compute grows. When model size is free to vary as well, the law picks out two fully trained loops as the best looped setting. In experiments, that setting reaches lower validation loss than the compute-optimal model without loops, with about twenty-seven percent fewer non-embedding parameters, at the price of only eight percent more training compute.

The conclusion holds within the regime that was studied: nearly six hundred runs over the model sizes, token counts and loop counts the authors swept, and a law fitted to them. The preference for full backpropagation is stated for the case where forward and backward loop counts are jointly optimized. The two-loop result is a statement about the best looped setting, compared against the compute-optimal non-looped model, at a small extra cost in training compute.

The intuitive way to make looped training cheaper is to cut the backward pass short. This scaling study points the other way. Spend the compute on fewer loops, train through all of them, and the gradients stay honest enough to pay for it.

Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation. By Xuyuan Xiong and seven others, from the University of Hong Kong and Yale University. Published on the twentieth of October. This is the long version: the whole paper rather than the abstract.

A looped transformer reuses one shared block of layers several times per token, so it gets more computational depth without more trainable parameters. Training it means backpropagating through every one of those loops, and that is expensive. The common shortcut is truncation: run all the loops forward, but send gradients back only through the last few. Huginn, Parcae and Iso-Depth all do some version of it. The catch is that a truncated gradient is biased, and nobody had measured whether the compute it saves pays for the looping benefit it gives up. Earlier scaling studies fixed the recipe first, often training half the loops. This study treats the number of forward loops and the number of backward loops as two separate dials, and asks which setting wins at a fixed compute budget.

Three pieces of vocabulary carry the rest. The model has a prelude that runs once, a recurrent block that runs T times, and a coda that runs once and makes the prediction. A setting is written as two numbers: T forward loops, of which the last B are trained. When B equals T, that is full backpropagation through time. One standard transformer is the setting one, one.

The second piece is the cost model. A forward pass costs two operations per parameter and token, a backward pass costs four. So each backward loop costs twice a forward loop, and the loop-dependent compute grows with T plus two B.

The third is the idea of effective parameters: the size of an ordinary, non-looped model that reaches the same loss on the same tokens. Looping is judged by how much effective size it buys.

The law starts from a standard size and data scaling law and replaces the parameter count with an effective count. The prelude and coda count once. The recurrent block counts several times, scaled by what the authors call the looping coefficient, a function of T and B.

That coefficient was designed to meet four requirements. One loop must recover the ordinary model. More loops of either kind must help. With no trained loop at all, the recurrent block adds no effective size. And the gain must saturate, on the argument that finite weights store finite information, so extra passes over them cannot grow capacity forever. An exponential form with just three fitted constants satisfies all four. Its ceiling works out to about three times the recurrent block, and the paper flags that ceiling as an extrapolation, not something observed.

The key fitted numbers are two exponents: one for forward loops, about zero point two four, and one for backward loops, about one point zero two. Trained loops matter far more than run loops.

From there the argument is short. Hold compute fixed, so T plus two B is constant. Then the only thing that changes is the looping coefficient, and maximising it is a convex problem. Whenever the backward exponent is more than twice the forward one, the unique answer is T equal to B, every loop trained. The factor of two is there because a backward loop costs double. The fitted exponents clear that bar by a wide margin.

The law is fitted on two hundred and thirty configurations from five hundred and ninety-two training runs, with a hundred and seventeen held back for testing. It holds up when whole model sizes or loop settings are left out, and beats the power-law gain of prior work in every evaluation, with about a third of its error when the sixteen-loop row is held out.

The most telling result comes before any law is fitted. Doubling forward loops without adding trained ones changed the loss by less than zero point zero zero seven nats in all twenty-four size and budget combinations tested. And at equal cost per token, two fully trained loops beat four loops with one trained in all twenty-four.

Then the direct test. Four loops fully trained, against eight loops with two trained: same model, same tokens, same compute. The fully trained setting had lower loss at all five sizes, though the margin shrinks as models grow.

With model size free, the law actually prefers no looping at all. But two fully trained loops come closest, worth nearly two untied copies of the block. Trained at their compute-optimal sizes with three seeds each, they stayed within zero point zero zero eight nats of the non-looped model while using twenty-seven to thirty-two percent fewer non-embedding parameters. Given eight percent more compute, the looped model pulled ahead.

The paper has a limitations section, and it names four. It studies dense models only, so whether this carries to looped mixture-of-experts models is open. It measures validation loss only, not downstream tasks or reasoning. Loops are capped at sixteen, so the ceiling and later switches to deeper settings are extrapolations. And the split between the recurrent block and the rest is fixed at each size.

The appendix adds two cautions the abstract does not. At a fixed size, the runs show four loops overtaking two at around two hundred tokens per parameter, well before the law's predicted switch. And against a genuinely untied model with its own block per loop, two shared loops trailed by roughly three to four times the gap the law predicted.

The idea to carry away is that a loop is only worth what you train. Running a block again without a gradient adds almost nothing, so truncation pays for depth you barely use. If you can choose the depth, choose less of it and train all of it. And if you want a looped model at all, two fully trained loops are where the law, and the runs, point.