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

Yuanyi Yan; Xinzhe Rao; Canyu Shen; Yang Chen; Yunlu Chen; Meng Tang; Teng Long; Vincent Tao Hu

cs.CVAug 29, 2026

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

Authors: Yuanyi Yan, Xinzhe Rao, Canyu Shen, Yang Chen, Yunlu Chen, Meng Tang, Teng Long, Vincent Tao Hu

Organizations: Huazhong University of Science and Technology · Tongji University · King Abdullah University of Science and Technology · University of California, Merced ·

Abstract

We aim to improve frozen flow-matching image generators by adding inference computation inside the denoiser, without changing model weights or the outer sampler. Existing generators usually spend extra test-time computation by increasing the number of sampling steps, which repeatedly evaluates the entire denoiser and couples quality gains to sampler cost. A key challenge is how to use extra computation inside a frozen transformer denoiser: the method must decide which tokens, layers, and sampling times receive repeated updates while preserving the original generation pipeline. We introduce a training-free looping framework that repeatedly applies selected transformer layers inside each denoising call. Dense and Sparse Token Loop vary the token scope; Sampling-Progress Gating and the loop layer range specify when and where looping is active; loop count and strength control the repeated updates; and Loop Guidance combines ordinary and looped vector-field predictions. Across two Scale-RAE model scales, loop variants improve primary and auxiliary quality metrics with competitive quality--efficiency trade-offs. Loop Guidance further improves both primary metrics across all three tested models; on Scale-RAE DiT2.4B, it raises GenEval from 0.4471 to 0.5691 and DPG-Bench from 0.7656 to 0.8053. Code will be released.

Figures & tables

Figure 1: Overview of training-free internal looping. Sampling-Progress Gating determines when looping is active. Dense Token Loop repeats the selected layer range for all tokens, while the optional Sparse Token Loop restricts repetition to selected tokens and reuses a complement residual initialized once. All weights remain frozen, and the outer sampler is unchanged.

Model	Method	GenEval $\uparrow$	DPG $\uparrow$	ImageReward $\uparrow$	CLIP $\uparrow$	OpenCLIP $\uparrow$	PickScore $\uparrow$	Time (s) $\downarrow$
	w/o Loop	0.4471	0.7656	0.0994	0.3204	0.3490	0.2122	1.3054
Scale-RAE DiT2.4B	Dense Token Loop (Ours)	0.5422	0.8007	0.4816	0.3302	0.3622	0.2165	2.1812
	Sparse Token Loop (Ours)	0.5082	0.7893	0.2798	0.3249	0.3555	0.2144	1.6884
	w/o Loop	0.5321	0.8003	0.4609	0.3285	0.3604	0.2175	1.6853
Scale-RAE DiT9.8B	Dense Token Loop (Ours)	0.5977	0.8178	0.7211	0.3344	0.3682	0.2221	4.0384
	Sparse Token Loop (Ours)	0.5857	0.8206	0.6365	0.3328	0.3663	0.2205	2.4410

Table 1: Main results across three RAE-based model configurations. w/o Loop denotes ordinary denoising. The best and second-best values for each model are shown in bold and underlined, respectively.

Model	Method	$\mathbf{g}_{\mathrm{lg}}$	GenEval $\uparrow$	DPG $\uparrow$	ImageReward $\uparrow$	CLIP $\uparrow$	OpenCLIP $\uparrow$	PickScore $\uparrow$
Scale-RAE DiT2.4B	w/o Loop	0.0	0.4471	0.7656	0.0994	0.3204	0.3490	0.2122
	Dense Token Loop (Ours)	2.8	0.5691	0.8053	0.5622	0.3306	0.3627	0.2171
Scale-RAE DiT9.8B	w/o Loop	0.0	0.5321	0.8003	0.4609	0.3285	0.3604	0.2175
	Dense Token Loop (Ours)	6.0	0.6432	0.8264	0.7975	0.3332	0.3661	0.2229
	w/o Loop	0.0	0.3829	0.7131	-0.3385	0.3076	0.3176	0.2055
RAEv2 SigLIP2-B (615M)	Dense Token Loop (Ours)	9.0	0.4532	0.7521	0.0460	0.3159	0.3297	0.2090

Table 2: Loop Guidance improves both primary metrics and all auxiliary metrics across the three evaluated generators. Baseline rows use

g_{\mathrm{lg}}=0

, and guided rows use the listed scales. Best values within each model are bolded.

Figure 2: Moderate Loop Guidance scales and loop counts capture most quality gains on Scale-RAE (DiT2.4B). Panel (a) varies the Loop Guidance scale, and panel (b) varies the loop count

K

