Resumo Técnico: JAGG (Jacobian-Aggregated Group Gradient)

Declaração do Problema

O Group Relative Policy Optimization (GRPO) emergiu como um algoritmo de aprendizado por reforço (RL) líder para alinhar modelos generativos com preferências humanas, eliminando com sucesso a necessidade de uma rede de valor separada em Large Language Models (LLMs). No entanto, estender o GRPO para modelos de difusão e flow matching (especificamente Diffusion Transformers ou DiTs) introduz um gargalo computacional severo.

No GRPO baseado em difusão, o processo de geração é tratado como um problema de decisão de múltiplos passos sobre uma trajetória de amostragem de $W$ passos temporais. O treinamento padrão requer uma passagem completa de forward e backward através da espinha dorsal (backbone) de alta capacidade do DiT para cada único passo na trajetória para calcular o gradiente da política. Para a geração de texto-para-imagem (T2I) de alta resolução, onde as trajetórias frequentemente abrangem dezenas de passos, isso resulta em custos de treinamento proibitivamente caros. Diferente dos modelos de linguagem, onde os passos de tokens são leves, cada passagem backward de um DiT moderno incorre no custo total de uma passagem forward de um transformer, tornando o requisito de backward por passo o componente dominante do custo de treinamento.

Embora métodos de aceleração de inferência sem treinamento (ex: $\Delta$-DiT, ScalingCache) tenham explorado com sucesso a observação de que os estados ocultos e as previsões de velocidade do DiT variam suavemente e quase linearmente ao longo da trajetória de denoising para pular avaliações intermediárias, permanecia uma questão aberta se essa mesma linearidade poderia ser aproveitada para reduzir o custo da passagem backward no treinamento de RL.

Metodologia: JAGG

Os autores propõem o JAGG (Jacobian-Aggregated Group Gradient), um método de aceleração matematicamente fundamentado que reduz o número de passagens backward da rede de $W$ para 2 para um grupo de $W$ passos de denoising consecutivos.

Insight Central

O método baseia-se no insight de que o Jacobiano do campo de velocidade de denoising $v_\theta(z, t)$ exibe uma linearidade aproximada nos inputs $(z, t)$, particularmente durante o regime de alto ruído do processo de denoising. Isso permite que os Jacobianos dos passos intermediários sejam precisamente aproximados via interpolação linear usando apenas os Jacobianos das extremidades ($J_0$ e $J_{W-1}$).

Fluxo de Trabalho Algorítmico

O JAGG processa um grupo de $W$ passos temporais através de quatro fases:

1. Passagens Forward: O modelo realiza passagens forward para todos os $W$ passos. No entanto, o rastreamento de gradiente é habilitado apenas para as duas extremidades (passo 0 e passo $W-1$). Os $W-2$ passos intermediários são executados usando torch.no_grad() para economizar memória e computação.
2. Computação do Gradiente Upstream: Para cada passo $j$, um proxy desconectado (detached) da saída do modelo é criado. A perda por passo $L_j$ (derivada do objetivo estilo PPO) é calculada e o gradiente upstream $s_j = \partial L_j / \partial \hat{v}_j$ é calculado via diferenciação automática. Crucialmente, este passo não envolve nenhuma passagem backward do transformer; o gradiente flui apenas através do passo SDE leve e do objetivo PPO.
3. Agregação Ponderada pelo Passo Temporal: Os gradientes upstream $s_j$ são agregados em dois sinais compostos, $s_{base}$ e $s_{corr}$, usando pesos de interpolação $\alpha_j$ derivados do passo temporal $t_j$:
   $$\alpha_j = \frac{t_0 - t_j}{t_0 - t_{W-1}}$$
   $$s_{base} = \sum_{j=0}^{W-1} (1 - \alpha_j) s_j, \quad s_{corr} = \sum_{j=0}^{W-1} \alpha_j s_j$$
   Esta agregação reconstrói efetivamente o sinal de gradiente completo usando a suposição de interpolação linear.
4. Backward Conjunto: Uma única chamada torch.autograd.backward é emitida nas duas saídas das extremidades $[\hat{v}_0, \hat{v}_{W-1}]$ com os gradientes agregados $[s_{base}, s_{corr}]$. Isso aciona exatamente um FSDP all-gather por módulo, substituindo os $W$ separadas passagens backward exigidas pelo treinamento padrão.

