Resumo Técnico: DiNa-LRM – Modelagem de Recompensa Latente Nativa de Difusão

1. Declaração do Problema

O alinhamento de modelos de difusão e de correspondência de fluxo com preferências humanas depende fortemente de funções de recompensa. As abordagens atuais de última geração utilizam predominantemente Modelos Visão-Linguagem (VLMs) como backbones de recompensa. Embora os VLMs ofereçam capacidades discriminativas robustas, eles apresentam duas limitações críticas para o alinhamento de modelos de difusão:

1. Sobrecarga Computacional e de Memória: Recompensas baseadas em VLMs são caras para consultar repetidamente durante a otimização, criando um gargalo para a eficiência do treinamento.
2. Incompatibilidade de Domínio: Os VLMs operam tipicamente no espaço de pixels, enquanto os geradores de difusão latente são treinados e otimizados no espaço latente. Essa incompatibilidade de pixel para latente complica o alinhamento, particularmente para métodos de gradiente de recompensa, e introduz sobrecarga desnecessária no sistema.

Por outro lado, embora os modelos de difusão tenham mostrado promessa como backbones discriminativos, tentativas anteriores de usá-los como recompensas frequentemente estiveram vinculadas a algoritmos específicos de otimização em nível de etapa ou focaram em estados intermediários ruidosos, sem estabelecer um framework de propósito geral comparável aos VLMs para pontuar amostras limpas.

2. Metodologia: DiNa-LRM

Os autores propõem o DiNa-LRM (Modelo de Recompensa Latente Nativa de Difusão), um framework que formula a aprendizagem de preferências diretamente sobre estados de difusão ruidosos dentro do espaço latente. O método consiste em três componentes principais:

A. Formulação de Preferência Nativa de Difusão

Em vez de avaliar amostras limpas ($x_0$) diretamente, o DiNa-LRM estende o modelo de Thurstone para estados de difusão ruidosos ($x_t$).

• Incerteza Calibrada por Ruído: Os autores reconhecem que os julgamentos de preferência sobre estados ruidosos são inerentemente menos certos à medida que o nível de ruído aumenta. Eles introduzem uma incerteza de comparação calibrada por ruído $\sigma_u^2(t)$ que escala com o nível de ruído de difusão $\sigma^2(t)$:
  $$\sigma_u^2(t) = k\sigma^2(t) + \sigma_u^2$$
  Esta formulação garante que o objetivo de recompensa permaneça consistente com a variedade ruidosa encontrada durante o pré-treinamento de difusão. Atua como um regularizador consciente da incerteza, prevenindo que gradientes de alta variância desestabilizem o treinamento em regimes de alto ruído.
• Objetivo de Treinamento: O modelo é otimizado usando uma função de perda de fidelidade sobre pares de preferência $(x_t^+, x_t^-)$, onde a probabilidade de comparação é modulada pela incerteza dependente do passo de tempo.

B. Arquitetura de Recompensa Latente

O DiNa-LRM é instanciado sobre um backbone de difusão latente pré-treinado (por exemplo, SD3.5-Medium) com uma cabeça de pontuação baseada em consultas leve.

• Extração de Recursos Multicamada: Recursos visuais e textuais são extraídos de múltiplas camadas do backbone.
• Adaptação Condicionada ao Passo de Tempo: Os recursos são modulados por incorporações de passo de tempo usando FiLM (Modulação Linear Específica por Recurso) para tornar a cabeça de recompensa explicitamente consciente do nível de ruído.
• Pontuação Q-Former com Portão: Um Q-Former com portão agrega esses recursos multicamada condicionados ao passo de tempo para produzir uma pontuação de recompensa escalar. Esta arquitetura opera inteiramente no espaço latente do VAE, evitando a decodificação no espaço de pixels.

C. Escalonamento no Tempo de Inferência via Ensemble de Ruído

Para melhorar a robustez e reduzir a sensibilidade a um passo de tempo de avaliação específico, o DiNa-LRM suporta ensemble de ruído.

• Em vez de pontuar uma amostra limpa em um único nível de ruído fixo, o modelo avalia a amostra em múltiplos passos de tempo $\{t_k\}$.
• Os recursos extraídos nesses diferentes passos de tempo são concatenados e processados pela mesma cabeça Q-Former para produzir uma pontuação agregada. Isso serve como um mecanismo de escalonamento no tempo de teste "nativo de difusão".

