轨迹即教师：通过能量导航蒸馏实现少步离散流匹配

用能量罗盘重塑离散流匹配的训练轨迹，让8步学生超越1024步教师

前置知识

离散流匹配 (Discrete Flow Matching, DFM)

DFM 是一种离散版的流匹配生成范式。它从一个无信息源序列（均匀随机 token 或 [MASK] token）出发，由速度场 v_θ 在 N 步迭代中把它逐步变换为有意义的文本 x_1。中间状态 x_t 处于源和目标之间的位置-级别混合。

DFM 是本文的基座方法，少步蒸馏的目标就是把它的多步生成压缩到几步。

少步蒸馏 (Few-Step Distillation)

蒸馏的核心思路是让学生模型模仿教师模型的多步生成轨迹，但在推理时只用很少的步数。FS-DFM 等方法使用四阶龙格-库塔 (RK-4) 积分器构造大步长目标，用半教师（学生参数的 EMA）作为轨迹估值器。学生训练信号完全来自这些轨迹，没有任何独立的监督信号。

本文聚焦于少步蒸馏中的关键瓶颈：轨迹本身的质量，而非学生容量。

RK-4 轨迹估计与盲跳跃

FS-DFM 用 RK-4 积分构造大步目标：依次计算 k_1, k_2, k_3, k_4 四个子步的 velocity logits。每个中间状态 x_mid 必须从 logits 经 CTMC 跳跃离散化成具体 token——即「盲跳跃」，对每个位置做随机抽样无任何质量评估，错误沿链路累积。

正是这三个盲跳跃使误差复利性放大，最终污染 v_RK-4 目标，这是本文要修正的核心症结。

CTMC (连续时间马尔可夫链) 跳跃

在离散流匹配中，把连续 velocity logits 转成离散 token 序列的过程由一个 Poisson-rate 过程驱动：每个位置 j 以 p_jump(j) 的概率发生转移，并从 velocity 的类别分布中采样一个新 token。这一过程是不可微的硬决策，无法用梯度反向传播来优化。

CTMC 跳跃的不可微性是为什么必须设计离散的候选选择策略（如 S2T 导航），而非做端到端梯度训练的原因。

能量模型 (Energy-Based Model)

能量模型 E_φ(x) 把一个序列映射为一个标量「能量」分数，分数越低表示质量越高。本文用 6 层 transformer + 注意力池化训练一个约 90M 参数的 E_φ，以 NCE 损失 + 速度近失负样本训练，使其能区分真实流状态与生成相关的损坏序列。

能量罗盘是本文的关键工具，承担「在多个候选跳转中挑出最好的那一个」的评分职能。

研究动机

现有的离散流匹配 (DFM) 通过数百至上千次前向传播迭代地把噪声 token 变成通顺文本，质量虽可与自回归 LLM 比肩，但速度慢了几个数量级。FS-DFM 等少步蒸馏方法虽然能把 1024 步压缩到 8 步，但蒸馏出来的学生质量仍受限于一个被忽视的瓶颈——训练轨迹本身。FS-DFM 在 RK-4 链中（目标为 $v_{\text{RK-4}} = \tfrac{1}{6}(k_1 + 2k_2 + 2k_3 + k_4)$）需要做三次盲跳跃：把 velocity logits 转成具体 token 序列时，对每个位置做一次不可逆的随机抽样，没有任何质量评估就被提交到下一步。一次坏抽样会污染后续 k_2、k_3 的 velocity 评估，最终 v_RK-4 目标吸收了三段累积漂移；学生只能忠实模仿这个被污染的目标，把错误学进权重里。当学生表现不佳时，业内通常归因于容量不足，但本文指出真正卡住性能上限的是轨迹质量。

本文的目标是本文要构建一个叫 TS-DFM (Trajectory-Shaped Discrete Flow Matching) 的框架，把 RK-4 中点处的盲跳跃替换为「导航式跳跃」——在每个中点处生成 K 个候选，由一个轻量级能量罗盘 E_φ 选出最连贯的那一个，最终让少步学生既能又快又显著超越少步基线 (FS-DFM)，又能反超 1024 步的全教师。具体可量化目标：170M 模型 8 步均匀源 GPT-2 困惑度从 87.6 降到 56.1，比 1024 步教师的 82.8 还低 32%，同时推理速度提升 128 倍。

