Zatom-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials

Zatom-1 ist das erste vollständig quelloffene Multimodal-Flow-Foundation-Modell, das durch eine gemeinsame generative Vortrainierung diskreter Atomtypen und kontinuierlicher 3D-Geometrien die generative und prädiktive Modellierung von 3D-Molekülen und Materialien vereint, dabei spezialisierte Baseline-Modelle in Genauigkeit und Geschwindigkeit übertrifft und einen positiven Transfer zwischen chemischen Domänen ermöglicht.

Das Wesentliche

Stell dir vor, du bist ein Architekt, der nicht nur Häuser, sondern ganze Universen aus Atomen entwerfen möchte. Bisher gab es für dieses Job zwei völlig getrennte Werkzeuge: Eines, das nur kleine, komplexe Moleküle (wie Medikamente) bauen konnte, und ein anderes, das nur riesige Kristallstrukturen (wie neue Materialien für Batterien) entwarf. Und das Schlimmste: Diese Werkzeuge konnten nicht miteinander reden. Sie lernten nicht voneinander.

ZATOM-1 ist wie ein genialer, neuer „All-in-One"-Architekt, der endlich beides kann.

Hier ist die Erklärung in einfachen Worten:

1. Das Problem: Zwei getrennte Welten

Bislang mussten Wissenschaftler zwei verschiedene KI-Modelle nutzen.

• Modell A lernte, wie man kleine Moleküle (wie Aspirin) aus dem Nichts erschafft.
• Modell B lernte, wie man riesige Kristalle (wie Solarzellen-Materialien) baut.

Das war ineffizient. Es war, als würde ein Koch, der nur Kuchen backen kann, versuchen, ein Steak zu grillen, und ein Grillmeister, der nur Steaks macht, versucht, einen Kuchen zu backen. Niemand lernte die allgemeinen Regeln des „Essens" (der Chemie) von beiden Seiten.

2. Die Lösung: ZATOM-1, der universelle Baumeister

ZATOM-1 ist das erste KI-Modell, das alles in einem einzigen Gehirn vereint. Es ist ein „Fundament-Modell" (Foundation Model). Das bedeutet, es lernt die grundlegenden Gesetze der Chemie, bevor es sich auf eine spezielle Aufgabe spezialisiert.

Stell dir vor, ZATOM-1 ist wie ein großer, lehrreicher Sommerkurs für Chemie-Studenten:

• Der Unterricht (Pretraining): Statt nur Rezepte auswendig zu lernen, bekommt das Modell Tausende von Beispielen gezeigt: „Hier ist ein Molekül, hier ist ein Kristall. Wie sind die Atome angeordnet? Wie halten sie zusammen?"
• Die Methode (Flow Matching): Anstatt das Modell zu zwingen, Schritt für Schritt zu raten (wie bei einem Puzzle, bei dem man ein Teil nach dem anderen einsetzt), nutzt ZATOM-1 eine Technik namens „Flow Matching". Stell dir vor, du hast ein Bild, das voller Rauschen (wie statisches Rauschen auf einem alten Fernseher) ist. ZATOM-1 lernt nicht, das Bild zu erstellen, sondern lernt, wie man das Rauschen glattzieht, bis das klare Bild (das Molekül oder Material) erscheint. Es ist wie das Entwirren eines verknoteten Seils: Das Modell lernt den Weg, an dem sich das Seil von einem Knäuel zu einer geraden Linie bewegt.

3. Warum ist das so schnell? (Der Turbo-Effekt)

Frühere Modelle waren oft wie ein langsamer Schachspieler, der jeden Zug einzeln berechnet. ZATOM-1 ist wie ein Supercomputer, der den gesamten Zugplan auf einmal sieht.

