Stepwise Credit Assignment for GRPO on Flow-Matching Models

Título: Atribuição de Crédito Passo a Passo para GRPO em Modelos de Fluxo (Flow-Matching)

1. Problema Identificado

O artigo aborda as limitações da aplicação de Aprendizado por Reforço (RL) em modelos de geração de imagens baseados em Flow-Matching (como o SD3.5), especificamente utilizando o algoritmo Flow-GRPO.

O problema central é a atribuição de crédito uniforme. No Flow-GRPO existente, a recompensa é calculada apenas com base na imagem final gerada e atribuída uniformemente a todos os passos de remoção de ruído (denoising steps) na trajetória. Isso ignora a estrutura temporal inerente à geração difusiva:

• Estrutura Hierárquica: Os passos iniciais determinam a composição, layout e estrutura de baixa frequência, enquanto os passos finais refinam detalhes e texturas de alta frequência.
• Recompensa de Erros Corrigidos: Uma trajetória pode cometer erros graves nos passos iniciais (ex: cor ou posição errada de um objeto) que são corrigidos nos passos subsequentes. Se a imagem final for boa, o Flow-GRPO padrão reforça erroneamente todos os passos, incluindo os passos iniciais defeituosos, o que pode levar a artefatos e instabilidade no treinamento.

2. Metodologia Proposta: Stepwise-Flow-GRPO

Os autores propõem o Stepwise-Flow-GRPO, uma abordagem que realiza uma atribuição de crédito granular baseada no ganho de recompensa em cada passo individual, em vez de apenas no resultado final.

Componentes Principais:

1. Estimativa de Recompensas Intermediárias (Tweedie's Formula):

   • Como os modelos de recompensa não podem avaliar diretamente estados ruidosos ($x_t$), o método utiliza a Fórmula de Tweedie para estimar a imagem limpa subjacente $\hat{x}_0(t) = \mathbb{E}[x_0|x_t]$ a partir de cada estado intermediário.
   • Para maior precisão, realizam uma sub-otimização determinística (vários passos ODE) a partir de $x_t$ para obter uma estimativa de alta qualidade antes de calcular a recompensa $r_t = R(\hat{x}_0(t), \text{prompt})$.

2. Atribuição de Crédito Baseada em Ganhos (Gain-Based Advantages):

   • Em vez de usar a recompensa final, calculam o ganho passo a passo: $g_t = r_{t-1} - r_t$.
   • Isso mede quanto cada passo específico melhorou (ou piorou) a qualidade da imagem em relação ao passo anterior.
   • Passos que aumentam a recompensa recebem reforço positivo; passos que a diminuem são penalizados.
   • Os ganhos são normalizados conjuntamente (across all steps) para preservar a magnitude natural dos ganhos iniciais (que são maiores e mais críticos para a composição).

3. SDE Inspirado em DDIM (Amostragem Melhorada):

   • O Flow-GRPO original converte a ODE determinística em uma Equação Diferencial Estocástica (SDE) para permitir exploração, mas isso introduz ruído visual que degrada a qualidade das imagens intermediárias avaliadas pelo modelo de recompensa.
   • Os autores introduzem uma nova SDE inspirada no DDIM (Denoising Diffusion Implicit Models). Esta formulação injeta estocasticidade controlada necessária para o gradiente de política, mas produz amostras intermediárias muito mais limpas, melhorando a qualidade do sinal de recompensa sem sacrificar a exploração.

3. Contribuições Chave

• Mecanismo de Atribuição de Crédito Passo a Passo: Substituição da recompensa de trajetória única por ganhos marginais por passo, alinhando o aprendizado com a estrutura temporal da geração de imagens.
• Teoria de Submodularidade: O método conecta-se teoricamente à otimização submodular adaptativa, onde a maximização gananciosa de ganhos marginais leva a políticas quase ótimas, especialmente quando os ganhos diminuem ao longo do tempo (como observado empiricamente).
• SDE de Alta Qualidade: Desenvolvimento de uma nova equação de difusão que equilibra a necessidade de estocasticidade para RL com a necessidade de imagens limpas para avaliação de recompensa.
• Eficiência de Amostra e Convergência: Demonstração de que o método aprende mais rápido e com menos iterações de treinamento do que o Flow-GRPO padrão.

4. Resultados Experimentais

Os experimentos foram conduzidos no modelo SD3.5-Medium utilizando os benchmarks GenEval (foco em composição e contagem) e PickScore (alinhamento humano).

• Eficiência de Amostra: O Stepwise-Flow-GRPO atinge recompensas mais altas em menos passos de treinamento em comparação ao Flow-GRPO padrão em todos os cenários testados.
• Convergência Temporal: Apesar de exigir computação adicional para estimar recompensas intermediárias, o método converge mais rápido em tempo de relógio (wall-clock time) devido à maior eficiência de amostragem (requer 2-3x menos iterações para atingir o mesmo desempenho).
• Qualidade Final:
  • No benchmark GenEval, o método superou o Flow-GRPO, especialmente em tarefas complexas como contagem de objetos, posicionamento espacial e ligação de atributos.
  • Em extensões de treinamento (400 horas de GPU), o método alcançou uma pontuação geral de 0.87, superando o Flow-GRPO (0.72) e aproximando-se de modelos autoregressivos de ponta como o GPT-4o (0.84).
• Estabilidade: O método demonstrou maior estabilidade ao treinar com modelos de recompensa complexos (como UnifiedReward-7b), onde o Flow-GRPO padrão tendia a divergir.
• Ablação: A normalização conjunta dos ganhos (preservando a magnitude temporal) provou ser superior à normalização passo a passo. O uso de 5 sub-passos ODE ($T'=5$) foi identificado como o melhor equilíbrio entre custo e qualidade da estimativa.

5. Significado e Impacto

Este trabalho representa um avanço significativo na aplicação de RL a modelos de geração de imagens difusivos/flow-matching. Ao reconhecer que a geração de imagens é um processo sequencial com contribuições temporais distintas, o Stepwise-Flow-GRPO resolve o problema de "recompensar erros corrigidos", permitindo que o modelo aprenda a tomar decisões corretas desde os primeiros passos de composição.

A combinação de uma atribuição de crédito mais inteligente com uma formulação de amostragem (SDE) aprimorada oferece um caminho robusto para melhorar o alinhamento texto-imagem, a coerência espacial e a qualidade estética, estabelecendo um novo padrão para o treinamento de modelos de geração via RL.