ConvergeFlow: Language Flow with Provable Convergence to Token Embeddings

AI summary

Overview

Research area: Continuous (embedding-space) flow-based and diffusion language models, with a theoretical focus on where sampling trajectories terminate.

Technical level: Advanced. The paper mixes a machine-learning systems contribution with formal convergence proofs (a theorem and three propositions).

One-sentence scope: ConvergeFlow is an embedding-space flow-based language model whose data predictor is constrained to the convex hull of vocabulary embeddings, and the paper proves the resulting flow converges to a valid token embedding so that tokens can be predicted directly without a cross-entropy-supervised decoder, backed by experiments on OpenWebText.

What This Paper Is About

Continuous diffusion and flow-based language models match discrete models in quality, but their learned trajectories are not guaranteed to end on a real token embedding, so existing systems (LangFlow, ELF, FLM) attach a separately trained, cross-entropy-supervised decoder to convert off-embedding states into tokens. The paper asks whether a flow-based LM's sampling trajectory can be made to converge directly to a valid token embedding, eliminating that decoder. ConvergeFlow answers yes by restricting the data predictor to convex combinations of vocabulary embeddings and training it purely with the mean squared error objective induced by flow matching.

Key Contributions

1. A flow-based LM with provable convergence to token embeddings. ConvergeFlow keeps a fully continuous formulation while encoding discrete token structure. Theorem 1 shows each token-level state converges in probability to a valid token embedding despite errors in the learned data predictor. The authors state that, to their knowledge, ConvergeFlow is the first flow-based LM with provable convergence to token embeddings.

2. A structured embedding-weighted data predictor. The data predictor is parameterized as a weighted average of vocabulary embeddings, with coefficients formed from a learned base weight function multiplied by an exact Gaussian kernel induced by the corruption process. Unlike LangFlow, the base weight is neither supervised nor interpreted as a token posterior; the resulting continuous data predictor is trained with the flow-matching MSE objective, so no cross-entropy loss is used anywhere.

3. Theory separating data-prediction accuracy from convergence. Proposition 2 gives a counterexample: a smooth, unconstrained data predictor can converge to the clean embedding while the flow it induces stays at distance at least 1 from every token embedding with probability at least some constant. Proposition 3 shows that if the convex-combination predictor converges to a token embedding, the weight vector converges to a one-hot vector.

4. Three sampling mechanisms plus competitive empirical results. The authors propose three sampling mechanisms controlling the trade-off between generative perplexity and entropy. On OpenWebText, ConvergeFlow reaches a generative perplexity of 33.17 at an entropy of 5.44 with a 130M-parameter model.

Main Findings

• Provable terminal convergence: Under the assumptions that every base weight is strictly positive, that log-weights are Lipschitz along the sampling trajectory with constant L̃, and that the time grid satisfies max (t_{k+1} − t_k)/(1 − t_{k+1})³ < δ for sufficiently small δ, Theorem 1 states that for each position i there exists some j_i in [V] such that x_{t_N}^{(i)} converges in probability to e_{j_i} as N goes to infinity. Nearest-neighbor decoding then returns the token.

• Accuracy alone is not enough: Proposition 2 constructs an unconstrained predictor μ_θ with μ_θ(x_t, t) converging to x* in probability as t goes to 1, yet there exists a constant c_lb > 0, independent of N, such that for every position i and all sufficiently large N, the probability that the flow output at t_N is at distance at least 1 from every token embedding is at least c_lb. Convergence of the predictor does not, by itself, force the flow onto the vocabulary.

• Weights collapse to a one-hot vector: Proposition 3 assumes the embeddings have equal norm and are pairwise separated, with max cosine similarity over j ≠ j′ bounded by 1 − ρ for some ρ > 0. If μ_θ^{(i)}(x_t, t) goes to e_j as t goes to 1, then w_θ^{(i)}(· | x_t, t) goes to the one-hot vector δ_j.

