Resumo Técnico: Diff-Instruct com Recompensa Difusa (Didr)

Declaração do Problema

Avanços recentes em geradores de texto para imagem (T2I) em um único passo permitiram a síntese em tempo real, mapeando ruído latente diretamente para imagens em uma única passagem forward. No entanto, alinhar esses geradores com preferências humanas por meio de Aprendizado por Reforço (RL) apresenta um desafio estrutural fundamental.

Métodos anteriores para RLHF em um único passo (por exemplo, Diff-Instruct*, Diff-Instruct++) combinam otimização de recompensa no espaço de imagem com correspondência de distribuição no espaço latente. Essas abordagens tipicamente aplicam sinais de preferência (recompensas) exclusivamente no ponto final da imagem limpa ($x_0$), enquanto impõem regularização Kullback-Leibler (KL) contra a distribuição de referência em toda a trajetória difusa ruidosa. Os autores identificam uma incompatibilidade estrutural crítica nesse paradigma: a recompensa atua apenas no estado terminal, enquanto o regularizador abrange todos os níveis de ruído.

Essa incompatibilidade leva a um modo de falha denominado dominação da recompensa terminal. Como o ruído forward mascara naturalmente os detalhes da imagem, a penalidade de regularização KL contra a divergência de modos enfraquece significativamente em altos níveis de ruído. Consequentemente, otimizadores exploram essa vulnerabilidade da trajetória, abandonando a distribuição de referência para colapsar em regiões de alta recompensa sem incorrer em penalidades KL suficientes. Isso resulta em pontuações de recompensa melhoradas, mas frequentemente à custa da fidelidade da imagem e da estabilidade distribucional.

Metodologia: Diff-Instruct com Recompensa Difusa (Didr)

Para resolver esse desalinhamento, os autores propõem o Didr, um framework de alinhamento livre de dados e em nível de trajetória derivado da minimização da KL Integral (IKL). A ideia central é propagar sistematicamente a distribuição de imagem limpa inclinada pela recompensa ótima de RLHF através de todos os níveis de ruído da trajetória de difusão.

1. O Alvo de Trajetória Princípiado

O método começa estabelecendo o alvo de alinhamento correto para imagens limpas sob RLHF regularizado por KL. A distribuição de imagem limpa ótima $q^*(x_0|c)$ é a densidade inclinada pela recompensa:
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
onde $q_0$ é a distribuição de referência, $r$ é a função de recompensa e $\tau$ é um parâmetro de temperatura.

Em vez de aplicar recompensas apenas em $x_0$, o Didr "eleva" esse alvo ótimo para a trajetória completa, difundindo $q^*$ através do processo forward de referência. Isso gera uma família de marginais inclinados pela recompensa $\{q^*_t\}_{t \in [0, T]}$ definidos em cada nível de ruído:
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
Os autores provam que minimizar a divergência KL Integral entre as marginais do gerador e essas marginais inclinadas pela recompensa compartilha o mesmo minimizador que o objetivo padrão de RLHF regularizado por KL. Isso garante que a recompensa e a regularização estejam equilibradas em cada nível de ruído, prevenindo a dominação da recompensa terminal.

2. A Pontuação de Recompensa Difusa (DRS)

Ao diferenciar o objetivo IKL, os autores derivam a função de pontuação alvo. A pontuação ótima decompõe-se em dois componentes:
$$s_{target} = s_{ref} + s_r$$
onde $s_{ref}$ é a função de pontuação de referência e $s_r$ é a Pontuação de Recompensa Difusa (DRS). A DRS atua como um termo de correção orientado pela recompensa definido em cada nível de ruído:
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
Ao contrário da orientação por classificador (que é uma aproximação de primeira ordem), a DRS retém a inclinação log-exponencial completa, propagando informações de preferência apenas onde a imagem limpa é identificável a partir da observação ruidosa.

3. O Proxy de Recompensa Difusa (DRP)

Como a DRS requer uma expectativa de posterior intratável $q(x_0|x_t, c)$, os autores introduzem o Proxy de Recompensa Difusa (DRP) para tornar o método prático. O DRP aproxima a expectativa de posterior usando cadeias de remoção de ruído de poucos passos diferenciáveis a partir do modelo de referência congelado.

