State Augmented Flows

Flow matching learns continuous-time transports from a simple source distribution to a target data distribution by training a velocity field on intermediate states along interpolating trajectories. However, because the model observes only the projected data-space state $(x_t,t)$, different transport trajectories can induce conflicting velocity directions at nearby intermediate states, leading to curved transport paths that limit the fidelity of low-NFE sampling. To address this limitation, we propose State Augmented Flows (SAF), a lightweight framework that lifts flow matching into a transient augmented state space. SAF transports an augmented state $(x_t,z_t)$, where auxiliary coordinates provide additional contextual information during transport while vanishing at the terminal time. This enlarges the state observed by the velocity model without changing the final sample space, and is compatible with existing coupling strategies and flow-matching architectures. Across MNIST, CIFAR-10, and ImageNet-32, SAF consistently improves over the rectified-flow baseline, with the largest gains in the low-NFE regime. These results suggest that augmenting the transported state provides a new and complementary mechanism for improving flow-based generative modeling.