Résumé Technique : AEGPO – Optimisation de Politique Guidée par l'Entropie Adaptative pour les Modèles de Diffusion

1. Énoncé du Problème

L'apprentissage par renforcement à partir de la rétroaction humaine (RLHF), en particulier l'optimisation de politique relative au groupe (GRPO), est devenu un paradigme standard pour aligner les modèles de diffusion et de flux (flow-matching) avec les préférences humaines. Cependant, les variantes existantes de GRPO souffrent de deux inefficacités fondamentales dues à leur dépendance à des heuristiques statiques :

1. Allocation Inefficace des Échantillons : Les méthodes standards assignent un nombre uniforme de déroulements (rollouts) à chaque prompt, ignorant la variation significative de la « valeur d'apprentissage » entre les différents échantillons. Certains prompts induisent des changements de politique minimaux, tandis que d'autres entraînent des améliorations substantielles.
2. Exploration Statique des Pas de Temps (Timesteps) : Les stratégies d'exploration (par exemple, le branchement dans les déroulements basés sur des arbres) suivent généralement des calendriers fixes et prédéfinis. Celles-ci ne parviennent pas à s'adapter à la nature dynamique du processus de débruitage, où les moments critiques pour l'exploration varient considérablement selon les prompts et les étapes d'entraînement.

Ces limitations entraînent un gaspillage de ressources computationnelles sur des échantillons de faible valeur et une exploration sous-optimale aux pas de temps non critiques, ce qui ralentit la convergence et limite la performance finale d'alignement.

2. Méthodologie : AEGPO

Les auteurs proposent l'Optimisation de Politique Guidée par l'Entropie Adaptative (AEGPO), un cadre plug-and-play qui exploite les signaux d'attention intrinsèques pour guider dynamiquement à la fois la sélection d'échantillons et l'exploration. L'idée centrale est que l'Entropie de l'Attention sert de puissant double signal proxy.

2.1 Signaux Fondamentaux

La méthode définit deux métriques intrinsèques dérivées des cartes d'attention du modèle :

• $\Delta$Entropie (Signal au niveau de l'échantillon) : Le changement relatif de la dispersion de l'attention entre la politique actuelle et la politique de base.
  • Définition : $\Delta\text{Entropie} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropie}_{\theta}(t) - \text{Entropie}_{\theta_{base}}(t)|$.
  • Fonction : Elle agit comme un proxy robuste de la valeur d'apprentissage de l'échantillon. Une $\Delta$Entropie élevée indique qu'un prompt force le modèle à réorganiser considérablement sa stratégie d'attention, impliquant un fort potentiel d'apprentissage.
• Entropie($t$) (Signal au niveau du pas de temps) : La dispersion absolue de l'attention à un pas de débruitage spécifique $t$.
  • Définition : La moyenne de l'entropie de Shannon des poids d'attention à travers les tokens textuels pour les caractéristiques d'image à l'étape $t$.
  • Fonction : Les pics d'Entropie($t$) identifient les pas de temps d'exploration critiques. Ces pics correspondent à des moments de haute incertitude ou d'intégration sémantique large (survenant souvent lors des phases initiales de structure et de raffinement tardif), où l'exploration produit la plus grande variance et diversité de récompense.

2.2 Stratégie Adaptative à Double Niveau

AEGPO utilise ces signaux pour réallouer le calcul sans augmenter le budget total de déroulement :

1. Allocation Adaptative Globale :

   • Le lot (batch) de prompts est divisé en paliers de « haute valeur » et de « faible valeur » basés sur la médiane de la $\Delta$Entropie.
   • Les prompts à haute valeur (haute $\Delta$Entropie) se voient assigner un budget de déroulement plus important ($r_{high}$), tandis que les prompts de faible valeur en reçoivent moins ($r_{low}$).
   • Contrainte : $r_{high} + r_{low} = 2r_{avg}$, garantissant que le calcul total reste constant.

