🤖 AI

NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation

Das Papier stellt NEAT vor, einen neuartigen, nachbarschaftsgeführten autoregressiven Set-Transformer, der auf den Datensätzen QM9 und GEOM-Drugs durch die Sicherstellung einer Permutationsinvarianz auf Atomniveau mittels einer ordnungsunabhängigen Trainingsstrategie eine state-of-the-art-Qualität bei der 3D-Molekülgeneierung mit erheblichen Geschwindigkeitsvorteilen erreicht.

Ursprüngliche Autoren: Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

Veröffentlicht 2026-05-07

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen einem Roboter beizubringen, wie man eine komplexe 3D-Struktur baut, etwa ein molekulares Modell aus verschiedenfarbigen Kugeln (Atomen), die durch Stäbe (Bindungen) verbunden sind. Das Ziel ist, dass der Roboter von Grund auf neue, gültige Strukturen erschafft und nicht nur bestehende kopiert.

Diese Arbeit stellt einen neuen Roboter-Architekten namens NEAT (Neighborhood-Guided, Efficient, Autoregressive Set Transformer) vor. So funktioniert er, einfach erklärt:

Das Problem: Die „Reihenfolge"-Verwirrung

Die meisten früheren Roboter versuchten, Moleküle zu bauen, indem sie einer strengen Reihenfolge folgten, wie Schüler, die nacheinander in ein Klassenzimmer eintreten. Sie mussten entscheiden: „Wer kommt zuerst? Wer kommt als Zweiter?"

Das Problem: In einem Molekül haben Atome tatsächlich weder einen „ersten" noch einen „zweiten" Platz. Sie existieren alle gleichzeitig in einem 3D-Cluster. Sie in eine Reihenfolge zu zwingen, ist wie der Versuch, eine Schneeflocke zu beschreiben, indem man ihre Kristalle in einer bestimmten Reihenfolge auflistet. Es ist willkürlich und verwirrend. Wenn der Roboter die Reihenfolge falsch wählt, zerfällt das Molekül.
Der alte Weg: Bisherige Methoden versuchten, dies zu lösen, indem sie eine „kanonische" (Standard-)Reihenfolge erfanden, was den Roboter jedoch verzerrte. Er lernte die Reihenfolge der Liste besser als die Form des Moleküls.

Die Lösung: Die „Nachbarschaft"-Strategie

NEAT ändert das Spiel. Anstatt einer Reihe betrachtet er das Molekül als eine Menge.

Wie es funktioniert: Stellen Sie sich vor, Sie bauen ein Haus Ziegel für Ziegel. Anstatt zu fragen: „Welcher Ziegel kommt als Nächstes in der Liste?", betrachtet NEAT die aktuelle Wand, die Sie gebaut haben, und fragt: „Wer sind die Nachbarn, die direkt neben dieser Wand stehen?"
Die Magie: NEAT betrachtet die „Grenze" der aktuellen Struktur. Er weiß, dass jedes Atom, das die aktuelle Struktur berührt, ein gültiger Kandidat ist, als Nächstes hinzugefügt zu werden. Ihm ist die Reihenfolge egal; ihm geht es nur um die Nachbarschaft.
Das Ergebnis: Dies macht den Roboter „permutationsinvariant". Es ist egal, ob Sie das Molekül von links, rechts oder auf dem Kopf stehend betrachten; der Roboter sieht dieselben Nachbarn und baut dieselbe gültige Struktur.

Der Bauprozess: Zwei Schritte gleichzeitig

Wenn NEAT beschließt, ein neues Atom hinzuzufügen, tut er zwei Dinge gleichzeitig:

Was ist es? Er rät den Atomtyp (z. B. „Das sollte eine Kohlenstoff-Kugel sein").
Wo ist es? Er verwendet eine spezielle mathematische Technik namens Flow Matching, um genau zu bestimmen, wo in 3D-Raum dieses Atom schweben sollte. Denken Sie daran wie an einen Flussstrom, der ein Blatt (das Atom) sanft von einem zufälligen Startpunkt zu seinem perfekten, stabilen Platz im Molekül schiebt.

Warum ist NEAT besonders?

Die Autoren behaupten, NEAT sei aus drei Gründen ein großer Fortschritt:

Es ist schnell: Da es nicht wie andere Roboter Tausende von verrauschten Schritten simulieren muss (was wie der Versuch ist, eine Statue zu schnitzen, indem man stundenlang an einem Steinblock herumbearbeitet), baut NEAT Moleküle viel schneller. Die Arbeit besagt, dass es 5- bis 16-mal schneller ist als seine Konkurrenten.
Es ist flexibel: Da es nicht auf einer festen Reihenfolge basiert, können Sie ihm ein „Starter-Kit" (ein Teil-Molekül) geben und bitten, die Arbeit zu beenden. Es kann die Struktur vervollständigen, egal wie Sie ihm das Starter-Kit übergeben. Andere Roboter scheitern oft, wenn das Starter-Kit nicht ihren spezifischen „Reihenfolge"-Regeln entspricht.
Es ist genau: Bei Standardtests (Datensätze namens QM9 und GEOM-Drugs) baut NEAT Moleküle, die chemisch gültig und einzigartig sind und die besten bestehenden Methoden erreichen oder übertreffen.

Der Haken (Einschränkungen)

Die Autoren sind ehrlich darüber, was NEAT noch nicht kann:

Größenlimits: Es kann nur Moleküle bauen, die so groß sind wie die, die es während des Trainings gesehen hat. Es kann nicht extrapolieren, um ein riesiges Molekül zu bauen, wenn es nur an kleinen gelernt hat.
Bindungs-Unklarheit: Es verlässt sich darauf zu wissen, welche Atome verbunden sind (Bindungen), um die Nachbarschaft zu bestimmen. Wenn Sie ihm eine seltsame Kristallstruktur geben, bei der die Bindungen unklar sind, könnte es verwirrt werden.
Keine „magischen" 3D-Regeln: Es verfügt nicht über eingebaute physikalische Regeln für 3D-Rotation (SE(3)-Äquivarianz), aber die Autoren fanden heraus, dass das zufällige Drehen der Trainingsdaten ausreichte, damit es trotzdem gut funktioniert.

Zusammenfassung

Kurz gesagt ist NEAT eine neue Art, Computern beizubringen, 3D-Moleküle zu entwerfen. Anstatt sie zu zwingen, einer starren, künstlichen Liste zu folgen, lehrt es sie, ihre unmittelbaren Nachbarn zu betrachten und natürlich zu bauen. Es ist schneller, flexibler und ebenso genau wie die aktuellen State-of-the-Art-Methoden.

Technische Zusammenfassung: NEAT – Neighborhood-Guided, Efficient, Autoregressive Set Transformer für die 3D-Molekülgenerierung

Problemstellung

Derzeitige computergestützte Workflows zur Entdeckung bioaktiver Moleküle stützen sich stark auf virtuelles Screening oder de-novo-Generierung. Während virtuelles Screening durch die endliche Größe chemischer Bibliotheken und steigende Rechenkosten begrenzt ist, bietet generative Modellierung eine vielversprechende Alternative. Die Generierung von 3D-Molekülstrukturen stellt jedoch einzigartige Herausforderungen dar. Moleküle sind inhärent dreidimensional, und ihre Eigenschaften werden oft durch die Konformation bestimmt.

Bestehende State-of-the-Art-Modelle zur 3D-Generierung nutzen überwiegend Diffusions- und Flow-Matching-Techniken. Diese Methoden leiden häufig unter einer quadratischen Skalierung bei Message-Passing-Updates, verfügen nicht über eine inhärente Anpassungsfähigkeit an variable Molekülgrößen (oft ist eine feste Atomzahl erforderlich, die bei der Initialisierung bestimmt wird), und haben Schwierigkeiten, dynamische Atominsertionen während des Denoisings zu handhaben.

Autoregressive Modelle bieten eine flexible Alternative, die variable Atomzahlen natürlich unterstützt und nur einen Encoder-Aufruf pro Token erfordert. Standard-autoregressive Transformer erfordern jedoch eine sequenzielle Anordnung von Tokens. In Molekülgraphen haben Atome keine inhärente sequenzielle Ordnung. Frühere Arbeiten haben dies durch die Verwendung kanonischer Atomordnungen adressiert, doch diese Ansätze sind nicht permutationsinvariant bezüglich der Atome. Sie verzerren die Vorhersage des nächsten Tokens zugunsten spezifischer Ordnungskonventionen, was die Fähigkeit des Modells einschränkt, zu generalisieren oder Aufgaben wie die Präfix-Vervollständigung durchzuführen, bei denen die Startatomreihenfolge möglicherweise nicht mit der kanonischen Ordnung des finalen Moleküls übereinstimmt. Umgekehrt erzeugen Graph Neural Network (GNN)-Ansätze, die Permutationsinvarianz gewährleisten, oft Moleküle mit geringer chemischer Validität.

Methodik

Die Autoren schlagen NEAT (Neighborhood-Guided, Efficient, Autoregressive Set Transformer) vor, das erste autoregressive Transformer-Modell zur 3D-Molekülgenerierung, das die Permutationsinvarianz von Atomen respektiert, ohne auf kanonische Atomordnungen zurückzugreifen.

Kernarchitektur

NEAT behandelt Molekülgraphen als Mengen von Atomen und nicht als Sequenzen. Die Architektur besteht aus drei Hauptkomponenten:

Daten-Augmentierungsmodul (Nachbarschaftsleitung): Anstelle sequenzieller Positionsencodings verwendet das Modell eine Sampling-Strategie, um Molekülwachstum nachzuahmen. Gegeben ein Molekülgraph wird ein verbundener Teilgraph (Quellmenge) gesampelt. Die Zielmenge wird als „Randknoten" (1-Hop-Nachbarn) der Quellmenge definiert. Dies liefert eine Menge gültiger Token-Fortsetzungen für den nächsten Schritt, unabhängig von der Reihenfolge, in der Atome hinzugefügt wurden.
Set-Transformer-Stamm: Die Quellmenge wird mittels eines Set-Transformers kodiert (basierend auf NanoGPT/GPT-2 mit bidirektionalem Attention und ohne Positionseingebettungen). Atomtypen werden linear eingebettet, und Positionen werden mittels Fourier-Features kodiert. Additives Pooling fasst diese Knotenrepräsentationen zu einer mengenbasierten Repräsentation ( $z_n$ ) zusammen. Dieses Design gewährleistet Permutationsinvarianz bezüglich der Reihenfolge der Atome in der Quellmenge.
Flow-Matching-Kopf: Die mengenbasierte Repräsentation $z_n$ $z_{n}$ wird verwendet, um den nächsten Token vorherzusagen, faktorisiert in Atomtyp und Position: $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$ $p (a_{n + 1}, x_{n + 1} ∣ z_{n}) = p (a_{n + 1} ∣ z_{n}) \cdot p (x_{n + 1} ∣ a_{n + 1}, z_{n})$ .
- Atomtyp: Modelliert als kategoriale Verteilung über eine lineare Schicht und Softmax.
- Atomposition: Modelliert mittels Conditional Flow Matching (CFM). Das Modell lernt ein Vektorfeld, um eine zufällige initiale Positionsverteilung zu den wahren Atompositionen zu transportieren.

Trainingsstrategie

Nachbarschaftsleitung: Während des Trainings wird das Modell nicht durch einen einzelnen nächsten Token überwacht, sondern durch die Menge aller gültigen Randknoten. Dies eliminiert die Notwendigkeit einer Kanonisierung.
Teacher Forcing: Die Vorhersage des Atomtyps und die Wiederherstellung der Position werden entkoppelt. Das Modell sagt den Typ voraus, und der Flow-Matching-Kopf lernt die Position, bedingt auf diesen Typ.
Augmentierung: Um die Robustheit gegenüber Fehlerakkumulation zu verbessern, wird während des Trainings Gaußsches Rauschen in einen Bruchteil der Atome der Quellmenge injiziert, und die Positionen werden relativ zur verrauschten Quellmenge neu zentriert.
Verlustfunktion: Der Gesamtverlust ist die Summe aus Cross-Entropy-Verlust für Atomtypen und Flow-Matching-Verlust für Positionen. Ein lineares Zuordnungsproblem wird gelöst, um Quell- und Zielverteilungen zu koppeln, Transportkosten zu minimieren und kreuzende Pfade zu vermeiden.

Inferenz

Die Generierung beginnt mit einem zufälligen initialen Atom. In jedem Schritt wird die aktuelle Menge von Atomen kodiert, der nächste Atomtyp vorhergesagt und eine zufällige initiale Position gesampelt. Das Flow-Matching-Vektorfeld wird dann integriert (mittels Euler- oder Euler-Maruyama-Verfahren), um die endgültigen 3D-Koordinaten zu bestimmen. Dieser Vorgang wiederholt sich, bis ein Stopp-Token vorhergesagt wird.

Hauptbeiträge

Permutationsinvariante autoregressive Generierung: NEAT führt eine neuartige Trainingsstrategie mittels Nachbarschaftsleitung und Set-Transformern ein, um eine 3D-Molekülgenerierung zu erreichen, die invariant gegenüber der Atomreihenfolge ist und die Einschränkungen kanonischer Ordnungen in autoregressiven Modellen überwindet.
Effizienz: Indem NEAT die Eingabe als Menge behandelt und einen einzigen Encoder-Aufruf pro Token verwendet, vermeidet es die quadratische Skalierung des diffusionsbasierten Message Passing sowie die Notwendigkeit mehrerer Integrationsstufen pro Atomgenerierungsschritt, wie sie in einigen Flow-Matching-Baselines vorkommt.
Fähigkeit zur Präfix-Vervollständigung: Die Ordnungsunabhängigkeit von NEAT ermöglicht eine hochwertige Präfix-Vervollständigung (Generierung von Molekülen um ein gegebenes Gerüst) ohne die Notwendigkeit, dass das Präfix mit einer spezifischen kanonischen Ordnung übereinstimmt, eine Aufgabe, bei der sequenzbasierte Modelle wie QUETZAL Schwierigkeiten haben.
State-of-the-Art-Leistung: Das Modell erzielt auf Standard-Benchmarks (QM9 und GEOM-Drugs) konkurrierende oder überlegene Ergebnisse im Vergleich zu sowohl autoregressiven als auch Diffusions-/Flow-Matching-Baselines.

Experimentelle Ergebnisse

Das Modell wurde auf den Datensätzen QM9 (kleine organische Moleküle) und GEOM-Drugs (arzneimittelähnliche Moleküle) evaluiert.

Generierungsqualität: Auf QM9 übertraf NEAT frühere autoregressive Modelle (G-SchNet, G-SphereNet, Symphony, Geo2Seq) hinsichtlich der xyz2mol-Validität und Einzigartigkeit. Es erreichte eine xyz2mol-Validität von 95,2 % und eine Einzigartigkeit von 91,6 %, was den zweiten Platz in Bezug auf Lookup-Validität und Einzigartigkeit einbrachte. Auf GEOM erreichte NEAT die höchste Lookup-Validität (99,7 %) und die zweitbeste xyz2mol-Validität (93,9 %) sowie Einzigartigkeit (91,6 %) unter autoregressiven Modellen und entsprach damit der Leistung von State-of-the-Art-Flow-Matching-Modellen wie FlowMol3.
Präfix-Vervollständigung: Bei einer Aufgabe, bei der 1.000 Moleküle aus 100 diversen Präfixen generiert wurden, übertraf NEAT QUETZAL signifikant. NEAT erreichte 97,8 % Lookup-Validität und 92,7 % xyz2mol-Validität, wohingegen QUETZAL lediglich 27,8 % bzw. 28,1 % erreichte. Dies demonstriert die Fähigkeit von NEAT, strukturelle Merkmale unabhängig von Atomordnungs-Konventionen zu generalisieren.
Laufzeiteffizienz: NEAT zeigt erhebliche Geschwindigkeitsvorteile. Es läuft etwa 5-mal schneller als SemlaFlow und 16-mal schneller als FlowMol3. In Worst-Case-Szenarien beträgt die Komplexität von NEAT $O(BN^3)$ (wobei $N$ die Atomzahl ist), während Diffusions-/Flow-Modelle als $O(BTN^2)$ skalieren (wobei $T$ die Integrationsstufen sind). Da $N$ für arzneimittelähnliche Moleküle typischerweise kleiner als $T$ ist, bietet NEAT einen praktischen Geschwindigkeitsvorteil.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass NEAT eine starke Grundlage für den Einsatz von Transformer-Architekturen in komplexen Arzneimitteldesign-Aufgaben bietet, speziell für die konditionale 3D-Molekülgenerierung. Durch die Beseitigung der Abhängigkeit von kanonischen Atomordnungen ermöglicht das Modell flexible Aufgaben wie die gerüstbasierte Generierung (Präfix-Vervollständigung), die für autoregressive Ansätze zuvor schwierig waren.

Die Autoren stellen fest, dass die aktuelle Architektur zwar nicht explizit SE(3)-äquivariant ist, die Verwendung von zufälliger Rotations-Augmentierung während des Trainings dies jedoch effektiv kompensiert, was zu nur geringen Varianzen in den Vorhersagen führt. Die Arbeit legt nahe, dass die Integration von NEAT mit einem SE(3)-äquivarianten Transformer eine vielversprechende Richtung für zukünftige Forschung ist.

Die Autoren betonen, dass, obwohl das Modell hochwertige 3D-Strukturen generiert, die Erzeugung von Molekülen nur ein Schritt in einem komplexen Pipeline ist, der umfangreiche Validierung, Synthese und Tests erfordert. Sie erkennen auch Einschränkungen an, wie etwa die Abhängigkeit von Kanteninformationen für die Nachbarschaftsleitung (was die Anwendbarkeit auf anorganische Kristalle einschränken kann) und das Fehlen von Extrapolation über die während des Trainings gesehene maximale Atomzahl hinaus.