Resumo Técnico: Reinforce Adjoint Matching (RAM)

Declaração do Problema

Modelos de difusão e de correspondência de fluxo (flow-matching) alcançam geração de alta fidelidade porque seu pré-treinamento depende de regressão supervisionada: uma amostra limpa é analiticamente ruidada, e o modelo regressa contra um alvo de forma fechada. No entanto, alinhar esses modelos com funções de recompensa (por exemplo, para correção composicional, renderização de texto ou preferência humana) via Aprendizado por Reforço (RL) no pós-treinamento interrompe essa estrutura eficiente.

Os métodos existentes de RL no pós-treinamento enfrentam obstáculos computacionais e estruturais significativos:

1. Métodos de Gradiente de Política: Tratam a remoção de ruído como um Processo de Decisão de Markov, exigindo simulações estocásticas de EDE (SDE) e otimização de limites inferiores variacionais. Isso introduz alta variância e requer muitas etapas de amostragem.
2. Métodos de Controle Ótimo Estocástico (Adjoint): Visam exatamente o objetivo de maximização de recompensa regularizada por KL, mas exigem varreduras adjuntas reversas ao longo de trajetórias de EDE armazenadas e gradientes de recompensa retropropagados através do processo de amostragem. Isso é numericamente instável e computacionalmente caro, particularmente para modelos de correspondência de fluxo onde os cronogramas de ruído divergem perto do final ruidoso.
3. Funções de Perda Surrogadas: Abordagens recentes substituem verossimilhanças intratáveis por surrogadas (por exemplo, ELBOs ou pontos finais contrastivos), sacrificando a estrutura de regressão limpa que permite a escalabilidade do pré-treinamento.

O desafio central é derivar um objetivo de RL no pós-treinamento que retenha a simplicidade e a escalabilidade do framework de regressão de pré-treinamento enquanto otimiza efetivamente as recompensas.

Metodologia: Reinforce Adjoint Matching (RAM)

Os autores propõem o Reinforce Adjoint Matching (RAM), um método que estende a estrutura de regressão de pré-treinamento para o pós-treinamento por RL sem exigir simulações de EDE, varreduras adjuntas reversas ou gradientes de recompensa.

Fundamentação Teórica

O método baseia-se na maximização de recompensa regularizada por KL formulada como um problema de controle ótimo estocástico.

1. Estrutura do Processo Ótimo: Sob regularização KL, o processo generativo ótimo inclina a distribuição de pontos finais limpos ($X_0$) em direção a amostras de alta recompensa ($p^*_0 \propto p_{ref}^0 \exp(r)$), mas deixa a lei condicional dos estados ruidosos dado um ponto final limpo inalterada. Especificamente, a trajetória de um $X_0$ fixo até um estado ruidoso $X_t$ permanece o mesmo kernel Gaussiano analítico usado no pré-treinamento.
2. Condição de Correspondência Adjoint: O controle ótimo $u_t$ satisfaz uma condição de ponto fixo onde o controle é igual ao gradiente escalado negativo da função de valor (o adjunto): $u_t = -\sigma_t \nabla_x V_t(x)$.
3. Decomposição REINFORCE: Para evitar calcular $\nabla r$ (que pode ser não diferenciável), os autores aplicam a identidade log-derivada (REINFORCE) à função de valor. Isso decompõe o adjunto em um termo de recompensa ponderado pela pontuação de ponte reversa e um termo de correção de custo de caminho.

O Objetivo RAM

O RAM aproxima o adjunto descartando a correção de custo de caminho (justificada pela baixa variância e estabilidade em escala) e utilizando a percepção estrutural de que o kernel de ruidagem permanece analítico.

1. Amostragem de Ponto Final: Em vez de simular uma EDE controlada, o RAM extrai um ponto final limpo $X_0$ da distribuição do modelo atual usando qualquer amostrador ODE pronto para uso.
2. Ruidagem Analítica: O $X_0$ amostrado é ruidado analiticamente para gerar estados de treinamento $X_t = (1-t)X_0 + t\epsilon$, exatamente como no pré-treinamento.
3. Alvo de Forma Fechada: O alvo de regressão para o campo de velocidade $v_\theta$ é derivado substituindo o termo de recompensa e a pontuação de ponte Bayesiana (que possui uma expressão de forma fechada para ruidagem linear) na condição adjunta.

