Trajetória como Professor: Flow Matching Discreto de Poucos Passos via Destilação Guiada por Energia

Resumo

Discrete flow matching gera texto transformando iterativamente tokens de ruído em linguagem coerente, mas pode exigir centenas de passagens diretas. A destilação utiliza a trajetória de múltiplos passos para treinar um estudante a reproduzir o processo em poucos passos. Quando o estudante apresenta desempenho inferior, a explicação usual é capacidade insuficiente. Argumentamos o oposto: a trajetória é o gargalo, não o estudante. Cada trajetória de treinamento é construída por uma cadeia de saltos estocásticos cegos, sem avaliação da qualidade da sequência; uma única decisão ruim em um ponto intermediário inicial se propaga pelos passos subsequentes, mas o estudante deve imitar o resultado. Trajectory-Shaped Discrete Flow Matching (TS-DFM) substitui esses saltos cegos por navegação guiada: uma bússola de energia leve avalia continuações candidatas em cada ponto intermediário, selecionando a mais coerente. Toda a modelagem ocorre apenas no treinamento; o custo de inferência permanece inalterado. Em modelagem de linguagem com 170M de parâmetros, o estudante modelado em 8 passos alcança perplexidade 32% menor do que o professor com 1.024 passos, sendo 128 vezes mais rápido, com ganhos consistentes entre distribuições de origem e três avaliadores de escala crescente. TS-DFM atinge a melhor perplexidade entre todas as bases de geração discreta com as quais comparamos, incluindo métodos treinados com 6 vezes mais dados ou que utilizam modelos 5 vezes maiores.

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.