🤖 machine learning

Scaling Categorical Flow Maps

This paper demonstrates the scalability of Categorical Flow Maps by training a 1.7B-parameter model on 2.1T tokens to generate high-quality text in just four steps, while also establishing a likelihood bound for evaluation and providing key insights into training challenges like loss weighting and time scheduling.

Original authors: Oscar Davis, Anastasiia Filippova, Pierre Ablin, Victor Turrisi, Amitis Shidani, Marco Cuturi, Louis Béthune

Published 2026-05-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Oscar Davis, Anastasiia Filippova, Pierre Ablin, Victor Turrisi, Amitis Shidani, Marco Cuturi, Louis Béthune

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: A New Way to Write Text

For a long time, the best way for computers to write text (like chatbots or story generators) has been Autoregressive (AR) models. Imagine a writer who builds a sentence one word at a time, like laying bricks. They write the first word, then the second, then the third. This is reliable, but it's slow because they have to wait for each brick to dry before laying the next one.

Recently, scientists tried a different approach called Diffusion. Imagine a sculptor starting with a block of marble covered in noise (static) and slowly chipping away the noise to reveal the statue. This is great for images, but for text, it's been tricky because text is made of distinct "blocks" (words), not smooth clay.

This paper introduces a new method called Categorical Flow Maps (CFMs). Think of this as a high-speed train that travels from a chaotic, noisy station directly to a specific, beautiful destination (a perfect sentence) in just a few stops, rather than chipping away slowly or laying bricks one by one.

The Problem: Scaling Up

Previous experiments with this "train" method worked well on small models (like a toy car), but no one knew if it could handle a massive, real-world train (a model with billions of parameters). The fear was that the math required to run the train would get too heavy and crash the system.

The Paper's Claim: The authors built a massive 1.7-billion-parameter model and proved that this "train" method works at a huge scale. They managed to generate high-quality text in as few as 4 steps (stops), whereas other methods might need hundreds.

How They Did It (The Recipe)

The authors didn't just turn on the machine; they had to tune the engine carefully. They used a two-stage process:

Stage 1: The Teacher (Pre-training)
First, they trained a standard model to understand how to move from noise to text. Think of this as a teacher learning the map. They tested over 350 different settings (like changing the fuel mix or the speed of the train) to find the perfect recipe. They found that mixing different time schedules and adjusting how much the model "listens" to its own mistakes was crucial.
Stage 2: The Student (Self-Distillation)
Once the teacher was ready, they used a technique called Self-Distillation. Imagine the teacher showing the student a shortcut: "Don't walk the whole path; just jump from the start to the finish."
- The student learns to predict the final destination directly from the noise, skipping the slow, step-by-step journey.
- The paper found that this "shortcut" allows the model to generate diverse, high-quality text in just 4 steps, while keeping the variety of words (entropy) high, so it doesn't sound robotic or repetitive.

The Results: Speed vs. Quality

The paper compares three main ways of generating text:

Autoregressive (The Brick Layer): Slow, but very accurate.
Standard Diffusion (The Sculptor): Can be fast, but often struggles with text quality.
Categorical Flow Maps (The High-Speed Train): The paper shows this method hits a "sweet spot."

Key Findings:

Speed: With the "shortcut" training, the model can produce text in 4 steps that is almost as good as the slow, brick-laying method.
Quality: The text generated is diverse and doesn't collapse into nonsense (a common problem where models just repeat the same word over and over).
Scoring: They created a new mathematical way to "grade" the model's confidence (likelihood), showing it performs competitively with other non-brick-laying methods.

The Challenges (The Bumps in the Road)

The authors are honest about the difficulties:

Memory Hunger: To run this "train," the computer needs to hold a massive map of every possible word at every position in the sentence. For a large model, this requires a lot of memory (they used 256 powerful GPUs to do it).
Tuning is Hard: They tried to use a "zero-shot" method (copying settings from a small model to a big one automatically), but it didn't work well. They still had to manually tune the big model, which is expensive and time-consuming.

Summary

This paper proves that Categorical Flow Maps are a viable, scalable alternative to the traditional "brick-by-brick" way of writing text. By training a massive model to take "shortcuts" from noise to text, they achieved high-quality generation in just a few steps, unlocking the potential for faster, more flexible language models.

Technical Summary: Scaling Categorical Flow Maps

Problem Statement

While autoregressive (AR) models dominate language modeling, continuous diffusion and flow matching models have demonstrated significant advantages in continuous modalities, including accelerated sampling and distribution tilting. Extending these benefits to discrete data has been a primary goal of recent research. Previous approaches, such as discrete diffusion, operate on discrete grids, limiting the application of established continuous techniques. Conversely, continuous-state processes (e.g., Categorical Flow Maps or CFMs) generate discrete data by evolving a probability distribution over a continuous simplex or $\mathbb{R}^d$ .

Despite recent demonstrations of CFMs achieving competitive sample quality in few-step regimes, their scalability remained unproven. It was unclear whether CFMs could scale to the billion-parameter, trillion-token regime required for modern Large Language Models (LLMs). A primary technical hurdle is the computational and memory overhead: for a vocabulary size $|V|$ and context length $L$ , CFMs require materializing $L \times |V|$ logit matrices at every step, which can be prohibitive when $|V| \approx 200,000$ and $L$ is large.

Methodology

The authors propose a two-stage training pipeline to scale CFMs to a 1.7B-parameter model trained on 2.1 trillion tokens.

