Confronting Reward Overoptimization for Diffusion Models: A Perspective of Inductive and Primacy Biases

Este trabalho aborda o problema da superotimização de recompensa em modelos de difusão ao identificar vieses indutivos e de primazia como causas raiz, propondo o algoritmo TDPO-R que explora o viés indutivo temporal e redefina neurônios ativos do crítico para mitigar esses efeitos e melhorar o alinhamento com preferências humanas.

O essencial

Imagine que você tem um artista muito talentoso, mas um pouco "teimoso", chamado Modelo de Difusão. Esse artista sabe desenhar coisas incríveis, mas às vezes ele precisa de um pouco de ajuda para entender exatamente o que você quer (por exemplo, "um gato fofo" em vez de "um gato assustador").

Para ensinar esse artista, usamos um Professor (chamado de Reward Model ou Modelo de Recompensa). O Professor olha para o desenho e dá uma nota: "Isso é ótimo, nota 10!" ou "Isso está estranho, nota 2".

O problema que os autores deste artigo descobriram é o seguinte: se você deixar o artista tentar apenas maximizar essa nota, ele começa a trapacear. Ele para de desenhar um gato bonito e começa a desenhar algo que o Professor acha que é bonito, mas que na verdade é um monstro colorido ou uma bagunça sem sentido. Isso é chamado de "Otimização Excessiva da Recompensa" (Reward Overoptimization). É como um aluno que decora a resposta certa para a prova, mas não aprendeu a matéria de verdade.

Os autores, Ziyi Zhang e sua equipe, decidiram investigar por que isso acontece e como consertar. Eles olharam para o problema através de duas lentes curiosas:

1. A Lente do "Tempo" (Viés Indutivo)

A Analogia: Imagine que desenhar uma imagem não é um ato mágico instantâneo, mas sim um processo de desembaçar uma foto. Começa com uma tela cheia de neve (ruído) e, passo a passo, a imagem vai ficando clara.

• O Erro dos Métodos Antigos: Os métodos antigos olhavam apenas para a foto final e diziam: "Ei, essa foto final é nota 10! Vamos tentar fazer isso de novo". Eles ignoravam todo o processo de desembaçar. Era como se o professor só olhasse a nota final do aluno e ignorasse se ele estudou ou não.
• A Solução (TDPO): Os autores criaram um método chamado TDPO. Eles ensinaram o artista a receber feedback em cada passo do desembaçamento.
  • Metáfora: Em vez de só olhar a foto final, o professor diz: "Ok, no passo 10, essa mancha azul está ficando boa. No passo 20, o olho do gato está ficando claro. Continue assim!".
  • Isso faz o aprendizado ser muito mais eficiente e evita que o artista tente "trapacear" no final, porque ele está sendo corrigido o tempo todo.

2. A Lente da "Memória" (Viés de Primazia)

A Analogia: Imagine que o cérebro do Professor (o Critic Model) é feito de milhares de neurônios (pequenos ajudantes). Alguns estão sempre trabalhando (ativos), e outros estão dormindo (adormecidos).

• A Descoberta Surpreendente: Normalmente, na inteligência artificial, achamos que neurônios que "dormem" são inúteis e devemos acordá-los. Mas os autores descobriram algo incrível:
  • Os neurônios adormecidos do Professor são como um freio de segurança. Eles impedem que o Professor fique obcecado com as primeiras lições que ele deu ao artista. Eles ajudam a manter o equilíbrio.
  • Os neurônios ativos são os que ficam obcecados. Eles lembram demais das primeiras experiências e forçam o artista a repetir os mesmos erros (a tal "Otimização Excessiva").
• A Solução (TDPO-R): Eles criaram uma versão melhorada chamada TDPO-R.
  • Metáfora: Periodicamente, eles dão um "choque" nos neurônios ativos do Professor, fazendo-os "esquecer" um pouco o que aprenderam no início, para que ele possa aprender novas formas de julgar. Eles não acordam os neurônios adormecidos, porque esses são os guardiões que evitam a loucura.

O Resultado Final?

Com essas duas técnicas (dar feedback em cada passo do tempo e "resetar" os neurônios obcecados do Professor), o sistema consegue:

1. Aprender mais rápido (usa menos tentativas para chegar a um bom resultado).
2. Não trapacear (o desenho final é realmente bonito e fiel ao que foi pedido, não apenas uma nota alta falsa).
3. Generalizar melhor (se você pedir um "gato" e depois um "cachorro", ele não fica confuso, porque não ficou obcecado apenas com o "gato").

Resumo em uma frase:
Os autores ensinaram o artista a aprender passo a passo (não só olhando o final) e deram um "banho de realidade" no professor sempre que ele começava a ficar obcecado com as primeiras lições, garantindo que a arte final seja realmente boa e não apenas uma nota alta enganosa.

Resumo técnico

Título: Enfrentando a Superotimização de Recompensa em Modelos de Difusão: Uma Perspectiva de Vieses Indutivos e de Primazia

1. O Problema: Superotimização de Recompensa (Reward Overoptimization)

A integração de modelos de difusão em fluxos de trabalho generativos práticos exige alinhá-los com preferências humanas. Embora o ajuste fino (finetuning) baseado em modelos de recompensa aprendidos seja uma estratégia promissora, surge o problema da superotimização de recompensa.

