Eingeladen für ICML 2025: Meta, die Universität von Cambridge und MIT stellen ein ganzatomares Diffusions-Transformer-Framework vor und ermöglichen erstmals die einheitliche Generierung von periodischen und aperiodischen Atomsystemen.

Das gemeinsame Forschungsteam von Meta FAIR, der Universität von Cambridge und dem Massachusetts Institute of Technology (MIT) hat den All-atom Diffusion Transformer (ADiT) vorgeschlagen, der die Modellierungsbarriere zwischen periodischen und nicht-periodischen Systemen gebrochen hat. Durch zwei Innovationen, nämlich die einheitliche latente Darstellung aller Atome und die latente Diffusion mit Transformer, wurde der Durchbruch erzielt, Moleküle und Kristalle mit einem einzigen Modell zu generieren.

Im vordersten Bereich der heutigen wissenschaftlichen Forschung und industriellen Anwendungen zeigt die generative Modellierung der dreidimensionalen Strukturen von Atomsystemen ein umwälzendes Potenzial und hat das Potenzial, die rückwärtsgerichtete Gestaltung von neuen Molekülen und Materialien grundlegend neu zu gestalten. Von der genauen Strukturvorhersage bis zur flexiblen bedingten Generierung haben die derzeit fortschrittlichsten Diffusionsmodelle und Flow-Matching-Modelle sich in Schlüsselaufgaben wie der Analyse von Biomolekülen, der Entwicklung neuer Materialien und der strukturbasierten Arzneimitteldesign hervorgetan und werden zu den Kernwerkzeugen für Forscher, um technische Engpässe zu überwinden.

Dennoch steht hinter diesem aufstrebenden Bereich ein Schlüsselproblem, das die technologische Entwicklung immer noch hemmt - die bestehenden Modelle fehlen an der Universalität zwischen verschiedenen Systemen. Obwohl alle Atomsysteme denselben physikalischen Prinzipien folgen, um ihre dreidimensionalen Strukturen und Wechselwirkungen zu bestimmen, werden die Modellierungen von kleinen Molekülen, Biomolekülen, Kristallen und ihren Komplexsystemen seit langem getrennt behandelt. Die meisten Diffusionsmodelle sind stark von den inhärenten Eigenschaften bestimmter Systeme abhängig und müssen in komplexen multiplikativen Mannigfaltigkeiten, in denen kategorische Daten (z. B. Atomtypen) und kontinuierliche Daten (z. B. dreidimensionale Koordinaten) miteinander verwoben sind, multimodale Generationen durchführen. Dies führt dazu, dass die Modelle zwischen verschiedenen Systemen schwer kompatibel sind.

Nehmen wir ein konkretes Szenario als Beispiel: Die De-novo-Generierung von kleinen Molekülen muss in zwei unabhängige Diffusionsprozesse aufgeteilt werden, nämlich den Atomtyp (kategorisch) und die dreidimensionale Koordinate (kontinuierlich). Obwohl das Entrauschungsmodell die kooperative Evolutionsregeln beider lernen muss, verringert die Verzerrung des Zwischenzustands häufig die Sampling-Effizienz. Bei der Modellierung von Biomolekülen muss eine zusätzliche Rotationsmannigfaltigkeit eingeführt werden, und die Atomgruppen werden als starre Körper behandelt. Der Diffusionsprozess von Kristallen und Materialien muss dagegen die periodischen Eigenschaften berücksichtigen und in einer gemeinsamen Mannigfaltigkeit, die aus mehrdimensionalen Parametern wie Atomtypen, Bruchkoordinaten und Gitterparametern besteht, ablaufen - diese Unterschiede machen die einheitliche Modellierung zwischen verschiedenen Systemen zu einer seit langem ungelösten Herausforderung in diesem Bereich.

