Resumo Técnico: Orientação Explícita de Crítico para Alinhamento de Modelos de Difusão

Declaração do Problema

O pós-treinamento de modelos de difusão para alinhamento com objetivos não diferenciáveis (por exemplo, preferência humana, correção composicional, renderização de texto) é cada vez mais crítico, pois o pré-treinamento isolado frequentemente produz resultados subótimos em propriedades centradas no usuário. Um desafio central neste domínio é a atribuição de crédito. Modelos de difusão geram amostras por meio de trajetórias de remoção de ruído em múltiplos passos, contudo, as recompensas são tipicamente definidas apenas na imagem final decodificada. Soluções existentes enfrentam limitações significativas:

• Métodos Relativos a Grupos (por exemplo, GRPO): Estes contornam críticos explícitos comparando amostras dentro de um grupo. Embora evitem o treinamento de uma função de valor, o sinal de recompensa resultante é grosseiro e em nível de trajetória, carecendo de atribuição de crédito granular para decisões intermediárias de remoção de ruído.
• Críticos no Espaço de Pixels: Métodos existentes de ator-crítico explícito frequentemente definem críticos no espaço de pixels, avaliando imagens "limpas" reconstruídas a partir de latentes ruidosos intermediários. Isso cria um descompasso de estado: a política opera sobre estados latentes ruidosos, enquanto o crítico avalia estados proxy que frequentemente estão fora da distribuição para codificadores de visão padrão (por exemplo, CLIP, BLIP). Esse descompasso leva a estimativas de valor frágeis, instabilidade na otimização e sobrecarga substancial de decodificação.
• Exploração de Recompensa (Reward Hacking): A otimização para recompensas verificáveis específicas frequentemente leva a soluções de atalho (por exemplo, fundos simplificados, padrões repetitivos) que maximizam a métrica-alvo enquanto degradam aspectos não medidos da qualidade da imagem.

Metodologia

Os autores propõem uma Estrutura de Ator-Crítico Latente Alinhada ao Estado que aborda essas questões integrando o crítico diretamente no processo de geração de difusão.

1. Crítico Latente Alinhado ao Estado

Em vez de operar no espaço de pixels, a estrutura aprende uma função de valor $V_\phi(z_t, t, y)$ diretamente sobre os estados latentes ruidosos $z_t$ visitados pela política.

• Arquitetura: O crítico é inicializado a partir do backbone de difusão pré-treinado (UNet ou DiT) e modificado com alterações arquitetônicas leves para prever um valor escalar.
  • UNet: O mapa de características é fragmentado em patches, e um token de resumo aprendível agrega informações por meio de uma cabeça de atenção e uma MLP.
  • DiT: Ramificações de crítico leves são inseridas em camadas intermediárias (por exemplo, camadas 7, 15, 23) usando tokens de consulta aprendíveis para agregar informações visuais.
• Alinhamento de Estado: Ao operar no mesmo espaço latente e compartilhar o prior generativo, o crítico está naturalmente ciente do passo de tempo e do nível de ruído, evitando problemas de deslocamento de distribuição típicos de críticos no espaço de pixels.

2. Técnicas de Estabilização de Treinamento

Para garantir uma otimização estável de ator-crítico, duas técnicas específicas são introduzidas:

• Condicionamento Consciente do Passo de Tempo: A cabeça de valor é explicitamente condicionada ao passo de tempo $t$ (via modulação estilo AdaLN), fornecendo acesso direto ao nível de ruído para melhorar a previsão de valor ao longo da trajetória.
• Pré-treinamento de Valor: Antes da otimização conjunta ator-crítico, uma fase curta de aquecimento é realizada onde apenas o crítico é atualizado usando perda de regressão contra alvos de retorno Monte Carlo. Isso produz estimativas de valor iniciais melhores, estabilizando as subsequentes atualizações PPO.

3. Direcionamento no Tempo de Inferência

Como o crítico é definido sobre os mesmos estados latentes que a política, ele pode ser reutilizado no tempo de inferência como um módulo de orientação. O gradiente da função de valor, $\nabla_{z_t} V_\phi$, aponta para regiões latentes previstas para gerar recompensas terminais mais altas. Esse gradiente é injetado na atualização do amostrador (modificando a previsão de ruído para DDIM ou velocidade para correspondência de fluxo), orientando efetivamente a trajetória de remoção de ruído sem exigir treinamento adicional.

4. Otimização Multi-Recompensa

