PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade, Manuel R. Ciosici, Yizhe Zhang, Irina Belousova

arXiv:2605.07924v1cs.LGcs.AIcs.CL

TL;DR

Discrete flow matching is slow because it relies on long refinement trajectories, and distillation can inherit errors from blind stochastic midpoint decisions. TS-DFM navigates those intermediate states with an energy compass during training, producing a few-step student that outperforms the 1,024-step teacher while preserving inference cost. At 170M parameters, it achieves 32% lower perplexity than the teacher at 8 steps and a 128× speedup, with state-of-the-art few-step performance.

  • Problem

    Discrete flow matching requires hundreds or thousands of refinement steps, while distillation can force students to imitate trajectories degraded by compounding blind stochastic-jump errors.

  • Method

    TS-DFM replaces blind midpoint jumps with candidate generation and energy-guided navigation, using a Sequence-to-Token policy and time threshold τ during training.

  • Results

    At 170M parameters, TS-DFM at 8 steps reaches 56.1 GPT-2 perplexity, 32% below the 1,024-step teacher and 128× faster, achieving state-of-the-art few-step performance.

  • Takeaways & Limitations

    Trajectory quality, rather than student capacity, sets the few-step DFM performance ceiling, while training-only shaping preserves the original inference procedure and cost.

  • Takeaways & Limitations

    The frozen compass may become less effective as the student’s flow-state distribution shifts, and fixed threshold τ schedules remain unexplored.

Abstract

from arXiv · show

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.

1 Introduction

Discrete flow matching can generate high-quality text but is slow because it refines sequences over hundreds or thousands of steps. TS-DFM treats training trajectories as the bottleneck and navigates intermediate states with an energy-guided selection procedure, improving few-step quality without adding inference cost.

  • Hundreds or thousands of refinement steps make discrete flow matching orders of magnitude slower than autoregressive language models.
  • Distillation trains a student from iterative trajectories, but blind stochastic midpoint decisions can compound errors that the student faithfully imitates.The resulting trajectory, rather than student capacity, is identified as the performance bottleneck.
  • TS-DFM generates multiple midpoint continuations and uses a lightweight energy compass with a Sequence-to-Token policy to select and refine the highest-quality candidate.Navigation activates after a time threshold τ, when sufficient linguistic structure has emerged.
  • 98.5% corruption-detection accuracy is achieved by the generation-aware energy compass on partially revealed flow states.
  • 56.1 GPT-2 perplexity at 8 steps is achieved by TS-DFM on 170M-parameter language modeling with a uniform source, 36% below FS-DFM and 32% below the 1,024-step teacher.The shaped student is also 128× faster, while training-only shaping leaves inference unchanged.

2 Related work

Prior work compresses iterative generation through distillation, self-distillation, remasking, or related consistency strategies. TS-DFM instead optimizes the trajectories that provide supervision, using an energy model to select higher-quality partial states.

  • Discrete diffusion and flow-based generation use iterative refinement, with quality degrading sharply when sampling steps are reduced.
  • Prior few-step methods include iterative self-distillation, inference-time remasking, and duality-based approaches, each with distinct training or inference trade-offs.SDTT requires multiple training rounds and a large base model, while ReMDM adds per-step inference overhead.
  • TS-DFM addresses what the student learns from rather than how distillation is performed, optimizing discrete trajectories rather than only changing the distillation mechanism.
  • An energy model scores partially revealed flow states and selects the highest-quality candidate during trajectory construction.This differs from prior uses focused on complete-sequence reranking or controlled generation.

3 Background and motivation

DFM distillation compresses long stochastic trajectories into few-step student predictions, but the student inherits errors from trajectory construction. In RK-4, un evaluated midpoint jumps can compound drift into the targets used for supervision.

  • Background: DFM transforms source sequences into coherent text through iterative velocity-field updates over N steps.The source may consist of uniform random tokens or mask tokens, and intermediate states mix source and target tokens according to α(t).
  • Few-step distillation: Distillation trains a student to reproduce long teacher trajectories in a few steps, using RK-4 targets for large step sizes.The frozen teacher supplies direct targets for small steps, while the EMA semi-teacher constructs multi-step targets for larger steps.
  • Trajectory construction: Each RK-4 midpoint is formed by a stochastic discrete jump without evaluating the resulting sequence’s quality.A Poisson-rate process determines position transitions, and categorical sampling selects replacement tokens.
  • Trajectory construction: A poor early midpoint changes later velocity evaluations, so three chained blind jumps accumulate drift in the final RK-4 target.The student receives no independent data or alternative error signal and therefore learns to match the drifted target.
  • Training overview: Figure 3 distinguishes direct frozen-teacher targets for small steps from semi-teacher RK-4 targets with guided midpoints for large steps.TS-DFM-specific components are shown separately from components inherited from FS-DFM.

