PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.29160v1cs.CVcs.AI

TL;DR

Frozen flow-matching generators raise the question of whether extra computation can be added inside their denoisers without changing the surrounding pipeline. This paper introduces training-free hidden-state looping, which improves quality across tested RAE-based generators and offers competitive efficiency, while remaining model-dependent.

  • Problem

    It remains unclear how to use extra computation inside frozen flow-matching denoisers when denoising time, layer depth, token routing, and guidance interact.

  • Method

    The method repeatedly applies selected transformer layers inside frozen denoisers, controlling token scope, sampling interval, layer range, update strength, and prediction combination.

  • Results

    Across three tested RAE-based generators, looping improves quality metrics, and at comparable generation time outperforms allocating computation only to additional outer sampling steps.

  • Takeaways & Limitations

    Internal layer looping provides a useful inference-time refinement axis, with Sparse Token Loop offering a lower-latency alternative to dense repetition.

  • Takeaways & Limitations

    The method is model-dependent: tested VAE-based generators show no consistent gain, and looping does not consistently repair counting or spatial composition.

Abstract

from arXiv · show

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.

Introduction

The paper introduces training-free internal transformer-layer looping for frozen flow-matching image generators, preserving model weights, conditioning, latent autoencoder, and outer sampler. It structures looping across token scope, sampling interval, layer range, and prediction-space control, with reported gains in quality and efficiency.

  • Motivation: Flow matching supports modern image generation, whose progress spans schedules, transformer backbones, guidance, inference-time scaling, and latent representations.
  • Method: The method repeatedly applies selected transformer layers inside a frozen denoiser without changing weights, text conditioning, the latent autoencoder, or outer sampler.
  • Method: Dense and Sparse Token Loop vary token scope, Sampling-Progress Gating controls the active interval, and the loop layer range selects repeated layers.
  • Results: Internal layer looping improves subject-level structure and reaches higher GenEval and DPG-Bench scores than increasing only outer sampling steps.
  • Method: Loop Guidance combines ordinary and looped vector-field predictions to control how the loop-induced correction is applied.
  • Evaluation: The study evaluates Scale-RAE and RAEv2 through ablations of loop count, layer range, sampling interval, sparsifier design, and guidance.

Related Work

Prior work combines flow matching with diffusion-transformer backbones and improves frozen generators through inference-time guidance or prediction modification. Looped-transformer research motivates repeated shared-layer computation, while this method applies training-free looping inside frozen flow-matching image denoisers and preserves the outer sampler.

  • Flow matching and diffusion-transformer design: Flow matching learns vector fields between base and data distributions and has been combined with diffusion-transformer backbones for image generation.The cited design directions include flow matching and diffusion-transformer architectures.
  • Inference-time control and guidance: Inference-time methods improve diffusion and flow models through schedule design, guidance, candidate selection, or prediction-space modification without changing generator weights.Examples include classifier guidance, classifier-free guidance, self-guidance, and autoguidance.
  • Looped transformers: Universal and recent looped transformers repeatedly apply shared layers to support iterative computation, algorithm learning, or latent reasoning.Training-free looped transformers also indicate that repeated layer evaluation can be useful without retraining.
  • Looped transformers: This method differs by performing training-free looping inside frozen flow-matching image denoisers while separating token, sampling-progress, and prediction-space controls.Its controls include Dense Token Loop, Sampling-Progress Gating, Sparse Token Loop with cached complement reuse, and Loop Guidance.

Method

The method adds training-free inference computation inside a frozen flow-matching denoiser by repeatedly applying selected transformer layers during chosen denoising calls. It controls repetition through token scope, sampling-progress gating, residual strength, and prediction-space Loop Guidance.

  • Framework: Loops repeatedly apply a selected transformer-layer range inside chosen denoising calls while leaving the frozen model and outer sampling procedure unchanged.The framework adds computation within one denoiser path rather than increasing the number of outer sampling steps.
  • Dense Token Loop: Dense Token Loop normalizes repeated residual updates with per-round strength η = λloop/K, where K is the loop count and λloop = 1 in experiments.With K = 1, the formulation recovers the ordinary layer output.
  • Sampling-Progress Gating: Sampling-Progress Gating activates looping from the initial noise state through a chosen endpoint send, after which calls follow the ordinary denoiser path.The active interval is Sloop = [0, send] over normalized sampling progress si = i/(N −1).
  • Sparse Token Loop: Sparse Token Loop repeatedly evaluates only a selected token subset while caching the complement residual once and updating both token groups each round.The selected fraction is ρsel, and synchronized loop depth preserves cross-group attention states.
  • Loop Guidance: Loop Guidance combines ordinary and looped vector-field predictions as vlg = vbase + glg(vloop −vbase), with glg controlling the loop-induced correction.Both predictions are required, so Loop Guidance costs more than a single looped path; glg > 1 extrapolates beyond the looped prediction.

