Gradient-Free Noise Optimization for Reward Alignment in Generative Models

Résumé Technique : ZeNO (Optimisation de Bruit d'Ordre Zéro)

Énoncé du Problème
Les méthodes existantes d'alignement par récompense pour les modèles génératifs (par exemple, les modèles de diffusion et les modèles de flux) reposent généralement sur l'optimisation de trajectoires stochastiques multi-étapes. Bien que efficaces dans ces contextes, de telles approches sont difficiles à étendre aux générateurs déterministes ou mono-étape (par exemple, les modèles de cohérence, les modèles de diffusion distillés), qui n'exposent pas naturellement la même structure de trajectoire. Une alternative naturelle consiste à optimiser directement le bruit d'entrée. Cependant, l'optimisation du bruit par gradient standard nécessite une rétropropagation à travers le générateur et la fonction de récompense. Cela limite l'applicabilité aux pipelines différentiables et peut s'avérer peu fiable lorsque les gradients reflètent des sensibilités locales spuriaires plutôt que des améliorations sémantiques significatives. De plus, de nombreuses fonctions de récompense pratiques (par exemple, la capacité de conception de protéines, l'alignement compositionnel texte-image) sont non différentiables ou boîtes noires, rendant les méthodes basées sur le gradient inapplicables.

Méthodologie : ZeNO
Les auteurs proposent ZeNO (Optimisation de Bruit d'Ordre Zéro), un cadre sans gradient qui formule l'optimisation du bruit comme un problème de contrôle par intégrale de chemin (PIC). L'idée centrale est de traiter l'évolution du bruit d'entrée comme un processus stochastique contrôlable, en l'optimisant à l'aide uniquement d'évaluations de récompense d'ordre zéro.

1. Formulation par Contrôle d'Intégrale de Chemin :
   La méthode modélise l'optimisation du bruit comme une Équation Différentielle Stochastique (EDS) contrôlée. L'objectif est de diriger la trajectoire du bruit vers des régions produisant des états terminaux à haute récompense tout en régularisant l'effort de contrôle. Le contrôle optimal $u^*(z, t)$ est dérivé sous une forme d'intégrale de chemin :
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   Crucialement, cet estimateur dépend uniquement des évaluations de coûts terminaux (récompenses) et ne nécessite pas de différentiation à travers le générateur $G_\theta$ ou la fonction de récompense $r$.

2. Dynamiques de Référence (Ornstein–Uhlenbeck) :
   Pour garantir que le bruit mis à jour reste compatible avec les générateurs pré-entraînés, ZeNO adopte le processus d'Ornstein–Uhlenbeck (OU) comme dynamique de référence. Contrairement aux EDS à explosion de variance, le processus OU préserve la distribution stationnaire gaussienne $N(0, I)$, qui correspond à l'a priori de la plupart des générateurs pré-entraînés. Ce choix relie la mise à jour du bruit aux dynamiques de Langevin ciblant implicitement une distribution inclinée par la récompense $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$.

3. Estimateur Linéarisé et Réduction de Variance :
   L'estimateur standard d'intégrale de chemin souffre d'une variance élevée due au poids exponentiel. Les auteurs introduisent une approximation linéarisée (Proposition 1) en soustrayant la récompense moyenne $\bar{r}$ comme ligne de base et en linéarisant le terme exponentiel sous l'hypothèse d'une faible variance de récompense :
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   Cet estimateur coïncide avec l'estimateur de fonction de score utilisé dans les Stratégies Évolutives Naturelles et surpasse empiriquement les variantes exponentielles.

4. Horizon Récedant Mono-étape :
   Pour éviter le coût computationnel des déroulements sur de longs horizons, ZeNO emploie une approximation d'horizon récedant mono-étape ($H=1$). À chaque itération, $N$ candidats de bruit perturbés sont échantillonnés à partir des dynamiques de référence, évalués par le modèle de récompense, et l'estimation de contrôle linéarisée est utilisée pour mettre à jour l'état actuel du bruit. Cela permet une mise à l'échelle au moment de l'inférence en augmentant le nombre de particules ($N$) ou d'itérations ($M$) sans modifier le générateur.

Contributions Clés

• Alignement Sans Gradient : ZeNO permet l'alignement par récompense pour les générateurs déterministes et mono-étape sans nécessiter de rétropropagation à travers le générateur ou le modèle de récompense.
• Compatibilité Boîte Noire : La méthode est applicable aux fonctions de récompense non différentiables et boîtes noires, un contexte où les méthodes existantes d'optimisation de bruit basées sur le gradient (par exemple, ORIGEN, ReNO) échouent.
• Lien Théorique : Le cadre établit un lien théorique entre l'optimisation de bruit d'ordre zéro et l'échantillonnage de Langevin ciblant une distribution inclinée par la récompense, fournissant une justification principielle pour la règle de mise à jour.
• Mise à l'Échelle au Moment de l'Inférence : ZeNO introduit un mécanisme flexible pour échanger le budget computationnel (via $N$ et $M$) contre une amélioration des performances de récompense.

Résultats Expérimentaux
Les auteurs évaluent ZeNO sur divers générateurs (SDXL-Turbo, DMD2, LCM, DFGAN) et fonctions de récompense (score esthétique, PickScore, GenEval, et capacité de conception de protéines).

• Génération d'Images : ZeNO améliore systématiquement les scores de récompense sur tous les générateurs et types de récompense testés. Il atteint des performances compétitives ou supérieures par rapport aux bases de référence à gradient (ORIGEN, $\Psi$-Sampler) et à l'échantillonnage Best-of-N. Notamment, sur le générateur LCM, les méthodes basées sur le gradient se sont dégradées en dessous de la base de référence non contrôlée, tandis que ZeNO a maintenu une amélioration, soulignant la robustesse des méthodes d'ordre zéro face aux sensibilités locales des gradients.
• Récompenses Non Différentiables : Dans la tâche GenEval (comptage et positionnement d'objets), où les récompenses sont non différentiables, ZeNO a surpassé Best-of-N sur les tâches de comptage, démontrant que l'optimisation itérative du bruit extrait plus de valeur du même budget computationnel que l'échantillonnage indépendant.
• Génération de Structures de Protéines : Appliqué à un générateur de squelette de protéines Riemannian MeanFlow (RMF) mono-étape, ZeNO a considérablement amélioré la capacité de conception (mesurée par le RMSD de cohérence interne). Il a atteint des niveaux de capacité de conception comparables à une base de référence RMF en 5 étapes en utilisant uniquement une seule évaluation de flux, démontrant sa capacité à compenser la réduction des étapes d'échantillonnage via l'optimisation du bruit.
• Diversité : L'analyse utilisant le Score Vendi indique que ZeNO améliore la récompense tout en préservant la diversité des échantillons, évitant l'effondrement vers un mode unique souvent observé dans les approches basées sur le gradient.

Signification et Revendications
L'article revendique que ZeNO fournit une solution pratique et théoriquement fondée pour aligner la classe croissante de modèles génératifs efficaces, déterministes et mono-étape avec des fonctions de récompense. Sa signification principale réside dans sa capacité à opérer dans des contextes où les gradients sont indisponibles ou peu fiables, tels que les modèles de récompense boîtes noires ou les tâches complexes de conception biologique. En formulant l'optimisation du bruit comme un problème de contrôle par intégrale de chemin, ZeNO offre un cadre évolutif et agnostique au générateur pour l'alignement au moment de l'inférence qui ne nécessite ni réentraînement ni modification du modèle génératif sous-jacent. Les auteurs notent une limitation : la méthode nécessite plusieurs évaluations de récompense par étape de mise à jour, ce qui peut être computationnellement exigeant si le modèle de récompense lui-même est coûteux.