WiT: Waypoint Diffusion Transformers via Trajectory Conflict Navigation

Hainuo Wang; Mingjia Li; Xiaojie Guo

Source-linked AI summary

WiT: Waypoint Diffusion Transformers via Trajectory Conflict Navigation

Hainuo Wang, Mingjia Li, Xiaojie Guo

arXiv:2603.15132v2Mar 16, 2026cs.CV

TL;DR

WiT addresses trajectory conflicts caused by the semantically entangled, non-discriminative pixel manifold in pixel-space Flow Matching. It routes generation through semantic waypoints from pre-trained vision models and uses spatially varying conditioning, achieving state-of-the-art performance among purely pixel-space models and a 2.2× training speedup over JiT-L/16 on ImageNet 256×256.

Problem
Pixel-space Flow Matching must learn mappings through an inherently entangled, non-discriminative pixel manifold, where semantically different targets can share dense neighborhoods and create trajectory ambiguity.
Method
WiT decouples semantic navigation from pixel-level generation by routing transport through semantic waypoints and conditioning the pixel generator with Just-Pixel AdaLN.
Results
WiT achieves state-of-the-art performance among purely pixel-space models on ImageNet 256×256 and provides a 2.2× training speedup compared with JiT-L/16.
Takeaways & Limitations
Explicit semantic grounding factors the noise-to-pixel transport path into semantic navigation and pixel-level texture synthesis.
Takeaways & Limitations
Directly estimating conditional variance in high-dimensional image space is empirically impractical, so trajectory conflict is assessed using inference-time proxies.

Abstract

from arXiv · show

While recent Flow Matching models avoid the reconstruction bottlenecks of latent autoencoders by operating directly in pixel space, the lack of semantic continuity in the pixel manifold severely intertwines optimal transport paths. This induces severe trajectory conflicts near intersections, yielding sub-optimal solutions. Rather than bypassing this issue via information-lossy latent representations, we directly untangle the pixel-space trajectories by proposing Waypoint Diffusion Transformers (WiT). WiT factorizes the continuous vector field via intermediate semantic waypoints projected from pre-trained vision models. It effectively disentangles the generation trajectories by breaking the optimal transport into prior-to-waypoint and waypoint-to-pixel segments. Specifically, during the iterative denoising process, a lightweight generator dynamically infers these intermediate waypoints from the current noisy state. They then continuously condition the primary diffusion transformer via the Just-Pixel AdaLN mechanism, steering the evolution towards the next state, ultimately yielding the final RGB pixels. Evaluated on ImageNet 256x256, WiT beats strong pixel-space baselines, accelerating JiT training convergence by 2.2x. Code will be publicly released at https://github.com/hainuo-wang/WiT.git.

1 Introduction

Pixel-space Flow Matching preserves fine-grained details but suffers trajectory conflict because raw pixels are semantically entangled. WiT introduces semantic waypoints to decouple semantic navigation from pixel generation and reports faster convergence with state-of-the-art pixel-space performance.

Motivation: Pixel-space Flow Matching avoids tokenizer compression artifacts but directly maps noise to a complex, entangled pixel distribution.This preserves fine-grained visual details while creating a difficult optimization problem.
Motivation: Trajectory conflict arises when paths to visually similar but semantically distinct endpoints converge in dense pixel-space neighborhoods.The shared noise-to-pixel mapping can therefore produce averaged, conflicting velocity fields.
WiT: WiT routes generation through discriminative semantic waypoints, forming separate noise-to-waypoint and waypoint-to-pixel mappings.This bipartite structure is intended to prevent conflicting paths from collapsing into averaged trajectories.
WiT: WiT projects pre-trained vision-model features onto a compact semantic manifold and uses these anchors to isolate semantic navigation from pixel-level generation.PCA reduces spatial redundancy and the associated regression burden.
WiT: WiT introduces Just-Pixel AdaLN, which uses dynamically predicted semantic waypoints for spatially varying guidance rather than uniform global conditioning.The mechanism continuously conditions the pixel-space transformer during generation.

2 Related Work

Related work spans diffusion and Flow Matching formulations, latent- and pixel-space generation, and representation alignment for reducing optimization conflict. These studies motivate addressing semantic entanglement directly in pixel-space transport.

Diffusion Models and Flow Matching: Diffusion models evolved from noise prediction toward velocity prediction and continuous-time Flow Matching formulations.Diffusion Transformers have helped scale these generative frameworks.
Generative Modeling in Pixel Space: Latent-space methods compress images into lower-dimensional manifolds because earlier pixel-space approaches were difficult to scale to high-resolution synthesis.The passage contrasts this paradigm shift with direct pixel-space generation.
Mitigating Optimization Conflict: Conditional Flow Matching must transport shared Gaussian noise toward many semantic classes, while semantically entangled pixels make paths converge or cross.The resulting regression objective can favor an averaged velocity field.

3 Methodology

WiT addresses trajectory conflict in pixel-space Flow Matching by introducing semantic waypoints that separate semantic navigation from pixel-level generation. Its lightweight waypoint predictor conditions a pixel generator with spatially varying guidance through Just-Pixel AdaLN.

