Resumo Técnico: TMPO para Alinhamento Diverso e Eficiente de Difusão

1. Declaração do Problema

A Aprendizagem por Reforço (RL) demonstrou potencial significativo no alinhamento de modelos de difusão e de correspondência de fluxo com tarefas downstream. No entanto, os métodos existentes dependem predominantemente da maximização de recompensa escalar, que os autores identificam como inerentemente buscadora de modos. Este objetivo de otimização faz com que os modelos concentrem a massa de probabilidade em poucas trajetórias de alta recompensa, levando a dois modos de falha críticos:

1. Degradação da Diversidade Generativa: O modelo colapsa em um conjunto estreito de composições e estilos (colapso de modo visual), ignorando a distribuição anisotrópica de saídas plausíveis aprendida durante o pré-treinamento.
2. Amplificação de Recompensas Não Confiáveis: Ao reforçar excessivamente características específicas que favorecem proxies de recompensa (por exemplo, precisão de OCR ou pontuações estéticas específicas), o modelo amplifica atributos espúrios, frequentemente às custas da qualidade geral da geração.

Estratégias de mitigação atuais, como regularização KL, falham em alterar a natureza fundamentalmente buscadora de modos do objetivo, enquanto as abordagens existentes de Rede de Fluxo Generativo (GFlowNet) frequentemente dependem de estimativas de fluxo de estado intratáveis ou de restrições locais que acumulam erros ao longo de horizontes de remoção de ruído longos.

2. Metodologia: Otimização de Política de Correspondência de Trajetória (TMPO)

O TMPO altera fundamentalmente o paradigma de otimização da maximização de recompensa escalar para a correspondência da distribuição de recompensa ao nível de trajetória.

2.1 Equilíbrio de Trajetória Softmax (Softmax-TB)

Em vez de maximizar a recompensa esperada, o TMPO visa corresponder as probabilidades de trajetória da política a uma distribuição-alvo induzida pelas recompensas.

• Objetivo: Para um grupo de $K$ trajetórias geradas a partir do mesmo prompt, o método corresponde as probabilidades da política $P_\theta(\tau)$ a uma distribuição de Boltzmann $q(\tau) \propto \exp(\beta R(\tau))$.
• Formulação sem Partição: Ao normalizar probabilidades e recompensas dentro do grupo de $K$ trajetórias, a função de partição global intratável é eliminada. O método define uma vantagem Softmax-TB:
  $$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$$
• Garantia Teórica: Os autores provam que minimizar este objetivo é equivalente a minimizar a divergência KL forward ($D_{KL}(q \parallel p)$) dentro do grupo. Diferentemente da KL reversa (usada em gradientes de política padrão), a KL forward é cobridora de modos: impõe uma penalidade infinita se a probabilidade da política $p_i$ se aproximar de zero para qualquer trajetória $i$ que tenha peso-alvo positivo $q_i$. Isso garante teoricamente a preservação da diversidade sem exigir regularização KL auxiliar.

2.2 Amostragem Árvore Estocástica Dinâmica

Para abordar o custo computacional de amostrar múltiplas trajetórias independentes para modelos de correspondência de fluxo em grande escala, o TMPO introduz uma estratégia de amostragem com estrutura de árvore:

• Compartilhamento de Prefixo: As trajetórias compartilham um prefixo ODE determinístico (etapas comuns de remoção de ruído) e ramificam-se em pontos estocásticos agendados dinamicamente.
• Ramificação Dinâmica: Os pontos de ramificação não são fixos; seguem um currículo que se desloca de regiões de alto ruído iniciais (diversidade de baixa granularidade) para regiões de baixo ruído posteriores (variações de alta granularidade) à medida que o treinamento progride. As posições de ramificação são perturbadas estocasticamente usando uma distribuição Beta para evitar o sobreajuste a uma geometria de árvore fixa.
• Eficiência: Esta abordagem reduz o cálculo redundante ao compartilhar a passagem direta para o prefixo compartilhado, diminuindo significativamente o tempo de treinamento enquanto mantém uma exploração eficaz.

3. Contribuições Principais

1. Identificação da Causa Raiz: O artigo identifica a natureza buscadora de modos da maximização de recompensa como a causa primária do hacking de recompensa e da degradação da diversidade no RL de difusão, propondo uma mudança para a correspondência de distribuições.
2. Objetivo Softmax-TB: Introdução de um objetivo de correspondência de distribuição ao nível de trajetória, sem partição, que herda as propriedades cobridoras de modos da divergência KL forward, preservando naturalmente a diversidade generativa.
3. Amostragem Árvore Estocástica Dinâmica: Um mecanismo de amostragem escalável que reduz o tempo de treinamento para modelos em grande escala ao compartilhar prefixos de remoção de ruído e ajustar dinamicamente os cronogramas de ramificação.

4. Resultados Experimentais

Os autores avaliaram o TMPO na base FLUX.1-dev em três tarefas de alinhamento: geração de imagens composicionais (GenEval), renderização de texto visual (OCR) e alinhamento de preferência humana (PickScore).

• Diversidade: O TMPO melhorou a diversidade generativa em 9,1% em média em comparação com os métodos mais avançados (Flow-GRPO, MixGRPO, TreeGRPO, GARDO). Foi o único método a manter consistentemente pontuações positivas de Distância Geométrica Média Logarítmica (LGMD), indicando diversidade saudável, enquanto as linhas de base frequentemente exibiam LGMD negativa (colapso de modo).
• Desempenho: O TMPO alcançou desempenho competitivo ou superior em métricas específicas da tarefa (precisão do GenEval, precisão do OCR) e pontuações de preferência humana (HPS-v2.1, ImageReward, PickScore).
• Eficiência: Ao utilizar compartilhamento de prefixo e amostragem em árvore, o TMPO reduziu o tempo de treinamento por iteração em até 27% em comparação com os métodos existentes.
• Compensações: A análise de Pareto demonstrou que o TMPO alcança a compensação ótima entre maximização de recompensa, preservação da diversidade e eficiência de treinamento.

5. Significado e Alegações

O artigo alega que o TMPO oferece uma solução fundamental para o dilema do hacking de recompensa no alinhamento de IA generativa. Ao reformular o objetivo para corresponder a distribuições induzidas por recompensa em vez de maximizar recompensas escalares, o TMPO impõe estruturalmente a preservação da diversidade sem depender de regularização externa ou reponderação heurística. O método demonstra que é possível alcançar altas pontuações de alinhamento enquanto se mantém a rica distribuição de saída multimodal aprendida durante o pré-treinamento, tudo isso reduzindo a sobrecarga computacional tipicamente associada a métodos de RL de múltiplas trajetórias. Os autores posicionam o TMPO como um framework escalável adequado para modelos de correspondência de fluxo em grande escala, com extensões potenciais para geração de vídeo e geração de ações robóticas.