⚡ electrical engineering

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

Cet article introduit JAGG (Jacobian-Aggregated Group Gradient), une méthode qui exploite la variation quasi linéaire des trajectoires de modèles de diffusion pour approximer les Jacobiens intermédiaires et agréger les gradients de groupe, réduisant ainsi le coût computationnel de l'entraînement par Group Relative Policy Optimization (GRPO) pour les Diffusion Transformers d'environ 2× sans perte de qualité significative.

Auteurs originaux : Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

Publié 2026-07-21

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'apprendre à un artiste numérique à mieux peindre des tableaux. Cet artiste n'est pas un humain ; c'est un programme informatique complexe appelé « modèle de diffusion ». Imaginez ce modèle comme un sculpteur qui part d'un bloc d'argile bruyant, rempli de statique, et qui retire progressivement le bruit, étape par étape, jusqu'à ce qu'une statue parfaite émerge. Pour apprendre à ce sculpteur à créer de l'art que les humains apprécient réellement, nous utilisons une technique appelée « Apprentissage par Renforcement » (Reinforcement Learning). C'est comme un jeu où le sculpteur tente différents mouvements de ciseau, reçoit un score de la part d'un juge (un modèle de récompense), puis apprend de ce score pour faire mieux la prochaine fois.

Le problème est que ce processus d'apprentissage est incroyablement coûteux. Chaque fois que le sculpteur fait un petit coup de ciseau (une « étape » dans le processus), l'ordinateur doit effectuer un calcul massif et gourmand en énergie pour déterminer exactement comment ajuster le cerveau du sculpteur. Si la statue nécessite 20 étapes pour être terminée, l'ordinateur doit effectuer ce calcul lourd 20 fois pour chaque leçon. C'est comme essayer d'apprendre une nouvelle danse en s'arrêtant après chaque mouvement pour demander une critique complète à un entraîneur, en réécrivant toute votre mémoire musculaire, puis en recommençant. Pour les images à haute résolution, cela prend tellement de temps et d'électricité qu'il est presque impossible d'entraîner ces modèles efficacement.

C'est là qu'intervient une nouvelle idée appelée JAGG. Les chercheurs derrière cet article se sont posé une question intelligente : « Avons-nous vraiment besoin de nous arrêter et de tout recalculer après chaque mouvement ? » Ils ont remarqué que dans les premières étapes désordonnées du processus de sculpture, les changements se produisent de manière très fluide et prévisible — presque comme une ligne droite. Si vous savez où la sculpture commence et où elle finit après quelques mouvements, vous pouvez en réalité deviner ce qui s'est passé au milieu sans effectuer le calcul lourd pour chaque étape.

L'article présente une méthode appelée JAGG (Jacobian-Aggregated Group Gradient). Au lieu de s'arrêter pour calculer la leçon après chaque étape, JAGG permet au modèle d'effectuer un groupe d'étapes (disons, quatre mouvements consécutifs) et ne réalise le calcul lourd de la « mise à jour du cerveau » que deux fois : une fois au tout début du groupe et une fois à la toute fin. Pour les étapes intermédiaires, il utilise un raccourci intelligent. Il suppose que le chemin entre le début et la fin est une ligne droite et comble les lacunes.

Les chercheurs ont prouvé mathématiquement que si le processus de sculpture est parfaitement fluide et linéaire, ce raccourci est en réalité parfait — il donne exactement la même réponse qu'en faisant le travail difficile à chaque fois. Dans le monde réel, où les choses ne sont pas parfaitement droites, ils ont constaté que ce raccourci fonctionne extrêmement bien durant les premières étapes « bruitées » de la création d'images. En utilisant cette méthode, ils ont pu réduire le temps nécessaire pour entraîner ces modèles d'environ 17 % à 18 % par étape.

Le meilleur dans tout cela ? La qualité des images n'en a pas souffert. Lorsqu'ils ont testé JAGG sur différents modèles (comme Flux et QwenImage), les images générées étaient tout aussi bonnes que celles créées par la méthode traditionnelle lente. En fait, dans certains cas, le raccourci a même aidé le modèle à apprendre plus stablement. L'article montre qu'en étant un peu plus astucieux sur le moment où effectuer le gros du travail, nous pouvons rendre l'entraînement de ces puissants artistes IA beaucoup plus rapide et moins coûteux, sans sacrifier la beauté de l'œuvre finale.

