Likelihood Matching for Diffusion Models

Résumé Technique : Appariement de Vraisemblance pour les Modèles de Diffusion

Énoncé du Problème
Les modèles de diffusion dominants actuels, en particulier les modèles génératifs basés sur le score (SGM), reposent sur l'appariement de score (SM - Score Matching) comme objectif d'entraînement principal. Bien qu'efficace empiriquement, le SM minimise la divergence entre une fonction de score paramétrée et la véritable fonction de score à différents niveaux de bruit. Théoriquement, cette approche ne minimise qu'une borne supérieure sur la log-vraisemblance de la distribution de données originale plutôt que la vraisemblance elle-même. Cette connexion indirecte peut entraîner une perte d'efficacité statistique par rapport à l'estimation du maximum de vraisemblance (MLE), qui est connue pour produire des estimateurs cohérents et efficaces. De plus, le SM standard se concentre exclusivement sur l'information du premier ordre (le score/gradient), négligeant potentiellement les informations structurelles d'ordre supérieur contenues dans les densités de transition.

Méthodologie
Les auteurs proposent l'appariement de vraisemblance (LM - Likelihood Matching), un cadre qui cible directement la vraisemblance de la distribution des données en exploitant l'équivalence entre la vraisemblance des données cibles et la vraisemblance de trajectoire du processus de diffusion inverse.

1. Fondement Théorique :
   La méthode repose sur la propriété selon laquelle la vraisemblance jointe de la trajectoire de diffusion inverse est équivalente à la vraisemblance de la distribution de données originale (à des constantes près liées au processus direct). Les auteurs établissent que maximiser la vraisemblance exacte de la trajectoire du processus inverse est équivalent à maximiser la log-vraisemblance des données.

2. Estimation Quasi-Maximum de Vraisemblance (QMLE) :
   Puisque les densités de transition inverse exactes sont intraitables, les auteurs les approximent à l'aide d'une approche de quasi-maximum de vraisemblance. Ils modélisent chaque densité de transition inverse comme une distribution gaussienne où la moyenne conditionnelle et la covariance sont appariées au véritable processus inverse.

   • Idée Clé : La moyenne et la covariance conditionnelles du processus inverse dépendent non seulement de la fonction de score ($\nabla \log q_t$), mais aussi de la fonction Hessienne ($\nabla^2 \log q_t$).
   • Paramétrage : Les auteurs paramètrent à la fois le score et l'Hessienne à l'aide de réseaux neuronaux (spécifiquement des U-Nets). Pour gérer la complexité computationnelle, l'Hessienne est souvent approximée par une structure diagonale plus de rang faible.

3. L'Objectif d'Appariement de Vraisemblance :
   La fonction d'objectif proposée, l'appariement de vraisemblance (LM), est dérivée en minimisant la quasi-log-vraisemblance négative des transitions inverses.

   • L'objectif unifie l'appariement de score et l'appariement de covariance.
   • Il incorpore un mécanisme de pondération de vraisemblance implicite via la matrice de covariance appariée $(I + \sigma^2 H_t)^{-1/2}$, qui ajuste automatiquement le poids du terme d'appariement de score en fonction de la courbure locale.
   • L'entraînement utilise un échantillonneur stochastique qui sélectionne aléatoirement des pas de temps le long de la trajectoire de diffusion pour calculer la perte, assurant une utilisation efficace de l'information évolutive de l'SDE (équation différentielle stochastique) directe.

4. Échantillonneur Stochastique :
   Une nouvelle procédure d'échantillonnage est introduite, utilisant à la fois les informations apprises du score et de l'Hessienne. L'étape inverse est définie comme :
   $$\tilde{Y}_{t-1} = \hat{\mu}_{t-1|t}(\tilde{Y}_t) + \hat{\Sigma}_{t-1|t}^{1/2}(\tilde{Y}_t) Z_t$$
   où la moyenne $\hat{\mu}$ et la covariance $\hat{\Sigma}$ sont des fonctions du score et de l'Hessienne estimés, respectivement.

