OmniMol: Transferring Particle Physics Knowledge to Molecular Dynamics with Point-Edge Transformers

Dieser Beitrag stellt OmniMol vor, einen hochmodernen, maschinell erlernten interatomaren Potenzial für kleine Moleküle, der eine Point-Edge-Transformer-Architektur nutzt und Wissenstransfer aus der Hochenergiephysik einsetzt, um mit minimalem Fine-Tuning und einzigartig schneller Inferenz hervorragende Leistung zu erzielen.

Das Wesentliche

Stellen Sie sich vor, Sie haben zwei völlig unterschiedliche Welten: die eine ist die chaotische, hochgeschwindigkeitsgetriebene Welt der Teilchenphysik (wo Wissenschaftler Atome zusammenstoßen lassen, um zu sehen, was herausfliegt), und die andere ist die komplexe, klebrige Welt der Molekülchemie (wo Atome zusammenkleben, um Medikamente, Materialien und Leben zu bilden).

Lange Zeit nutzten Wissenschaftler in diesen beiden Feldern völlig unterschiedliche Werkzeuge, um ihre Welten zu verstehen. Aber in diesem Papier stellen die Autoren OmniMol vor, ein neues Werkzeug, das versucht, den Teilchenphysik-Experten zu lehren, Chemie zu verstehen, indem es ein „Fundamentmodell" verwendet, das sie bereits entwickelt haben.

Hier ist die einfache Aufschlüsselung, wie sie es gemacht haben und was sie herausfanden:

1. Die „Meisterkoch"-Analogie

Stellen Sie sich das ursprüngliche Modell, genannt Omnilearned, als einen Meisterkoch vor, der Jahre damit verbracht hat, mit Teilchenjets zu kochen.

• Die Zutaten: In der Teilchenphysik ist ein „Jet" ein Sprühregen aus subatomaren Teilchen (wie Protonen und Neutronen), der aus einer Kollision herausfliegt.
• Die Fähigkeit: Dieser Koch lernte, Muster in diesen Sprühregen zu erkennen. Er weiß, wie Teilchen wechselwirken, wie sie sich gruppieren und wie man vorhersagen kann, was als Nächstes passiert. Er wurde auf eine Milliarde verschiedene Teilchensprühregen trainiert.

Nun fragten die Autoren: Kann derselbe Koch ein molekulares Mahl zubereiten?

• Die neuen Zutaten: Anstelle von subatomaren Teilchen sind die „Zutaten" Atome (wie Kohlenstoff, Sauerstoff, Wasserstoff) in einem Molekül.
• Die Herausforderung: Atome verhalten sich anders als subatomare Teilchen, aber sie teilen eine ähnliche Struktur: Sie sind einfach Punkte im Raum mit spezifischen Typen.

2. Der „Universalübersetzer" (Die Architektur)

Um dies zum Funktionieren zu bringen, bauten sie keinen neuen Koch von Grund auf neu. Sie nahmen den bestehenden „Meisterkoch" (Omnilearned) und gaben ihm ein neues Set an Werkzeugen:

• Der Punkt-Kanten-Transformer (PET): Stellen Sie sich vor, der Koch betrachtet einen Teller mit Essen. Anstatt nur auf eine Zutat nach der anderen zu schauen, ermöglicht ihm dieses Werkzeug, alle Zutaten gleichzeitig zu betrachten und zu sehen, wie jede einzelne mit jeder anderen zusammenhängt.
• Die „Physik-Bias": Dies ist die geheime Zutat. Das Modell hat ein eingebautes „Regelbuch", das ihm sagt: „Hey, diese beiden Teilchen/Atome sind nah beieinander, also sollten sie sich gegenseitig mehr Aufmerksamkeit schenken." Dies hilft dem Modell, sich auf die wichtigsten Beziehungen zu konzentrieren, ohne durch das Rauschen verwirrt zu werden.

3. Das Experiment: Feinabstimmung

Die Autoren nahmen dieses für Teilchen trainierte Modell und gaben ihm einen „Crashkurs" in Chemie mit einem Datensatz namens oMol (eine Sammlung von Millionen Molekülen).

• Das Ziel: Sie wollten, dass das Modell als maschinell gelerntes interatomares Potenzial (MLIP) fungiert. Auf Deutsch bedeutet dies, dass das Modell zwei Dinge für jede Gruppe von Atomen vorhersagen muss:
  1. Energie: Wie viel „Klebstoff" hält sie zusammen?
  2. Kraft: Wenn Sie ein Atom drücken, wie stark wird es zurückdrücken?

