FragmentNet: Adaptive Graph Fragmentation for Graph-to-Sequence Molecular Representation Learning

Der Artikel stellt FragmentNet vor, ein Graph-zu-Sequenz-Modell, das einen neuartigen adaptiven Tokenizer einsetzt, um Moleküle in chemisch valide Fragmente mit einstellbarer Granularität zu zerlegen, und zeigt, dass eine Vorverarbeitung auf dieser Fragmentebene die Leistung bei der nachgelagerten Eigenschaftsvorhersage im Vergleich zu herkömmlichen atomaren oder starren regelbasierten Ansätzen signifikant verbessert.

Das Wesentliche

Stellen Sie sich vor, Sie versuchen, einem Computer beizubringen, Chemie zu verstehen. Traditionell haben Wissenschaftler Computern auf zwei Hauptarten beigebracht, Moleküle zu betrachten, wobei beide Ansätze Mängel aufweisen:

1. Der „Atom-für-Atom"-Ansatz: Dies ist vergleichbar mit dem Versuch, einen Roman zu verstehen, indem man ihn Buchstabe für Buchstabe liest. Man sieht das „t", dann das „h", dann das „e", verpasst aber das Wort „the" vollständig. In der Chemie bedeutet dies, dass der Computer einzelne Atome sieht, aber Schwierigkeiten hat zu verstehen, wie sie sich zu funktionellen Teilen gruppieren (wie etwa ein Motor oder ein Türgriff eines Autos).
2. Der „Starre-Regel"-Ansatz: Dies ist vergleichbar mit der Verwendung eines Wörterbuchs, das nur vordefinierte, unveränderbare Wörter enthält. Wenn ein neuer Worttyp erscheint, kann das Wörterbuch damit nicht umgehen. In der Chemie bedeutet dies, feste Regeln zu verwenden, um Moleküle in Stücke zu zerteilen. Es funktioniert einigermaßen, ist aber unflexibel und kann sich nicht an die enorme Vielfalt der in der Natur vorkommenden chemischen Formen anpassen.

Einführung von FragmentNet: Der „Intelligente Lego"-Ansatz

Die Arbeit stellt FragmentNet vor, eine neue Methode, um Computern Moleküle beizubringen. Anstatt einzelne Atome zu betrachten oder starre Regeln zu verwenden, nutzt FragmentNet einen gelernten, adaptiven Tokenizer.

Stellen Sie sich ein Molekül als eine riesige, komplexe Struktur vor, die aus Lego-Steinen gebaut ist.

• Alte Methoden betrachteten entweder jeden einzelnen winzigen Plastiknippel auf den Steinen (Atome) oder versuchten, die Struktur in einige vordefinierte Kategorien zu zwängen.
• FragmentNet betrachtet die Struktur und lernt, die Steine selbstständig in sinnvolle Blöcke zu gruppieren. Es könnte entscheiden, dass eine bestimmte Ansammlung von Steinen ein „Rad" bildet, eine andere einen „Sitz" und eine weitere einen „Motor". Diese Blöcke sind die „Fragmente".

Wie es funktioniert (Die drei magischen Tricks)

1. Lernen des Gruppierens (Der adaptive Tokenizer):
   Das Modell errät nicht einfach, wie die Steine gruppiert werden sollen. Es studiert Millionen von Molekülen und lernt, welche Atomgruppen chemisch üblicherweise zusammenhängen. Es erstellt ein benutzerdefiniertes Wörterbuch, in dem ein „Token" nicht nur ein Buchstabe oder ein Atom ist, sondern ein chemisch valider Teil eines Moleküls (wie eine ganze funktionelle Gruppe). Dies ist vergleichbar damit, dem Computer beizubringen, dass „ing" eine Endung ist oder dass „car" ein Wortstamm ist, anstatt einfach nur „c-a-r" zu sehen.

2. Beibehaltung der Karte (Räumliche Positions-Encodings):
   Wenn man eine 3D-Lego-Burg in eine 1D-Wortliste (eine Sequenz) umwandelt, geht normalerweise die Information darüber verloren, wo die Teile relativ zueinander liegen. FragmentNet löst dieses Problem, indem es jedem Fragment ein spezielles „GPS-Tag" hinzufügt. Diese Tags sagen dem Computer: „Dieses Motorteil ist mit diesem Radteil verbunden, und sie sind drei Schritte vom Sitz entfernt." Dies stellt sicher, dass sich der Computer die Form des Moleküls merkt, selbst wenn es in eine Liste umgewandelt wurde.

