WiT: Waypoint Diffusion Transformers via Trajectory Conflict Navigation

1. Problem Statement

The paper addresses a fundamental limitation in pure pixel-space Flow Matching (FM) and Diffusion models. While recent architectures (like JiT) operate directly in pixel space to avoid the information loss and artifacts inherent in latent-space VAEs, they face a severe optimization challenge known as Trajectory Conflict.

• The Core Issue: The raw pixel manifold is "entangled" and non-discriminative. Visually similar but semantically distinct images (e.g., a cat and a dog) often map to overlapping regions in the high-dimensional noise space.
• The Consequence: During the denoising process, optimal transport paths for different classes converge or cross in dense local regions. To minimize regression loss (MSE), the neural network is forced to predict an averaged velocity field, resulting in:
  • Semantic Bleeding: Blurring of class-specific features.
  • Slow Convergence: The model struggles to learn the complex, unconstrained high-dimensional vector field from scratch.
  • Sub-optimal Solutions: Lower generation quality compared to latent-space models.

Existing solutions like Classifier-Free Guidance (CFG) amplify signals post-hoc but do not resolve the underlying geometric overlap of training trajectories.

2. Methodology: Waypoint Diffusion Transformers (WiT)

The authors propose WiT, a framework that decouples semantic navigation from pixel-level texture generation by introducing intermediate semantic waypoints. The core idea is to factorize the difficult "noise-to-pixel" mapping into two stable, decoupled stages: "noise-to-waypoint" and "waypoint-to-pixel."

Key Components:

1. Semantic Waypoints ($s_0$):

   • Instead of operating in the raw pixel space, WiT projects the target image into a low-dimensional, discriminative semantic manifold derived from pre-trained self-supervised vision models (specifically DINOv3).
   • PCA Projection: To reduce computational burden, features are compressed via Principal Component Analysis (PCA) to a compact dimension ($d=64$). This creates a "ground truth" semantic anchor that is highly separable by class.

2. Lightweight Waypoints Generator ($W_\psi$):

   • A small transformer (21M parameters) that takes the current noisy pixel state ($z_t$) and predicts the clean semantic waypoint ($\hat{s}_0$).
   • It is trained to minimize the velocity matching loss in the semantic space, effectively learning to navigate the noisy distribution toward the correct semantic cluster.

3. Pixel Space Generator ($G_\theta$) & Just-Pixel AdaLN:

   • The main generator is a large Diffusion Transformer (DiT) responsible for synthesizing high-fidelity pixels.
   • Just-Pixel AdaLN: A novel conditioning mechanism. Unlike standard AdaLN which uses global embeddings, this mechanism injects the predicted semantic waypoint $\hat{s}_0$ as spatially-varying modulation parameters (scale, shift, alpha) into the transformer blocks.
   • This allows the generator to be dynamically steered by the semantic blueprint at every spatial location, ensuring the pixel generation stays within the correct class manifold without "polluting" the token sequence.

Training Paradigm (Two-Stage):

1. Stage 1: Train the lightweight Waypoints Generator to predict semantic waypoints from noisy pixels.
2. Stage 2: Freeze $W_\psi$ and train the main Pixel Space Generator $G_\theta$. During inference, $W_\psi$ dynamically infers the waypoint at every ODE step, providing continuous spatial guidance to $G_\theta$.

3. Theoretical Motivation

The paper provides a theoretical justification based on Bayes-risk and variance decomposition:

• Total Variance Decomposition: The uncertainty in predicting the clean image $x$ from noise $z_t$ is decomposed as:
  $$\text{Var}(x | z_t) = \mathbb{E}_{s_0}[\text{Var}(x | z_t, s_0)] + \text{Var}_{s_0}(\mathbb{E}[x | z_t, s_0])$$
• Search Space Contraction: By explicitly predicting the semantic waypoint $s_0$, the model reduces the first term (residual variance). The generator no longer needs to resolve the massive ambiguity of the entire pixel manifold; it only needs to resolve the smaller variance within a specific semantic subspace.
• Result: This mathematically guarantees a reduction in the "irreducible error" and prevents the averaging of divergent trajectories.

4. Key Contributions

1. WiT Architecture: A novel generative paradigm that mitigates trajectory conflict in pixel-space Flow Matching by anchoring flows to low-dimensional semantic manifolds.
2. Just-Pixel AdaLN: A mechanism that enables spatially-varying semantic conditioning, decoupling semantic navigation from texture synthesis without requiring latent tokenizers.
3. Theoretical & Empirical Validation: Proves that semantic grounding shrinks the generative search space and demonstrates that this approach untangles overlapping transport paths.

5. Experimental Results

Evaluated on ImageNet 256×256, WiT achieves state-of-the-art (SOTA) performance among purely pixel-space models, rivaling and often surpassing heavy latent-space models.

• Performance Metrics:
  • WiT-L/16 (600 epochs): Achieves FID 2.22 and IS 303.3.
  • Comparison: Outperforms the baseline JiT-L/16 (FID 2.36) and surpasses the latent-space DiT-XL/2 (FID 2.27).
  • WiT-XL/16: Breaks the 2.0 FID barrier with FID 1.89.
• Training Efficiency:
  • WiT accelerates training convergence by 2.2× compared to JiT.
  • WiT-L/16 trained for only 265 epochs matches the performance of JiT-L/16 trained for 600 epochs.
• Qualitative Improvements:
  • Structural Coherence: Generates images with correct proportions and perspectives (e.g., animals, complex scenes) that often fail in unanchored pixel models.
  • Texture Fidelity: Preserves high-frequency micro-textures (feathers, wings) that are typically lost in VAE-based latent compression.

6. Significance

WiT represents a significant step forward in high-fidelity image generation. It demonstrates that:

1. Pixel-space generation is viable at scale without relying on lossy VAEs, provided the optimization landscape is regularized.
2. Semantic guidance is crucial: By introducing a "waypoint" mechanism, the model can navigate the complex pixel space efficiently, resolving the "trajectory conflict" that has plagued direct pixel-space diffusion.
3. Efficiency: It achieves superior results with significantly less training time and computational overhead compared to scaling up standard pixel-space baselines, offering a new blueprint for future generative models.