4 Method: navigation shaping

TS-DFM replaces blind midpoint jumps with navigation shaping: an energy compass selects coherent candidates, then velocity-based token refinement improves local decisions. Guidance activates only when flow states contain enough linguistic signal, while the student objective and inference procedure remain unchanged.

  • Overview: TS-DFM replaces blind stochastic midpoint jumps with an energy compass, Sequence-to-Token navigation, and a time threshold.These components shape training trajectories while building on FS-DFM.
  • Energy compass: The energy compass scores raw token sequences without time conditioning, treating sequence quality as intrinsic to content.It is trained on generation-aware negatives because random token replacements produce trivial, unrepresentative corruptions.
  • Energy compass: 98.5%–99.8% accuracy distinguishes real flow states from generation-relevant corruptions on held-out FineWeb-Edu.The energy objective combines noise contrastive learning with regularization and time-ordering terms.
  • Sequence-to-Token navigation: Navigate first selects the lowest-energy candidate from K sequence-level CTMC continuations, then refines individual tokens using velocity confidence.A safeguard accepts the refinement only when its energy does not exceed a tolerance.
  • Sequence-to-Token navigation: The sequence and token phases are complementary: at 64 steps, the full policy reaches 38.0 PPL versus 52.0 for sequence-only and 47.6 for token-only.Candidate diversity and the safeguard support meaningful selection and controlled refinement.
  • Navigated RK-4: Navigated RK-4 changes midpoint construction while preserving the semi-teacher’s velocity computation and the rest of the distillation setup.Navigated velocities and quality-checked midpoints are distinguished from their blind counterparts with separate notation.
  • Navigated RK-4: Quality-checked transitions can improve the target multiplicatively across the three RK-4 midpoints when each guided jump reduces error by a factor ρ < 1.The stated approximation is ∼ρ^3 relative to blind jumps.
  • Time-thresholded guidance: Navigation activates only at t ≥ τ, leaving source-dominated states to blind exploration and guiding states after linguistic structure emerges.Different RK-4 midpoints can cross the threshold within one estimation, and the student still minimizes KL divergence against the navigated target.

5 Experiments

Experiments evaluate TS-DFM across source distributions, evaluators, step budgets, candidate counts, training overhead, and model scale. Trajectory shaping consistently improves few-step perplexity, including results that surpass the 1,024-step teacher and other discrete-generation baselines.

  • TS-DFM surpasses the 1,024-step DFM teacher at 8 inference steps across both uniform and mask source distributions.
  • Uniform source: 56.1 GPT-2 PPL at 8 steps is 36% below FS-DFM and 32% below the 1,024-step teacher for the uniform source.
  • Mask source: 5.6× improvement over FS-DFM is reported for the mask source at 8 steps, while TS-DFM with DFM initialization beats the mask teacher at 8 and 16 steps.
  • Time threshold: τ=0.2 provides the best reported quality–diversity trade-off, while τ=0.02 improves raw perplexity but reduces entropy below 6.5.
  • State-of-the-art comparison: 56.1 GPT-2 PPL at 8 steps is 19% lower than Duo and 46% below SDTT, despite those methods using 6× more tokens or a 5× larger model.
  • Training cost: 2.2× training overhead is incurred at the default threshold, within a 2.0–2.4× range, while inference cost remains unchanged.
  • Scaling: 48.0 GPT-2 Large PPL at 1.3B parameters beats the 1,024-step teacher’s 57.5 with a 128× inference speedup.

6 Conclusion