Vor diesem Hintergrund hat das gemeinsame Forschungsteam von Meta Fundamental AI Research (FAIR), der Universität von Cambridge und dem MIT eine bahnbrechende Lösung vorgeschlagen - den All-atom Diffusion Transformer (ADiT).

Als ein auf Transformer basierendes einheitliches latentes Diffusionsframework liegt der Kernvorteil von ADiT darin, die Modellierungsbarriere zwischen periodischen und nicht-periodischen Systemen zu brechen. Durch die beiden Innovationen der einheitlichen latenten Darstellung aller Atome und der latenten Diffusion mit Transformer kann ADiT Moleküle und Kristalle mit einem einzigen Modell generieren. Seine Konstruktion führt fast keine induktiven Verzerrungen ein, sodass der Autoencoder und das Diffusionsmodell in der Trainings- und Inferenz-Effizienz weit über die traditionellen äquivarianten Diffusionsmodelle hinausgehen. Unter denselben Hardwarebedingungen wird die Zeit zur Generierung von 10.000 Stichproben von 2,5 Stunden auf weniger als 20 Minuten verkürzt. Noch bemerkenswerter ist, dass wenn die Modellparameter auf 500 Millionen erweitert werden, seine Leistung eine vorhersagbare lineare Verbesserung zeigt. Diese Eigenschaft legt die Schlüsselgrundlage für die Konstruktion eines universellen generativen chemischen Basismodells und markiert einen meilensteinartigen Schritt in Richtung Universalität und skalierbarer Anwendungen bei der Modellierung von Atomsystemen.

Die relevanten Forschungsergebnisse wurden unter dem Titel "All-atom Diffusion Transformers: Unified generative modelling of molecules and materials" in die ICML 2025 aufgenommen.

Highlights der Forschung:

* ADiT hat erstmals die Einheitlichkeit eines Generierungsmodells für periodische Materialien und nicht-periodische Molekülsysteme erreicht.

* ADiT basiert auf der einheitlichen latenten Darstellung aller Atome und der latenten Diffusion mit Transformer, was den Generierungsprozess effektiv vereinfacht und fast keine induktiven Verzerrungen aufweist.

* ADiT zeichnet sich durch hervorragende Skalierbarkeit und Effizienz aus, und seine Trainings- und Inferenzgeschwindigkeit übertrifft weit die äquivarianten Diffusionsmodelle.

Link zur Publikation: 

https://go.hyper.ai/27d7U

Datenmengen: Von periodischen bis nicht-periodischen, abdeckend experimentelle Daten aus verschiedenen Bereichen

In dieser Studie hat das Forschungsteam zunächst mehrere repräsentative Datensätze ausgewählt, um Experimente durchzuführen:

* MP20-Datensatz, der 45.231 metastabile Kristallstrukturen aus dem Materials Project enthält, mit maximal 20 Atomen in der Einheitszelle und 89 verschiedenen Elementen, der gut die periodischen Materialsysteme repräsentiert;

* QM9-Datensatz, der aus 130.000 stabilen organischen kleinen Molekülen besteht, mit maximal 9 schweren Atomen (C, N, O, F) und Wasserstoffatomen, der ein typischer Vertreter von nicht-periodischen Molekülsystemen ist;

* GEOM-DRUGS-Datensatz, der 430.000 große organische Moleküle mit maximal 180 Atomen enthält;

* QMOF-Datensatz, der 14.000 metallorganische Gerüststrukturen enthält.

Unter diesen entsprechen MP20 und QM9 jeweils verschiedenen Typen von Atomsystemen und bieten die Grundlage für das gemeinsame Training des Modells auf periodischen und nicht-periodischen Systemen. Das Forschungsteam hat die Daten wie in früheren Studien aufgeteilt, um die Fairness des Vergleichs mit anderen Modellen sicherzustellen. GEOM-DRUGS und QMOF erweitern weiter den Bereich des Modelltests und können die Generalisierungsfähigkeit des Modells umfassender überprüfen.

