Transformateurs de diffusion à tous les atomes : Modélisation générative unifiée des molécules et des matériaux

Les modèles de diffusion constituent l'outil standard pour la modélisation générative des systèmes atomiques en 3D. Cependant, pour différents types de systèmes atomiques, tels que les molécules et les matériaux, les processus génératifs sont généralement très spécifiques au système cible, malgré la physique sous-jacente étant la même. Nous présentons le All-atom Diffusion Transformer (ADiT), un cadre de diffusion latente unifié permettant de générer conjointement des matériaux périodiques et des systèmes moléculaires non périodiques à l'aide du même modèle : (1) Un autoencodeur mappe des représentations unifiées et atomiques complètes des molécules et des matériaux vers un espace d'embedding latent partagé ; et (2) Un modèle de diffusion est entraîné pour générer de nouveaux embeddings latents que l'autoencodeur peut décoder afin de produire de nouvelles molécules ou matériaux. Les expériences menées sur les jeux de données MP20, QM9 et GEOM-DRUGS démontrent que l'ADiT entraîné conjointement génère des molécules et des matériaux réalistes et valides, obtenant des résultats à l'état de l'art comparables à ceux des modèles spécifiques aux molécules et aux cristaux. L'ADiT utilise des Transformers standards avec un minimum de biais inductifs pour tant l'autoencodeur que le modèle de diffusion, ce qui entraîne une accélération significative lors de l'entraînement et de l'inférence par rapport aux modèles de diffusion équivariants. L'échelle d'ADiT jusqu'à 500 millions de paramètres améliore prévisiblement les performances, représentant une étape vers la création de modèles fondamentaux largement généralisables pour la chimie générative. Code open source : https://github.com/facebookresearch/all-atom-diffusion-transformer