← Back to projects

Training-Free Hidden-State Refinement for Flow-Matching Image Generators

Yuanyi Yan1, Xinzhe Rao1, Canyu Shen2, Yang Chen1, Yunlu Chen3, Meng Tang4, Teng Long5, Vincent Tao Hu1

  1. 1Huazhong University of Science and Technology
  2. 2Tongji University
  3. 3King Abdullah University of Science and Technology
  4. 4University of California, Merced
  5. 5University of Amsterdam

arXiv:2608.29160 [cs.CV], 2026

PDFarXivCodeBibTeX
Qualitative comparisons and quality-efficiency plots for Dense Token Loop, Sparse Token Loop, and Loop Guidance
Left: matched samples without looping, with Dense Token Loop, and with Loop Guidance. Right: internal looping improves GenEval and DPG-Bench more efficiently than adding outer sampling steps.

Abstract

We study how to spend additional inference-time computation inside a frozen flow-matching image denoiser, without changing its weights or outer sampler. The framework repeats selected transformer layers and exposes independent controls over token scope, layer depth, sampling progress, loop count, and prediction-space guidance. Across Scale-RAE and RAEv2 generators, the method improves prompt alignment and image-quality metrics while offering useful quality–latency trade-offs.

Key Results

+0.1220

GenEval

0.4471 → 0.5691 on Scale-RAE DiT2.4B

+0.0397

DPG-Bench

0.7656 → 0.8053 on Scale-RAE DiT2.4B

39.6%

Lower latency

Sparse versus Dense Token Loop on DiT9.8B

0

Training updates

Weights, conditioning, autoencoder, and sampler stay frozen

Method

Training-free internal looping with sampling-progress gating, dense or sparse token updates, and a selected transformer-layer range
Sampling-Progress Gating controls when looping is active. Dense Token Loop repeats the selected layer range for all tokens; Sparse Token Loop restricts repetition to selected tokens and reuses a cached complement residual. The outer sampler and all model weights remain unchanged.

Dense and Sparse Token Loop

Repeat a transformer-layer window for every token, or focus repeated computation on a subset while retaining full-token attention context.

Sampling-Progress Gating

Activate internal loops only during chosen denoising calls; the strongest shared setting concentrates them in the early half of sampling.

Layer Range and Loop Budget

Choose where repeated updates occur and distribute a total loop strength across a controllable number of inner evaluations.

Loop Guidance

Combine ordinary and looped vector-field predictions, using their difference as a prediction-space correction direction.

Results

Direct looping across three RAE-based generators

ModelMethodGenEval ↑DPG-Bench ↑ImageReward ↑Time (s) ↓
Scale-RAE DiT2.4Bw/o Loop0.44710.76560.09941.3054
Dense Token Loop (Ours)0.54220.80070.48162.1812
Sparse Token Loop (Ours)0.50820.78930.27981.6884
Scale-RAE DiT9.8Bw/o Loop0.53210.80030.46091.6853
Dense Token Loop (Ours)0.59770.81780.72114.0384
Sparse Token Loop (Ours)0.58570.82060.63652.4410
RAEv2 SigLIP2-B (615M)w/o Loop0.38290.7131-0.33850.8025
Dense Token Loop (Ours)0.37950.7226-0.25381.1023
Sparse Token Loop (Ours)0.38150.7152-0.29311.3572

Loop Guidance improves both primary metrics

ModelGenEvalDPG-Bench
Scale-RAE DiT2.4B0.4471 → 0.56910.7656 → 0.8053
Scale-RAE DiT9.8B0.5321 → 0.64320.8003 → 0.8264
RAEv2 SigLIP2-B0.3829 → 0.45320.7131 → 0.7521

Each metric reports the no-loop baseline followed by the loop-guided result.

Internal compute also improves the quality–latency trade-off: a 25-step Dense Token Loop reaches 0.5422 GenEval and 0.8007 DPG-Bench in 2.1812 seconds, outperforming the 50-step no-loop baseline (0.4742 and 0.7738) while remaining faster than its 2.4032-second runtime.

Qualitative Comparisons

Prompt- and seed-matched samples on Scale-RAE Qwen7B / DiT9.8B compare the ordinary path, Dense Token Loop, and Loop Guidance. Expand either benchmark gallery to inspect the full set shown here.

GenEval gallery · 8 comparisons
GenEval matched comparison 02_prompt085GenEval matched comparison 03_prompt099GenEval matched comparison 07_prompt184GenEval matched comparison 10_prompt269GenEval matched comparison 11_prompt311GenEval matched comparison 13_prompt410GenEval matched comparison 14_prompt510GenEval matched comparison 15_prompt552
DPG-Bench gallery · 9 comparisons
DPG-Bench matched comparison 01_prompt000DPG-Bench matched comparison 02_prompt055DPG-Bench matched comparison 03_prompt082DPG-Bench matched comparison 04_prompt181DPG-Bench matched comparison 05_prompt262DPG-Bench matched comparison 07_prompt453DPG-Bench matched comparison 09_prompt562DPG-Bench matched comparison 10_prompt672DPG-Bench matched comparison 12_prompt753

Scope and Limitations

Improvements are strongest for one- and two-subject prompts, where looping often repairs malformed structure and attribute binding. Aggregate GenEval2 does not improve, so hidden-state refinement alone does not solve difficult counting or crowded multi-object composition.

Reproducibility

The public repository contains the runtime, four backend adapters, exact presets, 2,418 ordered benchmark prompts, aggregate results, CPU-only tests, and generation, export, and aggregation utilities.

Browse the code and reproduction guide →

Citation

@misc{yan2026trainingfree,
  title={Training-Free Hidden-State Refinement for Flow-Matching Image Generators},
  author={Yan, Yuanyi and Rao, Xinzhe and Shen, Canyu and Chen, Yang and Chen, Yunlu and Tang, Meng and Long, Teng and Hu, Vincent Tao},
  year={2026},
  eprint={2608.29160},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2608.29160}
}