ADiT: Aufbau eines einheitlichen Generierungsmodells für Atomsysteme mit zwei Kernideen

Als latentes Diffusionsmodell basiert das Kernkonzept von ADiT auf zwei Schlüsselideen, um die einheitliche generative Modellierung von periodischen und nicht-periodischen Atomsystemen zu erreichen.

Die erste Schlüsselidee ist die einheitliche latente Darstellung aller Atome. Das Forschungsteam betrachtet sowohl periodische als auch nicht-periodische Atomsysteme als Atommengen im dreidimensionalen Raum und hat daher eine einheitliche Darstellungsmethode entwickelt, die sowohl die kategorischen Eigenschaften (z. B. Atomtypen) als auch die kontinuierlichen Eigenschaften (z. B. dreidimensionale Koordinaten) jedes Atoms enthält. Durch das Training eines Variational Autoencoders (VAE) für die Rekonstruktion aller Atome kann dieser Encoder Moleküle und Kristalle in einen gemeinsamen latenten Raum einbetten, was die Grundstruktur für die einheitliche Verarbeitung verschiedener Typen von Atomsystemen schafft.

Die zweite Schlüsselidee ist die latente Diffusion mit Transformer. Im latenten Raum, der vom VAE-Encoder aufgebaut wurde, hat das Forschungsteam einen Diffusion Transformer (DiT) eingeführt, um die generative Modellierung durchzuführen. Bei der Inferenz kann mit Hilfe der classifier-free guidance-Technik neue latente Variablen gesampelt werden, und diese latenten Variablen können dann durch den VAE-Decoder zu gültigen Molekülen oder Kristallen rekonstruiert werden, um so die Umwandlung vom latenten Raum in das reale Atomsystem abzuschließen.

Basierend auf diesen beiden Kernideen wird die experimentelle Methode von ADiT in zwei Phasen nacheinander durchgeführt.

In der ersten Phase bauen die Forscher einen Autoencoder für die Rekonstruktion auf. Durch die gemeinsame Rekonstruktion der All-Atom-Darstellung von Molekülen und Materialien mit VAE lernen und bauen sie einen gemeinsamen latenten Raum auf - dies ist die Voraussetzung für die einheitliche Modellierung verschiedener Atomsysteme und legt die Grundlage für den späteren Generierungsprozess.

In der zweiten Phase bauen die Forscher ein latentes Diffusionsgenerierungsmodell auf. Mit DiT werden neue Stichproben aus dem latenten Raum generiert, und diese Stichproben werden durch die classifier-free guidance decodiert zu gültigen Molekülen oder Kristallen. Der bedeutende Vorteil dieser latenten Diffusionskonstruktion besteht darin, dass die Komplexität der Verarbeitung von kategorischen und kontinuierlichen Eigenschaften auf den Autoencoder übertragen wird, sodass der Generierungsprozess im latenten Raum einfacher und skalierbarer wird, was die Effizienz und Anpassungsfähigkeit des Modells bei der Verarbeitung verschiedener Atomsysteme effektiv verbessert.

ADiT führt die generative Modellierung von chemischen Systemen in zwei Phasen durch.

ADiT führt in der Generierung von Kristallen und Molekülen an

Um die Leistungsvorteile von ADiT voll zu verdeutlichen, hat das Forschungsteam mehrere Baseline-Modelle ausgewählt, um gezielte Vergleiche durchzuführen. Im Bereich der Kristallgenerierung umfassen die Vergleichsobjekte äquivariante Diffusions- und Flow-Matching-Modelle wie CDVAE, DiffCSP, FlowMM, die auf multimodalen multiplikativen Mannigfaltigkeiten basieren, sowie nicht-äquivariante Diffusionsmodelle wie UniMat und das Zwei-Phasen-Framework FlowLLM. Im Bereich der Molekülgenerierung wird es mit äquivarianten Diffusionsmodellen wie GeoLDM und Symphony verglichen. Durch den systematischen Vergleich mit diesen fortschrittlichen Baseline-Modellen in diesem Bereich werden die Leistungsvorteile von ADiT deutlich.

