⚡ electrical engineering

AGG: Jacobian-Aggregated Group Gradient for Efficient GRPO Training of Diffusion Models

Questo articolo introduce JAGG (Jacobian-Aggregated Group Gradient), un metodo che sfrutta la variazione quasi lineare delle traiettorie dei modelli di diffusione per approssimare i Jacobiani intermedi e aggregare i gradienti di gruppo, riducendo così il costo computazionale dell'addestramento di Group Relative Policy Optimization (GRPO) per i Diffusion Transformer di circa 2 $\times$ senza una perdita significativa di qualità.

Autori originali: Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

Pubblicato 2026-07-21

📖 3 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di cercare di insegnare a un artista digitale come dipingere immagini migliori. Questo artista non è un essere umano; è un complesso programma per computer chiamato "modello di diffusione". Pensa a questo modello come a uno scultore che parte da un blocco di argilla rumorosa e piena di staticità e che lentamente scava via il rumore, passo dopo passo, finché non emerge una statua perfetta. Per insegnare a questo scultore come creare arte che piaccia davvero agli umani, usiamo una tecnica chiamata "Reinforcement Learning" (Apprendimento per Rinforzo). È come un gioco in cui lo scultore prova diverse mosse di scalpello, riceve un punteggio da un giudice (un modello di ricompensa) e poi impara da quel punteggio per fare meglio la volta successiva.

Il problema è che questo processo di apprendimento è incredibilmente costoso. Ogni volta che lo scultore compie una piccola scheggiatura (un "passo" nel processo), il computer deve eseguire un calcolo massiccio e vorace di energia per capire esattamente come regolare il cervello dello scultore. Se la statua richiede 20 passi per essere finita, il computer deve fare questa matematica pesante 20 volte per ogni singola lezione. È come cercare di imparare un nuovo ballo fermandosi dopo ogni singolo movimento per chiedere un critico completo all'allenatore, riscrivendo tutta la propria memoria muscolare e poi ricominciando da capo. Per le immagini ad alta risoluzione, questo richiede così tanto tempo ed elettricità che è quasi impossibile addestrare questi modelli in modo efficace.

È qui che entra in gioco una nuova idea chiamata JAGG. I ricercatori dietro questo articolo si sono posti una domanda intelligente: "Abbiamo davvero bisogno di fermarci e ricalcolare tutto dopo ogni singola mossa?". Hanno notato che nelle fasi iniziali, disordinate, del processo di scultura, i cambiamenti avvengono in modo molto fluido e prevedibile — quasi come una linea retta. Se sai dove inizia la scultura e dove finisce dopo alcuni movimenti, puoi effettivamente indovinare cosa è successo nel mezzo senza fare la matematica pesante per ogni singolo passo.

Il documento introduce un metodo chiamato JAGG (Jacobian-Aggregated Group Gradient). Invece di fermarsi a calcolare la lezione dopo ogni singolo passo, JAGG permette al modello di compiere un gruppo di passi (per esempio, quattro mosse di fila) e di eseguire la matematica pesante dell'aggiornamento del "cervello" solo due volte: una all'inizio del gruppo e una alla fine del gruppo. Per i passi nel mezzo, utilizza una scorciatoia intelligente. Assume che il percorso tra l'inizio e la fine sia una linea retta e riempie i vuoti.

I ricercatori hanno dimostrato matematicamente che se il processo di scultura è perfettamente fluido e lineare, questa scorciatoia è in realtà perfetta — fornisce esattamente la stessa risposta di chi fa il lavoro duro ogni volta. Nel mondo reale, dove le cose non sono perfettamente dritte, hanno scoperto che questa scorciatoia funziona incredibilmente bene durante le fasi iniziali di creazione dell'immagine, quelle più "rumorose". Utilizzando questo metodo, sono stati in grado di ridurre il tempo necessario per addestrare questi modelli di circa il 17% o 18% per passo.

La parte migliore? La qualità delle immagini non ne ha risentito. Quando hanno testato JAGG su diversi modelli (come Flux e QwenImage), le immagini generate erano buone quanto quelle create con il metodo lento e tradizionale. In alcuni casi, la scorciatoia ha persino aiutato il modello a imparare in modo più stabile. Il documento mostra che essendo un po' più intelligenti su quando fare il lavoro pesante, possiamo rendere l'addestramento di questi potenti artisti IA molto più veloce ed economico, senza sacrificare la bellezza dell'opera d'arte finale.

