Score-based diffusion models for accurate crystal-structure inpainting and reconstruction of hydrogen positions

Résumé technique : Modèles de diffusion basés sur le score pour l'inpainting et la reconstruction précises des structures cristallines et des positions de l'hydrogène

Énoncé du problème
La détermination précise des positions atomiques des atomes d'hydrogène dans les matériaux cristallins demeure un défi majeur en science des matériaux computationnelle. Bien que l'hydrogène soit abondant et crucial pour des applications allant du stockage d'énergie à l'industrie des semi-conducteurs, sa faible capacité de diffusion des rayons X et des électrons le rend difficile à localiser par diffraction de rayons X standard. Par conséquent, les bases de données cristallographiques historiques (telles que MC3D, dérivées de COD, ICSD et MPDS) contiennent souvent des structures où les positions de l'hydrogène sont soit omises, soit insérées par heuristique, soit déterminées avec une incertitude significative. Bien que les avancées expérimentales (par exemple, la cristallographie quantique) et la diffusion neutronique coûteuse puissent résoudre ces positions, les bases de données existantes dépendent encore largement de données incomplètes ou approximatives. Le problème central abordé est la reconstruction de ces positions d'hydrogène manquantes étant donné une structure cristalline « hôte » connue (le squelette non hydrogéné) afin de générer des structures cristallines précises, stables et complètes.

Méthodologie
Les auteurs proposent une approche générative combinant des modèles de diffusion basés sur le score avec des techniques adaptées de l'inpainting d'images en vision par ordinateur.

1. Architecture du modèle : L'étude utilise MatterGen de Microsoft, un modèle de diffusion basé sur le score pré-entraîné, conçu à l'origine pour générer de nouvelles structures cristallines. Les auteurs réentraînent et adaptent ce modèle pour la tâche spécifique d'inpainting (remplissage de données manquantes) plutôt que pour la génération complète.
2. Stratégies d'inpainting : Plusieurs approches sont comparées :
   • Baseline MatterGen : Le modèle original appliqué à la tâche avec un réseau et des types atomiques fixes.
   • pos-only : Un modèle réentraîné pour débruiter uniquement les positions atomiques, en gardant le réseau et les types fixes.
   • pos-only-RePaint : Le modèle réentraîné combiné avec l'algorithme RePaint, qui rééchantillonne de manière itérative les régions connues pour les harmoniser avec les régions complétées (inpainted).
   • pos-only-TD (Time-Dependent) : Une nouvelle adaptation du concept TD-Paint de la vision par ordinateur. Ce modèle est entraîné pour gérer des niveaux de bruit variables par site au sein d'une même structure. Il conditionne le processus de débruitage sur des positions connues « propres » (bruit nul) tout en prédisant de manière itérative les sites d'hydrogène manquants (bruit élevé). Cela permet au modèle d'utiliser directement la structure hôte connue sans ajouter de bruit à celle-ci.
3. Flux de travail (Workflow) :
   • Partir d'une structure hôte (sites d'hydrogène supprimés).
   • Initialiser aléatoirement $N_H$ sites d'hydrogène dans la maille élémentaire.
   • Appliquer le modèle de diffusion pos-only-TD pendant $N_{steps}$ itérations pour prédire les positions de l'hydrogène.
   • Effectuer une relaxation contrainte des sites d'hydrogène en utilisant un potentiel interatomique basé sur l'apprentissage automatique (MLIP), spécifiquement NequIP, tout en maintenant les atomes hôtes et la maille fixes.
   • Effectuer une relaxation complète de toutes les positions (incluant les atomes hôtes) en utilisant NequIP.
4. Métriques d'évaluation : Le succès est défini de deux manières :
   • Correspondance structurelle : La structure prédite correspond à la référence (issue de la base de données) dans les tolérances du StructureMatcher de pymatgen.
   • Correspondance LES ou structurelle (Lower Energy or Structural - LES) : La prédiction correspond soit à la structure, soit trouve une configuration qui est énergétiquement plus stable que la référence (selon la DFT ou le MLIP).

