Résumé Technique : DriftXpress

Énoncé du Problème

Les Modèles de Dérive (Drifting Models) représentent un changement de paradigme dans la modélisation générative en permettant une génération d'images en une seule étape, sans l'inférence itérative requise par les modèles de diffusion ou basés sur les flux. Au lieu d'un débruitage itératif, les Modèles de Dérive apprennent un champ d'attraction-répulsion dans l'espace des caractéristiques qui attire les échantillons générés vers la distribution des données et les repousse de la distribution actuelle du modèle.

Bien qu'efficaces, les Modèles de Dérive standards introduisent un compromis computationnel significatif : l'inférence est peu coûteuse (étape unique), mais l'entraînement est coûteux. Le goulot d'étranglement principal réside dans l'estimation répétée du champ d'attraction, qui nécessite le calcul des interactions exactes du noyau entre chaque échantillon généré et l'ensemble du support d'entraînement (ou de grands mini-lots) à chaque étape d'entraînement. Cela se traduit par des coûts d'entraînement élevés en temps réel, limitant l'évolutivité de l'approche. L'article s'interroge sur la possibilité d'améliorer ce compromis entraînement-inférence pour réduire les coûts d'entraînement sans sacrifier l'avantage de l'inférence en une seule étape.

Méthodologie : DriftXpress

Les auteurs proposent DriftXpress, une formulation accélérée des modèles de dérive basée sur des champs d'Espace de Hilbert à Noyau Reproduisant (RKHS) projetés. L'idée centrale est d'approximer la composante d'attraction du champ de dérive en utilisant un espace de caractéristiques de faible rang dérivé de la méthode de Nyström, tout en conservant la composante de répulsion exacte.

1. Projection RKHS via des Points Repères (Landmarks) Nyström

La méthode exploite l'approximation de Nyström pour projeter les interactions du noyau sur un sous-espace engendré par un ensemble fixe de points « repères » sélectionnés dans les données d'entraînement avant l'entraînement.

• Sélection des Points Repères : Un ensemble de points repères $U = \{u_1, \dots, u_r\}$ est échantillonné (par exemple, uniformément par classe) à partir des données d'entraînement. Ces points repères définissent un sous-espace $H_U \subset H_k$.
• Carte de Caractéristiques : Une carte de caractéristiques de Nyström $\phi(z)$ est construite en utilisant la matrice de noyau des points repères $K_{UU}$ et les évaluations du noyau entre un point $z$ et les points repères. Cela permet d'approximer la similarité du noyau $k(x, y)$ par un produit scalaire de faible rang $\phi(x)^\top \phi(y)$.

2. Attraction Projetée avec Résumés Mis en Cache

Dans la dérive standard, le champ attractif $V_p^+(x)$ est calculé comme une moyenne pondérée par le noyau normalisée des échantillons d'entraînement. DriftXpress remplace le noyau exact $k$ par le noyau projeté $k_U$.

• Précalcul : Au lieu de recalculer les interactions avec l'ensemble complet d'entraînement à chaque étape, DriftXpress précalcule deux résumés compacts sur l'ensemble du jeu de données d'entraînement :
  • $A_p = \sum \phi(y_j^+) (y_j^+)^\top$
  • $b_p = \sum \phi(y_j^+)$
• Évaluation Efficace : Pendant l'entraînement, la moyenne attractive $\mu_p^U(x)$ est calculée efficacement en utilisant ces résumés mis en cache et les caractéristiques projetées de la requête $\phi(x)$ :
  $$\mu_p^U(x) = \frac{A_p^\top \phi(x)}{\phi(x)^\top b_p + \epsilon}$$
  Cela réduit la complexité par étape, passant d'une dépendance à la taille de l'ensemble d'entraînement $N_+$ à une dépendance à la dimension des points repères $r$ (où $r \ll N_+$).

3. Répulsion Exacte

La composante répulsive $V_q^-(x)$, qui repousse les échantillons de l'lot généré actuel, est conservée exacte. Les auteurs observent que l'approximation de la répulsion entraîne un effondrement de mode sévère et une instabilité. Par conséquent, DriftXpress utilise une approche hybride : attraction projetée pour l'efficacité et répulsion exacte pour la stabilité.

