Résumé technique : Distance de Nouveauté de Transport (TNovD)

Énoncé du problème

L'avancement rapide de l'apprentissage automatique génératif pour la découverte de matériaux a dépassé le développement de métriques d'évaluation rigoureuses. Les approches existantes, telles que les métriques SUN (Stabilité, Unicité, Nouveauté), évaluent les structures générées sur une base instance par instance et ne parviennent pas à capturer les propriétés distributionnelles. Inversement, les métriques distributionnelles comme la distance de Fréchet Inception (FID) et son analogue moléculaire (FCD) reposent sur des hypothèses gaussiennes qui compressent les données en moyenne et variance, les rendant incapables de détecter la « mémorisation » (où un modèle se contente de reproduire les données d'entraînement). De plus, les méthodes actuelles peinent à évaluer simultanément à la fois la qualité (similitude statistique avec la distribution d'entraînement) et la nouveauté (dissimilitude par rapport aux échantillons d'entraînement spécifiques). L'article souligne qu'il est crucial de distinguer une généralisation de haute qualité d'une mémorisation dangereuse, particulièrement alors que les nouveaux modèles, comme les modèles de diffusion, deviennent plus sujets au surapprentissage.

Méthodologie

Les auteurs proposent la Distance de Nouveauté de Transport (TNovD), une métrique distributionnelle fondée sur la théorie du transport optimal (OT) et l'apprentissage contrastif. La méthodologie se compose de trois composantes principales :

1. Extraction de caractéristiques via un GNN équivariant :

   • Les structures cristallines sont intégrées dans un espace de caractéristiques à l'aide d'un réseau de neurones sur graphes (GNN) $SE(3)$-équivariant.
   • Le GNN est entraîné en utilisant une perte contrastive InfoNCE sur le jeu de données MP20.
   • Stratégie d'augmentation : Des paires positives sont créées en appliquant des rotations, des translations et des expansions de supercellules à la même structure cristalline. Les paires négatives consistent en des matériaux différents. Cet entraînement garantit que l'espace de caractéristiques est invariant aux transformations rigides et à la taille de la supercellule tout en distinguant les compositions chimiques distinctes.

2. Couplage par transport optimal :

   • La TNovD calcule le couplage entre les distributions de caractéristiques de l'ensemble d'entraînement ($P_D$) et de l'ensemble généré ($P_G$) en utilisant le cadre de la distance de Wasserstein.
   • Contrairement aux distances de Wasserstein standard, la TNovD introduit un seuil $\tau$ pour séparer la fonction de coût en deux régimes :
     • Régime de mémorisation ($x < \tau$) : Si un échantillon généré est trop proche d'un échantillon d'entraînement (distance $x < \tau$), il est lourdement pénalisé à l'aide d'une fonction de puissance $(\tau - x)^M$. Cela détecte les répliques exactes ou quasi exactes.
     • Régime de qualité ($x \ge \tau$) : Si un échantillon est suffisamment nouveau ($x \ge \tau$), le coût est linéaire $(x - \tau)$, mesurant à quel point l'échantillon s'écarte de la distribution d'entraînement.

3. Heuristiques des hyperparamètres :

   • $\tau$ (Seuil) : Déterminé de manière heuristique en analysant les distances du 1er et du 2ème plus proche voisin au sein de l'ensemble d'entraînement, approximant la condition selon laquelle un point est considéré comme « mémorisé ».
   • $M$ (Exposant) : Dérivé d'une condition d'équilibre où la dérivée de la TNovD par rapport à $\tau$ est nulle pour deux divisions d'une « bonne » distribution, garantissant que la métrique équilibre la pénalité pour la mémorisation et la pénalité pour la faible qualité.

Contributions clés

