KL-geodesics flow matching with a novel sampling scheme
This work addresses the problem of improving text generation quality for non-autoregressive language models, which offer speed advantages but struggle with dependency modeling.
The paper tackles the challenge of modeling complex dependencies in non-autoregressive text generation by proposing a conditional flow matching approach using KL-geodesics and a novel empirical sampling scheme. The hybrid inference method achieves superior performance compared to previous state-of-the-art discrete flow matching methods in both conditional and unconditional text generation experiments.
Non-autoregressive language models generate all tokens simultaneously, offering potential speed advantages over traditional autoregressive models, but they face challenges in modeling the complex dependencies inherent in text data. In this work, we investigate a conditional flow matching approach for text generation. We represent tokens as one-hot vectors in a \(V\)-dimensional simplex and utilize geodesics under the Kullback-Leibler (KL) divergence, which correspond to linear interpolation in logit space. We provide a theoretical justification that maximizing the conditional likelihood \(P_θ(x_1 \mid x_t, t)\) yields the exact flow matching velocity under logit interpolation. To address the suboptimal performance of basic inference, we propose a novel empirical sampling scheme that iteratively samples from the conditional distribution and introduces additional noise, significantly improving results despite lacking full theoretical underpinnings. Furthermore, we propose a hybrid inference method that combines the basic approach with the sampling scheme. This method demonstrates superior performance on both conditional and unconditional text generation experiments compared to previous SOTA method for discrete flow matching.