Chaitanya K. Joshi Xiang Fu Yi-Lun Liao Vahe Gharakhanyan Benjamin Kurt Miller Anuroop Sriram Zachary Ward Ulissi

Abstract

Diffusionsmodelle sind das Standardwerkzeug für die generative Modellierung von 3D atomaren Systemen. Allerdings sind die generativen Prozesse für verschiedene Arten von atomaren Systemen – wie Moleküle und Materialien – in der Regel hochspezifisch für das Zielobjekt, obwohl die zugrunde liegende Physik dieselbe ist. Wir stellen den All-atom Diffusion Transformer (ADiT) vor, einen einheitlichen latente Diffusionsrahmen, der es ermöglicht, sowohl periodische Materialien als auch nicht-periodische molekulare Systeme mit demselben Modell gemeinsam zu generieren: (1) Ein Autoencoder bildet einheitliche, all-atomare Darstellungen von Molekülen und Materialien auf einen gemeinsamen latenten Einbettungsraum ab; und (2) Ein Diffusionsmodell wird trainiert, um neue latente Einbettungen zu erzeugen, die der Autoencoder decodieren kann, um neue Moleküle oder Materialien zu sampeln. Experimente mit den Datensätzen MP20, QM9 und GEOM-DRUGS zeigen, dass das gemeinsam trainierte ADiT realistische und gültige Moleküle sowie Materialien erzeugt und dabei Stand der Technik vergleichbare Ergebnisse erzielt, die denen molekül- und kristall-spezifischer Modelle entsprechen. ADiT verwendet Standard-Transformers mit minimalen induktiven Vorurteilen sowohl für den Autoencoder als auch für das Diffusionsmodell, was zu erheblichen Beschleunigungen während des Trainings und der Inferenz im Vergleich zu äquivarianten Diffusionsmodellen führt. Die Skalierung von ADiT auf bis zu 500 Millionen Parameter verbessert die Leistung vorhersagbar, was einen Schritt in Richtung breit verallgemeinerbarer Grundlagenmodelle für die generative Chemie darstellt. Offener Quellcode: https://github.com/facebookresearch/all-atom-diffusion-transformer