Figure 3: Mid-to-late layer ranges and earlier loop-active sampling intervals give the strongest DPG-Bench settings on Scale-RAE (DiT2.4B). Panel (a) varies the loop layer range with the loop-active sampling-progress interval fixed; panel (b) varies the loop-active sampling-progress interval with the layer range fixed. The corresponding GenEval panels are provided in the appendix.

Loop-active interval	GenEval $\uparrow$	DPG $\uparrow$	Time(s) $\downarrow$
w/o Loop	0.4471	0.7656	1.3054
$[0,1]$	0.5597	0.7987	2.9957
$[0,0.5]$	0.5422	0.8007	2.1812

Table 3: For Dense Token Loop,

[0,1]

denotes full-schedule looping and

[0,0.5]

denotes Sampling-Progress Gating; the latter preserves most full-schedule Loop quality while reducing generation time on Scale-RAE.

Model	Sparsifier	Ratio	GenEval $\uparrow$	DPG-Bench $\uparrow$
	w/o Loop	–	0.4471	0.7656
	random	50%	0.5029	0.7951
	condition-aware	50%	0.5039	0.7846
Scale-RAE	attention-aware	50%	0.5008	0.7894
	random	70%	0.5140	0.7962
	condition-aware	70%	0.5175	0.7918

Table 4: Sparsifier comparison on Scale-RAE and RAEv2.

Model	Method	GenEval $\uparrow$	DPG-Bench $\uparrow$
	w/o Loop	0.4778	0.7038
PixArt-alpha	Dense Token Loop	0.4710	0.7170
	Loop Guidance	0.3928	0.6986
	w/o Loop	0.8459	0.8740
FLUX.2	Dense Token Loop	0.8482	0.8740
	Loop Guidance	0.8310	0.8720

Table 5: VAE transfer diagnostic on PixArt-alpha and FLUX.2.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Generator	Resolution	Precision	Steps	Dense range	Sparse range / ratio	$K/\lambda_{\mathrm{loop}}$	Progress	$g_{\mathrm{lg}}$ dense/sparse
Scale-RAE DiT2.4B	$224^{2}$	FP32	25	$[12,27]$	$[12,19]/70\%$	$4/1$	$[0,0.5]$	$2.8/\text{--}$
Scale-RAE DiT9.8B	$224^{2}$	FP32	25	$[12,27]$	$[12,27]/50\%$	$4/1$	$[0,0.5]$	$6.0/\text{--}$
RAE-SigLIP2-B proxy	$256^{2}$	BF16	50	$[15,22]$	$[15,22]/50\%$	$4/1$	$[0,0.5]$	$9.0/6.0$
PixArt-alpha	$1024^{2}$	BF16	20	$[14,19]$	–	$4/1$	$[0,0.5]$	$3.0/\text{--}$
FLUX.2	$1024^{2}$	BF16/4-bit	28	$[32,37]$	–	$4/1$	$[0,0.5]$	$3.0/\text{--}$

Appendix

Table 6: Complete loop configurations used for the main and transfer comparisons. Precision denotes parameter or weight precision. A dash indicates that Sparse Token Loop was not evaluated for that generator. The reported Loop Guidance scales are listed as dense/sparse; a dash indicates that the corresponding guided variant was not reported in the main comparison. Unguided loop variants are equivalent to

g_{\mathrm{lg}}=1

. The Scale-RAE DiT9.8B TF32 matmul policy is described in the text, and controlled Scale-RAE ablations vary only the named factor.

Figure 4: GenEval counterparts for the loop layer range and loop-active sampling-progress interval ablations on Scale-RAE (DiT2.4B). Mid-to-late layer ranges and earlier sampling-progress intervals give the strongest gains, consistent with the DPG-Bench results in the main paper.

Figure 5: Additional prompt-matched qualitative comparisons on Scale-RAE (DiT9.8B). Dense Token Loop uses loop layer range $[12,27]$ , a loop count of $K=4$ , and loop-active sampling-progress interval $s_{i}\in[0,0.5]$ ; Loop Guidance uses $g_{\mathrm{lg}}=6.0$ .

Panel of Figure 6 — Figure 6: GenEval2 qualitative correction study on Scale-RAE (DiT9.8B). Each row inside a panel shows the prompt followed by w/o Loop, Dense Token Loop (loop layer range $[12,27]$ , a loop count of $K=4$ , and sampling progress $s_{i}\in[0,0.5]$ ), and Loop Guidance ( $g_{\mathrm{lg}}=6.0$ ). Top left: few-subject prompts improve. Top right: counting prompts are not corrected. Bottom left: objects in crowded multi-object prompts stay intact. Bottom right: subjects become purer with weaker backgrounds.

