Pyramidal Patchification Flow for Visual Generation

1. Problem Statement

Diffusion Transformers (DiTs) and flow-based models have achieved state-of-the-art results in visual generation but suffer from high computational costs during the denoising process.

• Constant Token Budget: Standard DiTs use a fixed patch size (e.g., $2 \times 2$) across all timesteps. This results in a constant number of tokens ($L$) throughout the generation trajectory, leading to high computational complexity ($O(L^2)$ for attention) even at early, high-noise timesteps where fine-grained details are not yet necessary.
• Limitations of Existing Acceleration Methods:
  • Function Reduction: Methods like distillation or consistency models reduce the number of steps but often require complex training or compromise quality.
  • Pyramidal/Cascaded Approaches: Methods like Pyramidal Flow reduce token counts by operating on lower-resolution pyramid representations at high-noise stages. However, this introduces "resolution jumps" between stages, breaking trajectory continuity and requiring complex "re-noising" tricks to handle the resolution changes during inference.

2. Methodology: Pyramidal Patchification Flow (PPFlow)

PPFlow proposes a novel scheme to reduce the number of tokens at high-noise timesteps without changing the spatial resolution of the latent representation.

Core Concept

Instead of downsampling the image (changing resolution), PPFlow changes the patch size used for tokenization (Patchify) based on the noise level:

• High-noise timesteps (early stages): Use larger patches (e.g., $4 \times 4$). This reduces the number of tokens significantly while keeping the latent representation at full resolution.
• Low-noise timesteps (late stages): Use smaller patches (e.g., $2 \times 2$, standard DiT size) to capture fine details.

Architecture & Implementation

1. Pyramidal Patchification: The denoising trajectory is divided into $K$ stages (e.g., 2 or 3). Each stage $s$ has a specific patch size $p_s \times p_s$.
2. Shared DiT Blocks: The core Transformer blocks (Self-Attention, MLP) remain unchanged and share weights across all stages. This ensures the model learns a consistent velocity field.
3. Learned Projections:
   • Patchify: Separate linear projection matrices ($W_s$) are learned for each patch size to map the $p_s \times p_s \times C$ patch to a fixed token dimension $d$.
   • Unpatchify: Separate linear projection matrices are used to map tokens back to the latent space predictions.
4. Continuity: Because the latent representation resolution ($H \times W$) remains constant throughout the process, the approach satisfies the Continuity Equation. Unlike Pyramidal Flow, it does not require re-noising tricks or handling resolution jumps.

Training Strategies

• From Scratch: Trains the shared DiT blocks and stage-specific projections simultaneously.
• From Pretrained DiT: Copies weights from a standard pretrained DiT (e.g., SiT-XL/2).
  • Initialization: Projections for larger patches are initialized by averaging smaller patch weights; projections for Unpatchify are initialized by duplicating weights.
  • Efficiency: Requires only ~8.9% additional training FLOPs to fine-tune, making it highly efficient for adaptation.

3. Key Contributions

1. Novel Token Reduction Scheme: Introduces a method to dynamically reduce token counts at high-noise timesteps by varying patch sizes while maintaining full-resolution latent representations.
2. Elimination of Resolution Jumps: Unlike previous pyramidal approaches, PPFlow operates on full-resolution latents, ensuring trajectory continuity and removing the need for complex re-noising mechanisms during inference.
3. Parameter Efficiency: The core DiT architecture is shared; only lightweight linear projections are added for different patch sizes.
4. Flexible Training: Demonstrates effectiveness both when training from scratch and when fine-tuning large pretrained models (SiT, DiT, FLUX.1).

4. Experimental Results

The paper evaluates PPFlow on ImageNet (class-conditional) and text-to-image tasks (FLUX.1).

Class-Conditional Generation (ImageNet)

• Speedup: Achieves 1.60× to 2.02× inference speedup (denoising acceleration) compared to standard SiT/DiT baselines.
• Quality: Maintains comparable or better generation quality.
  • From Scratch (SiT-B/2): PPF-B-3 achieves an FID of 4.43 (vs. 4.46 baseline) with 2.04× speedup.
  • From Pretrained (SiT-XL/2): PPF-XL-3 achieves an FID of 2.23 (vs. 2.15 baseline) with 2.02× speedup, requiring only 7.1% extra training FLOPs.
• FLOPs Reduction: Reduces testing FLOPs by 37.8% (2-stage) and 50.6% (3-stage) for 256×256 images.

Text-to-Image Generation (FLUX.1)

• Applied to FLUX.1-dev across resolutions 512, 1024, and 2048.
• Achieves 1.61× to 1.86× speedup with comparable quality on benchmarks like GenEval, DPG-Bench, and T2I-CompBench.

Ablation Studies

• Comparison with Pyramidal Flow: PPFlow significantly outperforms standard Pyramidal Flow (which suffers from blurring and artifacts due to resolution jumps) and Lumina-Video methods.
• Training Consistency: PPFlow trains each patch-size model only on its corresponding timestep range, ensuring consistency between training and inference, unlike methods that train all patch sizes on all timesteps.

5. Significance

• Efficiency without Compromise: PPFlow offers a "free lunch" in terms of efficiency, accelerating the most expensive part of the generation process (early high-noise steps) without sacrificing the high-fidelity details captured in later steps.
• Simplicity: The method is easy to implement, requiring only modifications to the Patchify/Unpatchify layers while keeping the heavy DiT backbone identical.
• Scalability: The approach is proven to scale from small models (SiT-B) to massive text-to-image models (FLUX.1) and high resolutions (up to 2048×2048).
• Theoretical Soundness: By avoiding resolution jumps, it adheres to the continuity equation of flow matching, providing a theoretically cleaner alternative to cascaded or pyramidal generation methods.

In summary, PPFlow rethinks the tokenization strategy in Diffusion Transformers, proving that adaptive patch sizing is a highly effective lever for accelerating visual generation while preserving state-of-the-art image quality.