Beyond Pairwise Preferences: Listwise Reward-Aware Alignment for Diffusion Models

Resumo Técnico: Diffusion LAIR

Declaração do Problema
Os métodos atuais de otimização de preferências para alinhar modelos de difusão texto-para-imagem, como a Otimização Direta de Preferências (DPO) e suas variantes específicas para difusão (por exemplo, Diffusion-DPO, DSPO), dependem amplamente de supervisão por pares. Esses métodos reduzem os dados de treinamento a comparações binárias "vencedor-perdedor" para um determinado prompt. Os autores argumentam que essa redução é limitante por duas razões principais:

1. Estrutura de Dados: Conjuntos de dados de preferências práticos (por exemplo, Pick-a-Pic) frequentemente contêm múltiplas imagens candidatas para um único prompt, criando uma estrutura natural de lista que os métodos por pares ignoram.
2. Perda de Informação: Modelos de recompensa modernos fornecem pontuações contínuas. Reduzi-las a um único rótulo binário descarta a qualidade relativa das candidatas não selecionadas e ignora a magnitude das lacunas de recompensa entre as amostras.
   Além disso, as abordagens existentes de aprendizado por reforço (RL) online para alinhamento de difusão são computacionalmente caras, difíceis de paralelizar e sensíveis a mudanças de distribuição devido à sua natureza on-policy. Embora os métodos offline sejam mais eficientes, eles não aproveitaram plenamente o sinal mais rico disponível em dados de lista com pontuação de recompensa.

Metodologia: Diffusion LAIR
O artigo propõe o Diffusion LAIR (Recompensa Implícita Ponderada por Vantagem de Lista), um método de otimização de preferências offline que trata o alinhamento como um problema de lista consciente de recompensa.

• Representação de Dados: Em vez de pares independentes, o método processa um prompt $c$ com um conjunto de $N_c$ imagens candidatas $\{x^{(i)}_0\}$ e suas respectivas pontuações de recompensa contínua $\{r_i\}$.
• Pesagem por Vantagem: As pontuações de recompensa são convertidas em pesos de vantagem centralizados ($w_i$). Primeiro, uma probabilidade softmax $p_i$ é computada usando uma temperatura $\tau$. Em seguida, os pesos são centralizados contra uma linha de base uniforme:
  $$w_i = p_i - \frac{1}{N_c}$$
  Isso garante que $\sum w_i = 0$, atribuindo pesos positivos a amostras de alta recompensa e pesos negativos a amostras de baixa recompensa em relação à média do grupo.
• Objetivo de Otimização: O método otimiza um objetivo de regressão ponderado por vantagem na recompensa implícita. A recompensa implícita é definida como a melhoria na perda de remoção de ruído do modelo atual ($\theta$) sobre um modelo de referência fixo ($ref$):
  $$s_\theta(x_0, x_t, t, c, \epsilon) = \omega(\lambda_t) (l_{ref} - l_\theta)$$
  A função de perda é:
  $$L_{LAIR} = \mathbb{E} \left[ -\sum_{i=1}^{N_c} w_i s^{(i)}_\theta + \frac{\lambda}{N_c} \sum_{i=1}^{N_c} (s^{(i)}_\theta)^2 \right]$$
  O primeiro termo incentiva o modelo a aumentar a recompensa implícita para amostras de alta recompensa e diminuí-la para amostras de baixa recompensa. O segundo termo é uma penalidade quadrática que regulariza a magnitude da recompensa implícita, prevenindo desvios agressivos do modelo de referência.

Análise Teórica
Os autores fornecem uma caracterização teórica do objetivo:

1. Ótimo de Forma Fechada: O objetivo admite um ótimo de forma fechada único e limitado no espaço de recompensa implícita: $s^*_i = \frac{N_c}{2\lambda} w_i$. Isso esclarece que a força de regularização $\lambda$ controla diretamente a magnitude da atualização de preferência.
2. Limite Surrogado KL: Sob suposições específicas (aproximação de log-ratio e extensão de suporte total limitada), os autores derivam um limite para a divergência de Kullback-Leibler (KL) entre a distribuição induzida e a distribuição de referência: $D_{KL} \le \frac{\Delta_c}{\eta}$. Isso sugere que $\lambda$ controla a nitidez do viés de preferência, garantindo que a atualização permaneça conservadora.
3. Vantagem sobre Pares: Ao contrário de objetivos por pares (por exemplo, Bradley-Terry), que podem levar a recompensa implícita ao infinito para maximizar a margem entre um vencedor e um perdedor, a formulação de lista do LAIR atribui alvos finitos e relativos com base em todo o conjunto de candidatos, preservando a estrutura de classificação e o espaçamento de recompensas.

Resultados Experimentais
O método foi avaliado no Stable Diffusion 1.5 (SD1.5) e no Stable Diffusion XL (SDXL) usando o conjunto de dados Pick-a-Pic v2.

• Configuração: As imagens foram pontuadas offline usando o modelo de recompensa PickScore. Os modelos foram ajustados (fine-tuned) e avaliados em geração texto-para-imagem (HPD, Parti-prompts), geração composicional (GenEval) e edição de imagem baseada em instruções (InstructPix2Pix).
• Desempenho: O Diffusion LAIR superou bases fortes (Diffusion-DPO, DSPO, Diffusion-KTO, MaPO, InPO) na maioria das métricas.
  • No SD1.5, alcançou as pontuações mais altas em PickScore, HPS, Estética, CLIP e ImageReward para ambos os conjuntos de dados Parti-prompt e HPD.
  • No SDXL, superou igualmente as bases, incluindo MaPO e Diffusion-DPO.
  • Generalização: Notavelmente, apesar de ser treinado apenas com PickScore, o modelo mostrou melhorias significativas em outras métricas (por exemplo, ImageReward), sugerindo que o objetivo de lista melhora a qualidade geral das amostras além do modelo de recompensa específico utilizado.
  • Composicional e Edição: O método demonstrou ganhos consistentes nas taxas de vitória do GenEval (raciocínio composicional) e do InstructPix2Pix (edição de imagem) em comparação com bases por pares.
• Eficiência: O processo de treinamento foi significativamente mais eficiente do que as bases de RL online e comparável ou mais eficiente do que outros métodos offline, exigindo aproximadamente 24 horas de GPU para o SD1.5 e 140 horas de GPU A100 para o SDXL.

Significado e Alegações
O artigo afirma que o Diffusion LAIR oferece uma direção simples e eficaz para alinhar modelos de difusão, superando as limitações da redução por pares. Suas principais contribuições são:

1. Aprendizado de Lista: É o primeiro método a aproveitar diretamente grupos de imagens pontuadas por recompensa para otimização de preferências de difusão, utilizando a estrutura completa da qualidade relativa das amostras.
2. Clareza Teórica: Fornece um ótimo de forma fechada limitado, oferecendo insights sobre como a regularização controla a mudança de distribuição e a magnitude da atualização.
3. Superioridade Empírica: Demonstra forte desempenho em tarefas de geração, composição e edição, provando que a otimização de lista consciente de recompensa produz melhorias mais amplas do que os métodos por pares.

Os autores observam limitações, especificamente a dependência da qualidade do modelo de recompensa usado para o treinamento e as suposições idealizadas no limite teórico de KL. No entanto, eles postulam que o método é robusto a rótulos de preferência ruidosos comuns em conjuntos de dados anotados por humanos.