Paper control	JSON field	Shared CLI override	Semantics and accepted values
Outer sampler steps $N$	generation.num_ inference_steps	--outer-steps N	Positive integer; number of outer denoising steps.
Loop count $K$	num_loops	--loop-count K	Positive integer; inner updates at each active outer step.
Total loop strength $\lambda_{\mathrm{loop}}$	lambda_value	--lambda-loop VALUE	Positive scalar; the runtime applies $\lambda_{\mathrm{loop}}/K$ per inner round.
Loop-active sampling interval	start_frac , end_frac	--loop-active START END	$0\leq\texttt{START}\leq\texttt{END}\leq 1$ ; both endpoints are active.
Loop layer range $[a,b]$	block_indices	--loop-layers FIRST LAST	Inclusive, contiguous, ascending, nonnegative layer indices.
Token-loop base	token_operator	--token-loop none\|dense\|sparse	Ordinary path, Dense Token Loop, or Sparse Token Loop.

Appendix

Table 15: Released code-facing core hyperparameter interface. The JSON files in configs/ are the exact paper presets. A CLI override is applied to a copy and sets paper_preset=false in the resolved record, preventing a custom run from being mistaken for a reported configuration.

Sparse control	JSON field	Shared CLI override	Semantics and accepted values
Selector	selector	--selector NAME	all , random , condition_aware , attention_aware , or image_condition_split .
Selected-token fraction $\rho_{\mathrm{sel}}$	selection_ratio	--selection-ratio RATIO	$0<\rho_{\mathrm{sel}}\leq 1$ when required by the sparse selector.
Sparse-routing seed	selector_seed	--selector-seed SEED	Integer seed for reproducible token routing.
Eligible token domain	token_domain	--token-domain DOMAIN	backend_default , all_tokens , image_prefix_only , or image_vs_condition .

Appendix

Table 16: Sparse Token Loop routing controls. Unsupported backend/selector combinations fail validation before model loading.

Flow-Based Generative Modeling

Flow-based generative modeling learns probability transports from simple noise to data for fast sampling, density estimation, and controllable generation.

See all 538 papers

Oct 8, 2026

Dino Forcing Flow Models: Do not denoise what you can predict

Co-denoising pretrained representations such as DINO can substantially improve the training speed and quality of flow matching models, but it introduces a second denoising trajectory and requires carefully designed schedules. We propose a simpler alternative: predict the pretrained representation directly, then condition the model on its own prediction. This removes the need for a second ODE and any representation-specific denoising schedules, while retaining the benefits of representation guidance. Our approach converges substantially faster and achieves better generation quality as measured by FID score. On ImageNet, it outperforms the state of the art in latent space at 2x fewer epochs than prior methods; in pixel space, it improves FID over comparable prior methods by more than 20%. These results support a simple principle: do not denoise what you can predict. Our code is openly available at https://github.com/arijit-hub/dino_forcing.

Arijit Ghosh, Lucas Degeorge, Paul Couairon +3

ENPC, IP Paris · Ecole Polytechnique, IP Paris · AMIAD +1

Self-Conditioning Flow-Based Generative Modeling Flow Matching Representation Steering

Ask PDF

Oct 4, 2026

LiFT: Loop Flow Transformers

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.

Mohammad Mahdi Derakhshani, Pedro M. P. Curvo, Gertjan J. Burghouts +2

VISLab · University of Amsterdam · TNO, Intelligent Imaging +1

Diffusion Transformer Recurrent Transformers Diffusion-Based Image Generation

Ask PDF

Oct 17, 2026

CorrFlow: Self-Correcting Flow Matching

Flow-matching models generate samples through iterative velocity predictions, but typically do not revisit these predictions using the intermediate representations that produced them. We introduce CorrFlow, a self-correcting flow-matching framework that reuses this information through a lightweight refiner. Under rectified-flow interpolation, the error of the one-step clean prediction equals the realized velocity residual scaled by the noise level. Motivated by this identity, the refiner learns to predict residual corrections from the base model's clean prediction, conditioning, noise level and hidden states from the same forward pass. The correction is added to the base velocity before each sampling update, so subsequent predictions are made from the corrected state. This requires no additional base-model evaluations and only one lightweight refiner evaluation per step. Our theoretical analyses establish conditions under which feature reuse outperforms raw-input affine predictors and refiner training achieves greater residual-risk reduction than continued base-model training under a matched training compute budget. Experiments on class-conditional image generation, text-to-video generation and image-to-video generation show consistent improvements over the corresponding base models, including gains in visual quality and physical consistency. Project Page: https://fogradio.github.io/corrflow-project

Hangwei Zhang, Nathan Rouillé, Chuanxia Zheng

PVG, Nanyang Technological University · Télécom Paris, Institut Polytechnique de Paris

Flow Matching Flow-Based Generative Modeling Conditional Image Generation

Ask PDF

Abstract

Flow-Based Generative Modeling

Explore similar work

Dino Forcing Flow Models: Do not denoise what you can predict

LiFT: Loop Flow Transformers

CorrFlow: Self-Correcting Flow Matching