Résumé technique : Préconditionnement du gradient pour une génération guidée par la récompense efficace et fiable

Énoncé du problème

L'article aborde deux limitations critiques de l'optimisation de l'espace latent au moment de l'inférence pour les modèles génératifs en une seule étape (par exemple, les distillations de diffusion comme FLUX, SDXL-Turbo). Bien que l'optimisation du vecteur latent lors de l'inférence permette de maximiser la récompense sans réentraînement, les approches existantes souffrent de :

1. Triche à la récompense (Reward Hacking) : À mesure que l'optimisation progresse, le vecteur latent s'éloigne de l'a priori standard de bruit blanc gaussien. Cela amène le modèle à générer des artefacts irréalistes ou à « tricher » avec la fonction de récompense, dégradant la qualité des échantillons.
2. Inefficacité : Les méthodes actuelles nécessitent généralement plusieurs étapes de gradient par échantillon pour converger, prenant souvent plusieurs minutes par image, ce qui entrave le déploiement pratique.

Les solutions existantes reposent sur la régularisation (pénalités douces) pour encourager des propriétés gaussiennes (par exemple, norme $\ell_2$, kurtosis ou pénalités spectrales). Cependant, les auteurs soutiennent que la régularisation est intrinsèquement un mécanisme « doux » qui ne garantit pas le maintien de caractéristiques semblables à du bruit, échoue souvent à prévenir la triche à la récompense lorsque les optimiseurs trouvent des raccourcis, et introduit des hyperparamètres nécessitant un réglage.

Méthodologie

Les auteurs proposent le Préconditionnement du gradient, une méthode qui remplace la régularisation douce par une contrainte géométrique rigide sur la direction d'optimisation.

1. Le concept de base

Au lieu de pénaliser les écarts par rapport au bruit blanc gaussien, la méthode projette le gradient de récompense sur un ensemble réalisable de bruit blanc gaussien ($\mathcal{G}$) soigneusement conçu avant d'appliquer la mise à jour.

• Mécanisme : À chaque étape d'optimisation, le gradient de récompense $\nabla_x J$ est projeté sur $\mathcal{G}$. Le vecteur résultant devient la direction de mise à jour.
• Effet : Cela redéfinit chaque mise à jour dans une direction « alignée sur le bruit ». Cela favorise une ascension plus rapide de la récompense au sein du sous-espace compatible avec le bruit, tout en empêchant le vecteur latent de dériver vers des régions causant la triche à la récompense.
• Avantage : Contrairement à la régularisation, cela ne nécessite aucun hyperparamètre supplémentaire (comme un coefficient de pondération $\lambda$) et n'exige pas que le vecteur latent lui-même réside strictement dans l'ensemble, mais seulement la direction de mise à jour.

2. Conception de l'ensemble réalisable ($\mathcal{G}$)

L'ensemble réalisable est construit pour caractériser étroitement le bruit blanc gaussien dans le domaine spectral tout en admettant une projection sous forme fermée.

• Domaine spectral compact : Les auteurs définissent une application bijective $\mathcal{F}$ qui transforme le vecteur latent à valeurs réelles $x \in \mathbb{R}^N$ en un vecteur à valeurs complexes $y \in \mathbb{C}^{N/2}$, éliminant la redondance hermitienne des coefficients de la Transformée de Fourier Discrète (DFT).
• Contraintes par blocs : L'ensemble $\mathcal{G}_C$ dans le domaine spectral compact est défini en imposant que chaque bloc de taille $B$ satisfait des contraintes de norme $\ell_1$ et $\ell_2$ spécifiques correspondant aux attentes théoriques d'une distribution gaussienne complexe $\mathcal{CN}(0, I)$.
  • Contrainte $\ell_1$ : $\|y^{(p)}\|_1 = \sqrt{\frac{\pi}{2}} B$
  • Contrainte $\ell_2$ : $\|y^{(p)}\|_2^2 = B$