4. Die Ergebnisse: Schnell und überraschend gut

Das Papier fand einige aufregende Dinge heraus:

• Die „Few-Shot"-Superkraft: Normalerweise erfordert das Unterrichten eines Computers in Chemie massive Datenmengen. Aber da OmniMol mit dem „Wissen" der Teilchenphysik begann, lernte es Chemie sehr schnell. Selbst mit einer relativ kleinen Menge neuer Daten (wie 100.000 Molekülen) schnitt es fast so gut ab wie Modelle, die auf Millionen trainiert wurden. Es ist wie ein Meisterkoch, der eine neue Küche mit nur wenigen Rezepten erlernen kann, weil er bereits die Grundlagen von Geschmack und Hitze versteht.
• Geschwindigkeit: OmniMol ist unglaublich schnell. Während andere Modelle lange brauchen, um zu berechnen, wie sich ein Molekül bewegt, erledigt OmniMol dies im Handumdrehen. Die Autoren stellen fest, dass OmniMol pro Stunde Rechenzeit dreimal mehr Moleküle simulieren kann als einige seiner Konkurrenten.
• Der Kompromiss: Wenn sie riesige Datenmengen hatten (Millionen von Molekülen), schwächte sich der Vorteil des Starts mit Teilchenphysik-Wissen etwas ab. Dies deutet darauf hin, dass das „Teilchenphysik-Wissen" wie ein starker Vorsprung wirkt, aber wenn Sie genug Zeit und Daten haben, um ein Modell von Grund auf neu zu trainieren, ist dieser Vorsprung weniger wichtig.

5. Das große Ganze

Das Papier kommt zu dem Schluss, dass OmniMol das erste Mal ist, dass ein „Fundamentmodell", das für eine wissenschaftliche Disziplin (Teilchenphysik) entwickelt wurde, erfolgreich auf eine völlig andere (Chemie) übertragen wurde.

Sie bewiesen, dass wenn Sie ein intelligentes Modell haben, das versteht, wie Punkte im Raum in einem Feld interagieren, es angepasst werden kann, um zu verstehen, wie Punkte im Raum in einem anderen Feld interagieren, was Zeit und Rechenleistung spart.

Zusammenfassend: Die Autoren nahmen eine supersmart KI, die auf hochenergetischen Teilchenkollisionen trainiert war, passten ihr Gehirn an, um Atome statt Teilchen zu verstehen, und stellten fest, dass es zu einem blitzschnellen, hochpräzisen Werkzeug wurde, um vorherzusagen, wie Moleküle sich verhalten, insbesondere wenn Daten knapp sind.

Technische Zusammenfassung

Technischer Zusammenfassung: OmniMol

Problemstellung

Maschinelles Lernen (ML) hat die Darstellung und Simulation komplexer physikalischer Systeme transformiert, insbesondere in der Teilchenphysik und der Molekülchemie. Obwohl diese Domänen sich erheblich in ihren Energieskalen unterscheiden, teilen sie eine fundamentale Datenstruktur: variably große Mengen von Teilchen (oder Atomen) im Phasenraum, die effektiv strukturierte Punktwolken bilden.

Die primäre Herausforderung, die adressiert wird, ist die Entwicklung effizienter maschinell gelernter interatomarer Potentiale (MLIPs). Traditionelle Methoden wie die Dichtefunktionaltheorie (DFT) sind rechenintensiv, was großskalige und langfristige Molekulardynamik-Simulationen (MD) begrenzt. MLIPs zielen darauf ab, Potentialenergieflächen und Kräfte zu einem Bruchteil dieser Kosten zu approximieren. Das Training robuster MLIPs erfordert jedoch typischerweise massive Datensätze und erhebliche Rechenressourcen. Die Arbeit postuliert, dass ein auf Punktwolken in der Teilchenphysik (speziell Teilchenjets) vortrainiertes Fundamentmodell auf die Molekulardynamik übertragen werden könnte, was die Optimierung beschleunigen und die Genauigkeit in Daten-armen Regimen verbessern könnte.

Methodik

Architektur: Point-Edge Transformer (PET)

OmniMol wurde durch die Anpassung von Omnilearned entwickelt, einem Fundamentmodell, das ursprünglich für die Klassifizierung und Generierung von Teilchenjets in der Hochenergiephysik (HEP) konzipiert wurde. Die Kernarchitektur ist ein Point-Edge Transformer (PET), der lokale Aufmerksamkeit über $k$-nächste Nachbarn mit globalen All-zu-All-Transformer-Blöcken koppelt.

