Résumé Technique : Flash-GRPO

Énoncé du Problème

Les modèles de diffusion vidéo ont réalisé des progrès significatifs dans la génération de contenu réaliste, mais l'alignement de ces modèles avec les préférences humaines (qualité esthétique, adhérence au prompt, plausibilité physique) reste un défi critique. Bien que l'apprentissage par renforcement (RL), et plus spécifiquement l'optimisation de politique relative par groupe (GRPO), soit devenu un paradigme dominant pour cet alignement, il se heurte à un goulot d'étranglement computationnel sévère. Le GRPO standard nécessite une rétropropagation des gradients à travers des trajectoires complètes de débruitage spatio-temporel, exigeant des centaines de jours-GPU pour des modèles à grande échelle (par exemple, 14 milliards de paramètres).

Les méthodes d'efficacité existantes tentent d'atténuer ce coût par un échantillonnage des pas de temps (par exemple, des approches par fenêtre glissante). Cependant, les auteurs identifient que ces stratégies d'échantillonnage naïves compromettent fondamentalement l'optimisation. Elles présentent une instabilité d'entraînement sévère et échouent à atteindre le plafond de performance de l'entraînement sur trajectoire complète. L'article attribue cet échec à deux causes racines :

1. Estimation de l'avantage confondue par les pas de temps : Le mélange des pas de temps au sein des groupes d'avantage entrelace la variance de la récompense avec la difficulté intrinsèque des différents niveaux de bruit, obscurcissant le signal réel de la politique.
2. Mise à l'échelle des gradients dépendante du temps : Le gradient de politique contient intrinsèquement un facteur de mise à l'échelle dépendant du temps, dérivé de la discrétisation des équations différentielles stochastiques (SDE). Ce facteur varie de plusieurs ordres de grandeur à travers la trajectoire de diffusion, provoquant une domination des pas de temps précoces dans les mises à jour des paramètres, indépendamment de leur importance réelle, ce qui conduit à un déséquilibre d'optimisation.

Méthodologie : Flash-GRPO

Les auteurs proposent Flash-GRPO, un cadre d'entraînement en une seule étape conçu pour atteindre les performances d'une trajectoire complète tout en n'utilisant qu'un seul pas de temps par déploiement d'entraînement. La méthode résout les instabilités identifiées grâce à deux mécanismes fondamentaux :

1. Regroupement Iso-Temporel

Pour éliminer la variance confondue par les pas de temps, Flash-GRPO impose une cohérence temporelle au sein des groupes de prompts.

• Mécanisme : Pour un prompt donné, les $G$ déploiements d'un groupe partagent exactement le même pas de temps $t_k$ mais sont initialisés avec du bruit gaussien différent.
• Processus : Le groupe effectue une transition ODE-SDE en une seule étape au pas de temps assigné $t_k$ pour permettre l'exploration et le calcul des gradients. Tous les autres pas de temps de la trajectoire utilisent un échantillonnage déterministe d'équations différentielles ordinaires (ODE) pour garantir des générations de haute qualité et des signaux de récompense précis.
• Effet : Cela factorise le calcul de l'avantage, isolant la variance induite par la politique de la variance induite par le pas de temps. Cela garantit que les comparaisons de performance relative s'effectuent dans des conditions de débruitage identiques, découplant la performance de la politique de la difficulté du pas de temps.

2. Rectification Temporelle des Gradients

Pour résoudre l'hétérogénéité de l'échelle des gradients, les auteurs dérivent le gradient de politique explicite pour le processus de génération inverse.

• Dérivation : Il est démontré que le gradient est mis à l'échelle par un coefficient dépendant du temps $\lambda(t) = \frac{\sqrt{\Delta t}}{\sigma_t} + \frac{\sigma_t\sqrt{\Delta t}(1-t)}{2t}$. Dans un cadre en une seule étape avec des pas de temps divers, cela provoque des magnitudes de gradients extrêmement incohérentes.
• Mécanisme : Flash-GRPO introduit une Rectification Temporelle des Gradients, qui normalise explicitement le gradient par $1/\lambda(t)$.
• Effet : Cela force $\lambda(t) \to 1$ pour tous les pas de temps, assurant une contribution uniforme aux mises à jour des paramètres indépendamment de la position du pas de temps dans la trajectoire de diffusion. Cela élimine le biais induit par la discrétisation et stabilise la dynamique d'optimisation.

Contributions Clés

1. Identification des Sources d'Instabilité : L'article fournit des dérivations théoriques et une validation empirique pour deux causes spécifiques d'instabilité dans le GRPO vidéo en une seule étape : l'estimation de l'avantage confondue par les pas de temps et la mise à l'échelle des gradients dépendante du temps.
2. Cadre Flash-GRPO : Un cadre d'entraînement en une seule étape fondé sur des principes, combinant un regroupement iso-temporel pour une estimation précise de l'avantage et une rectification temporelle des gradients pour une optimisation équilibrée.
3. Évolutivité et Efficacité : Validation sur des modèles allant de 1,3 milliard à 14 milliards de paramètres, démontrant que la méthode élimine le compromis traditionnel entre efficacité et qualité, atteignant une qualité d'alignement de pointe avec des coûts computationnels substantiellement réduits.

Résultats Expérimentaux

Les auteurs ont évalué Flash-GRPO sur les modèles vidéo Wan2.1 (1,3 milliard et 14 milliards de paramètres) en utilisant le modèle de récompense HPSv3 et les métriques VBench.

• Efficacité d'Entraînement : Flash-GRPO réalise une accélération de 6× du coût d'entraînement par rapport aux méthodes standard sur trajectoire complète.
• Stabilité : Contrairement aux méthodes en une seule étape naïves (par exemple, Flow-GRPO-Fast1) qui souffrent d'un effondrement catastrophique de l'optimisation et d'une variance de récompense, Flash-GRPO maintient une croissance de récompense stable et monotone.
• Performance :
  • Sur le modèle 1,3B, Flash-GRPO a atteint la plus haute qualité esthétique (66,43) et la plus grande cohérence des sujets (98,70) parmi les méthodes comparées, surpassant à la fois Flow-GRPO et Flow-GRPO-Fast1.
  • Sur le modèle 14B, Flash-GRPO a démontré une stabilité constante et une croissance monotone, tandis que Flow-GRPO présentait une convergence plus lente et une instabilité.
  • Améliorations Qualitatives : Les comparaisons visuelles montrent une dynamique temporelle améliorée (réduction du scintillement), une fidélité visuelle accrue (détails plus riches) et une meilleure adhérence au prompt par rapport aux bases de référence classiques.

Signification et Revendications

L'article revendique que Flash-GRPO résout le compromis fondamental entre l'efficacité computationnelle et la qualité d'alignement dans les modèles de diffusion vidéo. En s'attaquant aux artefacts mathématiques spécifiques de l'optimisation en une seule étape (variance confondue et mise à l'échelle des gradients), la méthode permet à l'alignement vidéo RL de haute qualité de devenir à la fois pratique et évolutif. Les auteurs affirment que leur approche permet aux modèles d'atteindre le plafond de performance de l'entraînement sur trajectoire complète tout en opérant avec des budgets computationnels faibles, en faisant une alternative robuste pour l'alignement de modèles vidéo à grande échelle.