Résumé Technique : Diff-Instruct avec Récompense Diffusée (Didr)

Énoncé du Problème

Les récentes avancées des générateurs texte-vers-image (T2I) en une seule étape ont permis une synthèse en temps réel en mappant directement le bruit latent vers des images lors d'une seule passe avant. Cependant, l'alignement de ces générateurs avec les préférences humaines via l'apprentissage par renforcement (RL) présente un défi structurel fondamental.

Les méthodes précédentes pour le RLHF en une seule étape (par exemple, Diff-Instruct*, Diff-Instruct++) combinent l'optimisation de la récompense dans l'espace des images avec l'appariement de la distribution dans l'espace latent. Ces approches appliquent généralement les signaux de préférence (récompenses) exclusivement au point final d'image propre ($x_0$) tout en imposant une régularisation de Kullback-Leibler (KL) contre la distribution de référence sur toute la trajectoire de diffusion bruitée. Les auteurs identifient un décalage structurel critique dans ce paradigme : la récompense n'agit qu'à l'état terminal, tandis que le régularisateur s'étend sur tous les niveaux de bruit.

Ce décalage conduit à un mode de défaillance qualifié de domination de la récompense terminale. Parce que le bruit avant masque naturellement les détails de l'image, la pénalité de régularisation KL contre la divergence de mode s'affaiblit considérablement aux niveaux de bruit élevés. Par conséquent, les optimiseurs exploitent cette vulnérabilité de la trajectoire, abandonnant la distribution de référence pour s'effondrer sur des régions à haute récompense sans encourir de pénalités KL suffisantes. Cela se traduit par une amélioration des scores de récompense, mais souvent au détriment de la fidélité de l'image et de la stabilité distributionnelle.

Méthodologie : Diff-Instruct avec Récompense Diffusée (Didr)

Pour résoudre ce désalignement, les auteurs proposent Didr, un cadre d'alignement sans données et au niveau de la trajectoire, dérivé de la minimisation de la KL Intégrale (IKL). L'idée centrale est de propager systématiquement la distribution d'images propres biaisée par la récompense optimale RLHF à travers tous les niveaux de bruit de la trajectoire de diffusion.

1. La Cible de Trajectoire Principée

La méthode commence par établir la cible d'alignement correcte pour les images propres sous un RLHF régularisé par KL. La distribution d'images propres optimale $q^*(x_0|c)$ est la densité biaisée par la récompense :
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
où $q_0$ est la distribution de référence, $r$ est la fonction de récompense, et $\tau$ est un paramètre de température.

Au lieu d'appliquer les récompenses uniquement à $x_0$, Didr « élève » cette cible optimale à l'ensemble de la trajectoire en diffusant $q^*$ à travers le processus avant de référence. Cela génère une famille de marginales biaisées par la récompense $\{q^*_t\}_{t \in [0, T]}$ définies à chaque niveau de bruit :
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
Les auteurs prouvent que la minimisation de la divergence KL Intégrale entre les marginales du générateur et ces marginales biaisées par la récompense partage le même minimiseur que l'objectif standard de RLHF régularisé par KL. Cela garantit que la récompense et la régularisation sont équilibrées à chaque niveau de bruit, empêchant la domination de la récompense terminale.

2. Le Score de Récompense Diffusée (DRS)

En différenciant l'objectif IKL, les auteurs dérivent la fonction de score cible. Le score optimal se décompose en deux composantes :
$$s_{target} = s_{ref} + s_r$$
où $s_{ref}$ est la fonction de score de référence, et $s_r$ est le Score de Récompense Diffusée (DRS). Le DRS agit comme un terme de correction piloté par la récompense, défini à chaque niveau de bruit :
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
Contrairement à la guidance par classifieur (qui est une approximation du premier ordre), le DRS conserve l'ensemble du biais log-exponentiel, propageant l'information de préférence uniquement là où l'image propre est identifiable à partir de l'observation bruitée.

3. Le Proxy de Récompense Diffusée (DRP)

Puisque le DRS nécessite une espérance a posteriori intraitable $q(x_0|x_t, c)$, les auteurs introduisent le Proxy de Récompense Diffusée (DRP) pour rendre la méthode pratique. Le DRP approxime l'espérance a posteriori en utilisant des chaînes de débruitage à étapes courtes différentiables issues du modèle de référence figé.

