用能量导航优化生成路径,让模型用8步达到1024步教师的生成质量
Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation
- 用轻量级能量评分器引导每一步的生成选择,避免盲目跳跃
- 8步生成效果超越1024步教师,困惑度降低32%,速度提升128倍
- 适合追求高效高质文本生成的研究者和工程应用
离散流匹配通过逐步将噪声词元转化为连贯语言生成文本,但通常需要数百次前向传播。知识蒸馏利用多步轨迹训练学生模型在少数步骤内复现该过程。当学生表现不佳时,传统解释归因于容量不足。本文认为问题在于轨迹本身:训练轨迹由无质量评估的随机跳跃构成,早期错误会逐层传递,学生却必须模仿结果。为此提出轨迹塑形离散流匹配(TS-DFM),以轻量级能量导向器在每个中间点评估候选延续,选择最连贯的路径。所有塑形仅用于训练,推理成本不变。在1.7亿参数语言建模任务上,8步学生模型困惑度比1024步教师低32%,且速度提升128倍,跨不同数据分布和三个递增规模评估器均表现一致。TS-DFM在所有对比基线中达到最低困惑度,优于使用6倍数据或5倍大模型的方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。