Para mitigar a exploração de recompensa, a estrutura é estendida para otimização multi-recompensa. Um backbone latente compartilhado é usado com cabeças de valor separadas e leves para cada tipo de recompensa. Este design permite que recompensas heterogêneas (por exemplo, composicional vs. estética) se regularizem mutuamente enquanto compartilham a extração de características cara, impedindo o colapso em soluções de atalho comuns no treinamento de recompensa única.

Contribuições Principais

1. Crítico Latente Alinhado ao Estado: Uma estrutura que converte um modelo de difusão pré-treinado em sua própria rede de valor condicionada ao passo de tempo, permitindo previsão de valor em latentes ruidosos e direcionamento no tempo de inferência.
2. Estratégias de Estabilização: Identificação de condicionamento consciente do passo de tempo e pré-treinamento de valor como técnicas essenciais para estabilizar o treinamento explícito de ator-crítico em difusão.
3. Extensão Multi-Recompensa: Um mecanismo usando cabeças de valor separadas em um backbone compartilhado para otimizar conjuntamente recompensas complementares, mitigando a superotimização de proxy e a exploração de recompensa.
4. Superioridade Empírica: Desempenho consistente superior ao de métodos relativos a grupos (GRPO) e baselines anteriores de ator-crítico em benchmarks de recompensa única e multi-recompensa, tanto em backbones UNet quanto DiT.

Resultados Experimentais

O método foi avaliado em Stable Diffusion 1.5 (UNet) e Stable Diffusion 3.5M (DiT) através de diversos tipos de recompensa:

• Benchmarks de Recompensa Única: O método alcançou resultados state-of-the-art em métricas CLIP, HPSv2.1, PickScore, GenEval e OCR. Notavelmente, mostrou ganhos significativos em recompensas verificáveis (GenEval, OCR), onde sinais esparsos e discretos beneficiam-se mais de estimativas de valor por passo de tempo. Também demonstrou estabilidade onde o DDPO falhou em convergir no SD3.5M.
• Benchmarks Multi-Recompensa: Em tarefas de otimização conjunta (por exemplo, CLIP + HPS + GenEval), o método superou o GRPO, particularmente no GenEval, mantendo desempenho comparável em métricas de preferência.
• Direcionamento no Tempo de Inferência: Aplicar o crítico aprendido como sinal de orientação durante a amostragem melhorou ainda mais a qualidade de geração. No SD1.5, a orientação pelo crítico sozinha superou a amostragem Best-of-N (BoN); no SD3.5M, combinar orientação com BoN produziu os melhores resultados.
• Eficiência: O crítico latente foi significativamente mais rápido que críticos no espaço de pixels (CLIP/BLIP/HPS) tanto no cálculo de valor (10s vs. 28–35s) quanto em passos de treinamento (20s vs. 52–86s).
• Ablações:
  • Remover o pré-treinamento de valor degradou o desempenho, especialmente em recompensas esparsas.
  • Cabeças de valor separadas para tarefas multi-recompensa superaram designs de cabeça compartilhada.
  • Treinar sem Orientação Livre de Classificador (CFG) produziu a otimização mais estável e as recompensas finais mais altas, evitando a instabilidade e artefatos associados a escalas CFG elevadas.

Significado e Alegações

O artigo afirma estabelecer uma receita prática para aplicar aprendizado por reforço a modelos de difusão modernos. Seu significado reside em:

1. Resolução do Descompasso de Estado: Ao alinhar o crítico com o espaço de estados latentes da política, elimina-se a fragilidade e a sobrecarga da avaliação no espaço de pixels.
2. Habilitação de Treinamento Estável de Ator-Crítico: A combinação de alinhamento latente, condicionamento ao passo de tempo e pré-treinamento de valor torna o treinamento explícito de ator-crítico viável e estável para modelos de difusão, oferecendo uma alternativa mais eficiente em amostras aos métodos relativos a grupos.
3. Utilidade Dual: O crítico aprendido serve a um duplo propósito: otimizar a política durante o treinamento e orientar o processo de geração no tempo de inferência, proporcionando ganhos de desempenho além da fase de treinamento.
4. Mitigação da Exploração de Recompensa: A extensão multi-recompensa oferece um mecanismo prático para equilibrar objetivos concorrentes, reduzindo a prevalência de soluções de atalho no RL de difusão.

Os autores reconhecem limitações, notando o foco atual em difusão de imagem (não vídeo), a possível degradação de eficiência em configurações multi-recompensa comparada ao treinamento de recompensa única, e erros de discretização inerentes à amostragem SDE.