• Mécanisme : À partir d'un état bruité $x_t$, la méthode exécute $K$ chaînes de débruitage indépendantes de $S$ étapes en utilisant le modèle de référence figé pour générer des échantillons propres approximatifs $\{\hat{x}^{(k)}_0\}$.
• Estimation du Gradient : Le DRS est approximé comme une moyenne pondérée par softmax des gradients de récompense le long des trajectoires :
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  où les poids $\omega^{(k)}$ sont proportionnels à $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$.
• Efficacité : Cette approche permet aux gradients de récompense de se propager de manière stable en arrière à travers la trajectoire sans stocker les activations du modèle de référence ni nécessiter de données d'entraînement d'images.

4. Cadre d'Entraînement

Le processus d'entraînement alterne entre deux étapes :

1. Mise à jour de l'Assistant d'Enseignement (TA) : Un modèle de score TA ($s_\psi$) est entraîné par Matching de Score de Débruitage (DSM) pour suivre les marginales évolutives du générateur, servant de substitut au score implicite du générateur.
2. Mise à jour du Générateur : Le générateur ($g_\theta$) est mis à jour en utilisant le gradient de l'objectif IKL, qui compare le score TA à la somme du score de référence corrigé par CFG et de l'estimation DRP.

Contributions Clés

1. Identification de la Domination de la Récompense Terminale : L'article identifie et analyse formellement un mode de défaillance structurel dans les méthodes précédentes de RLHF en une seule étape, où les optimiseurs exploitent les niveaux de bruit pour s'effondrer sur des régions à haute récompense, violant l'équilibre RLHF intended.
2. Alignement au Niveau de la Trajectoire Principé : Les auteurs proposent Didr, un cadre qui dérive une cible au niveau de la trajectoire mathématiquement solide en propageant vers l'avant la densité d'images propres optimale RLHF. Cela garantit que la récompense et la régularisation sont équilibrées à chaque niveau de bruit.
3. Déduction Analytique du DRS : L'article déduit analytiquement le Score de Récompense Diffusée, montrant qu'il décompose exactement la cible de trajectoire en un score de référence plus une correction induite par la récompense.
4. Estimateur Pratique (DRP) : Les auteurs introduisent le Proxy de Récompense Diffusée, un estimateur sans données qui calcule le DRS via un débruitage à étapes courtes différentiable, permettant un entraînement stable sans données d'images.

Résultats Empiriques

Des expériences extensives démontrent que Didr surpasse systématiquement les bases de référence existantes en une seule étape :

• Dominance de Pareto : Sur l'architecture de base SDXL (1024×1024), la variante standard de Didr domine au sens de Pareto toutes les bases de référence d'alignement en une seule étape existantes dans le compromis PickScore–FID, surpassant simultanément les modèles multi-étapes comme FLUX-dev et SDXL-DPO sur les deux axes.
• Métriques de Préférence : La variante d'entraînement étendue (Didr_longer) atteint le PickScore le plus élevé (23,9) parmi les modèles en une seule étape évalués et améliore ImageReward de 8,9 % par rapport à la base de référence précédente la plus forte (Diff-Instruct*). Elle mène également en HPSv2.1 et en Score Esthétique.
• Généralisation : Didr se transfère avec succès à l'architecture de base Z-Image de 6 milliards de paramètres. En une seule étape de génération, il surpasse son modèle enseignant de 50 étapes sur les métriques de préférence, démontrant une évolutivité à travers les architectures.
• Améliorations Qualitatives : Les comparaisons visuelles montrent que Didr produit des détails fins plus nets et des couleurs plus naturelles par rapport aux méthodes en une seule étape précédentes, évitant la saturation excessive et les artefacts picturaux souvent observés dans les bases de référence optimisées par récompense.

Importance et Revendications

L'article revendique que la propagation de récompense au niveau de la trajectoire principée est une stratégie d'alignement efficace qui résout le désalignement inhérent des méthodes de récompense uniquement terminales. En propageant rigoureusement la distribution optimale RLHF à travers le processus de diffusion, Didr atteint le meilleur compromis entre alignement des préférences et fidélité de l'image parmi les générateurs en une seule étape.

Les auteurs soulignent que leur méthode ne nécessite aucune donnée d'entraînement d'images, reposant uniquement sur des invites textuelles et un modèle de référence figé. Ils positionnent Didr comme une solution permettant aux générateurs en une seule étape de correspondre ou de dépasser l'alignement des préférences des enseignants multi-étapes tout en maintenant des vitesses d'inférence en temps réel. Le travail suggère que le décalage structurel identifié dans les méthodes précédentes est une limitation fondamentale qui ne peut être résolue qu'en alignant l'ensemble de la trajectoire, et non seulement l'état terminal.