Loading post…
AppleProblem & solution
Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation
Summary
Discrete flow matching can parallelize text generation across positions, but reaching high quality may require hundreds or thousands of model evaluations. The paper argues that few-step distillation fails primarily because teacher trajectories are built from blind stochastic jumps: an early poor midpoint decision propagates, and the student is trained to imitate it. Its Trajectory-Shaped Discrete Flow Matching (TS-DFM) uses a lightweight energy compass during training to score candidate continuations at each midpoint and select the most coherent path, without adding inference cost. In experiments with 170M-parameter language models, an eight-step shaped student achieves 32% lower perplexity than a 1,024-step teacher and runs 128 times faster, with gains across source distributions and three evaluators; it also outperforms the compared discrete-generation baselines, including systems using more data or larger models.
Context
Autoregressive language models generate one token per forward pass, limiting throughput and increasing latency for long sequences. Diffusion language models parallelize across positions, but standard discrete diffusion may require hundreds to thousands of model evaluations. Distillation reduces the number of steps, yet the source argues that poor teacher trajectories, rather than insufficient student capacity, limit quality.
Approach / What changed
Trajectory-Shaped Discrete Flow Matching replaces blind stochastic jumps in the training trajectory with guided navigation. At each midpoint, a lightweight energy compass evaluates candidate continuations and selects the most coherent one. The shaping is applied only during training, so inference cost remains unchanged.
Takeaways
TS-DFM treats trajectory quality as the distillation bottleneck: early poor midpoint decisions can propagate through later steps and become targets for the student.
A lightweight energy compass scores candidate continuations at each midpoint during training, selecting more coherent transitions without increasing inference cost.
For a 170M-parameter language model, an eight-step shaped student has 32% lower perplexity than a 1,024-step teacher and is 128× faster.
FAQ
Why does the paper blame trajectories instead of student model capacity?
The teacher trajectories are created through blind stochastic jumps without evaluating sequence quality. A poor early decision propagates through subsequent steps, leaving the student to imitate a degraded result.
Does TS-DFM add computation when the model generates text?
No. The energy compass is used only while training to select coherent candidate continuations. The source states that inference cost is unchanged.
How does the shaped student compare with the teacher and other baselines?
At eight steps, it achieves 32% lower perplexity than the 1,024-step teacher and runs 128 times faster. It also achieves the best perplexity among the compared discrete-generation baselines, including methods using more data or larger models.
Source and freshness
- Source
- machinelearning.apple.com
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.