PaperlayerView exact revision on arXiv

Source-linked AI summary

LiFT: Loop Flow Transformers

Mohammad Mahdi Derakhshani, Pedro M. P. Curvo, Gertjan J. Burghouts, Jan-Willem van de Meent, Cees G. M. Snoek

arXiv:2610.05538v1cs.LGcs.AI

TL;DR

LiFT addresses mixed results from prior looped diffusion models by assigning recurrent steps distinct targets along a continuous path from an initial estimate to the flow-matching target. This lets a trained model vary inference depth, and on ImageNet LiFT-L/2 achieves a 3.34-point FID reduction versus dense DiT-XL/2 with fewer parameters and lower compute.

  • Problem

    Prior looped diffusion models had mixed evidence for useful rollouts beyond training depth and for outperforming dense models at equal compute.

    Source: §1 · p. 2

  • Method

    LiFT supervises each recurrent loop at its own point on a continuous path from the initial estimate to the flow-matching target.

    Source: §1 · p. 2

  • Results

    LiFT-L/2 lowers FID by 3.34 versus dense DiT-XL/2 on ImageNet 256 × 256, using approximately 60% fewer parameters, 32% fewer training FLOPs, and 52% fewer inference FLOPs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    A fixed LiFT checkpoint can adapt its inference computation by selecting recurrent depth, without adding parameters.

    Source: §6 · p. 14

    Evaluation is limited to class-conditional 256 × 256 image generation without classifier-free guidance, and transfer to other modalities or guided sampling remains open.

    Source: §6 · p. 14

Abstract

from arXiv · show

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.

LOOP IN DEPTH, FLOW IN TIME

Figure 1 compares LiFT and dense DiTs by generation quality and inference cost, highlighting where looping smaller models is competitive.

LOOP IN DEPTH, FLOW IN TIME

  • The figure plots FID against inference FLOPs per parameter.LiFT checkpoints use 50 integration steps; dense DiTs vary their step count.Source: Figure 1 · p. 1
  • LiFT outperforms larger dense DiTs at L/2 and XL/2 scales, while dense DiT remains stronger at B/2.The caption identifies the scale-specific comparison but supplies no numeric values.Source: Figure 1 · p. 1
  • Solid curves encode LiFT inference depth by color and parameter count by marker area; dotted curves show dense DiTs as steps vary.Both model families are compared in FID against inference FLOPs per parameter.Source: Figure 1 · p. 1

1 INTRODUCTION

LiFT addresses limits in looped diffusion models by assigning each loop a distinct point on a continuous correction path. This lets inference depth vary after training, and ImageNet results show improved generation at larger model scales.

1 INTRODUCTION

  • Earlier looped methods had mixed results: intermediate predictions trained toward the complete output could worsen beyond training depth, while MeanFlow DiT-B/4 did not beat dense models at equal compute.The cited prior results include FID worsening beyond training depth on ImageNet and quality peaking at 1.5 times training depth on video [12], and the MeanFlow comparison.Source: §1 · p. 2
  • LiFT gives each loop its own target along a straight path from the initial velocity estimate to the flow-matching target.Each loop is supervised at its matching point, with the final loop trained on the target itself.Source: §1 · p. 2
  • Because path progress is continuous, loop count can be chosen after training, making recurrent depth a second inference-compute axis alongside integration steps.The paper describes the two axes as complementary: steps follow the flow, while loops refine velocity estimates.Source: §1 · p. 2
  • At L/2 and XL/2, LiFT keeps improving as inference loops increase and reaches FIDs unavailable to dense models at lower inference cost; at B/2, dense remains stronger.The comparison is on ImageNet 256 × 256 [1]. Dense DiTs level off as integration steps grow to 250.Source: §1 · p. 2
  • LiFT-L/2 lowers FID by 3.34 versus dense DiT-XL/2 on ImageNet 256 × 256, with about 60% fewer parameters and 52% fewer inference FLOPs.The comparison uses the same training-token budget; LiFT also uses 32% fewer training FLOPs.Source: §1 · p. 2

2 BACKGROUND

The background establishes flow matching as the generative objective, describes the DiT prediction architecture, and explains how recurrent weight sharing separates executed depth from parameter count.

2.1 FLOW MATCHING

  • Flow matching learns a time-dependent velocity field transporting samples from a base distribution to the data distribution, and its target can be regressed without simulating an ODE.At inference, the learned field defines a neural ODE integrated from the base distribution toward the data distribution.Source: §2.1 · p. 3