Résultats clés
La méthode a été évaluée sur deux ensembles de données dérivés de la base de données MC3D :

• Jeu de données DFT : Structures avec des positions d'hydrogène connues qui ont été relaxées par la théorie de la fonctionnelle de la densité (DFT).

• Jeu de données EXP : Les structures expérimentales originales avant curation par la DFT.

• Comparaison des performances : Le modèle pos-only-TD a nettement surpassé la baseline MatterGen, le modèle pos-only et l'approche pos-only-RePaint.

  • RePaint a amélioré la variance mais a augmenté le coût de calcul (nécessitant environ 1000 étapes contre 300) sans augmenter significativement le taux de correspondance.
  • Le modèle pos-only-TD a atteint un taux de correspondance structurelle de 88 % sur un sous-ensemble où une méthode « pinball » basée sur la DFT atteignait environ 87 %, mais avec une approche d'optimisation globale plus efficace.

• Taux de succès :

  • Lors de la génération d'un échantillon unique par structure, la méthode a obtenu des taux de correspondance élevés.
  • En générant plusieurs échantillons (par exemple, $k=10$ ou $k=30$) et en sélectionnant la configuration la plus stable, le taux de correspondance LES a dépassé 97 % pour le jeu de données DFT+MLIP et 97,2 % pour le jeu de données EXP+MLIP.
  • En excluant les structures marquées comme « théoriques » dans les bases de données sources, le taux de succès sur le jeu de données EXP a dépassé 99 %.

• Stabilité et généralisation :

  • De nombreuses prédictions « non correspondantes » se sont révélées être énergétiquement plus stables que les structures de référence originales, suggérant que les références expérimentales dans les bases de données étaient sous-optimales ou incorrectes.
  • Le modèle a démontré une généralisation zero-shot à des structures comprenant jusqu'à 40 atomes par maille élémentaire, bien qu'il n'ait été entraîné que sur des structures allant jusqu'à 20 atomes. Bien que les performances diminuent légèrement pour les systèmes plus grands, le taux de correspondance LES reste supérieur à 90 %.

• Efficacité computationnelle : L'approche par diffusion est nettement plus rapide que la méthode pinball basée sur la DFT, qui souffre d'une explosion combinatoire et de coûts de calcul élevés.

Signification et revendications
L'article affirme que le transfert des techniques d'inpainting d'images (spécifiquement TD-Paint) aux sciences des matériaux permet une complétion des structures hôtes plus rapide et plus précise que les modèles de diffusion non conditionnés ou les approches traditionnelles basées sur la DFT.

• Curation de bases de données : La méthode fournit un outil robuste pour la curation et l'extension des bases de données de matériaux (comme MC3D) en reconstruisant de manière fiable les positions manquantes de l'hydrogène, permettant ainsi le criblage à haut débit de matériaux contenant de l'hydrogène.
• Découverte de configurations stables : L'approche ne se contente pas de reconstruire les structures connues ; elle identifie souvent des configurations plus stables que les structures rapportées, corrigeant potentiellement des erreurs historiques ou des insertions heuristiques.
• Transférabilité indépendante de l'hydrogène : Les modèles sont entraînés de manière « agnostique à l'hydrogène ». Les auteurs affirment que cette architecture peut être facilement transférée à d'autres tâches d'inpainting, telles que l'intercalation d'ions (ex: Li, Na) dans des matériaux de cathode ou la prédiction des sites d'arrêt des muons, sans nécessiter de réentraînement.
• Fondation pour des travaux futurs : L'article positionne cette méthode comme une base qui peut être affinée pour des structures plus grandes ou des applications spécifiques, plutôt que de nécessiter un réentraînement complet.

Les auteurs soulignent que leur définition du succès inclut à la fois la récupération structurelle et la découverte de configurations plus stables, présentant cette dernière non pas comme un échec de reconstruction, mais comme la prédiction réussie d'un état physiquement plus précis.