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

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.

💡 Research Summary

The paper investigates the application of flow matching (FM) – a continuous probability‑path framework underlying diffusion models – to binary data, where each datum is a vector of 0/1 bits. Recent successes in continuous domains have highlighted the simplicity and performance of signal‑space prediction (x‑prediction), as exemplified by the Just Image Transformer (JiT). The authors ask whether this paradigm remains robust when the target distribution is discrete.

Key Problem – Prediction‑Loss Mismatch
When x‑prediction is combined with the traditional velocity‑matching loss (v‑loss), the loss acquires a time‑dependent factor λ(t) = (1‑t)⁻² (see Eq. 3). As the integration time t approaches 1 (the point where the noisy latent is almost fully denoised), this factor diverges. The authors formalize gradient stability via the integrated second‑moment I = ∫₀¹E