• Mecanismo: A partir de um estado ruidoso $x_t$, o método executa $K$ cadeias independentes de remoção de ruído de $S$ passos usando o modelo de referência congelado para gerar amostras limpas aproximadas $\{\hat{x}^{(k)}_0\}$.
• Estimação de Gradiente: A DRS é aproximada como uma média ponderada por softmax dos gradientes de recompensa ao longo do caminho:
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  onde os pesos $\omega^{(k)}$ são proporcionais a $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$.
• Eficiência: Essa abordagem permite que os gradientes de recompensa se propaguem de forma estável de volta através da trajetória sem armazenar ativações do modelo de referência ou exigir dados de treinamento de imagem.

4. Framework de Treinamento

O processo de treinamento alterna entre duas etapas:

1. Atualização do Assistente de Ensino (TA): Um modelo de pontuação TA ($s_\psi$) é treinado via Correspondência de Pontuação de Remoção de Ruído (DSM) para rastrear as marginais em evolução do gerador, substituindo a pontuação implícita do gerador.
2. Atualização do Gerador: O gerador ($g_\theta$) é atualizado usando o gradiente do objetivo IKL, que compara a pontuação TA contra a soma da pontuação de referência corrigida por CFG e a estimativa DRP.

Contribuições Principais

1. Identificação da Dominação da Recompensa Terminal: O artigo identifica e analisa formalmente um modo de falha estrutural em métodos anteriores de RLHF em um único passo, onde otimizadores exploram níveis de ruído para colapsar em regiões de alta recompensa, violando o equilíbrio pretendido de RLHF.
2. Alinhamento em Nível de Trajetória Princípiado: Os autores propõem o Didr, um framework que deriva um alvo em nível de trajetória matematicamente sólido, propagando a densidade de imagem limpa ótima de RLHF para frente. Isso garante que a recompensa e a regularização estejam equilibradas em cada nível de ruído.
3. Derivação Analítica da DRS: O artigo deriva analiticamente a Pontuação de Recompensa Difusa, mostrando que ela decompõe exatamente o alvo da trajetória em uma pontuação de referência mais uma correção induzida pela recompensa.
4. Estimador Prático (DRP): Os autores introduzem o Proxy de Recompensa Difusa, um estimador livre de dados que computa a DRS via remoção de ruído diferenciável de poucos passos, permitindo treinamento estável sem dados de imagem.

Resultados Empíricos

Experimentos extensivos demonstram que o Didr supera consistentemente as bases existentes em um único passo:

• Dominação de Pareto: Na base SDXL (1024×1024), a variante padrão Didr domina Pareto todas as bases de alinhamento em um único passo existentes no trade-off PickScore–FID, superando simultaneamente modelos de múltiplos passos como FLUX-dev e SDXL-DPO em ambos os eixos.
• Métricas de Preferência: A variante de treinamento estendido (Didr_longer) alcança o maior PickScore (23,9) entre os modelos em um único passo avaliados e melhora o ImageReward em 8,9% sobre a base anterior mais forte (Diff-Instruct*). Também lidera em HPSv2.1 e Pontuação Estética.
• Generalização: O Didr transfere com sucesso para a base Z-Image de 6 bilhões de parâmetros. Em uma única etapa de geração, supera seu modelo professor de 50 passos nas métricas de preferência, demonstrando escalabilidade entre arquiteturas.
• Melhorias Qualitativas: Comparações visuais mostram que o Didr produz detalhes finos mais nítidos e cores mais naturais em comparação com métodos anteriores em um único passo, evitando a saturação excessiva e artefatos pictóricos frequentemente observados em bases otimizadas por recompensa.

Significado e Afirmações

O artigo afirma que a propagação de recompensa em nível de trajetória princípiada é uma estratégia de alinhamento eficaz que resolve o desalinhamento inerente dos métodos de recompensa apenas no ponto final. Ao propagar rigorosamente a distribuição ótima de RLHF através do processo de difusão, o Didr alcança o melhor trade-off entre alinhamento de preferência e fidelidade de imagem entre geradores em um único passo.

Os autores enfatizam que seu método requer nenhum dado de treinamento de imagem, dependendo exclusivamente de prompts de texto e de um modelo de referência congelado. Eles posicionam o Didr como uma solução que permite que geradores em um único passo igualem ou superem o alinhamento de preferência de professores de múltiplos passos, mantendo velocidades de inferência em tempo real. O trabalho sugere que a incompatibilidade estrutural identificada em métodos anteriores é uma limitação fundamental que só pode ser resolvida alinhando toda a trajetória, e não apenas o estado terminal.