Résumé Technique : JAGG (Jacobian-Aggregated Group Gradient)

Énoncé du Problème

Le Group Relative Policy Optimization (GRPO) est devenu un algorithme d'apprentissage par renforcement (RL) de premier plan pour aligner les modèles génératifs avec les préférences humaines, éliminant avec succès le besoin d'un réseau de valeur séparé pour les grands modèles de langage (LLM). Cependant, l'extension de GRPO aux modèles de diffusion et de flow matching (spécifiquement les Diffusion Transformers ou DiTs) introduit un goulot d'étranglement computationnel sévère.

Dans le GRPO basé sur la diffusion, le processus de génération est traité comme un problème de décision multi-étapes où la politique est optimisée sur une trajectoire d'échantillonnage de $W$ étapes temporelles. L'entraînement standard nécessite un passage avant (forward) et un passage arrière (backward) complet à travers l'architecture DiT à haute capacité pour chaque étape individuelle de la trajectoire afin de calculer le gradient de la politique. Pour la génération texte-vers-image (T2I) à haute résolution, où les trajectoires s'étendent souvent sur des dizaines d'étapes, cela se traduit par des coûts d'entraînement prohibitifs. Contrairement aux modèles de langage où les étapes de tokens sont légères, chaque passage arrière à travers un DiT moderne impose le coût complet d'un passage avant de transformer, rendant l'exigence de passage arrière par étape dominante dans la dépense d'entraînement.

Bien que les méthodes d'accélération de l'inférence sans entraînement (ex. $\Delta$ -DiT, ScalingCache) aient exploité avec succès l'observation selon laquelle les états cachés et les prédictions de vélocité des DiT varient de manière fluide et presque linéaire le long de la trajectoire de débruitage pour sauter des évaluations intermédiaires, il restait une question ouverte pour savoir si cette même linéarité pouvait être exploitée pour réduire le coût du passage arrière lors de l'entraînement par RL.

Méthodologie : JAGG

Les auteurs proposent JAGG (Jacobian-Aggregated Group Gradient), une méthode d'accélération mathématiquement fondée qui réduit le nombre de passages arrière du réseau de $W$ à 2 pour un groupe de $W$ étapes de débruitage consécutives.

Intuition Fondamentale

La méthode repose sur l'idée que le Jacobien du champ de vélocité de débruitage $v_\theta(z, t)$ par rapport aux paramètres du modèle $\theta$ présente une linéarité approximative dans les entrées $(z, t)$ , particulièrement durant le régime de bruit élevé du processus de débruitage. Cela permet d'approximer les Jacobiens des étapes intermédiaires via une interpolation linéaire en utilisant uniquement les Jacobiens des points extrêmes ( $J_0$ et $J_{W-1}$ ).

Flux Algorithmique

JAGG traite un groupe de $W$ étapes temporelles à travers quatre phases :

Passages Avant (Forward Passes) : Le modèle effectue des passages avant pour les $W$ étapes. Cependant, le suivi du gradient est activé uniquement pour les deux points extrêmes (étape 0 et étape $W-1$ ). Les $W-2$ étapes intermédiaires sont exécutées avec torch.no_grad() pour économiser la mémoire et la puissance de calcul.
Calcul du Gradient Amont (Up-stream Gradient) : Pour chaque étape $j$ , un proxy détaché de la sortie du modèle est créé. La perte par étape $L_j$ (dérivée de l'objectif de type PPO) est calculée, et le gradient amont $s_j = \partial L_j / \partial \hat{v}_j$ est calculé via la différenciation automatique. Crucialement, cette étape n'implique aucun passage arrière de transformer ; le gradient ne circule qu'à travers l'étape SDE légère et l'objectif PPO.
Agrégation Pondérée par l'Étape Temporelle : Les gradients amont $s_j$ sont agrégés en deux signaux composites, $s_{base}$ et $s_{corr}$ , en utilisant des poids d'interpolation $\alpha_j$ dérivés de l'étape temporelle $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$
Cette agrégation reconstruit efficacement le signal de gradient complet en utilisant l'hypothèse d'interpolation linéaire.
Backward Conjoint : Un seul appel torch.autograd.backward est émis sur les deux sorties des points extrêmes $[\hat{v}_0, \hat{v}_{W-1}]$ avec les gradients agrégés $[s_{base}, s_{corr}]$ . Cela déclenche exactement un "all-gather" FSDP par module, remplaçant les $W$ passages arrière distincts requis par l'entraînement standard.

Garantie Théorique

Les auteurs fournissent une preuve formelle (Théorème 1) stipulant que si le champ de vélocité $v_\theta(z, t)$ est strictement linéaire en $(z, t)$ , l'interpolation du Jacobien pondérée par $t$ est exacte, avec une erreur d'approximation nulle. Sous ces conditions, le passage arrière à deux points extrêmes récupère le gradient de groupe exact.

Routage Adaptatif ( $jagg\_frac$ )

Reconnaissant que l'hypothèse de linéarité se dégrade dans les étapes de débruitage tardives (faible bruit), les auteurs introduisent un paramètre de routage $jagg\_frac$ . Ce paramètre applique JAGG uniquement à la première fraction des groupes de débruitage (où l'hypothèse de linéarité du régime de bruit élevé tient, vérifiée empiriquement via la similitude cosinus des gradients) et utilise des passages arrière exacts par étape pour les groupes restants. Pour les modèles T2I testés, $jagg\_frac$ a été fixé à 0,6.