Principales Contributions

• Nouvel Objectif d'Entraînement : Le papier introduit un objectif d'entraînement basé sur la quasi-vraisemblance (LM) qui approxime la log-vraisemblance de la trajectoire inverse, combinant efficacement l'appariement de score et l'appariement de covariance avec une pondération de vraisemblance implicite.
• Échantillonnage Enrichi par l'Hessienne : Un échantillonneur stochastique est dérivé, exploitant à la fois les informations du score et de l'Hessienne, offrant une approximation plus précise de la densité de transition inverse que les échantillonneurs de type DDPM standard.
• Garanties Théoriques :
  • Convergence : Des garanties de convergence non asymptotiques sont fournies en distance de variation totale. La borne d'erreur est de l'ordre de $O(d^3 \log^{4.5} T / T)$ pour l'erreur de l'étape inverse, avec des termes supplémentaires pour les erreurs d'estimation du score et de l'Hessienne.
  • Cohérence : Le papier prouve la cohérence du quasi-MLE proposé sous un cadre paramétrique oracle, montrant que l'estimateur converge vers les vrais paramètres des données à mesure que la taille de l'échantillon et les étapes de diffusion augmentent.
• Mise en Œuvre Efficace : L'utilisation d'approximations de l'Hessienne de rang faible et de la formule de Sherman-Morrison-Woodbury permet un calcul efficace de la vraisemblance et des étapes d'échantillonnage.

Résultats Expérimentaux
Les auteurs valident la méthode sur des ensembles de données synthétiques et réels :

• Données Synthétiques : Sur des mélanges de distributions Gaussiennes et de t-distributions, le LM surpasse systématiquement l'appariement de score (SM) standard, particulièrement pour capturer des modes bien séparés. Sur une distribution "Swiss-roll" hautement enroulée (représentant des variétés courbes), le LM démontre une performance supérieure, validant l'avantage de la covariance informée par l'Hessienne.
• Estimation de Paramètres : Dans un cadre de mélange gaussien paramétrique, le LM a obtenu une erreur absolue moyenne (MAE) et une erreur standard plus faibles dans l'estimation des paramètres par rapport au SM, soutenant ainsi les affirmations de cohérence théorique.
• Jeux de Données d'Images (MNIST, CIFAR-10, CelebA) :
  • Le LM a obtenu une distance de Fréchet Inception (FID) plus faible et de meilleurs scores de log-vraisemblance négative (NLL) par rapport aux bases de référence SM.
  • Les performances s'améliorent avec le rang de l'approximation de l'Hessienne ($r$), atteignant un pic autour de $r=20-30$.
  • Compromis Vitesse-Qualité : Le LM a atteint un FID plus bas pour un temps de calcul comparable et a nécessité moins d'étapes d'échantillonnage inverse pour atteindre une haute qualité visuelle. Par exemple, sur MNIST, le LM a produit des chiffres reconnaissables après seulement 20 itérations inverses, alors que les échantillons SM restaient flous.

Signification et Revendications
Le papier affirme que l'appariement de vraisemblance offre une alternative plus rigoureuse sur le plan statistique que l'appariement de score en ciblant directement la vraisemblance des données via un cadre de quasi-vraisemblance. En incorporant l'information du second ordre (l'Hessienne), la méthode capture des densités de transition plus riches et améliore l'efficacité de l'échantillonnage, particulièrement dans les régimes à faible nombre d'étapes et sur des variétés de données complexes et courbes. Les auteurs soulignent que bien que le LM entraîne des coûts de calcul par étape plus élevés en raison de la modélisation de l'Hessienne, la réduction du nombre d'étapes d'échantillonnage requises et l'amélioration de la qualité de la génération (FID et NLL) valident son efficacité pour l'entraînement de modèles génératifs de haute fidélité. Ce travail comble le fossé entre l'efficacité statistique du MLE et le succès pratique des modèles de diffusion.