Riepilogo Tecnico: JAGG (Jacobian-Aggregated Group Gradient)

Definizione del Problema

Il Group Relative Policy Optimization (GRPO) è emerso come un algoritmo di reinforcement learning (RL) leader per l'allineamento di modelli generativi con le preferenze umane, eliminando con successo la necessità di una rete di valore separata nei Large Language Models (LLM). Tuttavia, estendere il GRPO ai modelli di diffusione e flow matching (specificamente i Diffusion Transformers o DiT) introduce un grave collo di bottiglia computazionale.

Nel GRPO basato sulla diffusione, il processo di generazione è trattato come un problema decisionale multi-step su una traiettoria di campionamento di $W$ timestep. L'addestramento standard richiede un passaggio forward e uno backward completo attraverso l'alto capacità del backbone DiT per ogni singolo timestep della traiettoria per calcolare il gradiente della policy. Per la generazione di testo-immagine (T2I) ad alta risoluzione, dove le traiettorie spesso coprono decine di step, ciò si traduce in costi di addestramento proibitivi. A differenza dei modelli linguistici dove gli step dei token sono leggeri, ogni passaggio backward di un moderno DiT comporta l'intero costo di un forward pass di un transformer, rendendo il requisito del backward per step l'elemento dominante della spesa di addestramento.

Sebbene i metodi di accelerazione dell'inferenza senza addestramento (ad es., $\Delta$ -DiT, ScalingCache) abbiano sfruttato con successo l'osservazione che gli stati nascosti del DiT e le predizioni di velocità variano in modo fluido e quasi lineare lungo la traiettoria di denoising per saltare le valutazioni intermedie, rimaneva aperto il quesito se questa stessa linearità potesse essere sfruttata per ridurre il costo del passaggio backward nell'addestramento RL.

Metodologia: JAGG

Gli autori propongono JAGG (Jacobian-Aggregated Group Gradient), un metodo di accelerazione matematicamente fondato che riduce il numero di passaggi backward della rete da $W$ a 2 per un gruppo di $W$ passi di denoising consecutivi.

Intuizione Centrale

Il metodo si basa sull'intuizione che il Jacobiano del campo di velocità di denoising $v_\theta(z, t)$ esibisce una linearità approssimativa rispetto agli input $(z, t)$ , particolarmente durante il regime di alto rumore del processo di denoising. Ciò consente di approssimare accuratamente i Jacobiani degli step intermedi tramite interpolazione lineare utilizzando solo i Jacobiani degli endpoint ( $J_0$ e $J_{W-1}$ ).

Workflow Algoritmico

JAGG elabora un gruppo di $W$ timestep attraverso quattro fasi:

Passaggi Forward: Il modello esegue passaggi forward per tutti i $W$ step. Tuttavia, il tracciamento del gradiente è abilitato solo per i due endpoint (step 0 e step $W-1$ ). I $W-2$ step intermedi sono eseguiti utilizzando torch.no_grad() per risparmiare memoria e computazione.
Computazione del Gradiente Upstream: Per ogni step $j$ , viene creato un proxy distaccato dell'output del modello. La perdita per step $L_j$ (derivata dall'obiettivo stile PPO) viene calcolata, e il gradiente upstream $s_j = \partial L_j / \partial \hat{v}_j$ viene calcolato tramite differenziazione automatica. Fondamentalmente, questo step non coinvolge alcun passaggio backward del transformer; il gradiente fluisce solo attraverso il leggero step SDE e l'obiettivo PPO.
Aggregazione Pesata per Timestep: I gradienti upstream $s_j$ vengono aggregati in due segnali compositi, $s_{base}$ e $s_{corr}$ , utilizzando pesi di interpolazione $\alpha_j$ derivati dal timestep $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$
Questa aggregazione ricostruisce efficacemente il segnale del gradiente completo utilizzando l'assunzione di interpolazione lineare.
Backward Congiunto: Viene emessa una singola chiamata torch.autograd.backward sugli output dei due endpoint $[\hat{v}_0, \hat{v}_{W-1}]$ con i gradienti aggregati $[s_{base}, s_{corr}]$ . Questo innesca esattamente un FSDP all-gather per modulo, sostituendo i $W$ separati passaggi backward richiesti dall'addestramento standard.