与已有工作不同的是，已有工作对少步离散生成的优化要么改蒸馏机制本身 (如 SDTT 多轮自蒸馏、ReMDM 推理时重掩码)，要么改蒸馏损失 (T3D 反向 KL)，但都把训练轨迹当作不可变的给定数据。TS-DFM 的独特切入点是：训练轨迹本身也是优化对象——可以在不动蒸馏算法的前提下，仅通过干预轨迹构造来突破性能上限。换言之，已有方法问「怎么蒸馏」，本文问「让学生学什么」，把对中间状态的评估从无 (blind) 变为有 (compass-guided)，并以序列-令牌两阶段策略把评估成本控制在可承受范围内。

核心方法

直觉上，TS-DFM 像给学生请了个「家教」——老师在讲题时会跳过关键步骤直接猜答案（盲跳跃），学生只能模仿这些猜测。TS-DFM 让老师每跳一步前先列出 K 个候选走法，由「能量罗盘 $E_\phi$」挑出最像正解的那个；选完之后再逐个检查候选里的每个 token，把明显写错的字改正过来。这两步分别解决「选错方向」和「选对了方向但个别字错了」两类问题，统称「Sequence-to-Token (S2T) 导航」。技术路线：用一个约 90M 的小 transformer 训一个能量罗盘 $E_\phi$ E_φ；在 RK-4 中点 t ≥ τ（默认 0.2）时启动导航，生成 K=5 个温度多样化的 CTMC 跳跃候选，一次性打分选最优，再用 velocity 置信度做逐位 token 修正；最后用一次能量评估做安全网。整个流程只发生在训练时，推理时罗盘不参与。

TS-DFM 与 FS-DFM 的本质区别在于：FS-DFM 在 RK-4 链的每个中点做一次盲随机抽样后即向前推进，把误差留作既成事实；TS-DFM 把这一硬决策拆成「先选候选、再修细节」的两阶段策略，由一个在「生成相关损坏」而非「通用噪声」上训练的能量罗盘来指导选择。这是首次把离散流轨迹的中间状态本身（而非仅端点）作为质量优化目标，并把这一干预与蒸馏算法正交化——任何基于轨迹的蒸馏框架都可以叠加 S2T 导航。

方法步骤详情

第一步训练并冻结能量罗盘 E_φ：用 (x_0, x_1) 对按标准插值构造正样本 x_t，NCE 损失让 E_φ 给正样本打更低能量；加 L_reg 锚定能量尺度，加 L_order 用 hinge 损失强制能量随时步 t 单调下降。第二步 RK-4 中点执行 Navigate：若 t < τ 退化为盲跳跃；否则 Phase 1 生成 K=5 个温度多样化候选，一次 batched 通过 E_ϕ 选最低能量 x_best；Phase 2 复用 velocity logits 抽出 argmax 与置信度，按 Δ_align ∈ {-1,0,+1} 调制 p_(i)^jump = clamp(p_base · exp(Δ_align · conf_i · β_adp · α_soft), 0, 0.95);最后用一次能量评估做安全网，仅当 E_ϕ(x_refined) ≤ E_best + ε_safe 才接受修正。第三步训练学生：L = KL(s_ψ ‖ v̂_RK-4) 与 FS-DFM 相同，v̂_RK-4 由导航后公式给出。训练相对 FS-DFM 增加 2.0-2.4× 开销。

技术新颖性

技术新颖性体现在三点：(i) 首次把离散流轨迹的中间中点视为可优化的对象而非固定监督信号，把 trajectory quality 提升为和 student capacity 并列的设计维度；(ii) 设计了 Sequence-to-Token 的两阶段导航，把离散引导分解到不同粒度——序列级用 K 次能量评估选全局最优 (5 次)，token 级用 velocity 置信度做免费修正 (0 次额外前向)，避免位置分解的 |V|·L ≈ 5×10^7 次能量调用；(iii) 用「速度近失 (velocity near-miss)」负样本构造方法训练罗盘：把 x_t 退回到 x_{t-h} 再让冻结教师前向重建得到 x̂_t，把 x̂_t 当作负样本——这模拟了 CTMC 跳跃时教师本身的重建错误，正是导航需要识别的失败模式。最终罗盘在四种生成相关损坏上都达到 98.2-100% 的判别准确率，bin 级能量严格单调。

