Resumo Técnico: AEGPO – Otimização de Política Guiada por Entropia Adaptativa para Modelos de Difusão

1. Definição do Problema

O Aprendizado por Reforço com Feedback Humano (RLHF), particularmente o Group Relative Policy Optimization (GRPO), tornou-se um paradigma padrão para alinhar modelos de difusão e de fluxo (flow-matching) com as preferências humanas. No entanto, as variantes existentes de GRPO sofrem de duas ineficiências fundamentais devido à sua dependência de heurísticas estáticas:

1. Alocação Ineficiente de Amostras: Métodos padrão atribuem um número uniforme de rollouts para cada prompt, ignorando a variação significativa no "valor de aprendizado" entre diferentes amostras. Alguns prompts induzem mudanças mínimas na política, enquanto outros impulsionam melhorias substanciais.
2. Exploração de Timestep Estática: Estratégias de exploração (por exemplo, ramificação em rollouts baseados em árvores) geralmente seguem cronogramas fixos e predefinidos. Estas falham em se adaptar à natureza dinâmica do processo de denoising, onde os momentos críticos para exploração variam significativamente entre prompts e estágios de treinamento.

Essas limitações resultam em desperdício de recursos computacionais em amostras de baixo valor e em uma exploração subótima em timesteps não críticos, retardando a convergência e limitando o desempenho final de alinhamento.

2. Metodologia: AEGPO

Os autores propõem o Adaptive Entropy-Guided Policy Optimization (AEGPO), um framework plug-and-play que utiliza sinais de atenção intrínsecos para guiar dinamicamente tanto a seleção de amostras quanto a exploração. A ideia central é que a Entropia de Atenção serve como um poderoso proxy de sinal duplo.

2.1 Sinais Principais

O método define duas métricas intrínsecas derivadas dos mapas de atenção do modelo:

• $\Delta$Entropia (Sinal de Nível de Amostra): A mudança relativa na dispersão da atenção entre a política atual e a política base.
  • Definição: $\Delta\text{Entropia} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropia}_{\theta}(t) - \text{Entropia}_{\theta_{base}}(t)|$.
  • Função: Atua como um proxy robusto para o valor de aprendizado da amostra. Uma $\Delta$Entropia alta indica que um prompt força o modelo a reorganizar significativamente sua estratégia de atenção, implicando alto potencial de aprendizado.
• Entropia($t$) (Sinal de Nível de Timestep): A dispersão absoluta da atenção em um passo de denoising específico $t$.
  • Definição: A média da entropia de Shannon dos pesos de atenção através dos tokens de texto para as características da imagem no passo $t$.
  • Função: Picos em Entropia($t$) identificam timesteps críticos de exploração. Esses picos correspondem a momentos de alta incerteza ou integração semântica ampla (frequentemente ocorrendo em estágios iniciais estruturais e de refinamento tardio), onde a exploração gera a maior variância e diversidade de recompensa.

2.2 Estratégia Adaptativa de Duplo Nível

O AEGPO utiliza esses sinais para realocar a computação sem aumentar o orçamento total de rollouts:

1. Alocação Adaptativa Global:

   • O lote de prompts é dividido em níveis de "alto valor" e "baixo valor" com base na mediana da $\Delta$Entropia.
   • Prompts de alto valor (alta $\Delta$Entropia) recebem um orçamento de rollout maior ($r_{high}$), enquanto prompts de baixo valor recebem menos rollouts ($r_{low}$).
   • Restrição: $r_{high} + r_{low} = 2r_{avg}$, garantindo que o total de computação permaneça constante.

2. Exploração Adaptativa Local:

   • Em vez de cronogramas de ramificação fixos, o método identifica dinamicamente os $K$ timesteps com os maiores picos de Entropia($t$) para cada rollout.
   • A ramificação (exploração) é realizada exclusivamente nesses timesteps de alta dispersão identificados.
   • Isso garante que a exploração seja concentrada nos momentos mais informativos da trajetória de denoising.

3. Principais Contribuições

• Descoberta de Sinal Duplo: O artigo estabelece a Entropia de Atenção como um proxy intrínseco de sinal duplo, onde a mudança relativa ($\Delta$Entropia) indica o valor de aprendizado da amostra e os picos absolutos (Entropia($t$)) identificam timesteps críticos de exploração.
• Framework AEGPO: Introdução de um framework adaptativo de duplo nível que melhora a seleção de amostras e a exploração de timesteps usando mecanismos leves guiados por entropia.
• Generalizabilidade: Demonstração de que o AEGPO é agnóstico ao modelo e agnóstico ao framework, aumentando com sucesso diversos métodos de otimização de política (DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT) e modelos base (FLUX.1-dev, SD3.5-M).

4. Resultados Experimentais

Os experimentos foram conduzidos nos datasets HPDv2.1 e Pick-a-Pic utilizando FLUX.1-dev e SD3.5-M como modelos base.

• Velocidade de Convergência: O AEGPO acelera significativamente a otimização da política. Ele alcança uma convergência 2$\times$ mais rápida no DanceGRPO e 5$\times$ mais rápida no DiffusionNFT em comparação com as variantes padrão, atingindo recompensas finais superiores em menos passos.
• Métricas de Desempenho:
  • No FLUX.1-dev, o BranchGRPO + AEGPO alcançou um novo estado da arte (SOTA) de HPS-v2.1 de 0.374 (subindo de 0.363) e Pick Score de 0.232 (subindo de 0.229).
  • No SD3.5-M, o DiffusionNFT + AEGPO melhorou o HPS-v2.1 para 0.343 e o ImageReward para 1.57.
• Trade-off Recompensa-KL: Os modelos AEGPO consistentemente dominam a fronteira de Pareto Recompensa-KL, alcançando recompensas de alinhamento significativamente mais altas para qualquer nível dado de desvio de política (divergência KL).
• Estudos de Ablação: Tanto a Alocação Global (G) quanto a Exploração Local (L) fornecem benefícios independentes, sendo que a abordagem combinada (G+L) produz o maior desempenho.
• Sobrecarga Computacional: O método introduz um overhead modesto (aumento de aproximadamente 11,1% no tempo de passo e aumento negligível de VRAM) devido à extração de mapas de atenção intermediários, o que é justificado pela redução drástica no total de passos de treinamento necessários para atingir o pico de desempenho.

5. Significância e Alegações

O artigo afirma que o AEGPO aborda as ineficiências centrais da amostragem estática em modelos de difusão para alinhamento. Ao mudar de uma alocação baseada em heurística para uma adaptação baseada em sinais intrínsecos, o método garante que os recursos computacionais sejam concentrados nas amostras mais informativas e nos momentos mais críticos do processo de geração.

Os autores enfatizam que o AEGPO é um módulo plug-and-play generalizável que não requer novas funções de perda ou modelagem de recompensa complexa. Em vez disso, ele utiliza a própria dinâmica de atenção interna do modelo para guiar a otimização, oferecendo um caminho mais eficiente para um alinhamento de alta qualidade em diversas arquiteturas de difusão e frameworks de RLHF. O trabalho sugere que sinais intrínsecos como a entropia de atenção são recursos subutilizados para melhorar a eficiência da otimização de política em modelos generativos.