軌跡を教師とする：エネルギー誘導蒸留による数ステップ離散フローマッチング

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

May 8, 2026

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

cs.AI

要旨

離散的フローマッチングは、ノイズトークンを反復的に変換して一貫性のある言語を生成するが、数百回のフォワードパスを必要とする場合がある。蒸留では、多段階の軌跡を利用して、そのプロセスをわずかなステップで再現する生徒モデルを訓練する。生徒モデルの性能が不十分な場合、通常は容量不足が原因とされる。本研究では逆の主張を行う。すなわち、ボトルネックは生徒モデルではなく軌跡にある。各訓練軌跡は、系列品質の評価を伴わない盲目的な確率的ジャンプの連鎖によって構築され、初期の中間地点での単一の誤った決定が後続のステップに伝播する一方で、生徒モデルはその結果を模倣しなければならない。軌跡整形型離散的フローマッチング（TS-DFM）は、これらの盲目的なジャンプを誘導ナビゲーションに置き換える。軽量なエネルギーコンパスが各中間地点で候補となる継続系列を評価し、最も一貫性のあるものを選択する。すべての整形は訓練時のみ行われ、推論コストは変わらない。1億7千万パラメータの言語モデリングにおいて、整形された生徒モデルは8ステップで、1024ステップの教師モデルよりも32%低いパープレキシティを達成し、かつ128倍高速であり、その利得はソース分布や3種類の評価者（規模が増大）において一貫している。TS-DFMは、6倍多くのデータで訓練された手法や5倍大きなモデルを用いた手法を含む、比較対象としたすべての離散生成ベースラインの中で最高のパープレキシティを達成する。

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.