**Figure 2**: During training, standard trajectory construction compounds errors through blind stochastic jumps. Navigation shaping selects among candidates via an energy compass, activating at t ≥ τ.

**Figure 3**: TS-DFM training overview. For small steps, the frozen teacher provides the target directly. For large steps, the semi-teacher constructs targets via RK-4 with each midpoint guided by energy compass E_φ.

**Figure 4**: The granularity dilemma. Sequence-level scoring vs Position-factorized vs Sequence-to-Token (ours).

**Figure 5**: Token-phase: Velocity-aligned token correction.

**Figure 6**: Generation trajectory (mask source, 16 steps).

实验结果

主要发现：(1) 170M 模型 8 步生成显著超越 1024 步教师：均匀源 TS-DFM (FS-DFM init) GPT-2 困惑度 56.1，比 FS-DFM 的 87.6 低 36%、比 1024 步教师的 82.8 低 32%，推理速度 128×。(2) 掩码源改进更惊人：TS-DFM (DFM init) 8 步 GPT-2 PPL 91.4，比 FS-DFM 的 516.6 好 5.6×，略好于 1024 步掩码教师的 93.6——解决了 FS-DFM 在掩码源上的遗留问题。(3) 与同类 SOTA 对比：8 步 TS-DFM 的 56.1 比 Duo 的 70.6 低 19%（6× 数据），比 SDTT 7 轮的 105.8 低 46%（5× 大的模型）。(4) 1.3B 扩大领先：8 步 TS-DFM 48.0 vs FS-DFM 81.5 vs 教师 57.5；同一 90M 罗盘无需重训即可指导 1.3B 学生。(5) GSM8K 迁移：1.3B 模型 4 步 pass@8 = 9.61，超过 256 步教师的 9.12，节省 64× 推理计算。

**Table 1**: Main results across source distributions for 1024-token unconditional generation.

**Table 2**: Effect of time threshold τ. Lower τ activates navigation earlier, improving perplexity but risking diversity loss.

**Table 4**: Few-step generation comparison at 8 and 16 steps.

**Table 5**: Scaling to 1.3B parameters (uniform source).

**Table 17**: Effect of candidate count K (S2T, 8 steps).

**Table 19**: TS-DFM vs diffusion language models in the few-step prefix-conditioned regime.

**Figure 1**: Generation quality across step budgets, evaluators, and source distributions.

查看结构化数据

任务	指标	本文	基线	提升
170M 均匀源 8 步生成	GPT-2 Large PPL ↓	56.1 (TS-DFM, FS-DFM init)	87.6 (FS-DFM) / 82.8 (1024 步 DFM 教师)	-36% vs FS-DFM, -32% vs 教师
170M 均匀源 16 步生成	GPT-2 Large PPL ↓	47.2	73.2 (FS-DFM) / 82.8 (教师)	-36% vs FS-DFM
170M 掩码源 8 步生成	GPT-2 Large PPL ↓	91.4 (TS-DFM, DFM init)	516.6 (FS-DFM) / 93.6 (1024 步教师)	-82% vs FS-DFM (5.6× 提升), -2.4% vs 教师
170M 均匀源 8 步 vs Duo	GPT-2 Large PPL ↓	56.1	70.6 (Duo, 524B tokens)	-19%, 同时少用 6× 训练数据
170M 均匀源 8 步 vs SDTT-7	GPT-2 Large PPL ↓	56.1	105.8 (SDTT 7 轮, 0.86B 模型)	-46%, 同时模型小 5×
1.3B 均匀源 8 步生成	GPT-2 Large PPL ↓	48.0	81.5 (FS-DFM) / 57.5 (1024 步教师)	-41% vs FS-DFM, -16% vs 教师
GSM8K 1.3B pass@8	pass@8 (无偏估计)	9.61 (TS-DFM, 4 步)	9.12 (教师 256 步) / 4.84 (教师 4 步)	+5.4% vs 256 步教师, +99% vs 4 步教师, 节省 64× 推理计算
Prefix-conditional 16 步生成	MAUVE ↑	0.707 (TS-DFM 0.17B)	0.390 (FS-DFM) / ≤0.005 (Dream 7B, LLaDA 8B)	+81% vs FS-DFM, 比 40× 大的未蒸馏模型高两个数量级

