Résumé Technique : Alignement des Modèles Génératifs à Peu d'Étapes par Amortissement de l'Inférence Variationnelle Basée sur des Échantillons

Énoncé du Problème

L'alignement des modèles génératifs à peu d'étapes (par exemple, GAN, modèles de cohérence, cartes de flux, modèles dérivants) avec des objectifs spécifiques, tels que des fonctions de récompense ou des préférences humaines, présente un défi majeur. Les cadres d'alignement existants reposent généralement sur des hypothèses structurelles restrictives que les générateurs à peu d'étapes violent souvent :

1. Vraisemblances Traitables : De nombreuses méthodes nécessitent une évaluation explicite de la vraisemblance sur les trajectoires de débruitage, ce qui est indisponible pour les modèles implicites comme les GAN ou les mappings en une seule étape.
2. Solveurs Spécifiques : Les approches basées sur le Contrôle Optique Stochastique (SOC) ou l'Apprentissage par Renforcement (RL) supposent souvent des solveurs d'Équations Différentielles Ordinaires (ODE) ou d'Équations Différentielles Stochastiques (SDE) spécifiques et des processus de débruitage séquentiels.
3. Spécificité du Modèle : Les formulations sont fréquemment adaptées à des familles de modèles spécifiques (par exemple, modèles de diffusion), échouant à se généraliser à des architectures à peu d'étapes diverses comme les VAE ou les cartes de flux.

Par conséquent, il manque un cadre d'alignement général capable d'affiner divers générateurs à peu d'étapes tout en préservant leur avantage central : une inférence rapide en une seule étape.

Méthodologie : FAV