1. Pretraining a Flow Model

The first stage involves training a standard flow matching model (a denoiser) to approximate the drift of the probability flow.

Architecture: A 1.7B-parameter transformer with AdaLN-Zero modulation, rotary position embeddings, and SwiGLU MLPs, following Lipschitz constraints for stability.
Objective: The model is trained to minimize a cross-entropy loss (Equation 2.10) between the predicted endpoint and the true data, effectively learning the vector field $v_\theta$ .
Design Space Exploration: Before scaling, the authors conducted an extensive hyperparameter sweep (350+ configurations) on a 400M proxy model. Key findings included:
- Time Scheduling: A mixture of linear and "error decoding" schedules ( $\alpha_t = (1-\lambda)t + \lambda\gamma_t$ ) with $\lambda \approx 0.75$ yielded optimal performance. Pure error decoding ( $\lambda=1$ ) caused performance collapse.
- Adaptive Loss Weighting: Reweighting the loss based on prediction error ( $w(x) = \text{sg}((\Delta(x) + c)^{-r})$ ) improved Generative Perplexity (Gen PPL) but required careful tuning to maintain distributional fidelity (MAUVE*).
- Random Unmasking: Training with a probability $p > 0$ of unmasking a random prefix was essential for generating coherent conditional text; without it, models produced out-of-distribution outputs.

2. Self-Distillation into a Categorical Flow Map (CFM)

The second stage transforms the pretrained flow model into a CFM capable of high-quality, few-step generation.

Mechanism: The model learns a flow map $X_{s,t}$ that transports probability mass from time $s$ to $t$ in a single step, rather than integrating the ODE over many steps.
Loss Functions: The authors evaluated Progressive Self-Distillation (PSD) and Stabilised Eulerian-logit Self-Distillation (ESD).
- PSD: Uses a continuous-time shortcut model approach.
- Divergence Choice: The authors compared Forward KL, Reverse KL, and Jensen-Shannon Divergence (JSD). They found that Forward KL and JSD improved token entropy and MAUVE* in the few-step regime, whereas Reverse KL caused instability and mode collapse.
- Stability: ESD training at the 1.7B scale proved unstable due to high-variance Jacobian estimates (computed via finite differences), leading the authors to focus on PSD results.
Sampling: The final CFM generates text by sampling from a prior and applying the learned flow map $X_{0,1}$ in a single pass (or few passes), significantly reducing inference time.

3. Likelihood Estimation

To evaluate CFMs on standard likelihood-based benchmarks (e.g., MCQA), the authors derived a Semi-discrete ELBO. This bound allows scoring answer choices without requiring the computationally expensive reverse integration of the ODE, which is particularly costly in high-dimensional discrete settings.

Key Contributions

Scalability Demonstration: The authors successfully trained a 1.7B-parameter flow model on 2.1T tokens and self-distilled it into a CFM, proving that CFMs can scale to the billion-parameter regime.
Hyperparameter Discovery: Through a systematic sweep on a 400M proxy, they identified robust settings for time scheduling, loss weighting, and input masking that transfer to larger scales.
Few-Step Generation: The self-distilled CFM generates diverse, high-quality text in as few as 4 inference steps (NFE) while maintaining near-data-level token entropy, outperforming the undistilled flow model in the low-NFE regime.
Likelihood Bound: The introduction of a semi-discrete ELBO enables the evaluation of CFMs on standard language modeling benchmarks, placing their performance in the same range as other non-autoregressive methods.
Limitations Analysis: The paper explicitly addresses the memory bottleneck of materializing $L \times |V|$ matrices and the challenges of zero-shot hyperparameter transfer for time-conditioned architectures.

Results

Quality vs. Steps: The self-distilled CFM achieves a Gen PPL of ~58.66 at 1024 steps (competitive with the AR baseline's 45.33) and a MAUVE* score of ~0.94. Crucially, at just 4 steps, the CFM achieves a MAUVE* of ~0.88, significantly outperforming the AR baseline (0.7) in the few-step regime.
Token Entropy: Self-distillation preserves high token entropy (diversity) even at low NFE, whereas undistilled models tend to collapse (low entropy) in the 1–2 step regime.
Benchmark Performance: On MCQA benchmarks (ARC-E, BoolQ, OBQA, etc.), the CFM achieves accuracies comparable to other non-autoregressive methods (e.g., Eso-LM) but remains below discrete diffusion methods trained directly on ELBO maximisation.
Conditional Generation: The self-distilled models show significant improvements in $\Delta$ PPL (conditional perplexity) in the few-step regime compared to the undistilled model.

Significance and Claims

The paper claims that Categorical Flow Maps are a viable path for discrete data generation at the scale of modern LLMs. By successfully scaling to 1.7B parameters, the authors demonstrate that the "flow map toolbox"—including guidance, distribution tilting, and reward fine-tuning—is now accessible for language modeling.

The work suggests that continuous-state generation can trade compute for sample quality, enabling high-quality text generation in significantly fewer function evaluations than traditional diffusion or AR models. While the authors note that likelihood estimation remains an open area for refinement (as the derived ELBO is a bound), the results indicate that CFMs can achieve competitive performance with discrete diffusion baselines while unlocking the benefits of continuous flow matching. The paper concludes that the scalability of CFMs opens new avenues for steering and fine-tuning language models beyond mere plausible output generation.