A função de perda resultante é:
$$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$$
onde $X_0 \sim p_\theta^0$, $\epsilon \sim \mathcal{N}(0, I)$, e $\text{sg}(\cdot)$ denota o operador de stop-gradient.

Mecanismos Chave de Eficiência

• Sem Simulações de EDE: Os estados de treinamento são gerados ruidando um único ponto final $K$ vezes com vetores de ruído independentes. Isso produz $K$ amostras de treinamento condicionalmente independentes por passagem frontal do modelo, fornecendo mais sinal de gradiente por etapa do que métodos que dependem de trajetórias únicas de EDE.
• Sem Gradientes de Recompensa: O método funciona com recompensas de caixa preta.
• Sem Varreduras Reversas: A natureza analítica do kernel de ruidagem elimina a necessidade de integração adjunta reversa.

Contribuições Principais

1. Extensão Estrutural: O artigo demonstra que a estrutura de regressão do pré-treinamento se estende ao pós-treinamento por RL sob regularização KL. O processo ótimo altera apenas a distribuição do ponto final, não a lei condicional de ruidagem.
2. Algoritmo RAM: Uma nova perda de consistência que corrige o alvo de pré-treinamento com informações de recompensa usando uma pontuação de ponte Bayesiana de forma fechada, evitando os gargalos computacionais dos métodos existentes.
3. Justificativa Teórica: Uma derivação mostrando que o ponto fixo do RAM corresponde ao ótimo regularizado por KL, com exatidão garantida para referências Gaussianas sob recompensas lineares e concordância de primeira ordem em geral.
4. Análise de Variância: Uma análise de estimadores de custo de caminho mostrando que descartar a correção de custo de caminho (como o RAM faz) reduz significativamente a variância em comparação com estimadores adjuntos de horizonte total ou locais, tornando-o o único estimador viável para pós-treinamento em escala de imagem.

Resultados Experimentais

Os autores avaliaram o RAM no Stable Diffusion 3.5 Medium (SD3.5M) em três tarefas:

1. Geração de Imagens Composicional (GenEval): Medindo a correção de objetos, atributos e relações espaciais.
2. Renderização Visual de Texto (OCR): Medindo a legibilidade do texto em imagens geradas.
3. Alinhamento de Preferência Humana (PickScore): Alinhamento com preferências humanas.

Desempenho:

• Maximização de Recompensa: O RAM alcançou as maiores recompensas de treinamento em todas as três tarefas, superando baselines incluindo Flow-GRPO, DiffusionNFT e AWM.
• Eficiência de Treinamento: O RAM atingiu a recompensa máxima do Flow-GRPO em 50× menos etapas para GenEval, 48× menos para OCR e 34× menos para PickScore.
• Qualidade da Imagem: Ao contrário de algumas baselines que sofreram com "hacking de recompensa" (degradação da qualidade estética para métricas de tarefa), o RAM manteve ou melhorou as métricas de qualidade da imagem (Aesthetic, DeQA, ImageReward, HPSv2) enquanto maximizava as recompensas de tarefa.
• Compatibilidade com CFG: O RAM permanece compatível com a Orientação Livre de Classificador (CFG) após o pós-treinamento, enquanto baselines como DiffusionNFT e AWM viram suas recompensas de tarefa cair significativamente quando o CFG foi aplicado.

Significado e Alegações

O artigo alega que o RAM representa uma mudança de paradigma no pós-treinamento por RL para modelos de difusão e de correspondência de fluxo. Ao aproveitar as propriedades estruturais específicas do problema de controle ótimo, o RAM restaura o paradigma de "regressão simples" que tornou o pré-treinamento escalável.

Os autores afirmam que:

• O RAM elimina a necessidade de simulações de EDE custosas e varreduras adjuntas reversas, que têm sido os principais gargalos para escalar o RL para geração de imagens de alta resolução.
• O método é "simples e escala", alcançando resultados state-of-the-art com significativamente menos etapas de treinamento e custos de computação por etapa comparáveis aos métodos existentes.
• A abordagem alinha com sucesso modelos com recompensas complexas (composicionalidade, texto, preferências) sem degradação visível da qualidade ou hacking de recompensa, demonstrando que a estrutura de pré-treinamento é robusta o suficiente para lidar com o alinhamento de pós-treinamento quando as condições teóricas corretas são atendidas.