TS-DFM treats training trajectories, rather than student capacity, as the bottleneck in few-step discrete flow distillation. It uses guided navigation during training and reports strong results across model sizes and source distributions, with limitations remaining in compass adaptation and threshold scheduling.

  • Conclusion: TS-DFM replaces blind RK-4 midpoint jumps with guided navigation using a lightweight energy compass and a two-phase Sequence-to-Token policy.The shaping occurs during training only.
  • Conclusion: At 170M parameters, TS-DFM achieves state-of-the-art few-step discrete generation quality on both uniform and mask sources.
  • Conclusion: At 1.3B parameters, the shaped student beats its 1,024-step teacher at 8 steps, delivering a 128× speedup with improved quality.
  • Conclusion: The PPL reduction over FS-DFM widens from 36% to 41% as model scale increases from 170M to 1.3B.
  • Limitations: The frozen compass may become less effective as the student’s flow-state distribution shifts, while adaptive threshold schedules remain unexplored.The threshold τ is fixed throughout training.

A.2 Training objective

The energy compass is trained as a compact sequence scorer with a three-part objective and generation-aware negatives. Training separately handles uniform and mask source distributions while checkpoint selection balances discrimination and temporal monotonicity.

  • Training objective: The total objective combines noise-contrastive, energy-regularization, and time-ordering losses with λreg = λorder = 1.0.The components target corruption discrimination, stable energy scale, and cross-timestep ordering.
  • Training objective: The noise-contrastive loss trains lower energy for real flow states than for corrupted negatives.
  • Negative construction: Generation-realistic negatives model teacher reconstruction mistakes and under-revealed states, while token-corruption negatives target wrong or repeated revealed tokens.Generation-realistic negatives account for approximately 45% of training negatives.
  • Checkpoint selection: Checkpoint selection equally weights within-timestep NCE accuracy and the fraction of correctly decreasing consecutive time-bin energies.Both discrimination and monotonicity are required for navigation.

B.1 Time-energy monotonicity

The compass is evaluated on whether energy decreases as more tokens are revealed along the flow. Both source-specific variants achieve perfect bin-level monotonicity, with high but imperfect sample-level smoothness and mild violations.

  • Criterion: As t increases and more tokens are revealed, the compass should assign monotonically lower energy: Eϕ(xt) < Eϕ(xt′) for t > t′.
  • Evaluation: The evaluation uses 200 held-out source-target pairs, 20 sampled timesteps, and the clean target, producing 4,200 evaluations per source distribution.
  • Results: Both source distributions show strictly decreasing mean energy across all 10 consecutive time-bin pairs.
  • Results: The uniform compass reaches 97.2% sample-level smoothness, while the mask compass reaches 94.0%.Violations occur in 44 and 74 of 200 samples, respectively.
  • Results: The largest wrong-direction energy gaps are 1.14 for the uniform compass and 1.40 for the mask compass, while bin-level monotonicity remains perfect.The violations are concentrated in a small number of outlier samples.

B.2 Generation-relevant corruption detection

The compass detects generation-relevant corruptions and supports navigation across sequence and token granularities. Sequence-level selection is affordable but can miss bad tokens, while token refinement adds local correction and an energy safeguard.

  • Corruption detection: The compass ranks real flow states below corrupted states using random replacement, frequency-matched replacement, token repetition, and time downstep.These corruptions represent content, repetition, and reduced-revelation failure modes.
  • Corruption detection: The uniform compass achieves 98.5% mean corruption-ranking accuracy, while the mask compass achieves 99.8%.The uniform variant has no strategy below 98%; the mask variant is perfect on three strategies and reaches 99.4% on token repetition.
  • Corruption detection: Both compasses handle wrong tokens, repetition, and reduced revelation well, indicating a general sequence-quality signal rather than a single-failure detector.
  • Corruption detection: Frequency-matched replacements perform nearly identically to random replacements, preventing reliance on unigram frequency cues.
  • Navigation design: Sequence-level navigation uses K energy calls but cannot correct individually poor tokens, whereas position-factorized guidance requires |V| × L calls per midpoint.
  • Navigation design: Phase 2 uses velocity alignment for token refinement at zero additional model-inference cost, with guidance intensity decaying from β0 at t=0 to β0/2 at t=1.
  • Energy safeguard: The Phase 2 safeguard accepts a refined sequence only when its energy remains within the allowed tolerance, otherwise retaining the best sequence.Velocity alignment is a proxy for, but not identical to, sequence-level energy.

C.6 Why both phases are necessary

