Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation
https://machinelearning.apple.com/research/trajectory-teacher-flow-matching📌 【Apple ML 研究】8 步打敗 1,024 步教師,蒸餾瓶頸竟然不在學生
TL;DR:Apple 團隊指出離散流匹配蒸餾失敗的元兇是訓練軌跡本身,而非學生模型容量不足。
當一個蒸餾後的學生模型表現不佳,直覺反應通常是「模型太小、學不會」。Apple ML 這篇論文提出相反的解讀:問題不在學生,而在教它的那條軌跡本身就是瞎子走出來的路。
🤔 離散流匹配的生成方式與它的蒸餾難題
離散流匹配(discrete flow matching)透過反覆迭代,把一串雜訊 token 逐步轉換成連貫的文字,但這個過程可能需要數百次前向傳遞,推理成本高昂。蒸餾的做法是利用教師模型走過的多步軌跡,訓練一個學生模型用更少步驟重現同樣的生成過程。論文作者觀察到:每一條訓練軌跡,其實是由一連串「盲目的隨機跳躍」串接而成,過程中完全沒有對序列品質做評估;只要早期某個中間點做出一個壞決策,這個錯誤就會沿著後續步驟一路傳遞下去,而學生模型卻被要求原樣模仿這條有缺陷的軌跡。
🧩 用「能量羅盤」導航,取代盲目跳躍
作者提出 Trajectory-Shaped Discrete Flow Matching(TS-DFM),核心思路是把教師軌跡中的盲目跳躍換成有引導的導航:在每個中間點,一個輕量級的「能量羅盤」(energy compass)會評估多個候選延續方向,選出其中最連貫的一個,再繼續往下走。值得注意的是,這種軌跡塑形完全發生在訓練階段,推理時的成本並未因此增加——學生模型在部署時的步數與速度不受影響,只有訓練用的教師軌跡品質變得更乾淨。
📊 8 步優於 1,024 步教師,還快 128 倍
在一個 170M 參數規模的語言模型上做語言建模實驗,經過軌跡塑形訓練的學生模型只用 8 步生成,困惑度(perplexity)就比用 1,024 步生成的教師模型低 32%,同時速度快了 128 倍。論文指出,這個優勢在不同來源分佈的資料,以及三種不同規模的評估器上都保持一致,顯示結果並非個別設定下的偶然。作者也將 TS-DFM 與其他離散生成類基準方法比較,結果顯示 TS-DFM 拿下所有比較對象中最佳的困惑度,即便對手是用了 6 倍訓練資料或 5 倍模型參數量的方法。
💡 軌跡品質可能比模型規模更關鍵
這篇論文最有意思的論點,是把「蒸餾效果不佳」的診斷方向從學生容量轉向教師軌跡的生成品質。如果一個壞的中間決策會像滾雪球一樣拖累整條軌跡,那麼與其一味放大學生模型或增加訓練資料,不如先確保用來教學生的軌跡本身是「乾淨」的。這也呼應了作者相關工作 FS-DFM 中提到的觀察:擴散語言模型雖然能平行生成,但要壓低步數同時維持品質,往往卡在生成路徑的穩定性,而不只是模型架構本身。
🎯 實務啟示
對於正在做模型蒸餾或探索少步數生成的團隊,這篇研究提示了一個容易被忽略的槓桿點:與其只調整學生模型的架構或訓練目標,先檢查教師軌跡的生成品質是否被有效評估與篩選,可能是更划算的投資方向。
🔗 來源
- 標題: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(Apple ML)
- 連結:https://machinelearning.apple.com/research/trajectory-teacher-flow-matching
#AppleML #DiffusionLanguageModel #KnowledgeDistillation #DiscreteFlowMatching #NLP #GenerativeAI #ModelCompression #LLM #TextGeneration #MachineLearning
原始資料 Apple ML · 收集於 2026-09-17
摘要原文
research area Methods and Algorithms , research area Speech and Natural Language Processing content type paper published September 2026 Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation Authors Amin Karimi Monsefi†**, Dominic Culver, Nikhil Bhendawade, Manuel R. Ciosici, Yizhe Zhang, Irina Belousova View publication Copy Bibtex 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 128× 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 6× more data or using 5× larger models. † The Ohio State University ** Work done while at Apple Related readings and updates. FS-DFM: Fast and Accurate Long Text Generation with Few-Step Diffusion Language Models October 13, 2025 research area Speech and Natural Language Processing conference ICLR Autoregressive language models (ARMs) deliver strong likelihoods, but are inherently serial: they generate one token per forward pass, which limits throughput and inflates latency for long sequences. Diffusion Language Models (DLMs) parallelize across positions and thus appear promising for language generation, yet standard discrete diffusion typically needs hundreds to thousands of model evaluations to reach high quality, trading serial depth… Read more Distillation Scaling Laws July 1, 2025 research area Methods and Algorithms , research area Speech and Natural Language Processing conference ICML We propose a distillation scaling law that estimates distilled model performance based on a compute budget and its allocation between the student and teacher. Our findings mitigate the risks associated with large-scale distillation by enabling compute-optimal allocation for both the teacher and student to maximize student performance. We provide compute-optimal distillation recipes for two key scenarios: when a teacher already exists, and when a… Read more Discover opportunities in Machine Learning. Our research in machine learning breaks new ground every day. Work with us
由 claude-code/sonnet 自動生成