2.2 DIFFUSION TRANSFORMERS

  • A DiT processes image or latent-patch tokens with transformer blocks conditioned on generative time and optional auxiliary information, then predicts the flow-matching velocity.The prediction head maps token representations back to the input geometry.Source: §2.2 · p. 3

2.3 LOOPED TRANSFORMERS AND COMPUTATIONAL DEPTH

  • Looped transformers repeatedly apply shared blocks, increasing executed depth while keeping trainable parameters fixed.The prelude–core–coda form embeds the input, recurrently transforms its state, and maps states to predictions.Source: §2.3 · p. 3
  • The loop count changes computation within a prediction independently of generative time, and can be changed at inference without changing trained parameters.Unique block count is LP + LR + LD, while executed block count is LP + KLR + LD.Source: §2.3 · pp. 3–4
  • Repeated image presentations count toward training data, but recurrent passes over the same tokens add computation without increasing the token count.Equal parameter or data budgets therefore need not imply equal training compute.Source: §2.3 · p. 4

3 LIFT

LiFT trains a recurrent transformer by assigning predictions along a continuous path from an initial estimate to the flow-matching target. Random depth coordinates train the progression, while a shared core and readout support variable inference depth.

3 LIFT

  • LiFT repeatedly applies a shared transformer core and trains its recurrence using a continuous reference path from the initial velocity estimate to the flow-matching target.Each core application receives a depth coordinate and is supervised at the corresponding path point.Source: §3 · p. 4

3.1 METHOD

  • The path connects a stop-gradient initial prediction to the sampled flow-matching target, and its conditional regression optimum moves linearly toward the marginal velocity.At depth s = 1, the objective reaches the ordinary flow-matching optimum.Source: §3.1 · p. 4
  • Training samples interior depth coordinates randomly and supervises each recurrent readout at its matching path value, while the final readout always targets the complete flow-matching task.The sampled grid is sorted and retains the endpoint sKtrain = 1.Source: §3.1 · p. 5
  • The trajectory objective distributes prediction across depth rather than requiring every readout to recover the endpoint immediately.An auxiliary regression loss encourages the initial prediction to provide a useful starting estimate.Source: §3.1 · p. 6
  • At inference, LiFT uses a uniform depth grid for the selected loop budget and sends the endpoint readout to the outer sampler.Changing the loop budget changes the computational grid while preserving its endpoint.Source: Figure 2 · p. 5, §3.1 · p. 6

3.2 ARCHITECTURE

  • The architecture uses a prelude, shared recurrent core, and shared prediction coda; the core receives depth conditioning and reinjects the initial state at every pass.The coda predicts from each state without changing the state passed to the next core application.Source: §3.2 · p. 6

4 EXPERIMENTS

The experiments test whether adjustable recurrent depth improves class-conditional ImageNet generation and how to allocate computation or choose models under resource budgets. LiFT improves on dense references at L/2 and XL/2, while dense B/2 remains stronger.

4.1 EXPERIMENTAL SETUP

  • The evaluation uses class-conditional ImageNet-1k at 256 × 256; within each scale, LiFT and dense references match executed depth and approximately match training computation.Models share a 3.2768 × 10^10 latent-patch-token training budget; LiFT varies shared-core size and training depth.Source: §4.1 · p. 7

4.2 Q1: GENERALIZATION BEYOND TRAINING DEPTH

  • L/2 R10 improves from FID 20.30 at two training loops to 10.95 at eight inference loops; large cores improve beyond training depth, while single-block cores plateau.The largest-core curves flatten after eight to sixteen inference loops.Source: Figure 3 · p. 8, §4.2 · pp. 8–9
  • L/2 R10 overtakes dense XL/2 at three inference loops, with FID 14.64 versus 15.54, slightly less inference compute, and 60% fewer parameters.Most other crossings cost more inference FLOPs than the dense models they beat.Source: §4.2 · p. 9

4.3 Q2: GENERATION QUALITY AT COMPARABLE COMPUTE

  • At B/2, dense remains stronger at every evaluated cost; at L/2 and XL/2, LiFT reaches FID 12.25 and 11.36 versus dense 16.17 and 14.69 at lower inference cost.LiFT uses 41% and 56% fewer parameters at L/2 and XL/2, respectively.Source: §4.3 · p. 9, Figure 4 · p. 10
  • 3.34 points lower FID: LiFT L/2 beats dense DiT-XL/2 while using approximately 60% fewer parameters, 52% fewer inference FLOPs, and 32% less training computation.Training computation was 61.98 versus 91.10 EFLOPs.Source: §4.3 · pp. 9–10, Figure 4 · p. 10

4.4 Q3: ALLOCATING COMPUTATION ACROSS DEPTH AND TIME

  • With 25 steps and four loops, LiFT L/2 reaches FID 13.16 versus dense L/2's 16.94 with 8% less inference computation, but the best reallocation depends on the starting configuration.At eight loops, increasing integration steps from 50 to 100 barely changes FID (10.95 to 10.96), while extending depth to 32 loops raises it.Source: §4.4 · pp. 10–11, Figure 5 · p. 11

4.6 QUALITATIVE COMPARISONS

  • Figure 7 pairs generations from LiFT L/2 R10 and dense L/2 as a qualitative complement to the metric comparisons.Source: §4.6 · p. 11

4.5 Q4: MODEL SELECTION UNDER RESOURCE BUDGETS

  • LiFT reaches every target FID with fewer parameters than dense DiTs and strict targets with less compute; at FID 20, compute- and parameter-minimizing choices differ.The compute-minimizing choice is L/2 R10 (2.96 TFLOPs per image); the parameter-minimizing choice is L/2 R5 (176M parameters).Source: §4.5 · p. 11, Figure 6 · p. 12

5 RELATED WORK

Prior work establishes recurrent depth as an inference-time resource and explores depth conditioning, generative recurrence, and intermediate supervision. LiFT distinguishes itself through depth-indexed velocity targets along a reference trajectory.

Recurrent computation and depth generalization

  • Recurrent models establish depth as an inference-time resource: Universal Transformers share transformations [5], and other work studies extrapolation beyond training depth and recurrence-count distributions.Source: §5 · p. 12

Depth conditioning and computational trajectories

  • Unlike LoopFormer’s endpoint-aligned trajectories [23], LiFT assigns intermediate velocity predictions distinct targets along the correction toward the flow-matching target.Source: §5 · pp. 12–13

Recurrence in generative transformers

  • LoopDiT shares LiFT’s prelude–core–output layout, but uses MeanFlow for its final prediction and retains a prefix schedule, whereas LiFT supervises intermediate velocity readouts and spans [0, 1].Source: §5 · p. 12

Intermediate supervision for generative prediction

  • Related generative methods study intermediate flow-matching supervision [13], recurrent-output losses [26], or intermediate-exit self-distillation [12]; LiFT instead indexes targets by each readout’s fraction of the correction.Source: §5 · pp. 12–13

Computational depth and generative time

  • Unlike models that carry representations across changing generative inputs [29], LiFT reinitializes recurrent state at each integration step and differentiates through its rollout; its finite depth-conditioned path also does not impose hidden-state equilibrium [31] [33] [34].Source: §5 · p. 14

6 CONCLUSION

LiFT’s depth-indexed supervision supports useful inference rollouts beyond training depth, and reallocating computation toward recurrence can outperform larger dense models with fewer parameters and less inference compute. The appendix reports extended experiments and finds that extrapolation gains vary across checkpoints and eventually diminish.

6 CONCLUSION

  • LiFT’s larger-core L/2 and XL/2 checkpoints can outperform larger dense DiTs with fewer parameters and less inference compute, although they trail same-scale dense models at training depth.Source: §6 · p. 14
  • A fixed checkpoint can adapt recurrent computation to available budgets without adding parameters, a flexibility the authors identify as potentially useful when parameter storage is constrained.Source: §6 · p. 14
  • The evaluation covers class-conditional 256 × 256 image generation without classifier-free guidance; transfer to other modalities or guided sampling remains open.Source: §6 · p. 14

A COMPLETE AND EXTENDED EXPERIMENTAL RESULTS

A.2 DEPTH-GENERALIZATION MATRICES

  • L/2 R10 lowers FID from 20.30 to 10.95 at four times training depth, and XL/2 R12 from 17.07 to 9.31 at eight times, but both regress slightly at still larger depths.Source: §A.2 · pp. 22–23, Table 4 · p. 22

A.3 JOINT ALLOCATION OF INFERENCE COMPUTATION

A joint sweep shows that inference quality depends on how computation is divided between recurrent loops and integration steps. The preferred allocation can depend on the quality metric, and gains diminish at larger budgets.

A.3 JOINT ALLOCATION OF INFERENCE COMPUTATION

  • Four loops at 25 integration steps improve on the two-loop training setting at lower cost, while sixteen loops at ten steps are worse and more expensive.Source: §A.3 · p. 23
  • At eight loops, raising integration steps from 25 to 50 lowers FID from 11.39 to 10.95, but doubling again to 100 leaves FID nearly unchanged at 10.96.Source: §A.3 · p. 23
  • A sweep of FID across integration steps and recurrent depth finds that extending the rollout from eight to 32 loops worsens FID at all four step budgets.Source: §A.3 · p. 24, Figure 9 · p. 24
  • At 25 integration steps, increasing depth from eight to 32 improves sFID from 4.36 to 3.96, while FID worsens from 11.39 to 12.14 and IS falls from 97.79 to 94.34.Source: §A.3 · p. 24

B ABLATIONS AND DIAGNOSTICS

  • The following ablations vary depth-coordinate sampling and prelude-loss weight in a shared B/2 setup, then examine inference-grid effects on a frozen L/2 checkpoint.Source: B ABLATIONS AND DIAGNOSTICS · p. 26

B.1 RANDOM VERSUS FIXED DEPTH COORDINATES

Randomizing training depth coordinates improves robustness away from the training depth. A small auxiliary prelude-loss weight also controls drift in the initial prediction and gives the best reported FID in the sweep.

B.1 RANDOM VERSUS FIXED DEPTH COORDINATES

Prelude-loss ablation

  • Without prelude supervision, anchor RMS reaches 3.39 against target RMS 1.30, while nonzero weights keep it near 0.93 and anchor MSE below 0.84; λ = 0.01 suffices.Source: §B.1 · p. 27, Table 9 · p. 27
  • λ = 0.01 has the lowest reported FID, 64.81 at eight inference loops, compared with 66.97 at zero weight, 65.32 at λ = 0.1, and 70.03 at λ = 1.Source: §B.1 · p. 27

B.3 PREDICTION TRAJECTORIES

The analysis tests whether intermediate predictions follow their assigned path toward the sampled target and whether the model’s own rollout follows a similar path. A two-loop-trained checkpoint maintains near-linear progression through 32 inference loops.

Analysis setup

  • The analysis compares intermediate readouts both with the supervised path toward the sampled target and with the chord between each rollout’s initial and final predictions.The paired evaluation uses a frozen LiFT L/2 R10 checkpoint trained with Ktrain = 2 and evaluates 512 validation examples at five generative times.Source: §B.3 · pp. 27, 29

Progress toward the supervised target

  • At uniform depth four, mean target-direction projections are 0.246, 0.493, 0.737, and 0.982; at depth 32 they remain 0.249, 0.495, 0.738, and 0.982.The measurements use coordinates s = 0.25, 0.5, 0.75, and 1, respectively.Source: §B.3 · p. 29, Figure 13 · p. 30

Geometry between the model’s own predictions

  • The two-loop-trained checkpoint’s own-chord projection tracks the depth coordinate closely through 32 loops, with mean absolute discrepancy rising from 0.0006 at K = 2 to 0.0051 at K = 32.The orthogonal residual changes from 6.07% of the final prediction norm at K = 2 to 8.68% at K = 32.Source: §B.3 · pp. 29–30, Figure 13 · p. 30

B.4 COMPUTATIONAL-GRID CONSISTENCY

The tested computational grids produce similar final velocities at fixed depth, while changing the number of recurrent applications produces larger variation. The measurements support a regular depth-conditioned progression that is mildly sensitive to the tested coordinate changes.

Endpoint agreement at fixed depth

  • Front- and back-loaded K = 4 grids change the endpoint by mean 1.89% and 1.82%, respectively, with mean cosine similarity above 0.9997.The three random grids have similar mean differences of 1.72–1.85%, but larger per-image time-averaged 95th percentiles of 4.26–4.61%, versus 2.58–2.66% for fixed nonuniform grids.Source: §B.4 · p. 30, Table 12 · p. 31

Changing the number of recurrent applications

Numerical precision

  • At intermediate times, grid-difference RMSE generally exceeds the BF16–FP32 precision baseline; the largest listed ratio is 7.5 for uniform depth 32 at t = 0.5.The precision baseline is pooled RMSE 0.0065–0.0150 on the same 16 examples at each time.Source: §B.4 · p. 32, Table 13 · p. 32

Interpretation

  • Across the trajectory and endpoint measurements, the learned progression is regular and only mildly sensitive to the tested intermediate-coordinate changes.Source: §B.4 · p. 32

C ANALYSIS OF THE LIFT OBJECTIVE

The appendix formalizes the LiFT objective under conditional-regression assumptions, derives training gradients with a detached anchor, and interprets the reference path through action and fixed-covariance Gaussian geometry.

