Resumo Técnico: Modelagem Boltzmann Consciente de Trajetória Auto-Distilada (TABOM)

1. Declaração do Problema

Modelos de Linguagem de Difusão (DLMs) oferecem vantagens sobre modelos autoregressivos, como consciência global superior e geração de tokens altamente paralela. No entanto, o pós-treinamento de DLMs via Ajuste Fino Supervisionado (SFT) padrão enfrenta uma discrepância crítica entre treinamento e inferência:

• Treinamento: O SFT padrão utiliza um objetivo de Limite Inferior Negativo da Verossimilhança (NELBO) com mascaramento aleatório uniforme, exigindo que o modelo reconstrua todos os tokens mascarados simultaneamente em uma única etapa. Isso instila um "viés indutivo uniforme", tratando todos os tokens como igualmente difíceis, independentemente de sua previsibilidade inerente.
• Inferência: A inferência de DLM segue uma trajetória "do fácil para o difícil" guiada por confiança e multi-etapa. Os tokens são desmascarados iterativamente, começando com aqueles que possuem a menor entropia preditiva (maior certeza) e progredindo para tokens mais incertos.

O Dilema:
Tentativas recentes de fechar essa lacuna usando trajetórias auto-distiladas (onde o modelo aprende a partir de seus próprios rastros de decodificação) produziram resultados limitados.

1. SFT Padrão sobre Verdade Terrena (SFT-GT): Melhora o desempenho no domínio, mas causa esquecimento catastrófico severo das capacidades fora de distribuição (OOD).
2. SFT Padrão sobre Dados Auto-Distilados (SFT-SD): Preserva as capacidades OOD (evitando o esquecimento), mas produz apenas ganhos marginais no domínio.
   O artigo postula que, embora as trajetórias auto-distiladas ofereçam uma barreira de otimização menor, a aplicação ingênua de objetivos NELBO padrão a elas falha em explorar a informação estruturada "do fácil para o difícil" embutida nas trajetórias, levando a um desempenho subótimo.

2. Metodologia: TABOM

Os autores propõem a Otimização Alinhada à Trajetória via Modelagem Boltzmann (TABOM), um framework de pós-treinamento que alinha o objetivo de treinamento com as dinâmicas de decodificação no momento da inferência.

2.1 Fundamentação Teórica: Distribuição Desmascarada Boltzmann

Os autores demonstram teoricamente que a distribuição desmascarada de inferência sob um cronograma de decodificação "do fácil para o difícil" pode ser modelada como uma distribuição de Boltzmann sobre a entropia preditiva ideal dos tokens.

• Proposição: A probabilidade de um conjunto desmascarado $U$ é proporcional a $\exp(-\beta \sum_{r \in U} H_{ideal}(x_r))$, onde $H_{ideal}$ é a entropia preditiva.
• Implicação: Esta formulação captura explicitamente o viés indutivo de que o modelo prefere desmascarar tokens "fáceis" (baixa entropia) primeiro. O objetivo é minimizar a divergência de Kullback-Leibler (KL) entre a distribuição prevista pelo modelo e esta distribuição alvo de Boltzmann.

2.2 Surrogado Tratável: Classificação Pares a Pares

A otimização direta da divergência de KL é intratável devido à necessidade de amostragem global e ao cálculo de uma função de partição sobre todos os estados desmascarados possíveis. O TABOM introduz um objetivo surrogado tratável baseado em Classificação Pares a Pares:

• Mecanismo: Em vez de maximizar a verossimilhança globalmente, o método impõe classificações relativas entre pares de estados desmascarados. Se o token $r$ é decodificado antes do token $s$ em uma trajetória (implicando que $r$ é "mais fácil"), a entropia prevista pelo modelo para $r$ deve ser menor do que para $s$.
• Função de Perda: Uma perda de dobradiça (hinge loss) pares a pares é aplicada sobre janelas locais de decodificação. Para um conjunto de tokens recém-desmascarados $\Delta$, a perda penaliza casos onde a entropia de um token "mais fácil" não é suficientemente menor do que a de um token "mais difícil":
  $$L_{rank} = \max(0, h_\theta(r) - h_\theta(s) + \gamma)$$
• Janela Local: Para evitar ruído ao comparar tokens com dificuldades vastamente diferentes (ex.: primeiro vs. último passo), a classificação é aplicada estritamente dentro de um tamanho de janela local fixo $W$ (ex.: $W=32$).