• Es braucht keine komplizierten Vorstufen (keine „Latent Spaces", die erst entschlüsselt werden müssen).
• Es kann in weniger als 4 Minuten 10.000 neue Moleküle oder Kristalle generieren. Andere Modelle brauchen dafür Stunden oder Tage.
• Es ist so effizient, dass es auf einer einzigen Grafikkarte läuft, die man auch für Videospiele nutzt.

4. Der magische Trick: Lernen durch Erfinden

Das Coolste an ZATOM-1 ist, dass es durch das Erfinden von neuen Dingen lernt, vorherzusagen, wie alte Dinge funktionieren.

• Die Analogie: Stell dir vor, du willst lernen, wie ein Auto funktioniert. Der beste Weg ist nicht, nur das Handbuch zu lesen. Der beste Weg ist, selbst ein Auto zu bauen. Wenn du weißt, wie man die Teile zusammenfügt, um ein Auto zu erschaffen, verstehst du automatisch, wie die Bremsen funktionieren oder warum der Motor heiß wird.
• ZATOM-1 wurde trainiert, neue Moleküle und Materialien zu erschaffen. Dadurch hat es so tiefes Verständnis entwickelt, dass es nun auch extrem gut darin ist, Eigenschaften vorherzusagen (z. B. „Wie stark ist dieses neue Material?" oder „Wie viel Energie speichert es?").

5. Der große Durchbruch: Chemie-Transfer

Früher dachte man, das Wissen über kleine Moleküle helfe nicht beim Verständnis von riesigen Kristallen. ZATOM-1 beweist das Gegenteil.

• Weil das Modell sowohl Moleküle als auch Kristalle gleichzeitig gelernt hat, kann es Wissen übertragen.
• Beispiel: Wenn das Modell gelernt hat, wie Atome in einem kleinen Molekül zusammenarbeiten, hilft ihm dieses Wissen, auch die Struktur eines riesigen Kristalls besser zu verstehen. Es ist, als würde ein Musiker, der Klavier und Geige spielt, durch das Spielen des Klaviers auch besser Geige spielen können, weil er die Musiktheorie in beiden Instrumenten versteht.

Zusammenfassung

ZATOM-1 ist wie ein Schweizer Taschenmesser für die Chemie.

• Es ist schnell (baut in Minuten, was andere in Stunden bauen).
• Es ist vielseitig (kann kleine Moleküle und große Materialien).
• Es ist klug (lernt durch das Erschaffen neuer Dinge, wie die Welt funktioniert).
• Es ist offen (jeder kann es nutzen, um neue Medikamente oder bessere Batterien zu entwickeln).

Mit diesem Werkzeug können Wissenschaftler jetzt viel schneller neue Entdeckungen machen, die unser Leben verbessern – von besseren Medikamenten bis hin zu umweltfreundlicheren Energiespeichern.

Technische Zusammenfassung

Hier ist eine detaillierte technische Zusammenfassung des Papers „ZATOM-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials" auf Deutsch:

1. Problemstellung

Die chemische Modellierung umfasst zwei Hauptdomänen: Moleküle (nicht-periodische Systeme) und Materialien (periodische Kristallstrukturen). Bisherige KI-Ansätze sind typischerweise auf eine dieser Domänen oder eine einzige Aufgabe (entweder Generierung oder Vorhersage) spezialisiert. Dies führt zu:

• Fehlender Wiederverwendbarkeit von Repräsentationen zwischen den Domänen.
• Ineffizienten Trainingsprozessen, da separate Modelle für jede Aufgabe benötigt werden.
• Komplexen Architekturen, die oft auf latente Räume (z. B. durch Autoencoder) angewiesen sind, was die Trainings- und Inferenzgeschwindigkeit verringert.

Das Ziel war es, ein einheitliches, end-zu-end-fähiges, Open-Source-Foundation-Modell zu entwickeln, das sowohl die generative Modellierung als auch die prädiktive Lernfähigkeit für 3D-Moleküle und -Materialien vereint.

2. Methodik: ZATOM-1

ZATOM-1 ist ein Transformer-basiertes Modell, das auf Multimodalem Flow Matching (Flow Matching, FM) trainiert wird.

A. Datendarstellung und Tokenisierung

Das Modell behandelt atomare Systeme auf fundamentaler Ebene durch fünf vereinheitlichte Modalitäten:

1. Atomtypen ($A$): Diskrete Modality (Elemente).
2. 3D-Koordinaten ($X$): Kontinuierliche Modality (in Ångström).
3. Bruchkoordinaten ($F$): Kontinuierlich im Bereich $[0, 1)$, wichtig für periodische Materialien.
4. Gitterlängen ($L_{len}$): Kontinuierlich (Kantenlängen der Einheitszelle).
5. Gitterwinkel ($L_{ang}$): Kontinuierlich (Winkel der Einheitszelle).

Für nicht-periodische Moleküle werden die Gitter- und Bruchkoordinaten maskiert, während für periodische Materialien die absoluten 3D-Koordinaten maskiert werden, um Rotationsinvarianz zu erzwingen.

B. Architektur

Das Herzstück ist ein standardisierter Transformer (genannt Trunk-based Flow Transformer oder TFT), der folgende Komponenten nutzt:

• Query-Key Normalization und Flash Attention für Stabilität und Geschwindigkeit.
• SwiGLU Feed-Forward-Netzwerke.
• Zweistufiger Trainingsprozess:
  1. Generatives Pretraining (Stage 1): Das Modell lernt eine gemeinsame Verteilung der Modalitäten durch Conditional Flow Matching. Es approximiert ein zeitabhängiges Geschwindigkeitsfeld, das Samples von einer einfachen Quellverteilung (Rauschen) zur Zielverteilung (echte Moleküle/Materialien) transportiert. Dies geschieht direkt im euklidischen Raum ohne latente Autoencoder.
  2. Prädiktives Finetuning (Stage 2): Die Gewichte des „Trunk" (Hauptteil des Transformers) werden eingefroren. Zusätzliche „Heads" (Köpfe) werden für Downstream-Aufgaben wie Eigenschaftsvorhersage, Energie- und Kraftberechnung trainiert.

C. Loss-Funktion

Der Gesamtverlust kombiniert Fehler für kontinuierliche Modalitäten (MSE für Koordinaten und Gitterparameter) und diskrete Modalitäten (Cross-Entropy für Atomtypen):
$$\mathcal{L}_{total} = \mathcal{L}_{metric}(X) + \mathcal{L}_{metric}(F) + \mathcal{L}_{metric}(L_{len}) + \mathcal{L}_{metric}(L_{ang}) + \lambda_{discrete} \cdot \mathcal{L}_{discrete}(A)$$

3. Schlüsselbeiträge

• Erstes einheitliches Foundation-Modell: ZATOM-1 ist das erste vollständig Open-Source-Modell, das sowohl generative als auch prädiktive Aufgaben für 3D-Moleküle und -Materialien in einem einzigen Netzwerk vereint.
• Skalierbarkeit und Geschwindigkeit: Durch den Verzicht auf latente Räume und die Nutzung von Flow Matching ist das Training und Sampling deutlich schneller als bei Diffusionsmodellen oder latenten Autoencodern.
• Positiver Transfer-Learning: Das Paper zeigt erstmals, dass das gemeinsame generative Pretraining auf Molekülen und Materialien die Vorhersagegenauigkeit für Eigenschaften in der jeweils anderen Domäne verbessert.
• Effizienz: Das Modell erreicht eine 12,5-fache Beschleunigung beim Sampling im Vergleich zu einem 500M-Parameter-Latent-Diffusion-Baseline und benötigt 3x weniger GPU-Stunden für das Training.

4. Ergebnisse

Die Evaluation erfolgte auf Standard-Datensätzen (QM9, GEOM-Drugs, MP20, Matbench).

• Generative Leistung:

  • Materialien (MP20): ZATOM-1 erreicht auf dem LeMat-GenBench-Benchmark Ergebnisse, die mit State-of-the-Art-Modellen (wie MatterGen oder ADiT) konkurrieren oder diese übertreffen, insbesondere in Bezug auf Stabilität und Einzigartigkeit (SUN-Metriken).
  • Moleküle (QM9 & GEOM-Drugs): Das Modell erreicht einen Validitätsgrad von ca. 95–99% und besteht fast alle physikalischen Sanity-Checks (PoseBusters), was einen neuen State-of-the-Art darstellt.
  • Geschwindigkeit: Die Inferenzzeit für 10.000 Proben liegt bei unter 4 Minuten auf einer einzigen NVIDIA A100 GPU (bei 300M Parametern), was eine Größenordnung schneller ist als vergleichbare Baselines.

• Prädiktive Leistung:

  • Eigenschaftsvorhersage: Auf QM9 erreicht ZATOM-1 State-of-the-Art-Ergebnisse bei Multi-Task-Vorhersagen (z. B. HOMO/LUMO-Energien, Dipolmomente).
  • Transfer-Learning: Modelle, die sowohl auf Molekülen als auch auf Materialien vortrainiert wurden, zeigen eine höhere Genauigkeit bei der Vorhersage von Moleküleigenschaften als Modelle, die nur auf Molekülen trainiert wurden.
  • Energie und Kräfte: Das Modell zeigt vielversprechende Ergebnisse bei der Vorhersage von Kräften für Materialien (MPtrj-Datensatz), die leicht über denen von Orb-v1 liegen.

• Skalierungsgesetze: Die Leistung verbessert sich vorhersehbar mit der Erhöhung der Parameterzahl (von 80M auf 300M), was auf die Eignung des Ansatzes für zukünftige Skalierung hindeutet.

5. Bedeutung und Ausblick

ZATOM-1 markiert einen Paradigmenwechsel in der Scientific Machine Learning (SciML) für die Chemie:

1. Vereinheitlichung: Es bricht die Silos zwischen molekularer und materialwissenschaftlicher Modellierung auf.
2. Effizienz: Durch die direkte Modellierung im euklidischen Raum ohne latente Kodierung werden Rechenressourcen drastisch gespart.
3. Foundation-Modell-Ansatz: Es demonstriert, dass generatives Pretraining als universelle Initialisierung für diverse Downstream-Aufgaben dient und positive Transfer-Effekte zwischen chemischen Domänen ermöglicht.

Das Modell und seine Gewichte sind Open Source verfügbar und legen den Grundstein für zukünftige Entwicklungen in der Entdeckung neuer Medikamente und Materialien durch KI.