4. Résumés Compositionnels (Partitionnement)

Pour gérer de grands jeux de données où le stockage des résumés pour toutes les classes en mémoire est impossible, DriftXpress emploie une stratégie de partitionnement (sharding). Les données d'entraînement sont partitionnées en fragments (par exemple, par classe), et des résumés séparés sont maintenus pour chaque fragment. Le champ attractif final est calculé en sommant les contributions de chaque fragment de manière séquentielle. Cela permet à la méthode de s'adapter à des jeux de données comme ImageNet sans matérialiser le tenseur de résumé complet dans la mémoire GPU, échangeant une légère réduction du débit contre la faisabilité mémoire.

Contributions Clés

1. Reformulation RKHS : Les auteurs reformulent les champs de dérive au sein du RKHS induit par le noyau de dérive, permettant une évaluation de faible rang et basée sur des points repères du champ d'attraction.
2. Algorithme DriftXpress : Ils introduisent un algorithme pratique qui remplace l'attraction exacte répétée par des résumés précalculés et mis en cache, et compose ces résumés à travers les fragments pour les jeux de données à grande échelle.
3. Garanties Théoriques : L'article fournit des bornes sur la distorsion du champ projeté, montrant que l'erreur est contrôlée par le résidu du noyau des points repères et la masse locale du noyau. Cela quantifie les conditions dans lesquelles l'approximation reste précise.
4. Validation Empirique : Des expériences extensives démontrent que DriftXpress préserve la qualité des échantillons tout en réduisant considérablement le temps d'entraînement.

Résultats

Les auteurs évaluent DriftXpress sur SVHN, CIFAR-10, CIFAR-100 et ImageNet, en le comparant aux Modèles de Dérive standards.

• Efficacité de l'Entraînement : DriftXpress réalise des accélérations substantielles en temps d'entraînement réel et en débit.
  • Sur CIFAR-10, il atteint une accélération de 6,63× en débit (de ~2 463 à ~16 342 images/sec) tout en maintenant des scores FID comparables ou légèrement meilleurs (5,52 contre 5,64).
  • Sur ImageNet, la variante partitionnée atteint une accélération de 2,64× (5 453 contre 2 058 images/sec) avec un FID comparable (9,21 contre 8,83).
• Convergence : DriftXpress atteint des régimes à faible FID plus rapidement en temps réel. Parce que le champ projeté utilise un résumé de l'ensemble complet des données plutôt que des interactions de mini-lots bruyantes, la dynamique d'entraînement est plus fluide, conduisant à une convergence plus rapide.
• Qualité des Échantillons : La méthode maintient une qualité d'échantillon comparable à la dérive standard. Les comparaisons qualitatives montrent que les images générées sont visuellement indiscernables de celles produites par la méthode standard.
• Sélection des Points Repères : Les expériences montrent que la sélection simple et aléatoire de points repères par classe est hautement compétitive et computationnellement efficace, surpassant souvent des stratégies de sélection plus complexes comme k-center.
• Inférence : Crucialement, le coût d'inférence reste identique à celui de la dérive standard (un seul passage avant), car l'approximation n'est appliquée que lors de l'estimation du champ d'entraînement.

Importance et Revendications

L'article revendique que DriftXpress repousse avec succès le compromis entraînement-inférence des modèles de dérive. En déplaçant la charge computationnelle de l'estimation du champ de la boucle d'entraînement vers une étape de prétraitement unique (et des recherches de faible rang efficaces), la méthode réduit le principal goulot d'étranglement des modèles de dérive — l'entraînement coûteux — sans compromettre l'avantage central de l'inférence en une seule étape.

Les auteurs soulignent que cette approche permet aux modèles de dérive de s'adapter à des jeux de données plus vastes (comme ImageNet) et à des tailles de lots plus importantes sans les coûts prohibitifs de mémoire et de calcul associés aux interactions exactes du noyau. Le travail suggère que le « coût d'entraînement » des modèles génératifs en une seule étape peut être considérablement découplé du « coût d'inférence », rendant la génération en une seule étape haute fidélité plus pratique. L'article conclut que, bien que l'approximation de la répulsion reste un défi ouvert, l'approche hybride d'attraction projetée et de répulsion exacte offre une voie robuste et évolutive vers l'avenir.