LiFT: Loop Flow Transformers

LiFT. Loop Flow Transformers. By Mohammad Mahdi Derakhshani and four others, from the University of Amsterdam and TNO. Published on the fourth of October. No code repository is linked to it.

A looped generative model scales computation by applying the same core over and over, rather than stacking more distinct layers. That raises a training question: what should each pass be asked to produce? One answer is to ask every recurrent step for the final prediction, so every step is judged against the finished output.

Loop Flow Transformers, or LiFT, take a different route. They are a family of looped generative models built by repeatedly applying a shared Diffusion Transformer core, with only light changes to the standard architecture. Each step is trained with a single regression target: a point on a straight path that runs from the model's initial estimate to the flow-matching target. Each loop is asked to get part of the way there, not all of it.

Those targets are indexed by a continuous depth coordinate. Because depth is continuous rather than a fixed count, a trained model can loop far beyond the depth it was trained at, without retraining, early exits, or any other modification. In the authors' experiments, those longer rollouts improve generation, so inference computation can grow without adding parameters.

On ImageNet at two hundred fifty-six by two hundred fifty-six, LiFT-L/2 reaches an FID three point three four points lower than the authors' dense DiT-XL/2 baseline. It does so with roughly sixty percent fewer parameters, thirty-two percent fewer training FLOPs, and fifty-two percent fewer inference FLOPs.

One straight path, indexed continuously, is what turns depth into a dial. More loops at inference means more computation, and the parameter count stays where it was.