• Une métrique distributionnelle unifiée : La TNovD est une métrique conçue pour évaluer conjointement la qualité et la nouveauté des modèles de génération de matériaux sans dépendre d'hypothèses gaussiennes. Bien que des travaux antérieurs [8] aient proposé une métrique distributionnelle basée sur un GNN, la TNovD se distingue en entraînant explicitement le GNN pour la détection de la nouveauté et en évitant les hypothèses gaussiennes pour mieux capturer la mémorisation.
• Détection de la mémorisation : En utilisant une fonction de coût non linéaire pour les petites distances, la TNovD pénalise explicitement la reproduction des données d'entraînement, une capacité qui fait défaut à la FID et à la FCD.
• Cadre agnostique au domaine : Bien qu'appliquée aux cristaux, les auteurs notent que le cadre est adaptable à d'autres domaines (ex: images, molécules) en remplaçant l'extracteur de caractéristiques.
• Apprentissage de caractéristiques robuste : L'utilisation de l'apprentissage contrastif avec des augmentations cristallographiques spécifiques (supercellules, rotations) crée un espace de caractéristiques qui respecte les symétries physiques et distingue la véritable nouveauté des variations structurelles du même matériau.

Résultats

Les auteurs ont validé la TNovD à travers plusieurs expériences :

1. Expériences de test (Validation MP20) :

   • Mémorisation : Lorsque les données de validation étaient progressivement remplacées par les données d'entraînement, la TNovD montrait une augmentation quasi linéaire, détectant avec succès la mémorisation.
   • Injection de bruit : L'ajout de bruit gaussien aux coordonnées atomiques ou le cisaillement du réseau entraînait une augmentation des scores TNovD, indiquant la détection de structures de faible qualité.
   • Invariance : La métrique est restée stable lors de la génération de supercellules, démontrant la capacité du GNN à gérer les variations de taille sans les pénaliser comme de la « nouveauté ».
   • Substitutions : Les substitutions aléatoires et intra-groupe d'atomes ont été correctement identifiées comme des matériaux de moindre qualité, les substitutions aléatoires étant plus lourdement pénalisées.

2. Jeu de données de substitution WBM :

   • Appliquée au jeu de données WBM (structures créées par substitution atomique itérative), la TNovD a montré qu'à mesure que les structures divergeaient de l'ensemble d'entraînement MP20, la composante de mémorisation chutait vers zéro, tandis que la composante de qualité augmentait légèrement, reflétant la croissance de la dissemblance.

3. Benchmarking de modèles génératifs :

   • Six modèles de génération de matériaux populaires (incluant MatterGen, CDVAE, DiffCSP et ADiT) ont été évalués sur 10 000 échantillons générés.
   • ADiT a obtenu le meilleur score global, porté par une haute qualité, bien que les auteurs émettent l'hypothèse qu'il puisse encore mémoriser des données d'une manière non totalement capturée par la métrique.
   • MatterGen a performé de manière similaire à la ligne de base du jeu de validation.
   • CDVAE a mal performé, montrant des pénalités de qualité élevées (faible qualité) et aucune mémorisation, indiquant un échec à correspondre à la distribution d'entraînement.
   • Chemeleon a présenté la composante de mémorisation la plus élevée.

Signification et Revendications

L'article affirme que la TNovD offre une « perspective complémentaire » nécessaire aux métriques existantes basées sur l'instance. Sa principale importance réside dans sa capacité à détecter la mémorisation tout en évaluant simultanément la qualité distributionnelle, comblant ainsi une lacune critique dans l'évaluation des modèles génératifs modernes.

Les auteurs notent modestement que :

• La TNovD n'est pas un objectif d'entraînement, car un générateur pourrait théoriquement « manipuler » la métrique ; c'est strictement un outil d'évaluation.
• La métrique ne capture pas directement la stabilité des matériaux (ex: énergie au-dessus de l'enveloppe) ou la synthétisabilité, bien qu'elle puisse être étendue avec des modèles d'énergie pré-entraînés dans des travaux futurs.
• Les résultats dépendent fortement du choix de la représentation des caractéristiques (GNN basé sur MP20) et des augmentations.
• Les faibles scores de mémorisation observés dans de nombreux modèles suggèrent que le principal défi pour les modèles génératifs de matériaux actuels est l'ajustement à la distribution de données correcte (qualité) plutôt que l'évitement de la mémorisation.

Ce travail établit un blueprint modulaire pour une évaluation holistique, soulignant qu'à mesure que les modèles génératifs progressent, la capacité à distinguer la découverte réelle de la simple réplication de l'ensemble d'entraînement devient primordiale.