Résumé Technique : Apprendre à échantillonner à partir de modèles de diffusion via l'apprentissage par renforcement inverse

Énoncé du Problème

Les modèles de diffusion sont devenus une norme pour la génération de données de haute dimension, pourtant la procédure d'échantillonnage demeure un goulot d'étranglement critique. Bien que le réseau de neurones de débruitage soit généralement fixe après l'entraînement, l'algorithme d'échantillonnage nécessite un réglage minutieux des hyperparamètres (par exemple, les schémas de bruit, les échelles de guidage, les profils de stochasticité). Actuellement, ce réglage est effectué via des recherches par grille empiriques coûteuses. Pour des modèles complexes comme ceux entraînés sur ImageNet-64, le coût computationnel d'une recherche exhaustive d'hyperparamètres peut approcher ou dépasser le coût de l'entraînement du débruiteur lui-même. De plus, les approches existantes d'apprentissage par renforcement (RL) ou de contrôle pour les modèles de diffusion nécessitent souvent des fonctions de récompense explicites ou un réentraînement du débruiteur, ce qui contredit l'objectif d'optimiser l'échantillonneur sans modifier le générateur pré-entraîné.

Méthodologie

Les auteurs proposent de formuler l'échantillonnage de la diffusion comme un problème de prise de décision séquentielle solvable via l'Apprentissage par Renforcement Inverse (IRL). L'innovation centrale est l'apprentissage d'une politique d'échantillonnage qui correspond aux distributions d'états cibles sans accès à des trajectoires d'experts, à des récompenses explicites ou à un réentraînement du débruiteur.

1. Formulation MDP

Le processus d'échantillonnage de la diffusion est formalisé comme un Processus de Décision de Markov (MDP) à horizon fini :

• État ($s_t$) : Défini comme $(x_t, \sigma_t)$, où $x_t$ est l'échantillon actuel et $\sigma_t$ est le niveau de bruit.
• Action ($a_t$) : Décisions de contrôle discrètes appliquées à la dynamique d'échantillonnage, telles que la sélection des échelles de guidage, les niveaux d'injection de bruit, ou la décision de re-bruiter (redémarrer) vers un niveau de bruit précédent.
• Dynamique : Le noyau de transition est induit par les dynamiques de diffusion discrétisées et l'action de contrôle sélectionnée. Le processus se termine lorsque le niveau de bruit atteint $\Sigma_0$ (un état absorbant).

2. Apprentissage par Renforcement Inverse basé uniquement sur l'état

Contrairement à l'IRL traditionnel qui suppose souvent des actions ou des trajectoires d'experts, cette méthode opère sous un régime basé uniquement sur l'état. L'"expert" est défini uniquement par les distributions marginales des échantillons de données aux niveaux de bruit sélectionnés.

• Objectif : Le but est de trouver une politique $\pi_\theta$ dont la mesure d'occupation d'état $\mu_\theta$ (la distribution marginale des états visités) correspond à la mesure d'occupation de l'expert $\mu_E$.
• Optimisation : Au lieu d'apprendre une fonction de récompense, les auteurs minimisent directement la divergence $f$ entre les mesures d'occupation de l'expert et de la politique :
  $$\min_\theta D_f(\mu_E \parallel \mu_\theta)$$
  Le papier explore spécifiquement les divergences de Kullback-Leibler (KL) et de Reverse-KL (rKL), notant que le choix de la divergence $f$ dicte le compromis entre la couverture de mode (diversité) et la recherche de mode (fidélité).

3. Gradient de Politique et Optimisation

Les auteurs dérivent un théorème de gradient de politique pour minimiser directement les divergences $f$. Le gradient est exprimé comme :
$$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$$
où $h_f$ est une fonction dérivée du générateur de la divergence $f$.

• Estimation du rapport de densité : Puisque $\mu_E$ et $\mu_\theta$ sont inconnus, le rapport de densité est estimé à l'aide d'un discriminateur entraîné pour distinguer les états de l'expert (données bruitées) des états générés par la politique.
• Stabilisation : Pour gérer la variance élevée des estimations de gradient, la méthode emploie un objectif de type PPO avec écrêtage (clipping), des corrections par échantillonnage préférentiel (importance sampling) et une base de moyenne de lot globale.

4. Stratégies d'Échantillonnage Évaluées

Le cadre est appliqué à trois interventions spécifiques au moment de l'inférence :