Les auteurs proposent FAV (Alignement des Modèles Génératifs à Peu d'Étapes par Inférence Variationnelle Basée sur des Échantillons), un cadre général qui découple l'alignement de la famille de générateurs et des dynamiques d'échantillonnage. FAV formule l'alignement comme un échantillonnage à partir d'une distribution inclinée par la récompense ancrée à une distribution de référence :
$$q^*(x) \propto p_{\text{ref}}(x) \exp(\beta \cdot r(x))$$
où $p_{\text{ref}}$ est une distribution de référence (par exemple, un générateur pré-entraîné ou des données empiriques) et $r(x)$ est la fonction objectif.

La méthode procède en trois étapes clés :

1. Approximation Basée sur des Échantillons via SVGD

FAV approxime la distribution cible $q^*$ en utilisant la Descente de Gradient Variationnel de Stein (SVGD). Au lieu de nécessiter une évaluation de densité, SVGD représente la distribution par des particules et les transporte itérativement vers la cible. Le champ de vitesse de Stein optimal $\phi^*$ est dérivé pour maximiser la diminution de la divergence KL :
$$\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \left[ k(x', x) \nabla_{x'} \log p_{\text{ref}}(x') + \beta \cdot k(x', x) \nabla_{x'} r(x') + \nabla_{x'} k(x', x) \right]$$
Ce champ comprend trois composantes :

• Alignement Prioritaire : Dirige les échantillons vers la distribution de référence $p_{\text{ref}}$.
• Guidage par Récompense : Pousse les échantillons vers les régions à haute récompense.
• Application de la Diversité : Empêche l'effondrement de mode via le terme de gradient du noyau.

2. Estimation de Score Non Paramétrique (KDE)

Un défi critique est que le score $\nabla \log p_{\text{ref}}$ est souvent intraitable pour les données empiriques ou les générateurs implicites. FAV résout cela en estimant le score de manière non paramétrique en utilisant l'Estimation de Densité par Noyau (KDE). En convoluant la distribution de référence avec un noyau gaussien, le score est approximé comme une moyenne pondérée des différences entre les échantillons de référence et les particules actuelles, contournant efficacement le besoin de vraisemblances explicites.

3. Amortissement par Régression de Point Fixe

Pour préserver la vitesse d'inférence rapide des modèles à peu d'étapes, FAV n'effectue pas de mises à jour SVGD itératives au moment de l'inférence. Au lieu de cela, il amortit le transport dans les paramètres du générateur $\theta$ via une régression de point fixe. Le générateur est entraîné pour minimiser la distance entre sa sortie et la cible "transportée" :
$$\mathcal{L}(\theta) = \mathbb{E}_{\epsilon} \left[ \| \text{stopgrad}(f_\theta(\epsilon) + \hat{\phi}^*_{q_\theta, q^*_\sigma}(f_\theta(\epsilon))) - f_\theta(\epsilon) \|_2^2 \right]$$
Ici, l'échantillon transporté agit comme une cible de régression stable. Cela permet au générateur d'apprendre à produire des échantillons alignés en un seul passage avant.

Extension aux Récompenses Boîte Noire : Pour les récompenses non différentiables, FAV emploie une estimation de gradient d'ordre zéro (lissage gaussien) pour approximer $\nabla r(x)$, permettant l'alignement même lorsque la fonction de récompense est une boîte noire.

Contributions Clés

1. Cadre d'Alignement Général : FAV est le premier cadre à aligner divers générateurs à peu d'étapes (incluant VAE, GAN, modèles de cohérence et cartes de flux) sans nécessiter de vraisemblances traitables, de solveurs ODE/SDE spécifiques ou de trajectoires multi-étapes.
2. Formulation VI Basée sur des Échantillons : En présentant l'alignement comme un échantillonnage à partir d'une distribution inclinée par la récompense utilisant SVGD et KDE, FAV découple la procédure d'alignement des dynamiques internes du générateur.
3. Préservation de l'Efficacité : Grâce à l'amortissement par régression de point fixe, FAV maintient l'efficacité d'inférence en une seule étape du générateur de base, évitant la surcharge computationnelle des processus d'échantillonnage itératif ou de distillation.
4. Compatibilité Boîte Noire : L'intégration de l'estimation de gradient d'ordre zéro permet à FAV d'optimiser des récompenses arbitraires de type boîte noire.

Résultats Expérimentaux

Les auteurs ont évalué FAV sur deux domaines principaux : la manipulation robotique et l'alignement de générateurs d'images.

1. Manipulation Robotique (RL Hors Ligne et Hors Ligne vers En Ligne)

• Benchmarks : 56 tâches de RL hors ligne (OGBench et D4RL) et 30 tâches de RL hors ligne vers en ligne.
• Performance : FAV a atteint des performances de pointe sur toutes les tâches, surpassant 10 baselines incluant des gradients de politique reparamétrés, l'Appariement Adjoint et des politiques génératives distillées (par exemple, FQL, IFQL).
• Découverte Clé : FAV a démontré des capacités supérieures d'extraction de politique et des performances de transition de l'apprentissage hors ligne vers en ligne, même sans modèles génératifs auxiliaires ou débruitage multi-étapes.

2. Alignement de Générateurs d'Images

• Échelles : Les expériences ont varié de ImageNet-256 (256x256) à la synthèse texte-à-image haute résolution (1024x1024).
• Modèles : Testé sur StyleGAN-XL, modèles dérivants, iMeanFlow et SANA-Sprint.
• Métriques : Évalué contre des scores esthétiques, des scores de préférence humaine (HPS) et des détecteurs de sécurité (NSFW).
• Performance :
  • FAV a constamment atteint la meilleure frontière de Pareto entre l'optimisation de la récompense et la diversité/qualité des échantillons, atténuant la sur-optimisation de la récompense mieux que des baselines comme DRaFT et l'Appariement Adjoint.
  • Efficacité d'Inférence : FAV a généré des échantillons alignés environ 180 à 280 fois plus rapidement que les méthodes d'alignement au moment de l'inférence (par exemple, Best-of-N, ReNO) tout en maintenant une qualité d'image compétitive.
  • Évolutivité : La méthode a réussi à passer de résolutions 256x256 à 1024x1024 et s'est généralisée à travers différentes architectures de modèles.

Signification et Revendications

L'article revendique que FAV fournit une solution unifiée et agnostique au modèle pour l'alignement des modèles génératifs à peu d'étapes. Sa signification réside dans :

• Suppression des Barrières Structurelles : Il élimine le besoin de vraisemblances traitables ou de trajectoires d'échantillonnage spécifiques, rendant l'alignement accessible à une classe plus large de modèles génératifs implicites.
• Équilibre entre Efficacité et Qualité : Il réconcilie avec succès le compromis entre la génération à haute récompense et l'efficacité computationnelle de l'inférence à peu d'étapes, un défi où les méthodes précédentes dégradaient souvent les performances ou nécessitaient une distillation complexe.
• Robustesse : Le cadre démontre une robustesse à travers des tâches diverses, allant du contrôle de faible dimension à la synthèse d'images haute résolution, et gère à la fois les fonctions de récompense différentiables et de type boîte noire.

Les auteurs concluent que FAV représente une approche forte et polyvalente pour l'extraction de politique en RL et l'alignement de récompense dans la génération d'images, offrant une alternative pratique aux méthodes plus complexes et chargées d'hypothèses.