Pixel-Space Flow Matching and Trajectory Conflict: Pixel-space Flow Matching can produce large conditional variance because semantically different targets share dense neighborhoods in the noise space.Under MSE, this ambiguity makes the optimal denoiser predict averaged states, causing gradient interference and limiting convergence.
Semantic Waypoints: WiT introduces a discriminative semantic waypoint to decouple semantic navigation from pixel-level texture generation.The waypoint partitions the transport path between the noise prior and raw pixels into more tractable semantic and residual pixel mappings.
Semantic Waypoints: WiT constructs semantic waypoints by projecting dense DINOv3 features onto a PCA manifold containing the top 64 principal components.The resulting low-dimensional representation is intended to provide a class-separable structural supervisory signal while reducing optimization burden.
Semantic Waypoints: A lightweight Waypoints Generator predicts the clean semantic waypoint from noisy pixel observations using a separately supervised semantic probability-flow objective.The generator is conditioned on timestep and class label, and its 21M-parameter design exploits the compressed semantic target dimension.
Semantic-Pixel Decoupled Architecture: The Pixel Space Generator receives predicted waypoints through Just-Pixel AdaLN, which combines global timestep-class conditioning with spatially varying semantic modulation.This decoupled architecture delegates semantic navigation to the waypoint generator while the primary transformer focuses on high-realistic spatial generation.

4 Experimental Validation

WiT is evaluated on class-conditional ImageNet 256 × 256 using FID and IS across model scales, training durations, and ablations. Results show faster convergence and stronger pixel-space generation, while waypoint dimensionality and injection strategy materially affect quality.

Experimental Setup: Experiments use ImageNet 256 × 256 with FID-50K and Inception Score, evaluating pixel-space models with a 50-step Heun solver.The waypoint generator uses ViT-S/16, while the primary generator matches JiT-Base and JiT-Large configurations.
Experimental Setup: WiT is configured at Base, Large, and Extra-Large capacity scales for architectural scaling comparisons.
Main Results: WiT consistently outperforms comparable pixel-space models across training stages, improving training efficiency and sample realism.
Main Results: 2.27× training speedup is achieved by WiT-L/16 at 265 epochs while matching JiT-L at 600 epochs, with 2.36 FID and 293.7 IS.
Main Results: 2.22 FID and 303.3 IS are reached by WiT-L/16 at 600 epochs, surpassing JiT-L/16 and DiT-XL/2 on FID.
Ablation Studies: d = 64 provides the best waypoint trade-off, whereas d = 128 yields 4.12 FID and d = 32 yields 5.11 FID.The intermediate dimension preserves structural variance while limiting optimization complexity and information loss.
Ablation Studies: Channel Concat performs worst among injection strategies, with 3.93 FID and 221.19 IS.

5 Quantitative Analysis of Trajectory Conflict

WiT frames trajectory conflict as conditional ambiguity in pixel-space denoising and analyzes how semantic conditioning contracts the search space. It then evaluates conflict proxies across generation trajectories against JiT.

Theoretical motivation: WiT models pixel-space trajectory conflict as ambiguity in the denoising target at a noisy state z_t.The analysis contrasts unconstrained p(x|z_t) with semantically constrained distributions involving an intermediate waypoint.
Theoretical motivation: The unconstrained prediction error is tied to Var(x|z_t), which becomes large when diverse images share the same noisy state.The paper connects this conditional variance to entangled pixel manifolds and empirically observed trajectory conflict.
Theoretical motivation: Oracle semantic conditioning decomposes total uncertainty into within-waypoint variance and between-waypoint variance.Equation 16 expresses the unconstrained variance as the sum of non-negative conditional and semantic components.
Theoretical motivation: Equation 16 guarantees E_oracle ≤ E_standard because the additional between-waypoint variance is non-negative.This formalizes search-space contraction under the oracle condition.
Method interpretation: WiT approximates oracle conditioning by predicting semantic constraints, leaving the pixel generator to resolve reduced residual variance.The paper presents this as motivation rather than a formal guarantee for the learned predictor.
Empirical validation: WiT and JiT are compared with pairwise directional conflict and CFG relative L2 distance over a 50-step Heun trajectory.The metrics use conditional, unconditional, and counterfactual velocities, with results summarized at midpoint and peak conflict.
Empirical validation: Table 4 reports approximately 1.62× higher trajectory-conflict performance for WiT after anchoring paths to low-dimensional semantic waypoints.The supplied passage states that JiT exhibits overlapping trajectories while WiT structurally untangles them.

6 Conclusion

WiT decouples semantic navigation from pixel-level texture synthesis by routing generation through compact semantic waypoints. On ImageNet 256 × 256, it achieves state-of-the-art performance among pure pixel-space architectures and speeds JiT training by 2.2×.

6 Conclusion: WiT factors the noise-to-pixel transport path through compact semantic waypoints extracted from pretrained vision features.A lightweight Waypoints Generator infers structural anchors during integration.
6 Conclusion: Just-Pixel AdaLN provides spatially varying guidance from predicted waypoints to the primary diffusion transformer.This conditions the transformer while separating semantic navigation from texture synthesis.
6 Conclusion: 2.2× training speedup over JiT is reported on ImageNet 256 × 256.The conclusion also reports state-of-the-art performance among pure pixel-space architectures and superiority to heavy VAE-compressed latent models.