Discrete Beckmann transport models enable faster language modeling and reasoning
Discrete Beckmann Transport Models for One-Step Language Modeling and Reasoning
Machine Learning
Summary
Language models usually generate text one word at a time, which can be slow. The authors introduce a new approach called Discrete Beckmann Transport Models that can transform this process into a single step or just a few steps. This method no longer relies on a previous teacher model to learn, making training simpler and potentially faster. It improves the quality and accuracy of generated language and reasoning compared to existing methods.
What this means in practice
- •For machine learning engineers: Speed up language model generation by using fewer inference steps without losing quality.
- •For natural language processing developers: Improve reasoning and language generation systems with one-step or few-step models that do not depend on teacher models.
Authors
Sophia Tang, Shiyi Wang
Abstract
Discrete diffusion and flow models are a promising alternative to autoregressive language models, but compressing many-step sampling into fewer steps typically requires distilling a pretrained teacher model. This caps the student at the teacher's quality and requires a costly two-stage training pipeline. We introduce Discrete Beckmann Transport Models (DBTM), built on a time-independent flow whose autonomous transport map provably carries any point in the ambient space to a fixed point on the vertices of the simplex in a single step. We show that this fixed-point property is characterized by a conservation equation whose residual can be minimized directly from data, removing the requirement for a teacher flow and time conditioning. Under this construction, a partially trained map corresponds to the flow truncated at finite time, so generation reduces to iterating one map until it reaches a fixed point. We further extend the map to a partial-context interpolant where additional function evaluations act as refinement steps rather than ODE integration steps. On language modeling and reasoning tasks, DBTM enables one- and few-step generation that improves quality and accuracy over discrete diffusion and continuous flow baselines.