2. Exploration Adaptative Locale :

   • Au lieu de calendriers de branchement fixes, la méthode identifie dynamiquement les $K$ meilleurs pas de temps présentant les pics d'Entropie($t$) les plus élevés pour chaque déroulement.
   • Le branchement (exploration) est effectué exclusivement à ces pas de temps identifiés à haute dispersion.
   • Cela garantit que l'exploration est concentrée sur les moments les plus informatifs de la trajectoire de débruitage.

3. Contributions Clés

• Découverte de Double Signal : L'article établit l'Entropie de l'Attention comme un proxy intrinsèque à double signal, où le changement relatif ($\Delta$Entropie) indique la valeur d'apprentissage de l'échantillon et les pics absolus (Entropie($t$)) identifient les pas de temps d'exploration critiques.
• Cadre AEGPO : Introduction d'un cadre adaptatif à double niveau qui améliore la sélection d'échantillons et l'exploration des pas de temps à l'aide de mécanismes légers guidés par l'entropie.
• Généralisabilité : Démonstration que l'AEGPO est indépendant du modèle (model-agnostic) et du cadre (framework-agnostic), améliorant avec succès diverses méthodes d'optimisation de politique (DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT) et des modèles de base (FLUX.1-dev, SD3.5-M).

4. Résultats Expérimentaux

Les expériences ont été menées sur les jeux de données HPDv2.1 et Pick-a-Pic en utilisant FLUX.1-dev et SD3.5-M comme modèles de base.

• Vitesse de Convergence : AEGPO accélère considérablement l'optimisation de la politique. Il atteint une convergence 2$\times$ plus rapide sur DanceGRPO et une convergence 5$\times$ plus rapide sur DiffusionNFT par rapport aux variantes standards, atteignant des récompenses finales supérieures en moins d'étapes.
• Métriques de Performance :
  • Sur FLUX.1-dev, BranchGRPO + AEGPO a atteint un nouveau score SOTA HPS-v2.1 de 0,374 (contre 0,363) et un Pick Score de 0,232 (contre 0,229).
  • Sur SD3.5-M, DiffusionNFT + AEGPO a amélioré le HPS-v2.1 à 0,343 et l'ImageReward à 1,57.
• Compromis Récompense-KL : Les modèles AEGPO dominent systématiquement les lignes de base sur la frontière de Pareto Récompense-KL, atteignant des récompenses d'alignement nettement plus élevées pour tout niveau donné de déviation de politique (divergence KL).
• Études d'Ablation : L'allocation globale (G) et l'exploration locale (L) apportent toutes deux des bénéfices indépendants, l'approche combinée (G+L) offrant la meilleure performance.
• Surcoût Computationnel : La méthode introduit un surcoût modeste (environ 11,1 % d'augmentation du temps par étape et une augmentation négligeable de la VRAM) dû à l'extraction des cartes d'attention intermédiaires, ce qui est justifié par la réduction spectaculaire du nombre total d'étapes d'entraînement nécessaires pour atteindre la performance de pointe.

5. Signification et Revendications

L'article affirme que l'AEGPO répond aux inefficacités centrales de l'échantillonnage statique dans l'alignement des modèles de diffusion. En passant d'une allocation basée sur des heuristiques à une adaptation basée sur des signaux intrinsèques, la méthode garantit que les ressources de calcul sont concentrées sur les échantillons les plus informatifs et les moments les plus critiques du processus de génération.

Les auteurs soulignent que l'AEGPO est un module généralisable et plug-and-play qui ne nécessite pas de nouvelles fonctions de perte ou de modélisation complexe de récompense. Au lieu de cela, il exploite la propre dynamique d'attention du modèle pour guider l'optimisation, offrant une voie plus efficace vers un alignement de haute qualité à travers diverses architectures de diffusion et cadres de RLHF. Ce travail suggère que les signaux intrinsèques tels que l'entropie de l'attention sont des ressources sous-utilisées pour améliorer l'efficacité de l'optimisation de politique dans les modèles génératifs.