• Conjecture-motivating factorization: Proposition 1 shows the posterior token distribution factorizes as a context-only posterior times an exact Gaussian kernel exp(−||x_t^{(i)} − α_t e_j||²/(2σ_t²)), normalized. The paper stresses this motivates the form of the parameterization only; the learned base weights are not required to estimate that context-only posterior.

• Embedding matrix is fixed: Because the target x* is itself defined by the embedding matrix, training the embedding matrix jointly with the data predictor under the flow-matching objective alone admits degenerate embedding-collapse solutions. The authors therefore use the pre-trained embedding matrix from LangFlow and keep it fixed during training and inference.

• OpenWebText results at dataset entropy: ConvergeFlow reports a generative perplexity of 33.17 at entropy 5.44 (model size 130M). The dataset entropy is 5.44. In the same table, the autoregressive Transformer (170M) reports 35.90 at 5.58; Duo (170M) 77.69 at 5.55; MDLM (170M) 104.85 at 5.63; SEDD Uniform (170M) 99.90 at 5.56; SEDD Absorb (170M) 105.03 at 5.62; FLM (179M) 62.23 at 5.33; LangFlow (130M) 60.09 at 5.43; ELF (105M) 65.30 at 5.40.

• Comparison to continuous flow baselines: The paper states that at the dataset entropy of 5.44, ConvergeFlow attains 33.17, whereas the lowest generative perplexity among the continuous flow-based LMs is approximately 60, and those models are evaluated at entropies below the dataset entropy.

• Two parameter-free decoding rules: After the final step, tokens can be read out either by nearest neighbor in embedding space, argmin over j of ||x_{t_N}^{(i)} − e_j||₂, or by argmax over j of the learned weights w_θ^{(i)}(j | x_{t_N}, t_N). Both require no separately trained terminal decoder.

• Generative perplexity definition: Perplexity is computed against a reference language model, the GPT-2 Large model, and the identity log PPL_gen = (1/L) KL(p_θ || p_ref) + (1/L) H(p_θ) motivates reporting entropy alongside it as a diversity measure; unigram entropy is used as a proxy for normalized sequence entropy. The entropy value 5.44 is marked on the trade-off figure as the entropy of the OpenWebText dataset.

Methodology in Plain English

The model starts from flow matching. Tokens are first mapped to vectors through an embedding matrix E in R^{V×d}, so a length-L sequence becomes an L×d target x*. Noise is mixed in along a linear Gaussian path, x_t = α_t x* + σ_t z with standard Gaussian z, and the network is trained to predict the clean target x* from the noisy x_t.

The training signal is the mean squared error implied by flow matching, with a schedule-dependent weighting (1 + α_t/σ_t)² on the squared error between x* and the predictor μ_θ(x_t, t). This is purely continuous supervision; no cross-entropy term appears.

The key design choice is what the predictor is allowed to output. Instead of letting it return any vector in R^{L×d}, each position's prediction is written as a convex combination of the vocabulary embeddings: μ_θ^{(i)}(x_t, t) = Eᵀ w_θ^{(i)}(· | x_t, t). The weight vector w is built by taking a learned base weight function f_θ^{(i)}(j | x_t, t) and multiplying it by the exact Gaussian kernel exp(−||x_t^{(i)} − α_t e_j||²/(2σ_t²)) from Proposition 1, then normalizing. Because the embeddings are fixed and the weights form a convex combination, the output is confined to the convex hull of the token embeddings, and the proof shows this is enough to force the flow onto an actual embedding at the end.

For generation, the data-prediction ODE is integrated with a first-order Euler method on a time grid, and self-conditioning feeds the previous step's prediction back in as an extra input, with an unconditional prediction used at the first step. At the end, the two parameter-free decoding rules above convert the terminal state into tokens. The paper proposes three sampling mechanisms for trading generative perplexity against entropy; two of them, schedules A and B, are described as time-adaptive guidance schedules, with the detailed definitions and the third mechanism falling in the portion of the manuscript not included here.