Experiments

Experiments show that internal looping improves quality across frozen RAE-based generators, with Loop Guidance delivering consistent primary- and auxiliary-metric gains. Ablations identify moderate guidance and loop counts, middle-to-late layers, and early sampling stages as effective settings, while transfer to VAE-based generators remains inconsistent.

  • Experimental setup: Experiments use frozen Qwen1.5B/DiT2.4B, Qwen7B/DiT9.8B, and RAE-SigLIP2-B generators, transferring a shared Scale-RAE configuration to the latter two models.Evaluation uses GenEval, PickScore, and ImageReward.
  • Unguided loop variants: Dense and Sparse Token Loop improve GenEval and DPG-Bench at both Scale-RAE scales, while Sparse Token Loop is generally more efficient.On RAEv2, both improve DPG-Bench and auxiliary metrics but slightly reduce GenEval.
  • Loop Guidance: Loop Guidance improves both primary metrics and auxiliary metrics across all three evaluated generators, including clear primary-metric gains on RAEv2.The guidance scale is not exhaustively tuned; each model uses one setting that performed reasonably well in trials.
  • Quality–efficiency trade-offs: 0.5422 GenEval and 0.8007 DPG-Bench are reached by 25-step Dense Token Loop in 2.1812 s, surpassing the 50-step no-loop baseline values of 0.4742 and 0.7738.Sparse Token Loop reaches 0.5082 GenEval and 0.7893 DPG-Bench in 1.6883 s versus 0.4770 and 0.7673 for a 35-step no-loop baseline.
  • Ablations: Moderate guidance scales and loop counts capture most gains, while middle-to-late layer ranges and early, higher-noise sampling stages outperform broader or later alternatives.The default settings are K = 4, layer range [12, 27], and loop-active interval si ∈[0, 0.5].
  • Transfer and limitations: Targeted condition-aware and attention-aware sparsifiers do not clearly improve upon random routing on Scale-RAE, and Loop Guidance lacks consistent gains across both primary metrics on PixArt-alpha and FLUX.2.The VAE-based transfer results suggest that benefits observed for RAE-based generators may not transfer directly.

Conclusion

The method adds inference computation inside frozen flow-matching denoisers through repeated selected-layer application without changing the generation pipeline. Looping improves quality across Scale-RAE scales, but its benefits are model-dependent and uneven across tasks and generator architectures.

  • Method: The framework repeatedly applies selected transformer layers inside each denoising call while preserving weights, conditioning, latent representations, and the outer sampler.Its controls organize looping by token scope, sampling progress, layer range, and prediction-space control.
  • Method: Dense Token Loop, Sparse Token Loop, Sampling-Progress Gating, and Loop Guidance instantiate the framework’s token, timing, layer, and prediction-space controls.Sparse Token Loop offers a lower-latency alternative to dense repetition.
  • Results: Across both Scale-RAE scales, looping improves primary and auxiliary quality metrics.The conclusion reports these gains across the evaluated Scale-RAE model scales.
  • Limitations: Looping most reliably corrects subject formation and attribute binding, but does not consistently repair counting or spatial composition on GenEval2.Its effect depends on layer location and sampling progress, and tested VAE-based generators show no consistent gain.
  • Limitations: Internal layer looping is a useful but model-dependent inference-time refinement axis.Future work should develop task-aware sparsifiers and adaptive loop control.

Method Pseudocode

Algorithms 1–4 define four inference-time operations for refining hidden states and vector-field predictions while keeping all model parameters frozen. Dense and sparse token loops repeatedly update selected layers, gating controls when looping occurs, and Loop Guidance blends ordinary and looped predictions.

  • Overview: All four inference-time operations preserve frozen model parameters.The algorithms summarize the core operations without changing model weights.
  • Dense Token Loop: Dense Token Loop repeatedly applies each frozen layer in a selected range for K iterations with per-iteration strength η = λloop/K.The refined state is propagated layer by layer after K repeated updates.
  • Sparse Token Loop: Sparse Token Loop samples m = ⌈ρsel|U|⌉ tokens uniformly from the valid set U and loops only on the selected subset.Non-selected tokens are initialized once while the selected tokens receive repeated updates.
  • Sampling-Progress Gating: Sampling-Progress Gating activates Dense Token Loop only when the normalized call index si = i/(N −1) lies in Sloop; otherwise it uses the ordinary path.The gate receives the layer range, loop count K, and total loop strength λloop as configuration.
  • Loop Guidance: Loop Guidance computes vlg = vbase + glg(vloop −vbase) by combining ordinary and gated-plus-looped vector-field predictions.The looped prediction uses Algorithms 1–3 within the same denoising call.

