Denoising, Fast and Slow: Difficulty-Aware Adaptive Sampling for Image Generation

1. Problem Statement

Modern diffusion and flow-based image generators (e.g., Diffusion Transformers, SiT) typically allocate computational resources uniformly across the image. At every denoising step, all spatial patches are updated with the same global timestep and the same number of function evaluations (NFE).

This uniform approach ignores the inherent heterogeneity of natural images:

• Easy regions: Large, low-frequency backgrounds or saturated areas are easy to denoise and require fewer steps.
• Hard regions: Thin structures, object boundaries, small text, and occlusion boundaries are ambiguous and require more refinement and context.

Treating all regions identically forces the model to waste compute on easy areas while under-allocating resources to difficult ones. Furthermore, existing attempts to introduce heterogeneous timesteps (e.g., Diffusion Forcing for video, Spatial Reasoning Models for images) suffer from a train-test gap. Naive patch-level sampling often results in training samples containing "clean" (low-noise) patches that do not exist at inference time (which starts from pure noise), leading to poor generalization.

2. Methodology: Patch Forcing (PF)

The authors propose Patch Forcing (PF), a framework that enables spatially adaptive image generation by assigning distinct noise trajectories to different image patches. The method consists of three core components:

A. Difficulty-Aware Training (The LTG Sampler)

To bridge the train-test gap, the authors introduce a specific timestep sampling strategy called Logit-Normal Truncated Gaussian (LTG).

• The Issue: Previous methods (like SRM) sampled the mean timestep uniformly, which inadvertently allowed some patches to be nearly fully denoised ($t \approx 1$) during training. This created "context leakage" where the model learned from information unavailable at inference.
• The Solution: The LTG sampler explicitly controls the maximum information available per training sample.
  1. Sample a maximum timestep $t_{max}$ from a Logit-Normal distribution.
  2. Sample individual patch timesteps $t_i$ from a Gaussian distribution centered at $t_{max}$ but truncated to the lower half ($t_i \le t_{max}$).
• Result: This ensures that no patch is ever fully denoised during training, forcing the model to learn to generate content even when context is limited, perfectly matching the inference distribution.

B. Learning Patch Difficulty

The model is augmented with a lightweight uncertainty head (predicting the standard deviation of the velocity field).

• Mechanism: The model predicts a per-patch uncertainty map alongside the denoising velocity.
• Interpretation: High uncertainty is interpreted as high difficulty. The authors validate that predicted uncertainty correlates strongly with validation loss (error) and that providing "future" context (from easier regions) reduces uncertainty in harder regions.

C. Adaptive Inference Samplers

Using the predicted difficulty map, PF introduces two adaptive sampling strategies that dynamically allocate compute:

1. Dual-Loop Sampler:
   • Outer Loop: Takes large timestep steps on low-uncertainty (easy) patches to advance them quickly.
   • Inner Loop: Uses the advanced states of easy patches as context to denoise high-uncertainty (hard) patches with smaller, more refined steps.
   • This creates a feedback loop where easy regions provide context for hard regions.
2. Look-Ahead Sampler:
   • Propagates confident (easy) patches forward to a future timestep $t_{ctx} = \min(\alpha \cdot t, 1)$.
   • These "future" states serve as high-fidelity context to guide the denoising of uncertain regions at the current timestep.

3. Key Contributions

1. Patch-Level Timestep Sampling: A novel training strategy (LTG) that solves the train-test mismatch in heterogeneous denoising by controlling the maximum information per sample rather than the average.
2. Difficulty-Aware Framework: The integration of a learned uncertainty head to identify hard vs. easy regions, enabling the model to self-generate contextual guidance.
3. Adaptive Samplers: The introduction of Dual-Loop and Look-Ahead samplers that dynamically allocate compute resources, advancing easy regions to help resolve difficult ones.
4. Orthogonality: The framework is shown to be orthogonal to existing techniques like Classifier-Free Guidance (CFG) and Representation Alignment (REPA), allowing for combined performance gains.

4. Results

The authors evaluated Patch Forcing on ImageNet 256x256 (class-conditional) and Text-to-Image synthesis.

• ImageNet Performance:

  • Baseline: Standard SiT (SiT-B/2) achieves an FID of 33.0.
  • Patch Forcing (Parallel): Simply using patch-level timesteps with parallel sampling improves FID to 27.9.
  • Adaptive Sampling: Adding the Look-Ahead sampler further reduces FID to 24.2 (SiT-B/2) and 9.8 (SiT-XL/2), significantly outperforming the baseline.
  • State-of-the-Art: The PFT-XL/2 + Look-Ahead + REPA combination achieves an FID of 6.7, surpassing previous SOTA methods.

• Text-to-Image (T2I):

  • A 1.2B parameter PFT trained on COYO data achieves competitive results on T2I-CompBench++ and GenEval.
  • Text Rendering: A notable qualitative and quantitative improvement is observed in text rendering. PFT models produce clearer, more accurate text compared to standard Flow Matching baselines, likely due to the ability to refine ambiguous text regions using context from clearer regions.

• Efficiency: The adaptive samplers achieve better quality at the same NFE budget compared to uniform Euler sampling.

5. Significance

This work fundamentally shifts the paradigm of diffusion generation from uniform processing to adaptive, difficulty-aware processing.

• Efficiency: It demonstrates that not all parts of an image require the same computational effort, offering a path to more efficient generation without sacrificing quality.
• Context Generation: It introduces a mechanism where the model generates its own internal context (by advancing easy regions) to solve hard regions, mimicking a "solve easy parts first, then use them to solve hard parts" reasoning process.
• Generalizability: The approach works across different model scales (B, L, XL) and modalities (class-conditional and text-conditional), suggesting that patch-level denoising schedules are a robust foundation for future adaptive image generation systems.

In summary, Patch Forcing proves that by acknowledging the varying difficulty of image regions and adapting the denoising schedule accordingly, one can significantly improve generation quality, text rendering, and computational efficiency.