Résumé Technique : Vers un contrôle plus général des modèles de diffusion via le Guidage de Jeffrey

1. Énoncé du Problème

Les modèles de diffusion offrent une flexibilité significative dans les tâches génératives, permettant un contrôle de la sortie lors de l'échantillonnage grâce à des mécanismes de guidage sans réentraînement du réseau de neurones sous-jacent. Bien que les techniques de guidage standard (par exemple, le guidage par classificateur) gèrent efficacement l'échantillonnage conditionnel pour des classes spécifiques, elles échouent souvent à traiter des objectifs plus complexes où la distribution cible est implicite ou définie uniquement par des heuristiques.

Les approches actuelles pour des tâches telles que l'équité (fairness), le contrôle de concepts et la génération compositionnelle reposent fréquemment sur des modifications heuristiques de la dynamique d'échantillonnage. Ces méthodes inclinent la distribution originale mais manquent d'une distribution cible explicitement définie, ce qui rend l'objectif réel difficile à caractériser ou à interpréter. Il existe un besoin pour un cadre principled capable de mettre à jour les distributions marginales vers une distribution cible prescrite tout en préservant la structure conditionnelle du modèle, permettant ainsi le contrôle de propriétés statistiques complexes (par exemple, l'indépendance des attributs ou l'appariement de distribution d'embeddings) que le conditionnement standard ne peut exprimer.

2. Méthodologie : Guidage de Jeffrey

Le papier propose le guidage de Jeffrey, un cadre qui exploite la règle de conditionnement de Jeffrey pour mettre à jour les modèles de diffusion. Contraなる aux approches basées sur Bayes qui conditionnent sur des événements spécifiques (conditionnement standard), la règle de Jeffrey permet de conditionner sur des prédicats complexes, tels que l'appariement d'une distribution marginale spécifique.

Fondement Théorique

Étant donné un modèle conjoint $p(x, y)$ où $x$ représente l'espace image et $y$ les attributs/caractéristiques, l'objectif est de mettre à jour la marginale $p(y)$ vers une distribution cible $p^\star(y)$ tout en préservant la structure conditionnelle $p(x|y)$. Sous la règle de Jeffrey, la distribution jointe mise à jour est définie comme :
$$p^\star(x, y) = p(x|y)p^\star(y)$$
Cette mise à jour minimise la perturbation de la distribution jointe originale (une projection d'information) tout en satisfaisant la contrainte marginale.

Le papier dérive une formulation de repondération pour la marginale sur $x$ :
$$p^\star(x) = p(x) \int p(y|x)\rho(y) dy$$
où $\rho(y) = p^\star(y)/p(y)$ est le rapport de densité. Dans les cas où les caractéristiques sont dérivées de manière déterministe à partir des images ($y = u(x)$), cela se simplifie en :
$$p^\star(x) = p(x)\rho(u(x))$$

Implémentation dans les Modèles de Diffusion

Pour appliquer cela aux modèles de diffusion basés sur le score, les auteurs dérivent une fonction de score guidée. Le score standard $\nabla_{x_t} \log p_t(x_t)$ est augmenté d'un terme de correction basé sur le rapport de densité :
$$\nabla_{x_t} \log \tilde{p}^\star_t(x_t) = \nabla_{x_t} \log p_t(x_t) + \nabla_{x_t} \log \int p(y|\hat{x}_0(x_t, t))\rho(y) dy$$
où $\hat{x}_0(x_t, t)$ est une estimation de l'échantillon propre obtenue via la formule de Tweedie.

Pour les mappages de caractéristiques déterministes (par exemple, les espaces d'embeddings), le guidage se simplifie en :
$$\nabla_{x_t} \log \tilde{p}^\star_t(x_t) = \nabla_{x_t} \log p_t(x_t) + \nabla_{x_t} \log \rho(u(\hat{x}_0(x_t, t)))$$

Le cadre introduit un hyperparamètre $\lambda$ pour mettre à l'échelle le terme de correction, contrôlant le compromis entre la qualité de l'échantillon et la fidélité à la distribution cible. Le rapport de densité $\rho$ est estimé à l'aide d'un classificateur binaire entraîné pour distinguer les échantillons provenant de la distribution cible des échantillons de la distribution actuelle du modèle.

3. Principales Contributions

Le papier apporte trois contributions principales :

1. Cadre de Guidage de Jeffrey : Une nouvelle méthode de guidage "plug-and-play" qui généralise le guidage par classificateur standard. Elle utilise la règle de Jeffrey pour faire correspondre les distributions de sortie générées à des distributions cibles arbitraires, permettant le contrôle d'objectifs auparavant inaccessibles au guidage standard.
2. Appariement de Distribution d'Embedding : Une preuve de concept démontrant la capacité de faire correspondre les distributions de sortie du modèle à une distribution cible dans un espace d'embedding (spécifiquement les embeddings Inception). Cette application réduit considérablement la distance de Fréchet Inception (FID) sans nécessordre de réentraînement.
3. Contrôle de la Distribution d'Attributs pour l'Équité : Une démonstration du contrôle des distributions d'attributs d'images pour atteindre des objectifs d'équité. La méthode impose avec succès la parité de genre (appariement de distribution) et décorréle les attributs (par exemple, "Mâle" et "Jeune") dans les visages générés, des tâches que le guidage par classificateur standard ne peut pas traiter directement.

4. Résultats Expérimentaux

Les auteurs ont évalué le guidage de Jeffrey sur les jeux de données CIFAR-10, FFHQ et CelebA-HQ.

Guidage d'Embedding (Réduction de la FID)

• Configuration : La méthode visait la distribution d'embedding Inception des données d'entraînement.
• Résultats : Sur CIFAR-10, la FID s'est améliorée de $\approx 5.88$ à $\approx 2.55$ (ensemble d'entraînement) et de $\approx 9.76$ à $\approx 6.71$ (ensemble de test). Sur FFHQ, la FID est passée de $\approx 20.13$ à $\approx 12.91$.
• Observation : Les auteurs notent que bien que la FID ait diminué drastiquement, les changements visuels aux images générées étaient minimaux, suggérant des limites de la FID en tant que métrique purement perceptuelle. Les meilleurs résultats ont été obtenus en appliquant le guidage uniquement aux dernières étapes de débruitage ($\delta=10$).

Applications d'Équité (Fairness)

• Parité de Genre (Appariement de Distribution) : Sur CelebA-HQ, la méthode a ajusté la proportion d'attributs "Mâle" d'environ $\approx 0.296$ à $\approx 0.522$ (ciblant 0.5), réduisant la Discrépance d'Équité (FD) de $0.288$à$0.032$. Ceci a été réalisé avec une dégradation minimale de la FID ($21.61$ contre $22.30$ pour la référence).
• Décorrélation d'Attributs : La méthode a réussi à décorréler les attributs "Mâle" et "Jeune", réduisant le coefficient de corrélation de Pearson de $-0.541$ à presque zéro ($0.005$) à$\lambda=3.0$. Cela a été accompli en réduisant le mode conjoint dominant (échantillons de femmes jeunes) tout en maintenant les distributions marginales relativement stables, avec une légère augmentation de la FID d'environ un point.
• Comparaison : Comparé à une référence d'échantillonnage ancestral combiné à un guidage standard, le guidage de Jeffrey a obtenu des métriques d'équité similaires avec des coûts de FID plus faibles et modifie moins d'échantillons, car il équilibre les directions de classe de manière continue le long de la trajectoire d'échantillonnage plutôt que de pousser entièrement les échantillons vers une seule classe.

5. Signification et Revendications

Le papier positionne le guidage de Jeffrey comme une extension principled du contrôle des modèles de diffusion. Sa signification réside dans :

• Généralisation : Il recontextualise le guidage par classificateur standard comme un cas particulier du conditionnement de Jeffrey, unifiant ainsi ces approches sous un cadre théorique plus large.
• Expressivité : Il permet de spécifier des distributions cibles explicites pour des objectifs complexes (comme l'indépendance des attributs) qui n'étaient auparavant réalisables que via des méthodes heuristiques ou implicites.
• Praticité : La méthode ne nécessite aucun réentraînement du modèle de diffusion, seulement l'ajout d'un terme de correction lors de l'échantillonnage et l'entraînement d'un estimateur de rapport de densité léger.

Les auteurs maintiennent un ton modeste, présentant leur travail comme une "preuve de concept". Ils reconnaissent que l'implémentation actuelle repose sur un classificateur pour estimer les rapports de densité, similaire à la façon dont le premier guidage par classificateur fonctionnait, et expriment l'espoir que les travaux futurs évolueront vers des variantes sans classificateur (classifier-free). Ils notent également que si la méthode offre une nouvelle puissance d'expression pour des applications bénéfiques (équité, atténuation de la mémorisation), elle porte aussi le risque de donner du pouvoir à des acteurs malveillants pour générer des contenus néfastes.