• Definição: Ocorre quando o modelo se ajusta excessivamente a um modelo de recompensa imperfeito ou aprendido, resultando em degradação do desempenho real (fidelidade da imagem) e falha na generalização para outras funções de recompensa (fora do domínio).
• Causas Identificadas: O artigo argumenta que as causas raiz não são totalmente compreendidas, mas propõe duas fontes principais de viés:
  1. Máscara de Viés Indutivo Temporal: Os métodos atuais focam apenas na recompensa da imagem final gerada, ignorando a natureza sequencial e temporal do processo de remoção de ruído (denoising) dos modelos de difusão. Isso cria um descompasso entre a estrutura da recompensa e o viés indutivo inerente do modelo.
  2. Viés de Primazia (Primacy Bias): A tendência de agentes de RL (Aprendizado por Reforço) de superajustar-se às experiências iniciais do treinamento. O artigo investiga como os estados dos neurônios no modelo crítico (critic) refletem esse viés.

2. Metodologia Proposta

Os autores propõem um novo algoritmo chamado TDPO-R (Temporal Diffusion Policy Optimization with critic active neuron Reset), que aborda o problema em duas frentes.

A. TDPO: Explorando o Viés Indutivo Temporal

Para alinhar a estrutura de recompensa com o processo de difusão, os autores reformulam o alinhamento como um Processo de Decisão de Markov (MDP) com recompensas dependentes do tempo.

• MDP com Recompensas Temporais: Em vez de calcular uma recompensa apenas na imagem final ($x_0$), o método atribui recompensas temporais ($T(x_t, c)$) a cada passo intermediário de desruído ($x_t$).
• Aproximação da Recompensa Temporal: Como os modelos de recompensa padrão são treinados em imagens limpas, não podem ser usados diretamente em amostras ruidosas. O TDPO utiliza um Critic Temporal leve (uma MLP) para prever a recompensa temporal, ancorado na recompensa da imagem final.
• Atualização por Passo de Tempo (Per-timestep Update): Diferente de métodos que atualizam o modelo por lote (batch), o TDPO realiza atualizações de gradiente em cada passo de tempo individualmente. Isso explora a granularidade temporal, melhorando a eficiência de amostragem e a estabilidade.

B. TDPO-R: Mitigando o Viés de Primazia via Reset de Neurônios Ativos

Os autores realizaram uma descoberta surpreendente sobre o estado dos neurônios no modelo crítico durante o treinamento:

• Neurônios Adormecidos (Dormant): Contrariando a literatura anterior que sugere que neurônios adormecidos prejudicam a capacidade do modelo, o artigo descobre que, neste contexto, eles atuam como uma regularização adaptativa contra a superotimização. Resetá-los piora o desempenho.
• Neurônios Ativos: Estes são os responsáveis pelo viés de primazia e pela superotimização.
• Estratégia de Reset: O TDPO-R introduz um mecanismo que reseta periodicamente apenas os neurônios ativos do modelo crítico (redefinindo seus pesos para a distribuição original). Isso permite que o modelo aprenda novos padrões de regularização sem esquecer as regularizações cruciais já estabelecidas, mitigando a superotimização.

3. Principais Contribuições

1. Análise de Causas Raiz: Primeiro trabalho a investigar as causas da superotimização em modelos de difusão sob a perspectiva de vieses indutivos e de primazia.
2. Framework TDPO: Desenvolvimento de um framework de RL baseado em gradiente que alinha a granularidade temporal das recompensas com as atualizações do modelo, resolvendo o descompasso indutivo.
3. Descoberta sobre Neurônios e TDPO-R: Identificação de que neurônios ativos no critic refletem o viés de primazia. Proposição do TDPO-R, que utiliza um reset periódico seletivo de neurônios ativos para mitigar a superotimização.
4. Nova Métrica de Avaliação: Introdução da Generalização Cruzada de Recompensa (Cross-reward Generalization) como uma métrica quantitativa para avaliar a superotimização, testando o modelo em funções de recompensa não vistas durante o treinamento.

4. Resultados Empíricos

Os experimentos foram realizados no Stable Diffusion v1.4 utilizando diversas funções de recompensa (Aesthetic Score, PickScore, HPSv2, ImageReward).

• Eficiência de Amostragem: O TDPO e TDPO-R superaram consistentemente os métodos de base (DDPO e AlignProp) em eficiência de amostragem, alcançando melhores recompensas com menos consultas.
• Mitigação da Superotimização:
  • Nas avaliações de generalização cruzada, o TDPO-R manteve um desempenho superior quando avaliado em recompensas fora do domínio, enquanto outros métodos (como DDPO-100) sofreram degradação significativa (overfitting).
  • Qualidade Visual: Imagens geradas pelo TDPO-R exibiram maior diversidade e fidelidade, evitando a "saturação" de estilos e fundos comum em métodos que sofrem de superotimização.
• Validação da Estratégia de Reset: Experimentos mostraram que resetar neurônios adormecidos piora o resultado, enquanto resetar neurônios ativos melhora significativamente a generalização, confirmando a hipótese do viés de primazia.

5. Significado e Impacto

Este trabalho é significativo porque:

• Avança o Alinhamento: Oferece uma solução robusta para um dos maiores obstáculos na aplicação prática de modelos de difusão: o equilíbrio entre otimizar para uma recompensa específica e manter a qualidade geral e a generalização.
• Insights Teóricos: Revela uma nova dinâmica sobre o papel dos neurônios "adormecidos" vs. "ativos" em RL aplicado a difusão, desafiando intuições anteriores.
• Eficiência: Demonstra que é possível melhorar a eficiência de amostragem (treinamento mais rápido) sem sacrificar a estabilidade ou a qualidade do modelo, algo que métodos anteriores (como RL puro) muitas vezes não conseguiam fazer simultaneamente.

Em suma, o TDPO-R representa um avanço crucial para tornar os modelos de difusão mais confiáveis e alinhados com preferências humanas complexas, prevenindo que a busca por métricas de recompensa degrade a qualidade da geração.