• Algorithme de projection : La projection sur l'intersection des sphères $\ell_1$ et $\ell_2$ est résolue sous forme fermée. L'algorithme consiste à trier les magnitudes au sein de chaque bloc et à résoudre une équation quadratique pour trouver les multiplicateurs de Lagrange optimaux.
• Complexité : La projection est mise en œuvre via la Transformée de Fourier Rapide (FFT) et un tri par blocs, résultant en une complexité de $O(N \log N)$. Dans les expériences avec le modèle FLUX, cela n'ajoute que 0,04 % au temps d'exécution.

Contributions clés

1. Cadre de préconditionnement du gradient : Une stratégie d'optimisation novatrice qui projette les gradients sur un ensemble réalisable de bruit, empêchant efficacement la triche à la récompense sans pénalités douces.
2. Caractérisation serrée du bruit : Un ensemble réalisable dans le domaine spectral défini par des contraintes par blocs $\ell_1$ et $\ell_2$ qui capture à la fois les statistiques gaussiennes marginales et l'absence de corrélation spatiale (propriétés du bruit blanc) plus étroitement que les régularisateurs marginaux ou spectraux antérieurs.
3. Projection efficace sous forme fermée : La dérivation d'un algorithme de projection en $O(N \log N)$ qui est computationnellement négligeable par rapport au passage avant du modèle génératif.
4. Fonctionnement sans hyperparamètres : La méthode élimine le besoin de régler les coefficients de régularisation ($\lambda$), qui sont une source d'instabilité dans les travaux antérieurs.

Résultats expérimentaux

La méthode a été évaluée sur des modèles texte-à-image en une étape (FLUX, SDXL-Turbo, SANA-Sprint, SD-Turbo) en utilisant quatre modèles de récompense : Aesthetic Score, PickScore, HPSv2 et ImageReward.

• Performance par rapport à l'état de l'art : Sur FLUX avec optimisation du Aesthetic Score, la méthode proposée a atteint un niveau de récompense comparable à la méthode de l'état de l'art basée sur la régularisation (MPGR) en utilisant seulement 30 % du temps réel (60 itérations contre 200 itérations).
• Récompense vs Qualité : La méthode a constamment atteint des récompenses cibles plus élevées (par exemple, Aesthetic Score ~8,9 contre ~7,1 pour les bases) tout en préservant ou en améliorant les métriques de préférence humaine hors échantillon (PickScore, HPSv2).
• Stabilité : Les résultats qualitatifs ont montré que la méthode produit de manière fiable des images réalistes alignées avec les invites textuelles, tandis que les méthodes basées sur la régularisation produisaient souvent des artefacts incohérents ou échouaient à maintenir la fidélité à l'invite.
• Diversité : L'analyse utilisant le Inception Score (IS) et le Vendi Score a confirmé que la méthode ne provoque pas d'effondrement de mode ; la diversité des échantillons est restée dans la variance des bases non optimisées.
• Ablation : Supprimer la projection latente (en ne conservant que la projection du gradient) a encore produit des améliorations significatives par rapport aux bases, mais la méthode complète (projection du gradient + projection latente) a fourni la meilleure stabilité et les meilleurs scores hors échantillon.

Importance et revendications

L'article revendique que, en résolvant à la fois le goulot d'étranglement de l'efficacité et le problème de la triche à la récompense, le préconditionnement du gradient rend l'optimisation de l'espace latent au moment de l'inférence substantiellement plus pratique pour les applications réelles.

Les auteurs soulignent que leur approche offre une alternative fondée sur des principes à la régularisation douce. En imposant une contrainte géométrique rigide sur la direction de mise à jour, ils garantissent que la trajectoire d'optimisation reste alignée avec l'a priori de bruit blanc gaussien. Cela assure que la capacité du modèle génératif est exploitée pour la maximisation de la récompense sans compromettre les propriétés statistiques requises pour une génération réaliste. La surcharge computationnelle négligeable et l'absence d'hyperparamètres font de cette méthode un remplacement direct pour les pipelines d'optimisation basés sur le gradient existants dans les modèles génératifs en une étape.