RiT 架构用原始输入旁路修复像素空间 Flow Matching 的 v-prediction 失效问题
What kills $v$-prediction? A Patch-wise PCA Perspective on Pixel-Space Flow Matching
做扩散或 Flow Matching 训练的工程和研究同学值得看这篇。它用 patch-wise PCA 给出了 v-prediction 失效的机理,提出的 RiT 不用改成 x-prediction 也能跑通像素空间训练,在 ImageNet 256×256 上验证有效。
一篇 arXiv 论文(编号 2611.12852)通过 patch-wise PCA 分析指出,像素空间 Flow Matching 中 v-prediction 训练困难的根源是带宽瓶颈:速度在次成分上主要由带噪原始输入决定,深层骨干网络需要额外容量保存这些信息,与主成分建模相互竞争。
论文用受控实验解释了 x-prediction 依然有效的原因:其参数化天然提供原始输入的旁路,学到的表示会自发集中在主方向上;在低秩观测下,x-prediction 相当于一个隐式条件均值解码器。基于此分析,作者提出 Raw-injected Transformer(RiT),用轻量级原始条件解码器加输出侧捷径,把语义建模与原始状态传播解耦。
RiT 恢复了 v-prediction 的有效性,并在 ImageNet 256×256 上取得有竞争力的成绩,说明预测目标的选择需要与架构带宽配合设计,而非简单换用 x-prediction。
What kills $v$-prediction? A Patch-wise PCA Perspective on Pixel-Space Flow Matching
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.
This report is categorized under AITOP's "Papers" desk. Within our cross-regional analytical framework, it directly maps to AI developer toolchains, coding agents, and engineering workflows. Key entities and technical topics under continuous tracking include Flow Matching, v-prediction, RiT, ImageNet.
Source attribution: curated from arXiv cs.LG (@Tianze Luo, Haofeng Huang, Yuan Yao).