Patch analysis explains challenges in pixel-space flow matching models
What kills $v$-prediction? A Patch-wise PCA Perspective on Pixel-Space Flow Matching
Computer Vision and Pattern RecognitionMachine Learning
Summary
High-quality image generation models often use complex two-step training with limited reconstruction quality. An alternative called pixel-space flow matching tries to train in one step but struggles with a method called v-prediction. The authors studied this issue by analyzing images in small patches to understand why v-prediction fails. They found that v-prediction requires processing noisy details that compete with main features, which is hard for standard models. Their new design, the Raw-injected Transformer, splits these tasks and makes v-prediction work well again.
What this means in practice
- •For image generation engineers: Improve pixel-space flow matching training for high-resolution images with better prediction stability and model design.
- •For computer vision developers: Use new model architecture to enhance end-to-end generative model training on challenging image datasets without two-stage processes.
Authors
Tianze Luo, Haofeng Huang, Yuan Yao
Abstract
Latent Diffusion Models dominate generative modeling but face two-stage training complexity and the inherent reconstruction limits of pre-trained autoencoders. Pixel-space Flow Matching offers a simplified end-to-end alternative yet often encounters training difficulties when using the standard $v$-prediction objective. Recent research attributes this difficulty to the off-manifold nature of noised targets, motivating a shift toward $x$-prediction. In this paper, we investigate this explanation through patch-wise Principal Component Analysis (PCA) and identify a bandwidth bottleneck. The velocity in minor components is largely determined by the raw noisy input, requiring deep backbones to preserve information that competes with principal-component modeling. Through controlled experiments and representation analysis, we further examine why $x$-prediction remains effective: its parameterization provides a raw-input bypass, while its learned input representation can spontaneously concentrate on principal directions. Under low-rank observations, our analysis of the induced flow characterizes $x$-prediction as an implicit conditional-mean decoder, reconstructing discarded components from the retained information. These observations motivate the Raw-injected Transformer (RiT), which combines a lightweight raw-conditioned decoder with an output-side shortcut to separate semantic modeling from raw-state propagation. Our method restores effective $v$-prediction and achieves competitive performance on ImageNet $256\times256$, providing a perspective on the interaction between prediction objectives and architectural bandwidth.