Stepwise Credit Assignment for GRPO on Flow-Matching Models

1. Problématique

L'article aborde les limitations de l'application de l'apprentissage par renforcement (RL), et plus spécifiquement de l'algorithme GRPO (Group Relative Policy Optimization), aux modèles de génération d'images basés sur le Flow-Matching (comme SD-3.5).

Les méthodes existantes, telles que Flow-GRPO, utilisent une attribution de crédit uniforme : elles attribuent la même récompense (ou pénalité) à chaque étape du processus de débruitage, basée uniquement sur la qualité de l'image finale. Cette approche présente deux défauts majeurs :

1. Ignorance de la structure temporelle hiérarchique : La génération d'images par diffusion suit une progression du grossier au fin. Les premières étapes déterminent la composition et la structure (basses fréquences), tandis que les dernières étapes affinent les détails et les textures (hautes fréquences). Traiter toutes les étapes de manière égale ignore cette hiérarchie.
2. Récompense d'erreurs corrigées : Une trajectoire peut commettre une erreur importante à un stade précoce (ex: mauvaise couleur d'objet) qui est ensuite corrigée par les étapes suivantes. Si l'image finale est bonne, Flow-GRPO récompense l'ensemble de la trajectoire, y compris l'étape erronée initiale, ce qui peut renforcer des comportements sous-optimaux et des artefacts.

2. Méthodologie : Stepwise-Flow-GRPO

Les auteurs proposent Stepwise-Flow-GRPO, une méthode qui attribue le crédit à chaque étape de débruitage individuellement, en se basant sur l'amélioration marginale de la récompense à chaque pas.

A. Estimation des récompenses intermédiaires

Puisque les modèles de récompense (Reward Models) ne peuvent pas évaluer directement les états bruités intermédiaires $x_t$, la méthode utilise la formule de Tweedie pour estimer l'image propre sous-jacente $\hat{x}_0(t) = E[x_0 | x_t]$ à partir de l'état bruité.

• Pour obtenir des estimations de haute qualité, ils effectuent un petit nombre de sous-étapes de débruitage déterministe (ODE) à partir de $x_t$ pour obtenir $\hat{x}_0(t)$.
• La récompense intermédiaire est calculée comme $r_t = R(\hat{x}_0(t), \text{prompt})$.

B. Attribution de crédit par gains (Gain-based Advantages)

Au lieu d'utiliser la récompense finale, l'algorithme calcule le gain $g_t$ à chaque étape, défini comme la différence de récompense entre deux étapes consécutives :
$$g_t = r_{t-1} - r_t$$

• Un gain positif indique que l'étape a amélioré l'image.
• Un gain négatif indique une détérioration.
• Ces gains sont ensuite normalisés de manière globale (sur toutes les étapes et toutes les trajectoires du groupe) pour préserver la structure temporelle naturelle (les gains sont généralement plus importants au début). Ces gains normalisés servent d'avantages pour la mise à jour de la politique via GRPO.

C. Équation Stochastique Différentielle (SDE) améliorée

L'article identifie que l'SDE utilisée par Flow-GRPO pour l'exploration RL introduit trop de bruit, dégradant la qualité des images intermédiaires et donc la fiabilité des récompenses.

• Les auteurs introduisent une nouvelle SDE inspirée de DDIM.
• Cette formulation permet d'injecter du bruit contrôlé pour l'exploration tout en produisant des échantillons intermédiaires beaucoup plus propres, ce qui améliore la qualité du signal de récompense sans sacrifier la stochasticité nécessaire aux gradients de politique.

3. Contributions Clés

1. Stepwise-Flow-GRPO : Introduction d'une attribution de crédit basée sur les gains marginaux par étape plutôt que sur la récompense finale, permettant d'apprendre quelles étapes spécifiques contribuent positivement à la génération.
2. Optimisation SDE (Style DDIM) : Développement d'une nouvelle équation de diffusion qui améliore la qualité des échantillons intermédiaires pour les modèles de récompense, tout en maintenant la compatibilité avec les marginales du Flow-ODE.
3. Théorie et Pratique : Lien théorique avec l'optimisation sous-modulaire adaptative, suggérant que la maximisation gloutonne de ces gains converge vers des politiques quasi-optimales.

4. Résultats Expérimentaux

Les expériences ont été menées sur des modèles SD3.5-Medium avec des récompenses variées (PickScore, ImageReward, UnifiedReward) et des benchmarks comme GenEval.

• Efficacité de l'échantillonnage (Sample Efficiency) : Stepwise-Flow-GRPO converge plus rapidement et atteint des récompenses finales plus élevées que Flow-GRPO standard. Il nécessite moins d'itérations d'entraînement pour atteindre un niveau de performance donné.
• Performance sur GenEval : La méthode surpasse Flow-GRPO, en particulier dans les tâches complexes nécessitant un raisonnement spatial, une liaison d'attributs et un comptage précis (ex: "quatre frisbees", "un chat à droite d'un ordinateur").
• Stabilité et Robustesse : La méthode démontre une stabilité supérieure lors de l'entraînement avec des modèles de récompense complexes (VLM comme UnifiedReward), là où Flow-GRPO a tendance à diverger.
• Efficacité temporelle (Wall-clock) : Bien que le calcul des récompenses intermédiaires ajoute une surcharge computationnelle par itération, la convergence plus rapide compense ce coût, résultant en un temps total d'entraînement réduit (20-40% plus rapide).
• Qualité visuelle : Les images générées sont plus réalistes, avec moins d'artefacts et une meilleure adhérence au prompt, notamment grâce à la correction des erreurs de composition dès les premières étapes.

5. Signification et Impact

Ce travail est significatif car il corrige une faille fondamentale dans l'application du RL aux modèles de génération d'images : l'ignorance de la dynamique temporelle du processus de débruitage.

• En passant d'une vision "boîte noire" (récompense finale) à une vision "boîte transparente" (récompense par étape), l'article permet un apprentissage plus fin et plus efficace.
• La méthode offre une alternative robuste aux approches concurrentes (comme TempFlow-GRPO) qui reposent sur des pondérations manuelles ou des sélections d'étapes arbitraires, en utilisant des gains dépendants des données qui s'adaptent naturellement à la structure de la tâche.
• Cela ouvre la voie à des modèles de génération d'images plus capables de gérer des instructions complexes et des compositions multi-objets, rapprochant les performances des modèles de diffusion des modèles autoregressifs de pointe.