Logit-KL Flow Matching: Non-Autoregressive Text Generation via Sampling-Hybrid Inference

TL;DR

This paper introduces Logit-KL Flow Matching, a novel non-autoregressive text generation method that uses flow matching with KL divergence geodesics, improving efficiency and performance over previous models.

Contribution

It applies conditional flow matching with KL divergence geodesics to NAR text generation and provides theoretical justification and a new sampling strategy for better results.

Findings

Improves perplexity and downstream metrics over prior NAR baselines.

Provides theoretical foundation linking likelihood maximization to flow matching velocity.

Develops a hybrid sampling scheme that enhances inference quality.

Abstract

Non-autoregressive (NAR) language models offer notable efficiency in text generation by circumventing the sequential bottleneck of autoregressive decoding. However, accurately modeling dependencies in discrete sequences remains challenging in this paradigm. In this work, we advance the field of NAR generation by applying conditional flow matching (CFM) methods grounded in geometrically principled interpolation, specifically leveraging Kullback-Leibler (KL) divergence geodesics, which correspond to linear interpolation in logit space. We rigorously establish that maximizing conditional likelihood in this setting precisely recovers the flow matching velocity field, supplying the theoretical justification for this approach in sequence modeling. To address practical performance gaps of basic inference, we propose a novel empirical sampling strategy that iteratively denoises and re-noises, along with a hybrid scheme that integrates our sampling method with basic procedure. Across unconditional and conditional text and code infilling, the approach improves perplexity and downstream metrics over prior NAR baselines under matched settings.