C.1 CONDITIONAL REGRESSION RESULT

  • For a fixed anchor and unrestricted predictions, the conditional regression optimum progresses from that anchor toward the population velocity field, reaching the ordinary flow-matching optimum at s = 1.Parameter sharing, finite capacity, and optimization can prevent simultaneous realization of these optima.Source: §C.1 · p. 32

C.2 GRADIENT FLOW WITH A DETACHED ANCHOR

  • The anchor is recomputed from the current model but detached within each update; trajectory gradients still traverse the recurrent computation, while the anchor evolves across updates.The construction does not require a fixed-point equation or an inner ODE solver, and the analysis does not guarantee reduced numerical integration error.Source: §C.2 · pp. 32–33

C.3 EUCLIDEAN MINIMUM ACTION

  • Under the stated endpoint and path assumptions, the unique Euclidean minimum-action path is the straight interpolation from b to u⋆; nonuniform grids select knots on that same path.The grid knots minimize the interpolated-path action, while the trajectory objective separately averages prediction errors with weights 1/K.Source: §C.3 · p. 33

C.4.1 GAUSSIAN PROBABILITY TRANSPORT

  • With covariance fixed, translating a Gaussian family along its mean path reduces transport-action minimization to minimizing the Euclidean action of the means.In this interpretation, the network passes the terminal mean to the sampler as its velocity prediction; it does not sample w or estimate posterior uncertainty.Source: §C.4.1 · pp. 33–34

C.4.2 FISHER GEOMETRY WITH FIXED COVARIANCE

  • For the fixed-covariance Gaussian mean submanifold, the Fisher metric is constant and its unique constant-speed energy minimizer is the straight reference path.The result requires covariance to remain fixed along every admissible path; allowing it to vary can produce geodesics outside this subfamily.Source: §C.4.2 · p. 34

C.4.3 GAUSSIAN KL PROJECTION

  • Fitting the predicted Gaussians to the reference Gaussians by KL divergence recovers the trajectory objective up to a positive constant when covariance is fixed.The equivalence needs no Gaussian sampling or variance parameter in the implementation.Source: §C.4.3 · p. 34

D REPRODUCIBILITY AND IMPLEMENTATION

D.1 DATA AND ARCHITECTURE

The implementation uses a shared ImageNet latent-data and transformer setup for dense and recurrent models, while LiFT trains intermediate readouts along its trajectory. The appendix specifies sampling, evaluation, and compute accounting for these comparisons.

D.1 DATA AND ARCHITECTURE

  • The models use class-conditional ImageNet-1k at 256 × 256, encoded as cached 4 × 32 × 32 latents, and share the same transformer block design.The 2 × 2 latent patches yield 256 tokens per image; LiFT uses four prelude blocks and no transformer coda in the main comparisons.Source: §D.1 · p. 35, Table 14 · p. 35

D.2 TRAINING PROTOCOL

  • Training uses 500,000 updates at global batch size 256; dense models use conditional flow matching, while LiFT uses trajectory supervision with prelude-loss weight λ = 0.01.LiFT samples interior depth coordinates and appends the endpoint s = 1.Source: §D.2 · p. 35

D.3 TRAINING AND SAMPLING PSEUDOCODE

  • LiFT training samples noise, generative time, and depth coordinates per example, supervises each recurrent readout, and backpropagates through predictions and recurrent states.At sampling, only the final readout advances the outer sample; the sampling algorithm uses Euler integration with T steps and Kinf inner loops.Source: §D.3 · pp. 35, 37

D.4 EVALUATION PROTOCOL

  • Evaluation generates 50,000 images per setting with EMA weights and Euler integration, without classifier-free guidance, and reports FID, sFID, and IS.FID and sFID use the published ImageNet reference statistics; depth sweeps hold integration steps at 50.Source: §D.4 · pp. 37–38

D.5 COMPUTE ACCOUNTING

  • Compute accounting includes recurrent applications in training and inference, while sampling requires only the final readout.Matching executed transformer depth approximately matches training computation; the estimate includes conditioning and readouts, and excludes tokenizer decoding from sampling cost.Source: §D.5 · p. 38, §D.5.1 · p. 38

D.5.1 COST OF TRAJECTORY SUPERVISION

  • Across matched-depth configurations, LiFT’s additional training cost is less than 0.2% of its same-scale dense reference under the stated accounting.The estimate includes conditioning and all trajectory readouts; the totals are approximately, not exactly, matched.Source: §D.5.1 · p. 39, Table 17 · p. 39

Related papers

Loading 2610.05538v1…