Wichtige architektonische Komponenten umfassen:

1. Input-Embeddings: Atome werden in einen Token-Raum eingebettet, der Positionsinformationen ($\vec{r}$), diskrete Atomnummern ($Z$) und zusätzliche Merkmale (Ladung, Spin) kombiniert.
2. Lokaler Attention-Block: Für jedes Atom wird eine lokale Nachbarschaft unter Verwendung der $K$-nächsten Nachbarn konstruiert ($K=15$ für Moleküle, verglichen mit $K=10$ für Jets). Paarweise physikalische Merkmale werden berechnet, einschließlich Distanztermen, inversen Potenzen der Distanz und gelernten Funktionen von Atom-Embeddings. Diese werden von einem kleinen lokalen Transformer verarbeitet, um einen lokalen Embedding-Vektor zu erzeugen.
3. Globale Aufmerksamkeit mit Interaktions-Bias: Der globale Self-Attention-Mechanismus integriert einen expliziten Bias, der von paarweisen physikalischen Merkmalen abgeleitet ist. Die Attention-Logits werden modifiziert als $A^*_{ij} = A_{ij} + B_{ij}$, wobei $B_{ij}$ ein von einem MLP eingebetteter Bias-Term ist. Dieser „Interaktionsmatrix-Aufmerksamkeits-Bias" injiziert paarweise physikalische Prioritäten direkt in den Transformer, lenkt das Netzwerk hin zu physikalisch sinnvollen Nachbarschaften, ohne die Ausdruckskraft zu beeinträchtigen.
4. Output-Köpfe: Der generative Kopf von Omnilearned wird für zwei Aufgaben umfunktioniert:
   • Kraftvorhersage: Ein permutationsäquivarianter Kopf, der kräfte pro Atom vorhersagt.
   • Energievorhersage: Ein Kopf, der Energiekorrekturen pro Atom vorhersagt, die summiert werden, um die gesamte molekulare Energie zu erhalten, wobei extensive Prioritäten erhalten bleiben.

Invarianz- und Erhaltungszwänge

Um physikalische Zwänge zu erfüllen, adressieren die Autoren zwei Anforderungen:

• Energieerhaltung: Kräfte werden nicht direkt vorhergesagt, sondern durch Rückwärtspropagierung der Energieausgabe berechnet ($\vec{F}_i = \nabla_{\vec{r}_i} E$). Dies gewährleistet eine exakte Energieerhaltung, erhöht jedoch die Rechenkosten während des Trainings (erforderlich ist eine doppelte Rückwärtspropagierung). Folglich wird diese Zwang nur auf die „kleine" Modellvariante angewendet.
• Rotationsäquivarianz: Die Standardarchitektur ist nicht inhärent äquivariant, da rohe Koordinatendifferenzen in MLPs eingespeist werden. Um dies zu beheben, führen die Autoren eine „äquivariante und konservative" Variante ein. Diese Version entfernt direkte Koordinatendifferenzterme aus den paarweisen Merkmalen und integriert stattdessen Winkelinformationen (Cosinus von Winkeln, die durch Vektoren zwischen benachbarten Atomen gebildet werden) in den lokalen Block. Diese Modifikation behält die Äquivarianz bei, während sie die mit dem Entfernen von Koordinatentermen verbundenen Leistungsverluste signifikant wiederherstellt.

Trainings- und Fine-Tuning-Strategien

Das Modell wird auf dem oMol-Datensatz feinabgestimmt (speziell die Teilmengen oMol-25, oMol-4M, oMol-100M und oMol-140M). Zwei Strategien werden untersucht:

1. LoRA (Low-Rank Adaptation): Die vortrainierten Gewichte des PET-Rückgrats werden eingefroren. Low-Rank-Adapter werden nur für die Transformer-Körper-Matrizen ($W_Q, W_K, W_V, W_O, W_{MLP}$) eingeführt, zusammen mit dem Training der molekularen Input-Encoder, des Bias-MLP und der Aufgaben-Köpfe. Eine „Embedding-Anpassung"-Schicht wird ebenfalls hinzugefügt, um gelernte Embeddings zu modifizieren.
2. Vollständiges Fine-Tuning: Alle Gewichte im Körper und in den Input-Encodern werden eingefroren und trainiert, während die Aufgaben-Köpfe von Grund auf neu trainiert werden.

Das Trainingsziel minimiert die Summe der mittleren absoluten Fehler (MAE) für Energien und Kräfte, wobei Kräfte stärker gewichtet werden ($\lambda_F = 10$).

