Résumé technique : Modèles de dérive à gradient de noyau

Énoncé du problème

La modélisation générative repose souvent sur des méthodes de transport qui transforment une loi a priori simple en une distribution de données complexe. Bien que performantes empiriquement, les approches standard (par exemple, les modèles de diffusion, l'appariement de flux) nécessitent généralement de résoudre des équations différentielles stochastiques (EDS) ou des équations différentielles ordinaires (EDO) via de nombreuses mises à jour séquentielles, rendant l'échantillonnage coûteux en calcul. Cela a stimulé l'intérêt pour la génération en une seule étape ou en quelques étapes, souvent obtenue par distillation de modèles préentraînés ou par l'entraînement direct de générateurs via des objectifs de carte de flux ou de cohérence.

Les modèles de dérive représentent une approche directe pour la génération en une seule étape. Ils entraînent un générateur en déplaçant les échantillons du modèle le long d'un champ d'attraction et de répulsion induit par un noyau, déplaçant ainsi le processus de raffinement du temps d'échantillonnage au temps d'entraînement. Cependant, la compréhension théorique de la dérive est actuellement limitée aux noyaux gaussiens. Dans le cas gaussien, le déplacement euclidien pondéré par le noyau coïncide avec une direction de matching de score lissée, assurant l'identifiabilité (c'est-à-dire qu'une dérive nulle implique que la distribution du modèle correspond à la distribution des données).

Pour les noyaux non gaussiens (par exemple, Laplace), qui fonctionnent souvent mieux empiriquement, la formulation standard combine une fonction de pondération par noyau avec un vecteur de déplacement euclidien fixe ($y-x$). L'article soutient que pour les noyaux non gaussiens, cette direction de déplacement ne s'aligne pas nécessairement avec la direction de l'augmentation la plus raide de la similarité (le gradient du noyau). Par conséquent, le champ de dérive devient biaisé vers les voisins lointains, brisant la propriété d'identifiabilité et obscurcissant l'objectif d'optimisation. De plus, la formulation standard est liée à la géométrie euclidienne, rendant incertaine l'application de la dérive à des données sur des variétés riemanniennes ou des espaces de données discrets.

Méthodologie : Dérive à gradient de noyau

Les auteurs proposent la Dérive à gradient de noyau, une reformulation où le noyau lui-même définit la direction du mouvement via son gradient, plutôt que de reposer sur un déplacement euclidien fixe.

1. Le champ de dérive

Au lieu du champ de déplacement original $V_{p,q_\theta}(x) = \mathbb{E}[k(x,y)(y-x)]$, la méthode proposée définit la dérive en utilisant le gradient du noyau :
$$V^\nabla_{p,q_\theta}(x) = \frac{\mathbb{E}_{y \sim p}[\nabla_x k(x, y)]}{\mathbb{E}_{y \sim p}[k(x, y)]} - \frac{\mathbb{E}_{x' \sim q_\theta}[\nabla_x k(x, x')]}{\mathbb{E}_{x' \sim q_\theta}[k(x, x')]}$$
Ce champ représente la différence entre l'attraction par gradient de noyau (vers les données) et la répulsion par gradient de noyau (auto-répulsion des échantillons du modèle).

2. Propriétés théoriques

• Interprétation du rapport de score : Sous des hypothèses de régularité standard, la dérive à gradient de noyau est exactement le gradient du logarithme du rapport des densités lissées par le noyau :
  $$V^\nabla_{p,q_\theta}(x) = \nabla_x \log \frac{\hat{p}_k(x)}{\hat{q}_{\theta,k}(x)}$$
  Cela permet de retrouver l'interprétation basée sur le score pour des noyaux généraux, et pas seulement pour les noyaux gaussiens.