局限与改进

作者明确承认两点：(1) 能量罗盘训练一次后即冻结，随着学生改善其流状态分布会偏移，可能降低罗盘有效性，迭代再训练可缓解但有协同适应风险；(2) 时间阈值 τ 在整段训练中是固定的，自适应或退火调度尚未探索。此外我自己观察到：(3) 在 τ=0.02 极端低阈值下生成熵会跌至 6.4 以下（<6.5 的多样性警戒线），出现过优化导致的重复，说明能量引导与多样性之间存在天然张力；(4) Mask 源上的 FS-DFM 初始化反不如 DFM 初始化（188.7 vs 91.4 PPL at 8 steps），暗示已被一次蒸馏锁定的速度模式可能束缚后续轨迹整形的能力；(5) 训练侧 2.0-2.4× 的额外开销对迭代实验是负担，对一次性的生产训练可接受但需在算力预算中预留；(6) 论文只在无条件 LM 和 GSM8K 上验证，未在更长文本摘要、代码补全等任务上检验轨迹整形是否依然鲁棒。

独立分析的弱点

独立分析的弱点：(1) 冻结罗盘的分布漂移问题在长训练中可能放大，因为学生学到的流状态越来越偏离 DFM 教师，但罗盘按早期分布训练——可周期性用学生最新中间状态对罗盘做增量自训练（类似 self-distillation 的 EMA 协同）。(2) 时间阈值 τ 的硬截断导致 t < τ 段无任何引导，但即便 t≈0.05 也有约 5% 位置被揭示；可改成 soft threshold 用 sigmoid 平滑过渡，避免边界阶跃。(3) 候选 K=5 在 GPU 内存吃紧时不够灵活；可探索 K=2 的轻量模式（附录 D.3 显示 PPL 仅比 K=5 高 9%），或对易混淆位置（命名实体、数字）给 $K=8$、明确位置（功能词）给 $K=1$。(4) 速度近失负样本需 9.1% 概率调用教师前向，可缓存或异步化教师重建，把额外开销压到 1.5× 以内。

未来方向

作者明确提出的方向：(a) 罗盘的迭代再训练 (compass refresh)，与学生协同进化；(b) 自适应或退火的 τ 调度，按训练进度或 batch loss 动态调整；(c) 探索把导航整形扩展到其他蒸馏框架 (SDTT 多轮、ReMDM 推理时整形)。基于成果可延伸的方向：(d) 把 S2T 思想推广到连续域扩散蒸馏，用 classifier guidance 替代离散能量罗盘；(e) 把能量罗盘用作推理时的可插拔 verifier，无需重训即可改善已有学生的输出 (附录 D 已证明推理时引导也能拿到 49% 改善)；(f) 把罗盘扩展到条件生成——例如用任务相关的负样本训练条件罗盘，实现可控文本生成；(g) 在指令微调、代码生成、对话等长上下文任务上验证，检验「每步权重高」的 reasoning 假设是否依然成立；(h) 与 RLHF/RLAIF 结合，让奖励模型作为罗盘的另一类生成相关负样本来源，进一步统一训练轨迹与人类偏好。

复现评估

开源情况：作者隶属 Apple 与 Ohio State 联合团队，论文未在正文明确给出代码仓库链接，需检查附录或后续 arXiv 版本。配置公开度：罗盘训练超参完整（AdamW β=(0.9, 0.999)、peak LR 1e-4、cosine decay、5000 步 warmup、weight decay 0.01、bf16、8 GPU）。数据：FineWeb-Edu 训练 + WikiText-103 验证均公开。算力：170M 实验单卡可跑 FS-DFM，TS-DFM 加上罗盘批处理需 4-8 张 A100/H100；1.3B 明显需多节点 8×H100。复现难度：中等偏上——主要门槛不在算法而在 (i) 罗盘训练的速度近失教师前向链路很慢 (ii) CTMC 跳跃数值稳定性调试（κ_t/(1-κ_t) 在小 t 处易爆）。建议先用 170M 均匀源复现表 1 的 56.1 PPL 做 sanity check，再叠加掩码源、1.3B、GSM8K。

轨迹即教师：通过能量导航蒸馏实现少步离散流匹配 Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation