Prediction-Loss Alignment for Sampler-Robust Flow Matching Training

The paper looks at a widely used training recipe for diffusion and flow-matching models, where the network predicts a clean sample that is then converted into a velocity for the loss. That conversion amplifies errors near the endpoints of the noise schedule, making training unstable and tied to a particular sampler. The authors propose aligning the prediction objective with the loss objective so that training stays robust regardless of which sampler is used at inference.