Aus den konkreten experimentellen Ergebnissen geht hervor, dass ADiT in den Aufgaben der Kristall- und Molekülgenerierung auf dem Stand der Technik (SOTA) liegt. Bei der Kristallgenerierung zeigt ADiT hervorragende Ergebnisse in den Schlüsselindikatoren wie Effektivität, Stabilität, Einzigartigkeit und Neuheit. Bei der Molekülgenerierung liegt ADiT an der Spitze in den Indikatoren der Effektivität und Einzigartigkeit von 10.000 gesampelten Molekülen.

Der gemeinsame Trainingsmechanismus von ADiT bringt auch signifikante Leistungsgewinne. Die experimentellen Daten zeigen, dass ADiT, das gleichzeitig auf den QM9- und MP20-Datensätzen trainiert wird, in den Aufgaben der Material- und Molekülgenerierung alle Versionen übertrifft, die nur auf einem einzigen Datensatz trainiert wurden.

Die Erweiterung der Modellgröße hat eine vorhersagbare Wirkung auf die Leistungssteigerung von ADiT. Wie in der folgenden Abbildung gezeigt, nimmt mit der Erhöhung der Anzahl der Parameter des DiT-Entrauschers von 32 Millionen (ADiT-S, blau) auf 130 Millionen (ADiT-B, orange) und schließlich auf 450 Millionen (ADiT-L, grün) selbst auf einem Datensatz mittlerer Größe von etwa 130.000 Stichproben der Diffusionstrainingverlust kontinuierlich ab und der Effektivitätsanteil steigt stetig an, was einen deutlichen Größeneffekt zeigt. Diese starke Korrelation zwischen Modellgröße und Leistung zeigt, dass durch die Erweiterung der Modellparameter und der Datenmenge ADiT möglicherweise weitere Durchbrüche erzielen kann.

Der Einfluss der Erhöhung der Anzahl der Entrauschungsparameter von ADiT auf den Trainingsverlust und die Generierungseffektivität

In Bezug auf die Effizienz zeigt ADiT im Vergleich zu äquivarianten Diffusionsmodellen einen deutlichen Geschwindigkeitsvorteil. Wie in der folgenden Abbildung gezeigt, hat ADiT, das auf einem Standard-Transformer basiert, bei der Generierung von 10.000 Stichproben auf einer NVIDIA V100 GPU eine weitaus bessere Skalierbarkeit in Bezug auf die Integrationsschritte als FlowMM und GeoLDM, die rechenintensive äquivariante Netzwerke verwenden. Selbst wenn die Parametergröße von ADiT-B 100 Mal größer ist als die der äquivarianten Baseline, ist seine Inferenzgeschwindigkeit dennoch schneller, was den Vorteil der Transformer-Architektur in Bezug auf die erweiterte Anwendbarkeit unterstreicht.

Der Zeitverlauf der Generierung von 10.000 Stichproben durch ADiT und äquivariante Diffusionsmodelle

Darüber hinaus wurde die Skalierbarkeit von ADiT auf größeren Systemen bestätigt. Auf dem GEOM-DRUGS-Moleküldatensatz, der 430.000 Moleküle mit maximal 180 Atomen enthält, zeigt ADiT im Vergleich zu den fortschrittlichsten äquivarianten Diffusions- und Flow-Matching-Modellen ähnliche Ergebnisse in den Indikatoren der Effektivität und PoseBusters. Bemerkenswerterweise basiert ADiT auf einer Standard-Transformer-Architektur, führt fast keine molekularen induktiven Verzerrungen ein und muss keine Atombindungen explizit vorhersagen, kann aber eine ähnliche Leistung wie die äquivari