3. Das „Lückentext"-Spiel (Maskierte Fragment-Modellierung):
   Um wirklich klug zu werden, spielt das Modell ein Spiel, das „Mad Libs" oder einem Kreuzworträtsel ähnelt.

   • Der Computer sieht ein Molekül, das aus Fragmenten besteht.
   • Es versteckt (maskiert) eines der Fragmente.
   • Es muss erraten, welches das fehlende Teil ist, basierend auf dem umgebenden Kontext.
   • Da es ganze Blöcke (Fragmente) statt einzelner Atome errät, lernt es die „Grammatik" der Chemie viel schneller. Es lernt, dass, wenn man ein „Rad" und einen „Sitz" sieht, das fehlende Teil wahrscheinlich ein „Motor" ist und nicht einfach nur ein zufälliger Plastikstein.

Was die Arbeit herausfand

Die Autoren testeten diese neue Methode gegen die alten „Atom-für-Atom"-Methoden bei mehreren Standard-Chemie-Tests (z. B. Vorhersage, wie gut sich ein Medikament in Wasser löst oder ob es die Blut-Hirn-Schranke überwinden kann).

• Das Ergebnis: Der „Intelligente Lego"-Ansatz (FragmentNet) gewann die meiste Zeit.
• Warum? Weil es den Kontext lernte. Durch das Training an ganzen Fragmenten verstand der Computer, dass bestimmte Atomgruppen zusammenarbeiten, was zu besseren Vorhersagen führte.
• Bonus-Feature: Die Arbeit zeigt auch, dass das Modell, da es diese Blöcke versteht, problemlos einen „Lego-Block" gegen einen anderen austauschen kann, um ein neues, gültiges Molekül zu erstellen. Dies ist vergleichbar damit, einem Auto den Motor zu entnehmen und einen anderen Motor einzusetzen, ohne dass das Auto auseinanderfällt.

Der Haken (Einschränkungen)

Die Arbeit ist ehrlich bezüglich ihrer Grenzen. Aufgrund von Budgetbeschränkungen wurde dieses Experiment auf einem einzigen Laptop (einem MacBook Pro) durchgeführt. Sie verwendeten einen relativ kleinen Datensatz (2 Millionen Moleküle) im Vergleich zu den Milliarden, die von massiven KI-Modellen verwendet werden. Außerdem testeten sie nur zwei Ebenen der „Blockigkeit" (sehr kleine Teile versus mittelgroße Teile).

Auf den Punkt gebracht

FragmentNet ist ein neues Werkzeug, das Computern beibringt, Chemie nicht durch das Starren auf einzelne Atome zu lesen, sondern indem es sinnvolle „Wörter" (Fragmente) erkennt und versteht, wie diese Wörter zusammengefügt werden, um einen Satz zu bilden. Dies macht den Computer zu einem viel besseren Schüler der Chemie und führt zu genaueren Vorhersagen darüber, wie Moleküle sich verhalten.

Technische Zusammenfassung

Technische Zusammenfassung: FragmentNet

Problemstellung

Das Lernen molekularer Repräsentationen hat sich traditionell darauf verlassen, Moleküle als einzelne Atome zu tokenisieren oder starre, regelbasierte Fragmentzerlegungen (z. B. BRICS) zu nutzen. Diese Ansätze stoßen auf erhebliche Einschränkungen:

1. Atom-Level-Tokenisierung erfasst oft keinen breiteren chemischen Kontext, was zu „negativem Transfer" führt, bei dem vortrainierte Modelle schlechter abschneiden als einfachere Baselines. Das Maskieren einzelner Atome kann chemisch inkonsistente Umgebungen schaffen, die das Erlernen von Bindungsregeln und Wechselwirkungen funktioneller Gruppen behindern.
2. Regelbasierte Fragmentierung fehlt es an Flexibilität und sie hat Schwierigkeiten, sich über diverse chemische Räume hinweg zu verallgemeinern.
3. Sequenzbasierte Methoden (z. B. SMILES-Tokenisierung) verlieren oft kritische topologische Informationen, die inhärent zu molekularen Graphen sind.