2.3 Objetivo Final

A perda total do TABOM combina:

1. Reconstrução Consciente de Trajetória: Uma perda de entropia cruzada padrão onde o modelo prevê tokens recém-desmascarados $\Delta$ dado o contexto de tokens previamente desmascarados (substituindo o mascaramento uniforme).
2. Regularização de Classificação Pares a Pares: A perda de dobradiça descrita acima para alinhar a paisagem de entropia com a trajetória "do fácil para o difícil".

$$L_{TABOM} = L_{reconstruction} + \lambda L_{rank}$$

3. Contribuições Principais

1. Insight Teórico: O artigo define formalmente a distribuição desmascarada de inferência como uma distribuição de Boltzmann sobre entropias preditivas, identificando o "viés indutivo uniforme" do NELBO padrão como a causa raiz da incompatibilidade entre treinamento e inferência.
2. Inovação Metodológica: O TABOM introduz um objetivo de classificação pares a pares que alinha a ordenação de certeza do modelo com trajetórias de decodificação observadas, sem exigir cálculos intratáveis de função de partição.
3. Resolução Empírica: O framework resolve o dilema "esquecimento vs. ganhos marginais". Ele alcança melhorias substanciais no desempenho no domínio, preservando ao mesmo tempo as capacidades fora de distribuição, mitigando efetivamente o esquecimento catastrófico.
4. Métrica Diagnóstica: Os autores introduzem o Score de Discriminação de Trajetória (TDS) para quantificar se um modelo preserva diferenças de incerteza ao nível de token ao longo da trajetória de decodificação.

4. Resultados Experimentais

Os experimentos foram conduzidos em dois DLMs (Dream-7B e LLaDA-8B) nos domínios de Geração de Código (HumanEval, MBPP) e Raciocínio Matemático (GSM8K, MATH500).

• Ganhos de Desempenho: O TABOM superou consistentemente as linhas de base (SFT-GT, SFT-SD, dInfer, T3D).
  • No Dream-7B (Código), o TABOM alcançou um ganho de +7,70 no HumanEval e +2,60 no MBPP, comparado a ganhos marginais para o SFT-SD (+1,00, +1,20).
  • Crucialmente, o TABOM manteve o desempenho OOD (ex.: GSM8K permaneceu alto em 81,73 versus a queda para 52,33 observada no SFT-GT).
• Robustez: Sob decodificação paralela de 2 tokens, o TABOM manteve alto desempenho, enquanto o SFT-GT sofreu degradação severa, destacando o alinhamento do TABOM com a distribuição de inferência natural do modelo.
• Estudos de Ablação:
  • A remoção da perda de classificação resultou em ganhos limitados, confirmando sua necessidade para alinhar a paisagem de entropia.
  • A aplicação da classificação globalmente (sem janelas locais) introduziu ruído e degradou o desempenho, validando a necessidade de restrições de janela local.
• Discriminação de Trajetória: A análise do TDS mostrou que o TABOM aumentou significativamente a variância da entropia preditiva entre os passos de decodificação (ex.: o TDS do HumanEval subiu de 0,035 no SFT-SD para 0,711 no TABOM), provando que o modelo aprendeu com sucesso a distinguir entre tokens fáceis e difíceis, ao contrário do SFT padrão que os trata uniformemente.

5. Significado e Alegações

O artigo alega que o TABOM representa uma mudança fundamental na forma como trajetórias auto-distiladas são utilizadas para DLMs:

• Além da Eficiência: Enquanto métodos anteriores baseados em trajetória focavam em compressão de passos de amostragem (aceleração), o TABOM usa essas trajetórias para aquisição genuína de conhecimento e aprimoramento de desempenho.
• Alinhamento em vez de Atalho: Ao modelar o comportamento de inferência como uma distribuição de Boltzmann e impor classificação baseada em entropia, o TABOM alinha a paisagem de otimização do modelo com suas dinâmicas reais de geração.
• Mitigação do Esquecimento Catastrófico: O método expande com sucesso a fronteira de conhecimento efetiva dos DLMs para novos domínios sem sacrificar capacidades pré-treinadas, uma falha comum do SFT padrão.

Os autores concluem que o TABOM fornece uma solução robusta para a discrepância entre treinamento e inferência em Modelos de Linguagem de Difusão, oferecendo um caminho para estratégias de pós-treinamento mais eficazes e eficientes que respeitam a estrutura inerente "do fácil para o difícil" dos processos generativos.