Wichtige Ergebnisse

Leistung auf oMol

• Vollständiges Training: Wenn auf großen Datensätzen (oMol-4M und oMol-100M/140M) trainiert, erreicht OmniMol eine wettbewerbsfähige Leistung mit dem Stand der Technik bei MLIPs. Beispielsweise erzielt das OmniMol-large-Modell auf oMol-140M einen Energie-MAE von 1,04 meV/Atom und einen Kraft-MAE von 13,59 meV/Å.
• Daten-armes Regime: Die signifikantesten Gewinne werden beobachtet, wenn die Trainingsdaten begrenzt sind. Wenn nur auf 100.000 Molekülen feinabgestimmt oder mit sehr wenigen Epochen (2 Durchläufe) über oMol-4M, übertreffen die vortrainierten OmniMol-Varianten Modelle, die von Grund auf neu trainiert wurden, erheblich.
  • Auf einer 100k-Teilmenge verbesserte das Vortraining den Energie-MAE um bis zu 29,4 % und den Kraft-MAE um 26,9 % für das mittlere Modell.
  • Mit nur zwei Epochen auf oMol-4M zeigte das mittlere Modell eine **54,6 %**ige Verbesserung im Energie-MAE und 56,9 % im Kraft-MAE im Vergleich zu seiner nicht vortrainierten Entsprechung.
• Äquivariante/Konservative Variante: Die äquivariante und konservative Modellvariante zeigt in daten-armen Regimen eine signifikant verbesserte Leistung (insbesondere für Kräfte), obwohl dieser Vorteil mit zunehmender Datensatzgröße abnimmt.

Skalierung und Inferenzgeschwindigkeit

• Skalierung: OmniMol folgt einer sauberen Potenzgesetz-Skalierung mit der Modellgröße und zeigt bis zu 1 Milliarde Parameter keine Anzeichen einer Sättigung, was mit jüngsten Erkenntnissen zu transformer-basierten MLIPs übereinstimmt.
• Inferenzgeschwindigkeit: Trotz großer Parameterzahlen demonstriert OmniMol einzigartige schnelle Inferenzgeschwindigkeiten aufgrund von Hardware-Optimierungen für Transformer. Auf einer A100-GPU für Systeme mit ~100 Atomen ist OmniMol-medium etwa 3x schneller als vergleichbare Graph Neural Network (GNN)-Baselines (eSEN-md-d und AllScAIP-md), während eine wettbewerbsfähige Genauigkeit beibehalten wird (nur ~0,7 meV/Atom höherer Energiefehler als AllScAIP-md).

Bedeutung und Behauptungen

Die Arbeit behauptet, die erste Demonstration eines disziplinübergreifenden Transfers für wissenschaftliche Punktwolken-Fundamentmodelle vorzustellen. Durch die Anpassung eines auf Hochenergiephysik-Teilchenjets vortrainierten Modells an die Molekulardynamik demonstrieren die Autoren, dass:

1. Disziplinübergreifender Transfer machbar ist: Ein Fundamentmodell, das für die Teilchenphysik entwickelt wurde, kann effektiv auf die Molekülchemie übertragen werden, was darauf hindeutet, dass die zugrunde liegenden Punktwolken-Strukturen über sehr unterschiedliche physikalische Skalen hinweg lernbare Merkmale teilen.
2. Induktive Verzerrung beschleunigt das Lernen: Das Vortraining wirkt als starke induktive Verzerrung. Ähnlich wie Äquivarianz hilft, wenn Daten knapp sind, ermöglicht die „bittere Lektion" des Vortrainings eine schnelle Optimierung und verbesserte Genauigkeit, wenn die Trainingsdaten begrenzt sind.
3. Effizienz: Der architektonische Transfer ermöglicht einzigartige schnelle Inferenzgeschwindigkeiten, was für Anwendungen kritisch ist, die eine schnelle Exploration von Designräumen erfordern, wie z. B. die Entdeckung kleiner Moleküle als Medikamente.

Die Autoren schließen, dass, obwohl die Studie sich auf MLIPs konzentriert, die Lehren bezüglich Punktwolken-Fundamentmodelle eine weitreichende Nützlichkeit in wissenschaftlichen Domänen haben könnten, in denen Systeme als ungeordnete Mengen interagierender Körper beschrieben werden. Sie beanspruchen keine universelle Überlegenheit gegenüber allen bestehenden Methoden in allen Regimen, heben jedoch die spezifischen Vorteile in daten-armen Szenarien und bei der Inferenzgeschwindigkeit hervor.