LiFT lets looped DiTs keep improving beyond training depth: L/2 and XL/2 beat larger dense DiTs with fewer parameters and less inference compute

LiFT gives each loop of a recurrent DiT its own regression target: a point on a straight path from the model's initial velocity estimate to the flow-matching target, indexed by a continuous depth coordinate, with the final readout still carrying the complete flow-matching task. Earlier looped diffusion models trained every intermediate prediction toward the complete output or relied on a detached full-depth teacher for self-distillation; LiFT instead distributes the remaining correction along a reference path, so successive loops perform successive stages of one correction. Validated on class-conditional ImageNet generation: L/2 R10 improves from FID 20.30 at its two-loop training depth to 10.95 at eight inference loops, XL/2 R12 from 17.07 to 9.31 at sixteen loops, and XL/2 R8 from 16.95 at three training loops to 9.50 at thirty-two loops.

The inference loop count is selectable after training, and loops and integration steps act as complementary axes of inference compute whose best split depends on the budget and starting configuration. Prior looped models worsened once inference exceeded training depth or showed no advantage over a dense model at equal compute; LiFT reports continued improvement under depth extrapolation and a joint allocation study over loops and steps. L/2 R10 at 25 steps and four loops reaches FID 13.16 at 7.397 TFLOPs, better than dense L/2 at 50 steps (16.94 at 8.069 TFLOPs); the same model at sixteen loops and ten steps gives FID 15.09 at 11.027 TFLOPs, both worse and more expensive.

At L/2 and XL/2 scales, reallocating inference compute toward recurrence yields operating points that beat larger dense DiTs with fewer parameters and less inference compute. Relative to a same-scale dense reference, LiFT trails at its training depth because shared blocks replace unique ones; the advantage emerges only once test-time compute is reallocated, not at the training operating point. LiFT L/2 R10 with eight loops reaches FID 10.95 at 28.24 TFLOPs, versus dense XL/2 at 250 steps with FID 14.30 at 59.31 TFLOPs: 3.34 FID points lower, roughly 60% fewer parameters, 52% fewer inference FLOPs, and about 32% less training computation (61.98 versus 91.10 EFLOPs).

The benefit of depth extrapolation depends on core capacity: single-block cores and the B/2 scale gain little, and the dense model stays stronger at every evaluated cost at B/2. The paper reports this scale dependence as an explicit scope condition rather than a universal claim. B/2 R1 moves only from FID 45.70 at eight loops to 45.58 at sixteen; the best LiFT B/2 result, FID 33.87, does not reach the 32.50 that dense B/2 attains with 25 steps at 1.150 TFLOPs, which the authors attribute to capacity, with at most 89M parameters for the same 32.8B training tokens.

The result applies to 256-resolution class-conditional ImageNet generation without classifier-free guidance, and is meant for deployment settings where a fixed checkpoint should adapt its computation to an available budget, particularly when parameter storage is more restrictive than computation; the authors also motivate budget-aware generation in video and generative dynamics models. Methodologically, LiFT splits a DiT into a prelude, a shared recurrent core, and a prediction-only coda, conditions the core on the depth coordinate through adaptive normalization and residual gates, and at inference reinitializes the prelude state at each integration step while using only the final readout.

Each result comes from a single training run and one 50,000-image sample set, and small FID differences near the plateaus were not tested for statistical significance; inference depths are prescribed, and the joint allocation study covers one checkpoint over a finite grid, leaving adaptive depth selection and broader optimization of inner and outer compute budgets open. Deployment implications require hardware evaluation, since the paper reports parameter counts and analytic FLOP estimates rather than device-level latency, energy consumption, or peak memory. Whether the same depth-scaling behavior extends to other modalities or guided sampling also remains open.