The two navigation phases address different trajectory-construction failures: global sequence selection chooses a coherent region, while token refinement corrects local weaknesses. Neither phase alone is sufficient, and training-time shaping adds overhead without changing inference.

  • Navigation procedure: The navigation procedure uses candidate exploration, token refinement, and an energy safeguard, with guidance applied at navigated midpoints.The supplied algorithm passages specify thresholding, candidate generation, token refinement, and safeguard acceptance.
  • Phase 1: sequence selection: Sequence selection alone can choose a globally low-energy candidate that still contains weak tokens, allowing local errors to compound across later midpoints.Its holistic score cannot isolate individually weak tokens.
  • Phase 2: token refinement: Token refinement alone lacks access to the energy landscape and may polish a locally coherent but globally incoherent candidate.Velocity confidence is a proxy for token quality, not a substitute for sequence-level energy.
  • Together: global selection + local correction: The combined policy first selects the best global continuation, then applies targeted token-level corrections within that region.The energy safeguard prevents accepted refinement from worsening the selected state.

D.1 Navigation policy comparison across step counts

Navigation policies show complementary behavior across step counts: sequence selection helps most at very low budgets, while the full Sequence-to-Token policy dominates as steps increase. Guided policies avoid the high-step degradation of unguided stochastic generation.

  • Low-step regime: At 8 steps, sequence selection provides the largest benefit, with the best policy reducing GPT-2 Large perplexity from 89.6 to 81.2.The corresponding LLaMA-2 7B reduction is from 58.2 to 52.2.
  • Higher-step regime: Starting at 16 steps, Sequence-to-Token consistently achieves the best or near-best perplexity across all three evaluators.Token refinement has enough steps to accumulate useful local corrections.
  • Higher-step regime: At 32 steps, Sequence-to-Token reduces GPT-2 Large perplexity by 39% relative to no guidance, from 70.6 to 43.1.It outperforms sequence-only at 53.3 and fine-only at 46.8.
  • Higher-step regime: At 64 steps, Sequence-to-Token reaches 38.0 GPT-2 Large perplexity, 49% below unguided generation and 20% better than fine-only.The comparison is in the regime most relevant to RK-4 midpoint construction.
  • Unguided degradation: Unguided GPT-2 Large perplexity worsens from 70.6 at 32 steps to 75.0 at 64 steps, whereas Sequence-to-Token improves from 43.1 to 38.0.The passage attributes the unguided deterioration to compounding stochastic-jump errors.
  • Token refinement behavior: Refinement becomes more selective as step count increases, replacing 7.8% of tokens at 8 steps versus 0.8% at 64 steps.Safeguard acceptance rises from 94.7% to 99.9%, and rejected refinements have positive energy changes.
  • Candidate count: Increasing candidate count has diminishing returns: K=5 reaches 66.2 LLaMA-3 8B perplexity, while K=8 provides no additional benefit at 70.6.K=2 reaches 88.7 and K=3 reaches 75.7, relative to 115.9 without guidance.
  • Cost: Inference-time guidance is costly, but the actual method applies shaping during training, leaving the distilled model’s inference speed unchanged.Inference-time guidance produces an 8.5× slowdown at K=5, unlike the training-only approach.

E.2 Midpoint diversity analysis

Midpoint stochastic jumps retain enough diversity for candidate selection to matter, and this diversity is largely stable across solver step counts. The analysis connects variable positions to the roles of global selection and local refinement.

  • Experimental setup: The diversity study launches 10 independent completions from each shared midpoint across 48 independent prefixes.It reports average diversity metrics across the completions.
  • Diversity across step counts: Midpoint diversity profiles are nearly identical at 8 and 64 steps, indicating that diversity is a property of flow structure rather than solver granularity.The measurements span seven midpoint times.
  • Diversity across midpoint times: At tmid=0.2, independent completions differ at approximately 147 of 1,024 token positions, or 14.4% edit distance.This is the activation threshold used in the main experiments.
  • Diversity across midpoint times: At tmid=0.3, completions still differ at approximately 98 positions, or 9.6% edit distance.The residual disagreement supplies candidate variation for selection.
  • Diversity across midpoint times: Even at tmid=0.5, approximately 32 positions vary across completions, while residual variation falls to approximately 13 positions at t=0.7 and 4 at t=0.9.The paper identifies these variable positions as consequential stochastic choices.
  • Implications for navigation shaping: The sequence phase selects globally across variable positions, while token refinement targets individual positions where stochastic jumps likely introduced errors.This division gives the two phases distinct roles in using midpoint diversity.
  • Scaling context: At 1.3B, trajectory shaping yields a 41% 8-step reduction over FS-DFM and beats the 1,024-step teacher at 8 steps with a 128× inference speedup.The same 90M compass is used at both model scales.
  • Scaling context: At 1.3B parameters, TS-DFM retains generation entropy from 7.0 to 7.6 across 4–32 steps, above the 6.5 repetition-risk threshold.The accompanying result states that quality gains therefore do not reflect entropy collapse.