Garantia Teórica

Os autores fornecem uma prova formal (Teorema 1) afirmando que, se o campo de velocidade $v_\theta(z, t)$ for estritamente linear em $(z, t)$, a interpolação do Jacobiano ponderada por $t$ é exata, com erro de aproximação zero. Sob estas condições, a passagem backward de duas extremidades recupera o gradiente de grupo exato.

Roteamento Adaptativo ($jagg\_frac$)

Reconhecendo que a suposição de linearidade degrada nos estágios finais (baixo ruído) do denoising, os autores introduzem um parâmetro de roteamento $jagg\_frac$. Este parâmetro aplica o JAGG apenas à primeira fração dos grupos de denoising (onde a suposição de linearidade de alto ruído se mantém, verificada empiricamente via similaridade de cosseno de gradientes) e utiliza passagens backward exatas por passo para os grupos restantes. Para os modelos T2I testados, o $jagg\_frac$ foi definido como 0.6.

Principais Contribuições

1. Identificação do Gargalo: O artigo identifica e caracteriza formalmente o gargalo computacional no treinamento de difusão GRPO causado pelas passagens backward por passo ao longo da trajetória de amostragem.
2. Framework JAGG: Propõe um framework de agregação de Jacobiano que reduz as passagens backward de $W$ para 2 por grupo, fundamentado em uma análise teórica rigorosa da linearidade do Jacobiano de velocidade.
3. Prova Teórica: Os autores provam que, sob condições de velocidade linear, a interpolação proposta é exata, fornecendo uma base teórica sólida para a aproximação.
4. Validação Empírica: Experimentos extensos em múltiplos backbones de difusão (Flux, QwenImage) e algoritmos de RL (DanceGRPO, Flow-GRPO-Fast) demonstram a generalidade e eficácia do método.

Resultados Experimentais

Os autores avaliaram o JAGG em benchmarks de T2I usando HPSv2.1, PickScore e CLIPScore como métricas.

• Eficiência: O JAGG alcançou um aumento de velocidade de tempo real (wall-clock speedup) consistente de 17,3% a 18,2% por passo de treinamento em configurações padrão de DanceGRPO (reduzindo o tempo de ~835s para ~691s no QwenImage e de ~279s para ~228s no Flux). Quando aplicado ao Flow-GRPO-Fast, apenas a etapa de atualização viu aumentos de velocidade de 29,9% a 30,0%.
• Qualidade: O método entregou esses aumentos de velocidade com degradação negligenciável na qualidade de geração. Em muitos casos, o JAGG acompanhou de perto as curvas de recompensa da linha de base (baseline) e igualou ou superou ligeiramente a linha de base na convergência.
• Ablação: Uma variante de ablação que realizava passagens backward nas extremidades mas sem a interpolação de Jacobiano (atribuindo a perda total às extremidades independentemente) apresentou desempenho consistentemente inferior. Isso confirma que a agregação da informação de gradiente intermediária via interpolação de Jacobiano é crítica para um treinamento estável.
• Qualidade do Gradiente: A análise empírica mostrou que em grupos de alto ruído (0–2), a similaridade de cosseno entre o gradiente agregado pelo JAGG e o gradiente exato foi alta (0,74–0,85), validando a suposição de linearidade. A similaridade caiu em grupos posteriores, justificando a estratégia de roteamento $jagg\_frac$.

Significância e Alegações

O artigo afirma que o JAGG representa um passo significativo para tornar o treinamento de difusão GRPO em larga escala prático. Ao explorar a quase-linearidade dos Jacobianos de velocidade no regime de alto ruído, o método reduz a carga computacional do treinamento de RL sem sacrificar a qualidade de geração.

Os autores enfatizam que o JAGG não é meramente um atalho heurístico, mas está fundamentado nas propriedades matemáticas da trajetória de denoising. O método é apresentado como ortogonal às técnicas existentes de aceleração de inferência, visando especificamente o gargalo da passagem backward exclusivo do RL de difusão. Os resultados sugerem que o JAGG permite uma adoção mais ampla do alinhamento baseado em RL para modelos generativos visuais ao reduzir significativamente a barreira de custo de treinamento.