3. Contribuições Principais

1. Formulação de Preferência Nativa de Difusão: O artigo estende o modelo de preferência de Thurstone para estados de difusão ruidosos, introduzindo uma incerteza calibrada por ruído que escala com o nível de ruído de difusão, fechando a lacuna distribucional entre o pré-treinamento e a aprendizagem de recompensa.
2. Escalonamento no Tempo de Inferência: Os autores propõem uma estratégia de ensemble de ruído que agrega recursos de múltiplos passos de tempo, fornecendo um mecanismo para escalonamento no tempo de teste que aproveita a trajetória global de difusão em vez de uma única estimativa pontual.
3. Validação Empírica: O trabalho demonstra que recompensas nativas de difusão podem alcançar desempenho competitivo com juízes baseados em VLMs fortes, operando nativamente no espaço latente.

4. Resultados Experimentais

Os autores avaliaram o DiNa-LRM em vários benchmarks e cenários de alinhamento:

• Precisão de Modelagem de Recompensa:

  • O DiNa-LRM supera substancialmente as linhas de base existentes baseadas em difusão (por exemplo, LRM-SD1.5, LRM-SDXL) em ImageReward, HPDv2, HPDv3 e GenAI Bench.
  • Alcança desempenho competitivo com recompensas baseadas em VLMs de última geração (por exemplo, HPSv3, UnifiedReward) em média, com uma lacuna notável se estreitando em conjuntos de dados específicos como HPDv2 e HPDv3.
  • Ensemble de Ruído: A aplicação de ensemble de ruído melhora ainda mais a precisão média (por exemplo, de 71,49% para 72,48% em média), validando o benefício de agregar evidências através de níveis de ruído.

• Alinhamento de Preferências (ReFL):

  • Em experimentos de alinhamento pós-treinamento usando ReFL (Aprendizagem com Feedback de Recompensa) no SD3.5-M, o DiNa-LRM demonstrou convergência mais rápida em pontuações proxy e métricas "douradas" (PickScore mantido de fora) em comparação com o HPSv3.
  • Eficiência: O DiNa-LRM reduz significativamente os custos de recursos em comparação com recompensas baseadas em VLMs. Reduz o uso de pico de VRAM em 51,4%, os FLOPs de cálculo de recompensa em 71,1% e os FLOPs da fase de otimização em 46,4%.

• Generalização:

  • A formulação generaliza entre diferentes backbones de difusão (SDXL, FLUX.1-Dev, Z-Image-Turbo), superando consistentemente as linhas de base anteriores baseadas em difusão.
  • Experimentos de transferibilidade mostram que um modelo de recompensa treinado no SD3.5-M pode alinhar efetivamente o SD3.5-L dentro do mesmo espaço latente.

5. Significância e Alegações

O artigo alega que o DiNa-LRM oferece uma alternativa prática e eficiente às recompensas de VLM no espaço de pixels para alinhamento de preferências. Sua significância principal reside em:

• Vantagem Estrutural: Ao operar nativamente no espaço latente, elimina a incompatibilidade de latente para pixel e a sobrecarga computacional da decodificação, tornando-o altamente compatível com métodos de alinhamento de gradiente de recompensa.
• Amigável à Otimização: Fornece um sinal diferenciável e eficiente que melhora a dinâmica de treinamento e reduz os requisitos de recursos sem sacrificar a qualidade do alinhamento.
• Escalabilidade: O mecanismo de ensemble de ruído fornece uma maneira nova e nativa de difusão para escalar o desempenho no tempo de inferência.

Os autores reconhecem limitações, observando que a recompensa está vinculada a espaços latentes específicos (limitando a transferência entre backbones sem retreinamento) e que a modelagem no espaço latente pode subestimar certos artefatos em nível de pixel, potencialmente levando a comportamentos de exploração de recompensa (por exemplo, injeção humana espúria ou deriva estilística) durante otimização de longo horizonte. Eles sugerem que trabalhos futuros devem focar em backbones unificados mais fortes e na incorporação de regularização leve no espaço de pixels.