Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Monsefi, Amin Karimi, Culver, Dominic, Bhendawade, Nikhil, Ciosici, Manuel R., Zhang, Yizhe, Belousova, Irina
Formato: Preprint
Publicado: 2026
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866910201855082496
author Monsefi, Amin Karimi
Culver, Dominic
Bhendawade, Nikhil
Ciosici, Manuel R.
Zhang, Yizhe
Belousova, Irina
author_facet Monsefi, Amin Karimi
Culver, Dominic
Bhendawade, Nikhil
Ciosici, Manuel R.
Zhang, Yizhe
Belousova, Irina
contents Discrete flow matching generates text by iteratively transforming noise tokens into coherent language, but may require hundreds of forward passes. Distillation uses the multi-step trajectory to train a student to reproduce the process in a few steps. When the student underperforms, the usual explanation is insufficient capacity. We argue the opposite: the trajectory is the bottleneck, not the student. Each training trajectory is built through a chain of blind stochastic jumps with no evaluation of sequence quality; a single bad decision at an early midpoint propagates through subsequent steps, yet the student must imitate the result. Trajectory-Shaped Discrete Flow Matching (TS-DFM) replaces these blind jumps with guided navigation: a lightweight energy compass evaluates candidate continuations at each midpoint, selecting the most coherent. All shaping is training-only; inference cost is unchanged. On 170M-parameter language modeling, the shaped student at 8 steps achieves 32% lower perplexity than the 1,024-step teacher while being 128x faster, with gains consistent across source distributions and three evaluators of increasing scale. TS-DFM achieves the best perplexity of any discrete-generation baseline we compare against, including methods trained on 6x more data or using 5x larger models.
format Preprint
id arxiv_https___arxiv_org_abs_2605_07924
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation
Monsefi, Amin Karimi
Culver, Dominic
Bhendawade, Nikhil
Ciosici, Manuel R.
Zhang, Yizhe
Belousova, Irina
Machine Learning
Artificial Intelligence
Computation and Language
Discrete flow matching generates text by iteratively transforming noise tokens into coherent language, but may require hundreds of forward passes. Distillation uses the multi-step trajectory to train a student to reproduce the process in a few steps. When the student underperforms, the usual explanation is insufficient capacity. We argue the opposite: the trajectory is the bottleneck, not the student. Each training trajectory is built through a chain of blind stochastic jumps with no evaluation of sequence quality; a single bad decision at an early midpoint propagates through subsequent steps, yet the student must imitate the result. Trajectory-Shaped Discrete Flow Matching (TS-DFM) replaces these blind jumps with guided navigation: a lightweight energy compass evaluates candidate continuations at each midpoint, selecting the most coherent. All shaping is training-only; inference cost is unchanged. On 170M-parameter language modeling, the shaped student at 8 steps achieves 32% lower perplexity than the 1,024-step teacher while being 128x faster, with gains consistent across source distributions and three evaluators of increasing scale. TS-DFM achieves the best perplexity of any discrete-generation baseline we compare against, including methods trained on 6x more data or using 5x larger models.
title Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2605.07924