Bestehende Strategien für das Maskierte Sprachmodellieren (MLM), die auf Graphen angewendet werden, maskieren häufig Atome, was die chemische Kohärenz unterbricht. Umgekehrt modellieren Methoden, die Subgraphen maskieren (z. B. SimSGT), die Wechselwirkungen zwischen diesen nicht explizit, was die Erfassung von Abhängigkeiten über große Distanzen begrenzt.

Methodik

Die Autoren stellen FragmentNet vor, ein Graph-zu-Sequenz-Modell, das die Lücke zwischen Graph-Topologie und Sequenzmodellierung durch adaptive, erlernte Tokenisierung überbrücken soll.

1. Adaptive, erlernte Tokenisierung

Im Gegensatz zu regelbasierten Methoden verwendet FragmentNet einen datengesteuerten Tokenizer, der molekulare Graphen in chemisch valide Fragmente mit einstellbarer Granularität zerlegt.

• Iteratives paarweises Zusammenführen: Der Tokenizer beginnt mit einzelnen Atomen und führt iterative Zusammenführungen verbundener Paare basierend auf einer erlernten Zusammenführungs-Historie durch, die aus dem Trainingskorpus abgeleitet wird.
• Granularitätskontrolle: Die Anzahl der Zusammenführungsiterationen ($T$) steuert die Tokengröße. Ein Molekül kann unter Verwendung der ersten $t$ Zusammenführungen ($t \le T$) tokenisiert werden, ohne neu trainiert werden zu müssen, was eine taskspezifische Optimierung der Granularität ermöglicht.
• Umgang mit hängenden Bindungen: Unterbrochene Bindungen werden durch „Dummy-Atome" (Ordnungszahl 0) dargestellt. Fragmente werden durch die Anzahl und Art der unterbrochenen Bindungen unterschieden (z. B. ein Kohlenstoff mit einer unterbrochenen Einfachbindung versus zwei).
• Einzigartigkeit: Um Stereoisomere und Tautomere zu unterscheiden, verwenden die Autoren den Weisfeiler-Lehman (WL) Graph-Hashing-Algorithmus, wodurch nicht-isomorphe Graphen eindeutige Hashes erhalten.

2. Hierarchischer Encoder (VQVAE + GCN)

Das Modell integriert Atom-Level- und Fragment-Level-Features mittels eines hybriden Encoders:

• VQ-VAE: Kodiert diskrete atomare Features in einen quantisierten latenten Raum.
• GCN: Aggregiert Features von benachbarten Knoten innerhalb der diskreten Fragmente, um strukturelle Beziehungen zu erfassen.
• Integration: Atom-Embeddings werden gemittelt, um Fragment-Repräsentationen zu bilden, die dann mit GCN-Ausgaben kombiniert werden, um komprimierte Fragment-Level-Feature-Embeddings zu erzeugen.

3. Chemisch bewusste räumliche Positionscodierungen (SPEs)

Um die molekulare Topologie beim Serialisieren von Graphen in Sequenzen zu erhalten, verwendet FragmentNet drei Arten von Positionscodierungen:

• Hop-basierte Codierung: Erfasst relative Verbundenheit über kürzeste Pfaddistanzen.
• WL Absolute Positionscodierung: Weist eindeutige Rollen-IDs basierend auf der Graphstruktur zu, um Isomere zu unterscheiden.
• Coulomb-Matrix-Codierung: Modelliert Wechselwirkungen basierend auf Abständen nach dem inversen Quadratgesetz und Atomladungen.
  Diese werden aggregiert, um einen umfassenden räumlichen Kontext für den Transformer bereitzustellen.

4. Maskierte Fragmentmodellierung (MFM)

Das Vor-Trainingsziel besteht darin, ganze chemisch valide Fragmente zu maskieren, anstatt einzelne Atome.

• Prozess: Ein Fragment wird durch ein [MASK]-Token ersetzt, und das Modell sagt das ursprüngliche Fragment unter Verwendung des Kontexts der unmaskierten Fragmente vorher.
• Vorteil: Dies erhält chemisch sinnvolle Kontexte, analog zur Rekonstruktion mehrwortiger Phrasen im NLP, und erleichtert das Erlernen von Bindungsregeln und funktionalen Beziehungen.
• Konfiguration: Die Autoren beschränken das Maskieren auf ein einziges Token pro Sequenz, um den Kontext zu erhalten; trainiert auf 2 Millionen Molekülen.

5. Architektur

