Discrete Flow Maps

Peter Potaptchik; Jason Yim; Adhi Saravanan; Peter Holderrieth; Eric Vanden-Eijnden; Michael S. Albergo

Source-linked AI summary

Discrete Flow Maps

Peter Potaptchik, Jason Yim, Adhi Saravanan, Peter Holderrieth, Eric Vanden-Eijnden, Michael S. Albergo

arXiv:2604.09784v2Apr 10, 2026stat.MLcs.LG

TL;DR

Autoregressive generation is sequential, and standard flow maps use Euclidean losses despite text distributions living on the probability simplex. This paper introduces Discrete Flow Maps, which use simplex-valued mean denoisers and cross-entropy or KL objectives; the method supports one- and few-step generation and surpasses prior state-of-the-art results.

Problem
Autoregressive generation is sequential, while standard flow-map objectives use Euclidean regression for discrete text distributions that lie on the probability simplex.
Method
Discrete Flow Maps reparameterize flow maps through a mean denoiser on the simplex and derive cross-entropy and KL-divergence training objectives.
Results
Discrete Flow Maps support one- and few-step generation with only minor performance degradation and surpass previous state-of-the-art performance in non-autoregressive language generation.
Takeaways & Limitations
The framework provides a geometrically consistent approach to discrete flow modeling while retaining trajectory-compressed generation and test-time steering or guidance.
Takeaways & Limitations
The main paper restricts its continuous-flow construction to a linear interpolant, although the framework extends to a broader class in the appendix.

Abstract

from arXiv · show

The sequential nature of autoregressive next-token prediction imposes a fundamental speed limit on large language models. While continuous flow models offer a path to parallel generation, they traditionally demand expensive iterative integration. Flow Maps bypass this bottleneck by compressing generative trajectories into single-step mappings, theoretically enabling the generation of full text sequences from noise in a single forward pass. However, standard formulations rely on Euclidean regression losses that are geometrically ill-suited for discrete data. In this work, we resolve this conflict with Discrete Flow Maps, a framework that reconciles trajectory compression with the geometry of the probability simplex. We recast standard flow map training for the discrete domain, aligning the training dynamics with the discrete nature of language. Empirically, this strict geometric alignment allows our method to surpass previous state-of-the-art results in discrete flow modeling.

1 Introduction

Autoregressive generation is inherently sequential, while continuous flow methods offer parallel generation but face a geometric mismatch when applied to discrete text. Discrete Flow Maps address this by parameterizing maps on the probability simplex and using likelihood-based objectives, enabling strong one- and few-step generation.

Autoregressive models generate text one token at a time, imposing linear computational cost and limiting efficient long-form or real-time synthesis.
Continuous flows and flow maps offer non-autoregressive generation, but standard flow-map training applies Euclidean L2 losses to probability distributions that live on the simplex.
The framework surpasses previous state-of-the-art performance in non-autoregressive language generation.
Discrete Flow Maps reparameterize flow maps through a mean denoiser that natively lies on the probability simplex.
The resulting objectives replace Euclidean regression with cross-entropy and KL-divergence losses aligned with discrete data geometry.
Discrete Flow Maps support one- and few-step language generation with only minor performance degradation, while enabling test-time steering and guidance.

2 Preliminaries

The preliminaries formulate language modeling on the probability simplex, review continuous flow transport and flow-map trajectory compression, and describe consistency training for direct time-to-time mappings. Standard flow maps use average-velocity parameterizations and consistency constraints to recover valid global trajectories.

2.1 Language Modeling on the Simplex: Language models represent vocabulary tokens as simplex vertices and learn distributions over discrete token sequences.
2.1 Language Modeling on the Simplex: Cross-entropy training over simplex-valued functions recovers conditional token probabilities because one-hot target expectations equal class-probability vectors.
2.2 Continuous Generative Flows: Continuous generative flows transport noise to data through an ODE whose trajectory endpoint follows the target data distribution.
2.2 Continuous Generative Flows: Numerical ODE integration requires many neural-drift evaluations, motivating flow maps and consistency models that compress trajectories into single- or few-step mappings.
2.3 Flow Maps and Trajectory Compression: A flow map directly transports a state between times, with its residual parameterization expressed through average velocity.
2.3 Flow Maps and Trajectory Compression: The diagonal loss matches average velocity to instantaneous drift, while semigroup, Lagrangian, and Eulerian constraints enforce consistency across distinct times.
2.3 Flow Maps and Trajectory Compression: Training the combined consistency objectives recovers the true flow map at optimality.

