On Feature Learning in Flow Matching under Different Couplings

Flow matching learns velocity fields from paired noise and data samples without simulating trajectories during training. We discuss feature learning of flow matching, focusing on the difference between the independent and optimal-transport (OT) couplings for the same source and target distributions. To understand the difference, we employ a theoretical model, where the data lie on an unknown low-dimensional subspace, with a non-Gaussian distribution along one unknown direction (feature) and Gaussian variation in the orthogonal directions. We show that, near the target endpoint, the nonlinear component remains nonzero under OT but vanishes under independent coupling, while both couplings have the same velocity component normal to the subspace. To examine how this difference changes feature learning, we analyze joint learning of the support and feature direction at a fixed flow time using a model with one cubic neuron. Although nonlinear feature learning becomes increasingly slow near the target endpoint under independent coupling, OT avoids this slowdown, presenting more eminent feature learning. We also empirically confirm this advantage of OT in joint training across flow times, using narrow windows near the endpoint or separate features for each time bin. These results suggest that OT coupling may benefit the learning of nonlinear structure in the target density near the endpoint.