Resumo Técnico: Pós-treinamento de Equilíbrio de Trajetória para Modelos de Linguagem Difusivos

Problema: Bloqueio de Trajetória em RL Difusivo

Modelos de Linguagem Difusivos (dLLMs) oferecem uma alternativa promissora aos modelos autoregressivos para raciocínio e geração de código. No entanto, os métodos existentes de pós-treinamento para dLLMs adaptam majoritariamente objetivos de Aprendizado por Reforço (RL) de maximização de recompensa (por exemplo, PPO, GRPO) projetados para modelos autoregressivos. Os autores identificam um modo de falha central nessa adaptação chamado bloqueio de trajetória.

Nos modelos difusivos, a recompensa é atribuída apenas à conclusão final desruidificada ($y$), enquanto a distribuição do modelo é induzida sobre toda a trajetória de desruidificação ($\tau$). Como as recompensas terminais são "indiferentes ao caminho" — significando que qualquer trajetória que atinja a mesma $y$ correta recebe a mesma recompensa —, as atualizações de gradiente de política amostradas tendem a reforçar as trajetórias específicas que acabam sendo amostradas no início do treinamento. Isso cria um ciclo de retroalimentação auto-reforçador onde a massa de probabilidade colapsa sobre um conjunto estreito de caminhos de desruidificação. Consequentemente, mesmo que um modelo aprenda a produzir uma resposta correta, ele perde a capacidade de gerar um conjunto diversificado de soluções válidas (modos) sob amostragem repetida. Os autores formalizam isso via Teorema 1, provando que a cobertura de modos terminais é limitada superiormente pela diversidade de trajetória; assim, o bloqueio de trajetória limita diretamente a cobertura de soluções do modelo.

Metodologia: TraFL (Equilíbrio de Fluxo de Trajetória)

Para abordar o bloqueio de trajetória, os autores propõem TraFL, um objetivo de pós-treinamento fundamentado no princípio de equilíbrio de trajetória das Redes de Fluxo Generativo (GFlowNets). Em vez de maximizar a recompensa diretamente, o TraFL treina a política para corresponder a uma distribuição alvo inclinada pela recompensa ancorada em um modelo de referência congelado.

O Objetivo

A distribuição alvo $p^*(\tau | x)$ é definida como:
$$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$$
onde:

• $p_{\text{ref}}$ é um modelo de referência congelado.
• $r(x, y)$ é a recompensa terminal.
• $\beta$ controla a nitidez da recompensa.
• $Z(x)$ é um termo de normalização dependente do prompt (função de partição).

O objetivo de treinamento minimiza o resíduo quadrático entre o log-probabilidade do modelo atual e o alvo:
$$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$$
$$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$$

Implementação Prática para Difusão

Como os dLLMs não expõem log-probabilidades exatas de trajetória, os autores introduzem dois componentes-chave para tornar o TraFL prático:

1. Surrogado de Nível de Sequência Compatível com Difusão: Eles derivam um surrogado de log-probabilidade tratável ($\log \hat{p}_\theta(y|x)$) baseado em um limite inferior de reconstrução mascarada. Isso envolve amostrar um nível de máscara $l$, corromper $l$ tokens e calcular a média das log-probabilidades dos tokens reconstruídos. Um fator de $1/l$ é usado para normalizar a pontuação independentemente do número de tokens mascarados.
2. Normalização Aprendida Dependente de Prompt: Diferentemente de métodos que estimam $Z(x)$ via softmax sobre um lote (o que pode induzir concentração), o TraFL aprende uma função de partição $Z_\phi(x)$ via uma pequena MLP. Este termo recebe sinais de gradiente a cada etapa, independentemente de quais conclusões são amostradas, prevenindo a falha de normalização "local ao lote" observada em outras abordagens de correspondência de distribuição.

Contribuições Principais

1. Identificação do Bloqueio de Trajetória: O artigo identifica formalmente e prova que objetivos de maximização de recompensa em dLLMs levam a um colapso da diversidade de trajetória, o que, por sua vez, limita a cobertura de modos distintos de soluções corretas.
2. Objetivo TraFL: A proposta de um objetivo de equilíbrio de trajetória que ancora a política a um modelo de referência, separando o papel da recompensa (inclinando conclusões terminais) do modelo de referência (alocando massa de probabilidade sobre trajetórias plausíveis).
3. Surrogado Prático e Normalização: O desenvolvimento de um surrogado de nível de sequência compatível com difusão e um termo de normalização aprendido que permite treinamento estável sem probabilidades exatas de trajetória.
4. Superioridade Empírica: O TraFL é o único método de pós-treinamento avaliado que melhora consistentemente sobre o modelo base em todas as configurações de comprimento de benchmark (GSM8K, MATH-500, HumanEval, MBPP) e orçamentos de amostragem.

Resultados

Os autores avaliam o TraFL em benchmarks de raciocínio matemático e geração de código contra o modelo base (LLaDA-8B-Instruct) e baselines fortes (ESPO, JUSTGRPO).

• Desempenho In-Distribution: O TraFL melhora o Pass@k em todos os orçamentos de amostragem ($k \in \{1, 2, 4, 8, 16\}$) e temperaturas. Notavelmente, ele melhora o Pass@1 (precisão de amostra única) enquanto também escala melhor com orçamentos de amostragem aumentados, ao contrário de baselines que frequentemente sacrificam a precisão de amostra única para ganhos em múltiplas amostras.
• Generalização em Dados Retidos: Os ganhos transferem-se para benchmarks fora da distribuição. No Minerva Math, o TraFL permanece acima do modelo base. No LiveCodeBench, o TraFL é o método mais forte em todas as divisões de dificuldade (Fácil, Médio, Difícil), superando significativamente as baselines em problemas médios e difíceis.
• Análise de Diversidade: Usando um protocolo de LLM como juiz, os autores mostram que os ganhos do TraFL estão associados a uma cobertura mais ampla de estratégias de solução distintas (por exemplo, diferentes abordagens de prova ou decomposições algorítmicas) em vez de apenas afiar a pontuação de um único modo. Baselines como o ESPO, embora melhorem a precisão, não mostram ganhos significativos de diversidade sobre o modelo base, consistente com a hipótese de bloqueio de trajetória.

Significado e Alegações

O artigo alega que o TraFL fornece uma solução principiante para a falha de "busca de modo" do pós-treinamento padrão por RL para modelos difusivos. Ao mudar da maximização de recompensa para o equilíbrio de trajetória, o método preserva a diversidade de caminhos de desruidificação necessária para cobrir múltiplas soluções válidas.

Os autores afirmam que o TraFL é o único método avaliado que melhora sobre o modelo base em todas as configurações de comprimento de benchmark testadas. Eles enfatizam que as melhorias são robustas ao aumento do orçamento de amostragem e generalizam para conjuntos de dados retidos, sugerindo que o objetivo de equilíbrio de trajetória desacopla efetivamente o sinal de recompensa do mecanismo de alocação de caminho. O trabalho defende objetivos de correspondência de distribuição sobre a pura maximização de recompensa para pós-treinamento de dLLM, oferecendo um framework prático que evita os modos de falha de concentração inerentes às abordagens padrão de gradiente de política.

O artigo permanece modesto quanto às limitações, notando que o método depende de um surrogado aproximado de nível de sequência em vez de probabilidades exatas de trajetória e usa recompensas binárias grosseiras (correspondência exata ou aprovação de teste), o que pode limitar a eficiência de aprendizado em cenários com sucessos raros.