3 Discrete Flow Maps

Discrete Flow Maps reparameterize flow maps with a mean denoiser on the probability simplex, replacing Euclidean regression with distribution-aware objectives. The framework uses cross-entropy and KL-based consistency losses to preserve valid discrete geometry during training.

Mean Denoiser Parametrization: The method replaces unconstrained velocity parameterization with a mean denoiser that explicitly resides on the probability simplex.Standard flow maps use velocities in R^K, which need not remain on the simplex.
Mean Denoiser Parametrization: The mean denoiser is guaranteed to lie on the simplex because it is formed from expectations of one-hot vectors and their weighted convex combinations.This guarantees geometrically valid probability-distribution outputs.
Mean Denoiser Parametrization: Softmax parameterization ensures the learned denoiser is always a valid probability distribution, while the framework also extends to general convex sets.A suitable link function can map outputs into the convex hull of the data support.
Training Objectives: Training uses cross-entropy for the diagonal denoiser and KL divergence for consistency objectives, replacing Euclidean losses with simplex-compatible alternatives.The consistency objectives include semigroup, Lagrangian, and Eulerian formulations.
Training Objectives: Semigroup, Lagrangian, and Eulerian consistency losses enforce the corresponding flow identities, with forward-KL training targets defined as valid probability distributions.The semigroup target is a convex combination on the simplex; Lagrangian and Eulerian objectives enforce their respective consistency rules.

4 Algorithmic Details

The algorithmic design extends discrete flow maps through flexible schedules, conditional block generation, guidance, and stabilized simplex-aware objectives.

Interpolants and schedules: Time reparameterization distributes denoising progress more evenly without changing the noise-to-data path or its endpoints.The schedule is chosen so the probability of identifying the final argmax increases approximately linearly in reparameterized time.
Interpolants and schedules: Position-dependent schedules reveal earlier positions sooner, creating an autoregressive bias within an otherwise parallel model.
Conditional generation and guidance: Conditional flow generation keeps a variable-length context fixed while modeling the continuation distribution p1(x | c).
Conditional generation and guidance: Block generation produces future-token blocks in parallel, appends them to the context, and repeats this process for subsequent blocks.Mixed attention uses causal attention for the fixed context, bidirectional attention within each block, and cached key–value states for context reuse.
Conditional generation and guidance: Classifier-free guidance combines conditional and unconditional drifts, and the distilled flow map inherits this mechanism for one-step or few-step sampling.The guidance strength is controlled by ω ≥ 0.
Conditional generation and guidance: Guided sampling still terminates at valid discrete tokens because it recovers the conditional-data support contained in simplex vertices.
Loss implementation: Stable logit-space rearrangements avoid ill-conditioned correction ratios near the simplex boundary without changing the target distribution.Detached loss weighting additionally rescales student-logit gradients to stabilize optimization without changing the optimum.

5 Experiments

Experiments evaluate Discrete Flow Maps on LM1B and OpenWebText across accelerated sampling regimes, finding strong few-step performance and guidance-dependent diversity trade-offs.

Setup: DFMs are evaluated on LM1B and OpenWebText with sequence lengths of 128 and 1024, respectively.LM1B uses bert-base-uncased tokenization, while OpenWebText uses the GPT-2 tokenizer.
Setup: Training uses 1M diagonal steps followed by off-diagonal consistency training for 200k LM1B steps and 100k OpenWebText steps.
Results: Across both datasets, DFMs outperform the compared baselines in few-step generative perplexity while generally preserving diversity.ESD outperforms PSD at 2 and 4 NFEs on both entropy and generative perplexity, but exhibits mode collapse at 1 NFE.
Results: Consistency training produces a few-step sampler that significantly outperforms the model trained on the diagonal alone.
Guidance: Classifier-free guidance is evaluated for blocks of 256 tokens conditioned on prompts, across guidance scales ω ∈ {0.0, 0.5, 1.0, 1.5, 2.0}.The experiment uses four blocks of 256 tokens, each with 1024 steps.
Guidance: As ω increases beyond standard conditional generation, both generative perplexity and entropy decrease.

6 Related Work

Related work spans continuous representations of language, discrete diffusion, flow maps, and categorical flow maps, while this paper centers exact simplex-aware objectives.