Experimental Details

Experiments use fixed prompts, seeds, checkpoints, and resolutions across baseline–loop comparisons, evaluating Scale-RAE configurations and transfer models with pinned metrics and standardized generation settings. Ablations examine sampling progress, loop count, layer range, sparsification, and Loop Guidance under controlled DiT2.4B protocols.

  • Evaluation protocol: The standard protocol evaluates all 553 GenEval, 1,065 DPG-Bench, and 800 GenEval2 prompts, generating one image per prompt with seed 42.GenEval and DPG-Bench support primary comparisons, while GenEval2 supports additional qualitative analysis; compared settings share prompt, noise seed, checkpoint, and resolution.
  • Models and checkpoints: Two official Scale-RAE configurations are evaluated: Qwen1.5B/DiT2.4B and Qwen7B/DiT9.8B, both retaining FP32 parameters and model guidance 1.0.Architecture transfer additionally uses a 615M-parameter RAE-SigLIP2-B/DDT proxy checkpoint rather than the unreleased official RAEv2 text-to-image checkpoint.
  • Compute and implementation: All generation runs use NVIDIA H100 GPUs, with two GPUs for resident Scale-RAE Qwen7B/DiT9.8B workers and one GPU for other generators.The Qwen7B/DiT9.8B checkpoint uses device_map=auto for layer placement rather than tensor parallelism; FLUX.2 additionally uses model CPU offload.
  • Ablation design: Controlled DiT2.4B ablations vary active progress intervals, loop counts, layer ranges, sparsifier routing, and Loop Guidance while holding other settings fixed.The main ablation protocol uses 25 outer steps, K = 4 where applicable, layers [12, 27] or [12, 19], λloop = 1, progress [0, 0.5], and glg = 1; moderate loop counts and guidance provide the strongest quality–time balance or gains.
  • Sampling settings: Twenty-five outer sampling steps are the default efficiency setting, while unlooped Scale-RAE timing comparisons also test 15, 35, and 50 steps.The timing averages 1,618 generation records per setting and uses the same pinned evaluators.

Other Experimental Results

The experiments show that carefully chosen middle-layer looping and Loop Guidance improve both primary benchmarks, while qualitative corrections are strongest for subject-level structure. GenEval2 is a limitation: neither looping variant improves its aggregate score.

  • Layer-Range Sweep: A middle RAE layer range improves both primary benchmarks, whereas extreme early and late ranges collapse.Several other middle ranges raise DPG-Bench while slightly lowering GenEval.
  • Loop Guidance: Loop Guidance substantially improves both DiffusionBench RAE-SigLIP2-B scores, with glg = 9 maximizing GenEval and glg = 7.5 maximizing DPG-Bench.The sweep fixes layer range [15, 22] and loop count K = 4; increasing glg generally strengthens both metrics.
  • Subject-Level Corrections: Looping most consistently corrects subject-level appearance by forming incomplete shapes, removing spurious parts, and binding attributes more cleanly to intended entities.The review covers prompt-matched pairs on GenEval and DPG-Bench, whose prompts are dominated by one or two subjects.
  • GenEval2 Findings: 0.1801 to 0.1736: Dense Token Loop lowers GenEval2, while Loop Guidance reaches 0.1441; neither variant improves the aggregate score.The results use official soft-TIFA-style geometric-mean scores over 800 prompts.

Reproducibility

The release provides a backend-neutral runtime, model adapters, exact configurations, benchmark utilities, and a composable interface for token loops and Loop Guidance. Reported ablations are rescored with pinned evaluators and exclude incompatible artifacts to preserve evaluation integrity.

  • Code and Data Supplement: The Code and Data Supplement includes the backend-neutral loop runtime, model adapters, exact JSON configurations, benchmark snapshots, utilities, and CPU-only tests.It also provides aggregate values used by the paper tables; pretrained weights and several generated or evaluator artifacts are excluded.
  • Code and Data Supplement: Public benchmark inputs and checkpoints are obtained separately through public identifiers in configs/.
  • Composable method and hyperparameter interface: The released resolver treats the token-loop base and Loop Guidance as orthogonal, allowing guidance to compose with either Dense Token Loop or Sparse Token Loop.Scale-RAE and RAEv2 expose both token-loop bases with or without Loop Guidance.
  • Composable method and hyperparameter interface: Exact paper presets are stored in configs/, while CLI overrides create non-paper-preset records to prevent custom runs from being mistaken for reported configurations.
  • Evaluator and artifact integrity: Every reported Scale-RAE ablation row is rescored from its bound image grid with pinned evaluators, while incomplete or incompatible artifacts are excluded.GenEval and DPG-Bench use fixed evaluator procedures, including native resolution and a local mPLUG checkpoint for DPG-Bench.
Loading 2608.29160v1…