MAR-GRPO: Stabilized GRPO for AR-diffusion Hybrid Image Generation

Resumo Técnico: MAR-GRPO

1. Problema Identificado

O artigo aborda os desafios de aplicar Aprendizado por Reforço (RL), especificamente o algoritmo GRPO (Group Relative Policy Optimization), a modelos híbridos Autoregressivos (AR) e de Difusão (conhecidos como MAR - Masked Autoregressive).

Embora os modelos MAR combinem a geração sequencial de tokens (AR) com cabeças de difusão leve para refinar latentes e gerar imagens de alta qualidade, a otimização via RL online enfrenta instabilidades fundamentais:

• Gradientes Ruidosos: A cabeça de difusão, sendo estocástica, introduz variância significativa nos gradientes. Como o modelo AR depende da saída da difusão para calcular a verossimilhança, essa variância se propaga, criando um sinal de recompensa não estacionário.
• Colapso de Diversidade e Saturação Precoce: A aplicação direta do GRPO padrão (end-to-end) leva a uma otimização instável, onde a qualidade da imagem degrada após um certo número de iterações (saturação precoce) e ocorre o colapso da diversidade de amostras.
• Desacoplamento de Escala: Existe uma incompatibilidade entre a grande escala do transformador AR e a natureza leve da cabeça de difusão, exacerbando a instabilidade durante o ajuste fino.

2. Metodologia Proposta (MAR-GRPO)

Os autores propõem o MAR-GRPO, um framework de RL estabilizado que introduz três componentes principais para mitigar o ruído da difusão e melhorar a dinâmica de treinamento:

1. Expectativa de Múltiplas Trajetórias (MTE - Multi-Trajectory Expectation):

   • Em vez de estimar a direção de otimização baseada em uma única trajetória de difusão (que é ruidosa), o método amostra múltiplas trajetórias de difusão (com diferentes sementes) a partir do mesmo latente AR.
   • Calcula-se a expectativa (média) sobre essas trajetórias para estimar a verossimilhança do token. Isso reduz a variância do gradiente induzida pela estocasticidade da difusão, fornecendo uma direção de atualização mais estável para o modelo AR.

2. Seleção de Tokens Baseada em Incerteza (Top-k%):

   • Aplicar a média de múltiplas trajetórias a todos os tokens seria computacionalmente custoso e poderia suavizar excessivamente gradientes informativos.
   • O método estima um mapa de incerteza token a token (calculando a variância entre as múltiplas trajetórias).
   • A otimização de múltiplas trajetórias é aplicada seletivamente apenas aos top-k% de tokens mais incertos (geralmente regiões com semântica ambígua ou detalhes finos). Os tokens com baixa incerteza (estáveis) continuam sendo otimizados com a trajetória padrão, preservando a eficiência e a nitidez dos gradientes.

3. Seleção de Tokens Consciente de Consistência (Consistency-Aware):

   • Nem todos os tokens gerados durante o rollout contribuem positivamente para o resultado final.
   • O método avalia a consistência das previsões do AR ao longo das etapas de mascaramento. Tokens cujas previsões mudam drasticamente ou não convergem para a imagem final são filtrados (excluídos) do cálculo da perda de otimização. Isso evita que o modelo aprenda com ruídos ou tokens que não refletem a intenção final.

Estratégia de Treinamento: O estudo também revela que congelar a cabeça de difusão durante o treinamento do RL (otimizando apenas o transformador AR) é crucial para estabilizar o processo, removendo a não estacionariedade causada pela atualização contínua dos parâmetros de difusão.

3. Principais Contribuições

• Primeiro Framework GRPO para MAR: É a primeira proposta de otimização baseada em GRPO especificamente adaptada para a arquitetura híbrida AR-Difusão.
• Identificação da Fonte de Instabilidade: Demonstra empiricamente e analiticamente que a cabeça de difusão é a principal fonte de variância de gradiente e instabilidade em modelos MAR sob RL.
• Mecanismo de Correção de Ruído: Introduz a técnica de Expectativa de Múltiplas Trajetórias combinada com seleção seletiva de tokens, que reduz o ruído sem sacrificar a eficiência computacional (devido à natureza leve da cabeça de difusão).
• Estratégia de Filtragem de Tokens: Propõe uma métrica de consistência para filtrar tokens que degradam a qualidade do treinamento.

4. Resultados Experimentais

Os experimentos foram conduzidos em modelos base como NOVA e Harmon, utilizando benchmarks de preferência humana e compreensão espacial.

• Qualidade Visual e Estabilidade: O MAR-GRPO supera consistentemente o GRPO padrão e modelos pré-RL. Enquanto o GRPO padrão sofre de colapso de diversidade e degradação de qualidade após ~250 iterações, o MAR-GRPO mantém uma melhoria estável e atinge picos de desempenho mais altos.
• Métricas de Preferência Humana:
  • HPS (Human Preference Score): Aumento significativo (ex: de 26.76 para 29.35 no modelo NOVA).
  • ImageReward e Aesthetic Score: Melhorias consistentes, indicando imagens mais esteticamente agradáveis e alinhadas com o prompt.
• Compreensão Espacial e Composicional:
  • No benchmark T2I-CompBench, o método demonstrou melhorias em relações espaciais (2D e 3D), contagem de objetos e fidelidade estrutural.
  • O GRPO padrão frequentemente introduzia distorções espaciais e texturas não naturais, enquanto o MAR-GRPO preservou a fidelidade estrutural.
• Eficiência: A adição de múltiplas trajetórias e seleção de tokens introduz um custo computacional negligenciável, pois a cabeça de difusão é leve e a seleção é seletiva.

5. Significado e Impacto

Este trabalho é significativo porque resolve uma barreira fundamental na aplicação de RL a modelos de geração de imagem híbridos (AR-Difusão). Ao estabilizar a dinâmica de treinamento, o MAR-GRPO permite que esses modelos alcancem seu potencial máximo em termos de fidelidade ao prompt, qualidade estética e compreensão espacial, superando as limitações de métodos anteriores que sofriam com instabilidade e colapso.

O artigo estabelece um novo padrão para o ajuste fino de modelos MAR via RL, sugerindo que a gestão cuidadosa da estocasticidade da difusão é tão crítica quanto a própria arquitetura do modelo. O código e os resultados abrem caminho para futuras pesquisas em geração unificada de texto e imagem com reforço.