• Identifiabilité : Si le noyau $k$ est caractéristique (l'application de lissage est injective), alors une dérive nulle ($V^\nabla = 0$) implique que les densités lissées sont égales, ce qui implique à son tour que les distributions originales sont égales ($p = q_\theta$). Cela restaure la garantie théorique d'identifiabilité pour les noyaux non gaussiens.
• Descente de KL lissée : La dynamique correspond au flot de gradient de Wasserstein de la divergence de Kullback-Leibler (KL) entre les distributions lissées, fournissant une interprétation variationnelle de l'objectif d'entraînement.

3. Extension aux variétés et aux données discrètes

Comme les gradients de noyau sont des vecteurs tangents intrinsèques, le cadre s'étend naturellement au-delà de l'espace euclidien :

• Variétés riemanniennes : La dérive est définie en utilisant le gradient riemannien $\nabla_M$. Les échantillons sont transportés via l'application exponentielle ($\text{Exp}_x$), respectant la géométrie intrinsèque du support des données (par exemple, les sphères).
• Données discrètes : En considérant les distributions catégorielles comme des points sur un simplexe de probabilité muni de la métrique de Fisher-Rao, la méthode mappe le problème vers le quadrant positif d'une hypersphère. Cela permet une génération en une seule étape de données discrètes (par exemple, séquences d'ADN, molécules) sans distillation.

Contributions clés

1. Reformulation : Introduction de la dérive à gradient de noyau, remplaçant les déplacements euclidiens fixes par des directions de gradient induites par le noyau.
2. Unification théorique : Démonstration que la dérive est un processus de matching de score lissé pour n'importe quel noyau différentiable, récupérant la dérive gaussienne comme cas particulier.
3. Preuve d'identifiabilité : Démonstration que l'identifiabilité vaut pour tout noyau caractéristique, résolvant l'écart théorique pour les cas non gaussiens.
4. Généralisation géométrique : Extension du cadre aux variétés riemanniennes et aux données discrètes via la géométrie de Fisher-Rao, permettant une génération en une seule étape sur des sphères et pour des variables catégorielles.

Résultats expérimentaux

Les auteurs ont évalué la méthode dans trois contextes :

1. Expériences synthétiques :

   • Sur des cibles en damier et en rouleau suisse sur des variétés sphériques ($S^2$) et hyperboliques ($H^2$), la formulation par gradient a constamment surpassé la formulation de base (basée sur le déplacement) en termes de distance de Wasserstein tranchée et de précision des tuiles.
   • Des études d'ablation sur les noyaux de Matérn ont montré que, bien que les noyaux gaussiens ($\nu \to \infty$) aient donné des performances similaires pour les deux formulations, la direction du gradient était cruciale pour les noyaux non gaussiens (par exemple, Laplace, $\nu=0.5$).

2. Données géospatiales (Événements terrestres) :

   • Évaluation sur des distributions réelles de volcans, tremblements de terre, incendies et inondations sur $S^2$.
   • Les variantes Gradient Laplace Sphérique et Spectrale ont obtenu les meilleurs scores de Discrepancy Moyenne Maximale (MMD) sur les quatre jeux de données, surpassant la base Laplace euclidienne vanille et d'autres variantes de gradient.

3. Génération de données discrètes :

   • ADN promoteur : Génération de séquences de longueur 1024. Le modèle Gradient Sphérique Gaussien a atteint une corrélation 6-mer de 0,90, réduisant l'écart avec les bases distillées par flux (qui ont atteint 0,96) tout en restant sans enseignant.
   • Génération de molécules (QM9) : Génération de chaînes SMILES. Le modèle Gradient Sphérique Laplace a atteint 38,9 % de validité et 44,1 % d'unicité, surpassant considérablement la base Laplace euclidienne (22,0 % de validité) et approchant les performances des modèles de flux distillés (par exemple, PairFlow à 44,3 % de validité), malgré d'être une méthode en une seule étape et sans enseignant.

Importance et affirmations

L'article affirme que la dérive à gradient de noyau fournit une extension fondée sur des principes des modèles de dérive au-delà du cas euclidien gaussien. En alignant la direction de dérive avec le gradient du noyau, la méthode :

• Restaure les garanties théoriques (identifiabilité) pour les noyaux non gaussiens.
• Offre une interprétation cohérente basée sur le score pour des noyaux généraux.
• Se généralise naturellement à des domaines non euclidiens (variétés) et à des données discrètes sans nécessiter de procédures de distillation complexes.

Les auteurs positionnent cela comme une voie flexible et sans enseignant vers une génération de haute qualité en une seule étape. Ils reconnaissent des limites, notant que la méthode est actuellement sensible au choix des hyperparamètres (spécifiquement la largeur de bande du noyau et la régularité) et se concentre principalement sur les noyaux de Matérn. Ils suggèrent que des travaux futurs devraient explorer des classes de noyaux plus larges et une analyse systématique de la stabilité de l'entraînement pour réduire davantage l'écart avec les modèles distillés les plus avancés.