LiFT: Loop Flow Transformers

2026-10-04

AI Summary

This study addresses the high computational costs incurred when diffusion models rely on increasing parameters or complex architectures to enhance performance. To this end, we propose the Recurrent Flow Transformer, which iteratively applies a shared DiT core and introduces a novel single-step regression training scheme based on continuous depth indexing. This design enables test-time compute scaling without retraining, supporting both early exit and infinite-loop inference beyond the trained depth. Evaluated on ImageNet 256×256, our method reduces FID by 3.34 points while decreasing parameter count by 60% and lowering training and inference FLOPs by 32% and 52%, respectively. These results demonstrate an efficient and scalable approach to generative modeling.

Abstract

We introduce Loop Flow Transformers (LiFT), a family of looped generative models that scales computation by repeatedly applying a shared Diffusion Transformer (DiT) core, with only light changes to the standard architecture. Rather than asking every recurrent step for the final prediction, LiFT trains each step with a single regression target: a point on a straight path from the model's initial estimate to the flow-matching target. Because we index these targets by a continuous depth coordinate, a trained model can loop far beyond its training depth with no retraining, early exits, or other modifications. In our experiments, these longer rollouts improve generation, so inference computation can grow without adding parameters. On ImageNet at 256x256, LiFT-L/2 achieves an FID 3.34 points lower than our dense DiT-XL/2 baseline while using approximately 60% fewer parameters, 32% fewer training FLOPs, and 52% fewer inference FLOPs.

Problem

• generative models
• Diffusion Transformer
• inference scaling
• parameter efficiency
• flow matching

Innovation

• Loop Flow Transformers
• Diffusion Transformer
• Flow Matching
• Weight Sharing
• Continuous Depth