1. Guidage de type Classifier-Free Adaptatif : Apprentissage d'une échelle de guidage dépendante de l'état $\omega(x, \sigma)$.
2. Injection de Stochasticité : Apprentissage d'un niveau d'injection de bruit dépendant de l'état $\gamma_{EDM}(x, \sigma)$.
3. Re-bruitage (Échantillonnage par Redémarrage) : Apprendre quand revenir à un niveau de bruit précédent pour échapper à de mauvaises générations locales.

Contributions Clés

1. Formulation MDP pour l'Échantillonnage : Le papier présente l'échantillonnage de la diffusion comme un MDP pour adapter les échantillonneurs au moment de l'inférence sans réentraîner le débruiteur.
2. Correspondance Directe des Marges d'État : Il dérive un objectif de gradient de politique qui minimise les divergences $f$ entre les mesures d'occupation d'état, contournant ainsi le besoin de modélisation de récompense explicite ou de données d'actions d'experts.
3. Compromis Interprétables : Le cadre expose un compromis contrôlable entre la diversité des échantillons (Recall) et la fidélité (Precision) via le choix de la divergence $f$ (KL vs rKL).
4. Contrôle Dépendant de l'État : Les politiques apprises agissent comme des contrôles au moment de l'inférence qui interviennent sélectivement (par exemple, en appliquant un guidage élevé uniquement à des niveaux de bruit spécifiques ou pour des types d'échantillons spécifiques) plutôt qu'en appliquant des calendriers globaux fixes.

Résultats Expérimentaux

La méthode a été évaluée sur CIFAR-10, FFHQ et ImageNet-64 en utilisant des modèles EDM pré-entraînés.

• Performance : Les politiques apprises égalent ou dépassent les performances des meilleures bases hand-tuned (trouvées via recherche par grille) à travers les métriques FID, Precision et Recall.
  • Guidage Adaptatif : A amélioré le compromis qualité-diversité par rapport au guidage fixe. La divergence KL a amélioré le Recall (diversité), tandis que la rKL a amélioré la Precision (fidélité).
  • Injection de Stochasticité : Les profils dépendants de l'état appris ont récupéré les motifs qualitatifs des schémas optimaux d'EDM (faible stochasticité à la fin de la trajectoire) et ont surpassé les bases fixes sur des jeux de données complexes comme ImageNet.
  • Re-bruitage : A réduit considérablement le nombre d'évaluations de fonctions (NFE) sur ImageNet tout en maintenant ou en améliorant le FID par rapport aux stratégies de redémarrage standard.
• Efficacité Computationnelle :
  • Entraînement : Un seul cycle d'entraînement remplace la recherche par grille exhaustive. Sur ImageNet-64, cela réduit le budget d'optimisation jusqu'à 9$\times$ (de ~1437 EFLOPs pour une recherche par grille de 3 hyperparamètres à ~161 EFLOPs pour l'entraînement IRL).
  • Inférence : La politique ajoute un overhead modeste de 16% à 40% en FLOPs par rapport au débruiteur seul, ce qui est considéré comme négligeable compte tenu de l'élimination du réglage manuel.
• Résultats d'Ablation :
  • Conditionnement : Les politiques conditionnées par l'état complet $(x, \sigma)$ surpassent nettement celles conditionnées uniquement par le niveau de bruit $\sigma$, et sont presque aussi performantes que celles avec des étiquettes de classe explicites $(x, \sigma, c)$, suggérant que la politique apprend implicitement des caractéristiques discriminantes de classe à partir du contenu de l'image.
  • Choix de la Divergence : La KL et la rKL se sont révélées être les plus efficaces pour contrôler le compromis précision-rappel, tandis que les divergences symétriques (TV, Hellinger) produisaient des résultats moins interprétables.

Signification et Revendications

Le papier affirme que ce cadre fournit une alternative systématique et automatisée au réglage manuel des hyperparamètres pour l'échantillonnage de la diffusion. En traitant l'échantillonnage comme un problème d'IRL, il permet l'optimisation de stratégies d'échantillonnage complexes et dépendantes de l'état qui sont difficiles à concevoir manuellement.

Les auteurs soulignent que la méthode est particulièrement précieuse pour les modèles à grande échelle où le coût de l'évaluation par essai dans une recherche par grille est prohibitif. L'approche ne modifie pas le modèle génératif lui-même mais apprend plutôt une "politique de contrôle" optimale pour la trajectoire d'inférence, offrant une voie pratique pour améliorer la qualité et la diversité des échantillons sans le fardeau computationnel du réentraînement ou de la recherche exhaustive. Le papier reste modeste, notant que les bénéfices dépendent du jeu de données et que le cadre actuel repose sur des espaces d'actions discrets, suggérant des travaux futurs dans les espaces d'actions continus et une initialisation plus robuste.