Technical Summary: Latent-Kernel Discrete Flow Maps for Few-Step Generation

Problem Statement

Discrete diffusion and flow-matching models generate text by denoising a sequence over many steps. To maintain computational efficiency, these models typically factorize the transition probability across token positions, deciding every token independently given the current sequence state. While this factorization makes individual steps cheap, it creates a fundamental bottleneck for few-step generation.

When a target sequence requires strong correlations between positions (e.g., subject-verb agreement or matching brackets), independent updates commit to tokens separately. This often results in locally consistent but globally inconsistent combinations (e.g., "dog barks" vs. "dogs bark"). Correcting these mismatches requires many additional denoising steps, negating the efficiency gains of few-step generation.

Existing solutions to accelerate generation fall into two categories:

1. Likelihood models (e.g., MDLM, SEDD): Train strong factorized denoisers from scratch but require many steps at inference to resolve correlations.
2. Few-step accelerators (e.g., Di4C, SDTT, ReDi): Distill or rectify a trained "teacher" model into a few-step student. However, these methods inherit the teacher's quality ceiling and cannot surpass the performance of the model they imitate.

The core challenge is whether a model can express correlated steps natively without relying on a teacher, thereby breaking the factorization assumption at its root while maintaining low computational cost.

Methodology: Latent-Kernel Discrete Flow Maps (LKF)

The authors propose Latent-Kernel Discrete Flow Maps (LKF), a from-scratch flow-map model that introduces intra-step correlation through a mixture of factorized components tied by a single shared latent variable.

The Model Architecture

The finite-time transition kernel $\Psi_{t \to s}(x_s | x_t)$ is defined as a mixture of $M$ factorized components:
$$\Psi_{t \to s}(x_s | x_t) = \sum_{k=1}^{M} w_k(x_t, t, s) \prod_{i=1}^{L} P_i(x_s^i | x_t, k, t, s)$$
Where:

• $k \in [M]$ is a shared latent component index sampled once per sequence.
• $w_k$ are router weights predicting the probability of component $k$.
• $P_i$ are factorized token distributions for position $i$ conditioned on the latent $k$.

Key Mechanism:

1. Shared Latent: At the start of a denoising trajectory (or each step), a single latent $k$ is drawn from the router $w(\cdot | x_t)$.
2. Conditional Independence: Given $k$, tokens are still sampled independently across positions. However, because $k$ is shared across the entire sequence, marginalizing over $k$ induces correlations between positions.
3. Efficiency: The model uses a late latent injection architecture. A shared transformer trunk processes the input to create a latent-agnostic representation. Only the final $L_k$ blocks are replicated for each of the $M$ components, conditioned on $k$. This keeps the computational cost of marginalization low ($O(M \cdot L_k)$ rather than $O(M \cdot \text{depth})$).

Training and Inference

• Training: The model is trained by marginalizing the latent $k$ exactly. The objective is the log-likelihood of the mixture:
  $$\log \Psi = \log \sum_{k} \exp(\log w_k + \ell_k)$$
  where $\ell_k$ is the log-probability assigned by component $k$ to the target tokens. This forces the router to learn a mixture where the latent captures sequence-level correlations. A router regularizer ensures the mixture components are utilized (preventing collapse to a single component) and that the latent carries sufficient information.
• Inference (Best-of-M): To leverage the heterogeneity of the mixture, the authors employ a Best-of-M decoding strategy. $M$ parallel rollouts are generated, each freezing a different component $k$. The final output is selected based on the model's own sequence-level mixture negative ELBO (NELBO). This selection requires no external reward model or teacher.

Key Contributions

1. Native Correlated Steps: Introduction of LKF, a flow-map kernel that expresses correlated steps natively via a mixture of $M$ factorized components tied by a single shared latent, eliminating the need for a teacher model.
2. Theoretical Analysis: An exact information decomposition showing that the total correlation (TC) a rank-$M$ step can carry is bounded by $(L-1)H(k|xt)$. The authors prove that the approximation error of the model is lower-bounded by the uncaptured total correlation of the target distribution.
3. Synthetic Validation: Experiments on synthetic datasets (hidden-agreement and parity) demonstrate that LKF captures information and improves likelihood monotonically with $M$ for representable correlations, while correctly failing to capture correlations that require exponential mixture rank (parity), aligning with theoretical hardness results.
4. Empirical Performance: On LM1B and WikiText-103 benchmarks, LKF achieves state-of-the-art few-step generative perplexity. At $M=8$, it improves generative perplexity by 2.1× to 3.3× over likelihood baselines (MDLM, SEDD) and surpasses distilled/rectified accelerators (Di4C, ReDi) without using any teacher.

Experimental Results

• Benchmarks: Evaluated on LM1B and WikiText-103.
• Performance:
  • LM1B: The $M=8$ LKF model achieves a generative perplexity (gen-PPL) of 105 at 32 function evaluations (NFE), compared to 304 for MDLM and 117 for the best rectified accelerator (ReDi).
  • WikiText-103: The $M=8$ model reaches a gen-PPL of 75, significantly outperforming MDLM (228) and SEDD (247).
• Diversity: Unlike rectified accelerators that often suffer from entropy collapse (producing low-diversity text), LKF maintains high sample entropy while improving likelihood.
• Ablations:
  • The gain is attributed to the mixture structure itself, not just reranking. Replacing the $M$ components with $M$ noise draws of a single component results in a significant performance drop (e.g., 35% higher PPL at $M=8$).
  • Best-of-M selection is crucial; a single-commitment policy (commit-M) yields results similar to the baseline, confirming that the mixture provides candidate diversity that selection exploits.
  • Router Efficiency: The router acts as a mixing weight rather than a quality signal; the selection is most effective when based on the model's own likelihood (running evidence) rather than the router's initial prediction.

Significance and Claims

The paper claims that LKF represents a shift from "teacher-dependent" acceleration to "from-scratch" few-step generation. By modeling the transition kernel as a low-rank mixture, LKF allows a single step to place mass on correlated completions (e.g., subject-verb agreement) with the same sampling complexity as a factorized model.

The authors emphasize that:

• The method does not inherit a teacher's quality ceiling.
• The gain scales with the mixture size $M$.
• The approach is theoretically grounded, with bounds on correlation capacity and approximation error.
• It achieves the best few-step generative perplexity among baseline methods on standard language modeling benchmarks, offering a viable path to efficient, high-quality discrete generation without the overhead of distillation or rectification.