Binary Flow Matching: Prediction-Loss Space Alignment for Robust Learning

TL;DR

This paper extends flow matching to binary data, identifying the importance of aligning prediction loss with the signal space to ensure robust training and eliminate gradient issues.

Contribution

It formalizes prediction-loss alignment as essential for flow matching, proving that x-loss alignment improves robustness and provides practical guidelines for discrete data modeling.

Findings

Aligning prediction loss with signal space removes singular gradient weighting.

x-loss alignment enables robust training without heuristic timestep schedules.

Topology-dependent differences between probabilistic and geometric losses are revealed.

Abstract

Flow matching has emerged as a powerful framework for generative modeling, with recent empirical successes highlighting the effectiveness of signal-space prediction ($x$-prediction). In this work, we investigate the transfer of this paradigm to binary manifolds, a fundamental setting for generative modeling of discrete data. While $x$-prediction remains effective, we identify a latent structural mismatch that arises when it is coupled with velocity-based objectives ($v$-loss), leading to a time-dependent singular weighting that amplifies gradient sensitivity to approximation errors. Motivated by this observation, we formalize prediction-loss alignment as a necessary condition for flow matching training. We prove that re-aligning the objective to the signal space ($x$-loss) eliminates the singular weighting, yielding uniformly bounded gradients and enabling robust training under uniform timestep sampling without reliance on heuristic schedules. Finally, with alignment secured, we examine design choices specific to binary data, revealing a topology-dependent distinction between probabilistic objectives (e.g., cross-entropy) and geometric losses (e.g., mean squared error). Together, these results provide theoretical foundations and practical guidelines for robust flow matching on binary -- and related discrete -- domains, positioning signal-space alignment as a key principle for robust diffusion learning.