OP-GRPO: Efficient Off-Policy GRPO for Flow-Matching Models

Título: OP-GRPO: GRPO Off-Policy Eficiente para Modelos de Flow-Matching

1. O Problema

Os modelos de Flow-Matching (como SD3.5, Flux e Wan) têm se destacado na geração de imagens e vídeos devido à sua eficiência de amostragem e fidelidade visual. Para alinhar esses modelos com preferências humanas, utiliza-se frequentemente o GRPO (Group Relative Policy Optimization), uma técnica de aprendizado por reforço (RL).

No entanto, o GRPO padrão enfrenta duas limitações críticas em modelos de Flow-Matching:

1. Baixa Eficiência de Amostragem (On-Policy): O GRPO segue um paradigma estritamente on-policy. Isso significa que, a cada iteração de treinamento, novas trajetórias (amostras) são geradas e, ao final do passo, são descartadas. Mesmo amostras de alta qualidade não podem ser reutilizadas, resultando em um custo computacional elevado e tempo de treinamento longo (milhares de horas de GPU).
2. Instabilidade em Tarefas Difíceis: Quando a tarefa é complexa, o modelo raramente gera trajetórias bem-sucedidas, fazendo com que as recompensas decaiam para zero. Isso causa o colapso dos advantages (vantagens) e o desaparecimento dos sinais de gradiente, estagnando o aprendizado.

A aplicação direta de métodos off-policy (que reutilizam dados antigos) é desafiadora devido ao deslocamento de distribuição (distribution shift). Reutilizar trajetórias geradas por políticas antigas distorce o mecanismo de clipping do GRPO, podendo levar à instabilidade do treinamento. Além disso, observa-se que os passos finais do processo de denoising (remoção de ruído) tornam-se numericamente instáveis para estimativas off-policy.

2. Metodologia: OP-GRPO

O OP-GRPO é o primeiro framework off-policy de GRPO projetado especificamente para modelos de Flow-Matching. A abordagem integra três componentes principais para mitigar as limitações acima:

A. Replay Buffer de Alta Qualidade

• O método mantém um buffer de replay ($\mathcal{B}_{off}$) que armazena trajetórias de alta qualidade coletadas durante fases anteriores.
• Seleção Ativa: Durante o rollout, o buffer é atualizado mantendo apenas as trajetórias com as maiores recompensas para cada condição de entrada ($c$), garantindo diversidade e evitando a degeneração para amostras triviais.
• Decaimento de Recompensa: Para evitar que amostras muito antigas (demasiado off-policy) dominem o treinamento, aplica-se um decaimento nos valores de recompensa das trajetórias no buffer, favorecendo amostras mais recentes.
• Híbrido: No treinamento, uma parte das amostras é gerada on-policy (pela política atual) e outra parte é substituída por amostras do buffer.

B. Correção de Importância em Nível de Sequência

• Um desafio central é corrigir o viés introduzido pela reutilização de dados sem quebrar o mecanismo de clipping do GRPO (que limita a magnitude da atualização da política).
• Uma correção ingênua (aplicar importance sampling passo a passo) faria com que o termo de clipping medisse a diferença entre a política atual e a política antiga de muitos passos atrás, resultando em um alto índice de amostras descartadas (mais de 40%).
• Solução: O OP-GRPO propõe uma correção de importância em nível de sequência. O termo de clipping continua a medir a mudança em relação à política de referência imediata ($\pi_{old}$), preservando a estabilidade. A correção de distribuição é aplicada como um termo multiplicativo global na função de objetivo:
  $$\mathcal{J}_{\text{OP-GRPO}}(\theta) = \mathbb{E} \left[ \frac{P_{\pi_{\text{old}}}(\tau)}{P_{\pi_{\text{off}}}(\tau)} \cdot f(r, \hat{A}, \theta, \epsilon, \beta) \right]$$
  Isso reduz drasticamente a taxa de amostras off-policy descartadas (de >40% para ~11.8%), permitindo atualizações estáveis e eficientes.

C. Truncamento de Trajetória (Trajectory Truncation)

• Problema Numérico: O artigo demonstra teórica e empiricamente que, nos passos finais do denoising (quando o ruído $\sigma \to 0$), a distribuição de transição torna-se quase determinística. Isso faz com que os pesos de importância se tornem mal condicionados (valores extremamente grandes ou pequenos), causando instabilidade numérica.
• Solução: Adota-se uma estratégia de truncamento. As trajetórias off-policy são utilizadas apenas até um passo de truncamento ($t_{off}$). Os passos subsequentes (de $t_{off}$ até 0) são re-gerados pela política atual (on-policy). Isso elimina a instabilidade numérica dos passos finais sem prejudicar significativamente a qualidade perceptual da imagem/vídeo.

3. Contribuições Principais

1. Primeiro Framework Off-Policy para Flow-Matching: Estende o GRPO para um cenário off-policy, permitindo a reutilização de trajetórias históricas e aumentando a eficiência de amostragem.
2. Correção de Distribuição Eficiente: Desenvolveu um mecanismo de correção de importância em nível de sequência que preserva a estabilidade do clipping do GRPO, reduzindo a taxa de rejeição de amostras úteis.
3. Análise Teórica e Prática de Estabilidade: Identificou e mitigou a instabilidade numérica nos passos finais de denoising através do truncamento de trajetórias.
4. Eficiência de Treinamento: Demonstra que é possível atingir desempenho superior ou comparável com uma fração significativa dos passos de treinamento necessários para métodos on-policy.

4. Resultados Experimentais

Os experimentos foram realizados em modelos de geração de imagem (SD3.5-M) e vídeo (Wan2.1-1.4B) em três tarefas: Geração Composicional, Renderização de Texto Visual e Alinhamento de Preferência Humana.

• Eficiência de Treinamento: O OP-GRPO alcançou o desempenho final do Flow-GRPO (base on-policy) utilizando apenas 34.2% dos passos de treinamento em média. Em tarefas de vídeo, essa redução foi de 30.1%.
• Desempenho: O OP-GRPO não apenas acelerou a convergência, mas também atingiu métricas finais superiores em várias tarefas (ex: melhoria média de 1.17% nas métricas de alinhamento e qualidade).
• Generalização: O modelo treinado com OP-GRPO demonstrou melhor generalização em métricas não vistas durante o treinamento (como PickScore e ImageReward) e em benchmarks mais complexos (T2I-CompBench++), indicando que o mecanismo off-policy não causa overfitting excessivo.
• Estabilidade: Estudos de ablação confirmaram que remover a correção de importância ou o truncamento leva à divergência ou instabilidade do treinamento.

5. Significado e Impacto

O trabalho OP-GRPO representa um avanço significativo na otimização de modelos generativos baseados em Flow-Matching. Ao resolver o gargalo da ineficiência de amostragem do GRPO on-policy, o método permite:

• Redução de Custos: Diminuição drástica no tempo de treinamento e uso de GPU para alinhar modelos de última geração.
• Escalabilidade: Facilita o treinamento de modelos em tarefas complexas onde a geração de amostras de alta qualidade é rara, pois permite que o modelo "reaproveite" e aprenda com esses poucos exemplos valiosos.
• Novo Paradigma: Estabelece as bases para o uso de aprendizado off-policy em modelos de difusão e fluxo contínuo, abrindo caminho para futuras aplicações em áudio, 3D e outras modalidades.

Em resumo, o OP-GRPO transforma o treinamento de modelos de Flow-Matching de um processo caro e ineficiente em um processo mais rápido, estável e escalável, sem sacrificar a qualidade da geração.