Principales Contributions

Identification du Goulot d'Étranglement : Le papier identifie et caractérise formellement le goulot d'étranglement computationnel dans l'entraînement du GRPO de diffusion causé par les passages arrière par étape temporelle à travers la trajectoire d'échantillonnage.
Cadre JAGG : Il propose un cadre d'agrégation de Jacobien qui réduit les passages arrière de $W$ à 2 par groupe, fondé sur une analyse théorique rigoureuse de la linéarité du Jacobien de la vélocité.
Preuve Théorique : Les auteurs prouvent que sous des conditions de vélocité linéaire, l'interpolation proposée est exacte, fournissant une base théorique solide pour l'approximation.
Validation Empirique : Des expériences étendues sur plusieurs backbones de diffusion (Flux, QwenImage) et algorithmes de RL (DanceGRPO, Flow-GRPO-Fast) démontrent la généralité et l'efficacité de la méthode.

Résultats Expérimentaux

Les auteurs ont évalué JAGG sur des benchmarks T2I en utilisant HPSv2.1, PickScore et CLIPScore comme métriques.

Efficacité : JAGG a obtenu une accélération constante du temps de calcul (wall-clock speedup) de 17,3 % à 18,2 % par étape d'entraînement sur des configurations DanceGRPO standards (réduisant le temps de ~835s à ~691s sur QwenImage et de ~279s à ~228s sur Flux). Appliqué à Flow-GRPO-Fast, l'étape de mise à jour seule a vu des accélérations de 29,9 % à 30,0 %.
Qualité : La méthode a délivré ces accélérations avec une dégradation négligeable de la qualité de génération. Dans de nombreux cas, JAGG a étroitement suivi les courbes de récompense de référence et a égalé ou légèrement dépassé la référence à la convergence.
Ablation : Une variante d'ablation effectuant des passages arrière sur les points extrêmes mais sans interpolation de Jacobien (attribuant la perte complète aux points extrêmes indépendamment) a systématiquement sous-performé. Cela confirme que l'agrégation de l'information de gradient intermédiaire via l'interpolation de Jacobien est critique pour un entraînement stable.
Qualité du Gradient : Une analyse empirique a montré que dans les groupes de bruit élevé (0–2), la similitude cosinus entre le gradient agrégé par JAGG et le gradient exact était élevée (0,74–0,85), validant l'hypothèse de linéarité. La similitude a chuté dans les groupes ultérieurs, justifiant la stratégie de routage $jagg\_frac$ .

Signification et Revendications

Le papier affirme que JAGG représente une étape significative vers la rendre pratique l'entraînement à grande échelle du GRPO de diffusion. En exploitant la quasi-linéarité des Jacobiens de vélocité dans le régime de bruit élevé, la méthode réduit la charge computationnelle de l'entraînement par RL sans sacrifier la qualité de génération.

Les auteurs soulignent que JAGG n'est pas simplement un raccourci heuristique mais qu'il est ancré dans les propriétés mathématiques de la trajectoire de débruitage. La méthode est présentée comme orthogonale aux techniques existantes d'accélération de l'inférence, ciblant spécifiquement le goulot d'étranglement du passage arrière propre au RL de diffusion. Les résultats suggèrent que JAGG permet une adoption plus large de l'alignement par RL pour les modèles génératifs visuels en abaissant significativement la barrière du coût d'entraînement.