Principled RL for Flow Matching Emerges from the Chunk-level Policy Optimization

Resumo Técnico: RL Principiada para Flow Matching via Otimização de Política em Nível de Chunk

Declaração do Problema
Avanços recentes em modelos de flow matching texto-para-imagem (T2I) pós-treinamento têm aproveitado a Otimização de Política Relativa em Grupo (GRPO) para alinhar imagens geradas com preferências humanas. No entanto, a GRPO sofre de uma limitação crítica: atribuição imprecisa de vantagem. Na GRPO padrão, a recompensa final obtida de uma imagem gerada é atribuída uniformemente a cada etapa da trajetória de geração. Isso depende da suposição de que um resultado final superior implica uma política superior em cada etapa intermediária. Os autores demonstram que essa suposição é frequentemente violada; uma trajetória que produz uma recompensa final melhor pode conter etapas intermediárias com políticas inferiores em comparação com uma trajetória com uma recompensa final pior. Essa má atribuição introduz sinais de otimização ruidosos, levando à instabilidade no treinamento. Embora modelos de recompensa conscientes do processo possam teoricamente resolver isso, eles exigem grandes quantidades de dados de preferência para imagens intermediárias ruidosas, que atualmente não estão disponíveis.

Metodologia: Otimização de Política de Agrupamento em Chunk (GCPO)
Para abordar a atribuição imprecisa de vantagem sem introduzir recompensas de processo, os autores propõem a Otimização de Política de Agrupamento em Chunk (GCPO). Esta abordagem desloca o paradigma de otimização de política do nível de etapa para o nível de chunk.

1. Agregação em Nível de Chunk: Em vez de tratar cada etapa de geração como uma ação independente, a GCPO agrega etapas consecutivas em um "chunk de ação" coerente. A política é otimizada sobre esses chunks, tratando cada chunk como uma única ação atômica. Essa agregação suaviza as flutuações de gradiente causadas pela má atribuição de vantagens em etapas individuais.
2. Agrupamento Guiado por Dinâmica Temporal: Os autores observam que o flow matching exibe dinâmicas temporais distintas, onde diferentes etapas contribuem desigualmente para a imagem final. Eles utilizam a distância relativa $L_1$ no espaço latente ($L1_{rel}(x, t)$) para quantificar a taxa de mudança. Etapas com dinâmicas temporais semelhantes (por exemplo, taxas semelhantes de mudança latente) são agrupadas no mesmo chunk, enquanto etapas com dinâmicas distintas são separadas. Uma estratégia adaptativa é empregada, analisando recursivamente derivadas da distância relativa para particionar trajetórias em chunks significativos.
3. Razão de Importância Principiada: O objetivo de otimização é reformulado para usar uma razão de importância em nível de chunk. Essa razão é calculada com base na probabilidade de todo o chunk, normalizada pelo tamanho do chunk, garantindo estabilidade.
4. Amostragem Ponderada (Opcional): Para aprimorar ainda mais o desempenho, é introduzida uma estratégia opcional de amostragem ponderada. Com base em estudos de ablação mostrando que regiões de alto ruído (etapas iniciais) contribuem mais para o alinhamento de preferências, mas sofrem de instabilidade, essa estratégia atribui pesos de amostragem mais altos a chunks com valores médios mais elevados de $L1_{rel}$.

Principais Contribuições

• Pioneirismo em RL em Nível de Chunk: O artigo introduz a primeira abordagem de aprendizado por reforço em nível de chunk para flow matching pós-treinamento em geração T2I.
• Mitigação de Problemas de Atribuição de Vantagem: Ao tratar etapas consecutivas como uma única ação, a GCPO mitiga efetivamente os efeitos adversos da atribuição imprecisa de vantagem inerente à GRPO em nível de etapa, levando a um treinamento mais estável.
• Integração de Dinâmicas Temporais: O método integra agrupamento guiado por dinâmica temporal com uma razão de importância em nível de chunk principiada, alinhando o processo de otimização com a estrutura intrínseca do flow matching.
• Ganhos de Desempenho: Experimentos extensivos demonstram que a GCPO alcança desempenho superior tanto no alinhamento de preferências quanto em benchmarks padrão T2I, com ganhos relativos de até 43% sobre as bases GRPO.

Resultados Experimentais
Os autores avaliaram a GCPO usando o modelo FLUX.1 Dev no conjunto de dados HPDv2.1, comparando-o com bases como Dance-GRPO e Flow-GRPO.

• Alinhamento de Preferências: A GCPO alcançou melhorias significativas nas métricas intra-domínio (HPSv3) e extra-domínio (ImageReward), superando as bases em até 43%.
• Benchmarks Padrão: Nos benchmarks GenEval e DPG, a GCPO demonstrou o desempenho mais forte, com ganhos relativos até três vezes maiores que as bases.
• Estudo com Usuários: Em um estudo de preferência humana com 40 prompts, variantes da GCPO foram preferidas em relação ao Dance-GRPO em 72,5% dos casos.
• Estudos de Ablação:
  • Estratégia de Agrupamento: O agrupamento guiado por dinâmica temporal superou o agrupamento de tamanho fixo, confirmando a importância de alinhar a otimização com estruturas temporais intrínsecas.
  • Amostragem Ponderada: Embora essa estratégia tenha melhorado o alinhamento de preferências, mostrou efeitos mistos nos benchmarks padrão T2I, ocasionalmente desestabilizando estruturas de imagem em regiões de alto ruído.

Significado e Alegações
O artigo alega que a GCPO representa um avanço principiado em RL para flow matching ao abordar a questão fundamental da atribuição imprecisa de vantagem por meio de uma mudança na granularidade de otimização. Os autores argumentam que a transição da otimização em nível de etapa para a otimização em nível de chunk, guiada pelas dinâmicas temporais inerentes ao processo de geração, fornece uma solução robusta que não requer os dados indisponíveis para modelos de recompensa de processo. O trabalho destaca que a otimização de política em nível de chunk é uma direção promissora para estabilizar e melhorar o pós-treinamento de modelos de fluxo generativos. Os autores notam modestamente limitações, como a possibilidade de melhorias adicionais combinando recompensas heterogêneas entre diferentes chunks, e reconhecem os riscos gerais associados à IA generativa avançada, incluindo o potencial de uso indevido na criação de deepfakes.