E.4 Comparison to diffusion language models in the few-step regime

In the few-step prefix-conditioned regime, undistilled diffusion language models collapse despite much larger parameter counts, whereas distilled TS-DFM improves on FS-DFM. TS-DFM reaches substantially better perplexity and MAUVE without diversity collapse.

  • Undistilled diffusion models: Across 4–16 steps, Dream and LLaDA have perplexity above 400, entropy of 0.42–2.31, and MAUVE near 0.005.Their outputs are described as essentially indistinguishable from random in this regime.
  • Distilled-model comparison: TS-DFM outperforms FS-DFM on perplexity at every step count, reducing PPL by 29% at 4 steps, 25% at 8 steps, and 20% at 16 steps.The reported pairs are 97.07 →68.95, 75.78 →57.11, and 67.42 →53.97.
  • Distilled-model comparison: At 16 steps, TS-DFM improves MAUVE from 0.390 to 0.707, a 1.8× improvement over FS-DFM.The gain places generations closer to the natural-text distribution under the reported metric.
  • Diversity: TS-DFM entropy remains in the 7.70–7.86 range, indicating that its quality gains occur without diversity collapse.The only reported exception is a marginal FS-DFM MAUVE lead at 4 steps.
  • Parameter-efficiency comparison: At 16 steps, TS-DFM reaches PPL 54 and MAUVE 0.71, while Dream and LLaDA remain at PPL ≥432 and MAUVE ≤0.005.TS-DFM uses 0.17B parameters versus 7B and 8B for the undistilled comparisons.
  • Overall comparison: The comparison separates distilled from undistilled few-step behavior: large undistilled models collapse, while TS-DFM extends the distilled family’s advantage with 20–29% PPL reductions.The supplied summary also reports a near-doubling of MAUVE at higher step counts.

E.5 Generalization to mathematical reasoning at scale

TS-DFM transfers trajectory shaping from 170M-parameter unconditional language modeling to 1.3B-parameter GSM8K reasoning without task-specific modification. At 4 and 8 steps, it surpasses the 256-step teacher on pass@8 while using substantially less inference compute.

  • TS-DFM at 4 steps reaches 9.61 pass@8 and at 8 steps reaches 9.17, both above the 256-step teacher’s 9.12.
  • 32×–64× less inference compute accompanies TS-DFM’s higher pass@8 accuracy than the full-step teacher.
  • TS-DFM retains 92–93% of the teacher’s full-step quality on pass@32, reaching 25.47 at 4 steps versus the teacher’s 27.52.
  • At 8 steps, TS-DFM improves over the same-step teacher by +61%, +35%, and +20% on pass@8, pass@16, and pass@32.
  • TS-DFM outperforms FS-DFM at every reported step-count and pass@k setting, with the largest gains at the lowest budgets.
  • The same energy-navigated RK-4 procedure transfers to GSM8K at 1.3B parameters without methodological modification, a task-specific compass, or per-task tuning.

F.1 Generated samples

The generated samples and trajectory visualization provide qualitative inspection of TS-DFM outputs, source settings, inference steps, and progressive token revelation. The shown text is often imperfect, while the trajectory moves from masks toward increasingly contextual completions.

  • The sample metadata records source distribution, generation method, and inference-step count for each output.
  • The qualitative outputs contain fragmented, repetitive, and semantically incoherent passages across medical, dental, plant, and educational topics.
  • The displayed TS-DFM samples include uniform-source generation at 8 steps and mask-source generation at 16 steps.
  • Figure 6 uses gray blocks for unrevealed positions, blue bold tokens for newly revealed positions, and black text for previously revealed tokens.
  • The mask-source trajectory first reveals scattered function and content words, then fills context-dependent tokens before resolving remaining ambiguities.
Loading 2605.07924v1…