Garanzia Teorica

Gli autori forniscono una prova formale (Teorema 1) affermando che, se il campo di velocità $v_\theta(z, t)$ è strettamente lineare in $(z, t)$ , l'interpolazione del Jacobiano pesata per $t$ è esatta con errore di approssimazione nullo. Sotto queste condizioni, il passaggio backward a due endpoint recupera il gradiente di gruppo esatto.

Routing Adattivo ( $jagg\_frac$ )

Riconoscendo che l'assunzione di linearità degrada nelle fasi di basso rumore (tardive) del denoising, gli autori introducono un parametro di routing $jagg\_frac$ . Questo parametro applica JAGG solo alla prima frazione dei gruppi di denoising (dove l'assunzione di linearità ad alto rumore è valida, verificata empiricamente tramite la similarità del coseno tra i gradienti) e utilizza passaggi backward esatti per i gruppi rimanenti. Per i modelli T2I testati, $jagg\_frac$ è stato impostato a 0.6.

Contributi Chiave

Identificazione del Collo di Bottiglia: Il documento identifica e caratterizza formalmente il collo di bottiglia computazionale nell'addestramento di diffusione GRPO causato dai passaggi backward per timestep attraverso la traiettoria di campionamento.
Framework JAGG: Propone un framework di aggregazione del Jacobiano che riduce i passaggi backward da $W$ a 2 per gruppo, basato su un'analisi teorica rigorosa della linearità del Jacobiano della velocità.
Prova Teorica: Gli autori dimostrano che, sotto condizioni di velocità lineare, l'interpolazione proposta è esatta, fornendo una solida base teorica per l'approssimazione.
Validazione Empirica: Esperimenti estesi su molteplici backbone di diffusione (Flux, QwenImage) e algoritmi RL (DanceGRPO, Flow-GRPO-Fast) dimostrano la generalità ed l'efficacia del metodo.

Risultati Sperimentali

Gli autori hanno valutato JAGG su benchmark T2I utilizzando HPSv2.1, PickScore e CLIPScore come metriche.

Efficienza: JAGG ha ottenuto un aumento della velocità di tempo reale (wall-clock speedup) costante del 17,3% - 18,2% per step di addestramento su configurazioni standard di DanceGRPO (riducendo il tempo da ~835s a ~691s su QwenImage e da ~279s a ~228s su Flux). Applicato a Flow-GRPO-Fast, la sola fase di aggiornamento ha visto accelerazioni del 29,9% - 30,0%.
Qualità: Il metodo ha fornito questi aumenti di velocità con una degradazione trascurabile della qualità della generazione. In molti casi, JAGG ha seguito da vicino le curve del reward del baseline e ha eguagliato o leggermente superato il baseline alla convergenza.
Ablazione: Una variante di ablazione che eseguiva i passaggi backward sugli endpoint ma senza l'interpolazione del Jacobiano (assegnando la perdita completa agli endpoint indipendentemente) è stata costantemente inferiore alle prestazioni. Ciò conferma che l'aggregazione delle informazioni del gradiente intermedio tramite l'interpolazione del Jacobiano è critica per un addestramento stabile.
Qualità del Gradiente: L'analisi empirica ha mostrato che nei gruppi ad alto rumore (0–2), la similarità del coseno tra il gradiente aggregato da JAGG e il gradiente esatto era alta (0,74–0,85), validando l'assunzione di linearità. La similarità è diminuita nei gruppi successivi, giustificando la strategia di routing $jagg\_frac$ .

Significato e Rivendicazioni

Il documento sostiene che JAGG rappresenti un passo significativo verso la praticità dell'addestramento su larga scala della diffusione GRPO. Sfruttando la quasi-linearità dei Jacobiani di velocità nel regime di alto rumore, il metodo riduce il carico computazionale dell'addestramento RL senza sacrificare la qualità della generazione.

Gli autori sottolineano che JAGG non è meramente un' scorciatoia euristica, ma è fondato sulle proprietà matematiche della traiettoria di denoising. Il metodo è presentato come ortogonale alle esistenti tecniche di accelerazione dell'inferenza, mirando specificamente al collo di bottiglia del passaggio backward unico della diffusione RL. I risultati suggeriscono che JAGG permetta un'adozione più diffusa dell'allineamento tramite RL per i modelli generativi visivi, abbassando significativamente la barriera del costo di addestramento.