Die serialisierten Fragment-Embeddings, angereichert mit SPEs und einem Molecular Descriptor CLS-Token (abgeleitet aus RDKit-Deskriptoren), werden von einem Transformer-Encoder verarbeitet. Ein Kopf zur Eigenschaftsvorhersage verwendet Max-Pooling über die Sequenz für nachgelagerte Aufgaben.

Hauptbeiträge

1. Neuartiger erlernter adaptiver Tokenizer: Eine Methode zur Zerlegung molekularer Graphen in chemisch valide Fragmente unter Beibehaltung der strukturellen Konnektivität, die eine einstellbare Granularität ermöglicht.
2. Räumliche Positionscodierungen: Eine Reihe von Codierungen (Hop, WL, Coulomb), die die Topologie molekularer Graphen in einem sequenzkompatiblen Format erfassen und eine effektive Graph-zu-Sequenz-Modellierung ermöglichen.
3. Empirische Studie zur Granularität: Ein Nachweis, dass die Tokenisierungsgranularität eine kritische Designentscheidung ist. Die Arbeit zeigt, dass die Fragment-Level-Tokenisierung in Kombination mit MFM-Vor-Training die Atom-Level-Tokenisierung bei der Mehrzahl der Eigenschaftsvorhersageaufgaben übertrifft.

Ergebnisse

Das Modell wurde auf MoleculeNet- und Malaria-Benchmarks unter Verwendung von Scaffold-Splitting (80-10-10) evaluiert.

• Auswirkung des Vor-Trainings: FragmentNet, das mit MFM vor-trainiert wurde, schnitt konsistent besser ab als nicht vor-trainierte Modelle.
• Fragment vs. Atom: Mit MFM-Vor-Training übertraf die Fragment-Level-Variante (100 Zusammenführungsiterationen) die Atom-Level-Variante (0 Zusammenführungsiterationen) bei 5 von 7 Datensätzen (BBBP, Tox21, ToxCast, BACE, ESOL, Lipo, Malaria). Ohne Vor-Training schnitt die Atom-Level-Tokenisierung oft besser ab, was darauf hindeutet, dass die Vorteile einer gröberen Tokenisierung spezifisch durch Vor-Training freigesetzt werden.
• Interpretierbarkeit: Attention-Karten zeigten chemisch intuitive Muster, wie etwa Attention-Köpfe, die sich bei der Löslichkeit (ESOL) auf Hydroxylgruppen oder bei der antimalarischen Aktivität auf Chinazolin-Kerne konzentrierten, was mit bekannten Pharmakophoren übereinstimmt.
• Fragmentaustausch: Der erlernte Tokenizer ermöglichte ein Fragmentaustausch-Modul, um chemisch valide Analoga zu generieren (z. B. Modifikation von Ibuprofen) ohne Substruktur-Matching, was die Nützlichkeit bei der molekularen Bearbeitung demonstriert.

Bedeutung und Behauptungen

Die Arbeit geht davon aus, dass die Tokenisierungsgranularität ein entscheidender Hebel zur Verbesserung molekularer Repräsentationen ist. Durch den Wechsel von der Atom-Level- zur Fragment-Level-Modellierung adressiert FragmentNet die bei der Atom-Level-Maskierung üblichen Probleme des negativen Transfers und erfasst höherstufige strukturelle Motive.

Die Autoren betonen, dass ihr Ansatz „chemisch informiert" ist, Sequenzlängen verkürzt und die Rechenkosten im Vergleich zu Standard-Transformer-Modellen senkt. Trotz des Trainings auf einer bescheidenen Einrichtung (ein einzelner Laptop mit 2 Millionen Molekülen und einem kleinen Vokabular) zeigte das vor-trainierte Fragmentmodell erhebliche Verbesserungen gegenüber nicht vor-trainierten Varianten.

Die Arbeit etabliert, dass adaptive, erlernte Tokenisierung in Kombination mit maskierter Fragmentmodellierung eine gangbare und effektive Strategie für das Lernen molekularer Repräsentationen ist, die verbesserte nachgelagerte Leistung und eine verbesserte chemische Interpretierbarkeit bietet. Die Autoren räumen Einschränkungen hinsichtlich des Umfangs ihrer Experimente (einzelner Laptop, kleiner Datensatz) ein und schlagen vor, dass zukünftige Arbeiten die optimale Granularität für spezifische Aufgaben untersuchen und auf größere Modelle und Datensätze skalieren sollten.