Continuous and discrete generative models: Continuous flow and diffusion language models differ mainly in how discrete data is represented in continuous space, including embeddings and simplex-based manifolds.
Flow maps: Compared with concurrent flow map language models, the paper shares diagonal training but additionally develops and emphasizes an Eulerian perspective.
Flow maps: Compared with concurrent categorical flow maps, the approach places probability-simplex geometry and exact cross-entropy and KL losses at the framework’s center.The related work is described as relying instead on composite loss bounds or inexact objectives.
Continuous and discrete generative models: Discrete diffusion models use discrete-time or continuous-time Markov chains on discrete state spaces and have produced substantial speed-ups at scale.
Flow maps: The paper reparameterizes flow maps with a mean denoiser that preserves the data’s convex constraints for discrete modeling.

A Proofs

The appendix generalizes the main linear schedule to differentiable time schedules with specified endpoint and positivity constraints, recovering the main-paper case as a special instance.

General schedules: All results in the appendix specialize to the main paper when α_t = 1 − t and β_t = t.
General schedules: General schedules α_t and β_t are C1, with α_t > 0 for t < 1 and endpoint constraints α_0 = 1, β_0 = 0, α_1 = 0, and β_1 = 1.

A.1 Mean Denoiser

The mean denoiser is parameterized through a simplex-valued predictor and an affine flow-map form, with general interpolant identities recovering the linear schedule as a special case.

The derivation obtains the stated mean denoiser by normalizing the integral term and conditioning on It=x.
The flow map uses a simplex-valued predictor ψs,t: R^K → ΔK−1.
The equivalent residual form expresses Xs,t(x) as x plus (t−s)vs,t(x).
For αt=1−t and βt=t, the general parameterization recovers the formulation used in the main paper.
The trajectory equation is obtained by differentiating the interpolant and taking a conditional expectation.

A.2 Flow Map Identities for Mean Denoiser

The appendix derives semigroup and consistency identities for general interpolants, then reformulates them in logit space to support optimization while recovering the linear-interpolant case.

General flow maps satisfy the semigroup composition identity Xs,t=Xu,t∘Xs,u.
The coefficient derivation assumes u↦βu/αu is non-decreasing when establishing the relevant identity.
For αt=1−t and βt=t, the coefficients reduce to the linear-interpolant identities.
Logit Consistency: The derivation converts consistency identities into logit-space equations involving ψs,t and its time derivative.
The Eulerian derivation uses the trajectory derivative and Jacobian of the predictor along the vector field.

A.3 Training objectives for Mean Denoiser

Training combines cross-entropy for the diagonal predictor with consistency losses based on semigroup, Lagrangian, and Eulerian formulations, while retaining the linear schedule as a special case.

Diagonal Loss: The diagonal predictor is trained with cross-entropy under the general interpolant.
Diagonal Loss: The predicted flow map and vector field use affine forms based on Γs,t, Ξs,t, and the diagonal predictor.
Consistency via Semigroup Loss (PSD): The semigroup loss enforces consistency through KL distillation from a teacher combining intermediate predictions.
The Lagrangian and Eulerian losses derive consistency teachers from the logit-space and Eulerian identities, respectively.
Reduction to the Linear Schedule: Under αt=1−t and βt=t, the general objectives recover the linear-interpolant coefficients and identities.

B Experimental Details

The experiments use fixed optimization settings, objective-specific stabilization choices, and a modified noise schedule that spends more time near low- and high-noise regions.

Training Details: Training uses batch size 512, 2500 warm-up steps, learning rate 3×10−4, and Adam with β1=0.9 and β2=0.999.
Diagonal: Diagonal training uses the adaptive loss with r=0.5 and c=0.01.
PSD: PSD training prioritizes the diagonal-loss gradient and projects out conflicting distillation-gradient components.Without projection, optimization was reported to collapse toward degenerate solutions at the expense of the diagonal loss and overall performance.
ESD: ESD omits gradient surgery and uses learnable loss weighting dependent on both s and t.
Noise Schedule: The modified schedule uses β̃(t)=λβ(t)+(1−λ)t with λ=0.9.This combination spends more time in low- and high-noise regions than the argmax schedule.

C Example Generations

The examples compare LM1B generations using 2 versus 1024 function evaluations and show OWT block generations under multiple CFG settings. OWT samples are organized into four blocks of 256 tokens, with color changes marking block boundaries.

Figure 5 compares LM1B generations produced with 2 and 1024 function evaluations.
Figure 6 shows OWT generation with CFG ω = 0.0 using block sampling across 4 blocks of size 256.A change in text colour denotes a new block.
The displayed OWT generations contain extended text samples spanning multiple generated blocks.
Figures 7–9 show OWT block-sampling generations with CFG scales ω = 1.0, 2.0, and 3.0.Each example generates 4 blocks of size 256, with color changes marking new blocks.