# Prediction--Loss Alignment for Sampler--Robust Flow Matching Training

arXiv: https://arxiv.org/abs/2602.10420 · HTML: https://mlsanity.com/paper/2602.10420/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 16 of its 25 references that are in the map.
- First posted on arXiv 11 Feb 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 16 of 16)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design](https://mlsanity.com/paper/2402.04997/) (2024)
- [Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation](https://mlsanity.com/paper/2303.00848/) (2023)
- [Efficient Diffusion Training via Min-SNR Weighting Strategy](https://mlsanity.com/paper/2303.09556/) (2023)
- [Continuous diffusion for categorical data](https://mlsanity.com/paper/2211.15089/) (2022)
- [Variational Flow Matching for Graph Generation](https://mlsanity.com/paper/2406.04843/) (2024)
- [Soft Graph Transformer for MIMO Detection](https://mlsanity.com/paper/2509.12694/) (2025)

## Built on by (most cited first, 4 of 4)

- [Soft Graph Diffusion Transformer for MIMO Detection](https://mlsanity.com/paper/2605.00449/) (2026)
- [CAT-Flow: Curvature-Adaptive sTeps for Flow Matching](https://mlsanity.com/paper/2609.01746/) (2026)
- [TACD: Distilling Efficient Text-to-Motion Models via Terminal Amplification Control](https://mlsanity.com/paper/2610.02867/) (2026)
- [Bernoulli Flow Models: Self-Consistent Generative Modeling for Binary Data](https://mlsanity.com/paper/2610.11362/) (2026)

## Abstract (arXiv)

Recent work has popularized a practical recipe in diffusion and flow matching: predict the clean signal x, convert it to a velocity, and train through a velocity-space loss. The conversion contains a singular endpoint amplification and therefore appears prone to unstable optimization, yet recent systems obtain strong empirical results with this recipe. We investigate this tension through the integrability of the pre-optimizer stochastic-gradient second moment. Under stated initialization conditions, the moment diverges under Uniform sampling; boundary-suppressing sampling can restore integrability under an additional upper-growth condition. We then show that prediction--loss alignment eliminates this conversion-induced source of non-integrability. Under a uniform moment bound, alignment yields a finite second moment for every timestep density, including Uniform sampling. Controlled experiments across continuous and binary settings reproduce the predicted sampler-dependent instability and show that aligned objectives remain trainable across the tested samplers. These results reconcile pointwise amplification with sampler-dependent empirical success and support alignment as a principled route to more robust flow-matching training.