Why This Matters

The paper removes a structural inconsistency in continuous flow-based language modeling: if the flow is guaranteed to land on a token embedding, then the cross-entropy-trained decoder that re-introduces discrete supervision becomes unnecessary, and the training objective and the inference procedure finally agree. This is what allows the standard continuous-diffusion toolkit (classifier-free guidance, self-conditioning, few-step ODE solvers, distillation) to be applied cleanly to language, and it supports bidirectional context, global planning, parallel token updates, and iterative revision rather than strict left-to-right generation.

Real-world applications implied by parallel, bidirectional, iterative generation (the paper itself does not report application-specific evaluations):

• Latency-sensitive text generation, where all positions are refined in parallel instead of one token at a time.
• Controllable or constrained generation, where intermediate states can be revised and guided during sampling.
• Infilling and text editing, where surrounding context on both sides informs the generated span.
• Distillation or few-step sampling pipelines for deployment, which the continuous formulation makes available.

Industry relevance: at 130M parameters, ConvergeFlow reports a generative perplexity of 33.17 at the dataset entropy of 5.44 on OpenWebText, against 60.09 for LangFlow at 5.43 (130M), 62.23 for FLM (179M), and 65.30 for ELF (105M). A quality-diversity trade-off that is explicitly controllable at this scale is directly relevant to anyone choosing between autoregressive and parallel-decoding language models for deployment.

Future Directions

• Joint embedding and predictor learning. The authors freeze the pre-trained LangFlow embedding matrix because the flow-matching objective alone permits degenerate embedding-collapse solutions. Removing that dependence is an open problem flagged by the paper itself.
• Scaling. Evaluation is reported only on OpenWebText with 130M-parameter models. Whether the convergence guarantee and the generative-perplexity advantage persist at larger scale is untested in the provided content.
• Full characterization of the sampling mechanisms. The paper promises three mechanisms for controlling the perplexity-entropy trade-off, but the definitions of the third mechanism and the details of schedules A and B are outside the included content; systematic ablations of each are a natural next step.
• Exploiting theory for efficiency. Because convergence holds for any sufficiently fine time grid under the stated conditions, adapting that grid (adaptive step counts, higher-order solvers, or distillation) is a natural way to reduce sampling cost while retaining the guarantee.

Target Audience

This paper is most useful to researchers working on diffusion and flow-based language models who want to understand why prior continuous LMs needed a cross-entropy decoder and how a structural constraint on the data predictor removes that need. It also suits machine-learning theorists interested in convergence guarantees for sampling trajectories over discrete-supported distributions, and practitioners evaluating parallel-decoding language models who care about explicit control over the quality-diversity trade-off. The empirical sections are readable with intermediate background; the theorem and propositions require comfort with flow matching, probability, and Lipschitz-style regularity arguments.

Authors’ abstract

Recent advances in continuous diffusion and flow-based language models (LMs) have achieved performance competitive with discrete LMs. However, existing continuous frameworks still rely on decoders supervised with cross entropy (CE) because the flow trajectories are not guaranteed to terminate at valid token embeddings. Motivated by this limitation, we introduce \textbf{ConvergeFlow}, an embedding-space flow-based LM, which constrains the data predictor to the convex hull of token embeddings and trains it solely with the mean squared error objective induced by flow matching. Under suitable regularity conditions, we prove that the resulting flow converges to valid token embeddings despite errors in the data predictor, enabling direct token prediction without a CE-supervised decoder. We further develop three sampling mechanisms for controlling the trade-off between the generative perplexity and entropy. Experiments on OpenWebText demonstrate that ConvergeFlow achieves performance competitive with existing continuous and discrete diffusion LMs. These findings demonstrate the potential of the flow-based paradigm for language modeling. Our code is available at https://github.com/Na-Li66/ConvergeFlow.