Tokenization for Molecular Foundation Models

Technische Zusammenfassung: Tokenisierung für molekulare Fundamentmodelle

Problemstellung
Die genaue Vorhersage chemischer Eigenschaften ist entscheidend für Branchen, die von der Energiespeicherung bis hin zur pharmazeutischen Forschung reichen. Während Transformer-Architekturen die natürliche Sprachverarbeitung (NLP) revolutioniert haben, stößt ihre Anwendung auf molekulare Fundamentmodelle auf einen grundlegenden Engpass: die Tokenisierung. Aktuelle molekulare Modelle verlassen sich überwiegend auf eine „Atom-weise“ Tokenisierung, bei der SMILES-Strings (Simplified Molecular Input Line Entry System) mittels fester Vokabulare in Atomebene-Token zerlegt werden.

Die primäre Einschränkung dieses Ansatzes ist die Unfähigkeit, die gesamte OpenSMILES-Spezifikation vollständig abzudecken. Atom-weise Tokenizer behandeln in Klammern gesetzte Atome (die Isotope, Chiralitätszentren, Ladungen und explizite Wasserstoffzahlen kodieren) als einzelne, unteilbare Token. Um jede mögliche Permutation dieser Merkmale abzudecken, wäre ein Vokabular von über 28 Billionen Token erforderlich. Infolgedessen nutzen bestehende Modelle Vokabulare von weniger als 3.000 Tokens, was zu erheblichen Deckungslücken führt. Beim Auftreffen auf neuartige Atome in Klammern greifen diese Closed-Vocabulary-Tokenizer auf ein generisches Unknown-Token [UNK] zurück, was kritische chemische Informationen wie Chiralität oder spezifische Isotopenzusammensetzungen potenziell verschleiern kann. Darüber hinaus leiden bestehende Open-Vocabulary-Alternativen (wie BPE-basierte Verfahren) oft unter Ambiguität, wobei unterschiedliche chemische Entitäten (z. B. eine Schwefel-Kohlenstoff-Bindung vs. das Element Scandium) in dasselbe Token konfundiert werden.

Methodik
Die Autoren schlagen ein neues Tokenisierungs-Framework vor, Smirk, sowie eine komprimierte Variante, Smirk-GPE, die darauf ausgelegt sind, eine vollständige Abdeckung der OpenSMILES-Spezifikation zu erreichen und gleichzeitig die Recheneffizienz aufrechtzuerhalten.

1. Smirk-Tokenisierung: Dieses Schema verwendet eine zweistufige, zeichenbasierte Dekomposition von SMILES-Strings basierend auf den Glyphen, die in der OpenSMILES-Spezifikation definiert sind.

   • Stufe 1: Dekomposition in Atome (z. B. OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]).
   • Stufe 2: Dekomposition von Atomen in Klammern in konstituierende Glyphen (z. B. [C@@H] $\rightarrow$ [ C @ @ H ]).
   • Dieser Ansatz unterscheidet zwischen mehrdeutigen Sequenzen (z. B. Sc als Bindung vs. [Sc] als Scandium), indem er die Klammern und internen Symbole als distinkte Token behandelt. Das resultierende Vokabular ist auf 165 Tokens fixiert, erfordert kein Training und garantiert, dass jedes OpenSMILES-kodierte Molekül ohne die Verwendung eines [UNK]-Tokens tokenisiert werden kann.

2. Smirk-GPE (Glyph Pair Encoding): Um die durch die vollständige Dekomposition von Atomen in Klammern verursachte erhöhte Sequenzlänge (Fertilität) zu adressieren, implementierten die Autoren Smirk-GPE. Diese Variante wendet eine Byte-Pair-Encoding (BPE)-ähnliche Kompressionsstrategie spezifisch auf die Glyph-Tokens an. Im Gegensatz zu Standard-BPE, das Strings zusammenführt, lernt Smirk-GPE Merge-Regeln auf Token-IDs, wodurch sichergestellt wird, dass chemisch bedeutsame Zusammenführungen (z. B. die Kombination einer Schwefel- und einer Kohlenstoff-Glyphe) keine Ambiguität mit Atomsymbolen (z. B. Scandium) erzeugen.

3. Evaluationsrahmen:

   • Intrinsische Metriken: Die Autoren evaluierten Tokenizer anhand von Fertilität (mittlere Sequenzlänge), normalisierter Entropie (Kompressions-Effizienz), Token-Imbalance und der Häufigkeit des [UNK]-Tokens.
   • Kostengünstiger Proxy: Da das Training vollständiger Transformer-Modelle für jeden Tokenizer rechenintensiv ist, nutzten die Autoren n-Gramm-Modelle als Proxy. Sie trainierten n-Gramm-Modelle auf 1,6 Milliarden SMILES-Strings und maßen die Cross-Entropy-Loss sowie den Informationsverlust (via KL-Divergenz), um die Downstream-Performance zu schätzen.
   • Extrinsische Validierung: Um den n-Gramm-Proxy zu validieren, prätrainierten die Autoren 18 Encoder-only RoBERTa-Modelle (von Grund auf) unter Verwendung von 11 verschiedenen Tokenizern und drei molekularen Kodierungen. Diese Modelle wurden auf sechs Regressions- und sieben Klassifikationsaufgaben aus MoleculeNet und tmQM feinjustiert.

Kernergebnisse

• Abdeckung: Smirk und Smirk-GPE sind die einzigen evaluierten Tokenizer, die eine 100%ige Abdeckung der OpenSMILES-Spezifikation erreichen und somit die Verwendung des [UNK]-Tokens eliminieren. Im Gegensatz dazu emittieren bestehende chemiespezifische Tokenizer (einschließlich SPE, APE und verschiedener BPE-Varianten) das [UNK]-Token mit einer nicht vernachlässigbaren Häufigkeit (bis zu ~50 % beim tmQM-Datensatz).
• Informationsverlust: Tokenizer mit begrenzter Abdeckung zeigen einen erheblichen Informationsverlust, insbesondere bei Datensätzen, die reich an Übergangsmetallen und Stereochemie sind (z. B. tmQM). Beispielsweise verursacht der MoLFormer-Tokenizer einen Verlust von 40,3 nats/Molekül auf tmQM aufgrund unbekannter Tokens, während Smirk-Varianten diese Degradation abmildern.
• Performance-Korrelation: Die Studie fand eine starke lineare Korrelation zwischen n-Gramm-Metriken (Cross-Entropy und Informationsverlust) und der Downstream-Performance von Transformer-basierten Modellen. Dies validiert die Nutzung von n-Gramm-Modellen als kostengünstigen Proxy zur Evaluierung der Tokenizer-Qualität.
• Downstream-Auswirkungen:
  • Smirk zeigte einen positiven Effekt auf die Pretraining-Qualität und die Downstream-Performance auf dem tmQM-Datensatz.
  • Auf MoleculeNet-Aufgaben (die von kleinen organischen Molekülen dominiert werden) performte Smirk ähnlich wie die Standard-Atom-weise Tokenisierung.
  • Tokenizer mit schlechter Abdeckung (SPE/APE) beeinträchtigten sowohl das Pretraining als auch die Downstream-Performance im Vergleich zur Baseline negativ.
  • Es wurde festgestellt, dass die Wahl der molekularen Kodierung (SMILES vs. SELFIES) im Vergleich zur Wahl des Tokenizers einen vernachlässigbaren Einfluss hat.

Bedeutung und Ansprüche
Das Paper argumentiert, dass ein Fundamentmodell für die Chemie das gesamte Spektrum des chemischen Raums kodieren muss, um zu verhindern, dass kritische Merkmale verschleiert werden. Die Autoren behaupten, dass aktuelle Tokenizer unbeabsichtigt Informationen auf Atomebene (wie die Chiralität in Cisplatin oder spezifische Isotope) verschleiern, was zu einem potenziell signifikanten Informationsverlust führt, der nicht nur theoretischer Natur ist, sondern klinisch und industriell relevante Moleküle tatsächlich beeinflusst.

Die Bedeutung dieser Arbeit liegt in:

1. Robustheit: Der Nachweis, dass Open-Vocabulary-Tokenizer (Smirk/Smirk-GPE) eine robuste Abdeckung des chemischen Raums bieten und den Informationsverlust durch Unknown-Tokens verhindern.
2. Effizienz: Die Etablierung von n-Gramm-Modellen als zuverlässiger, kostengünstiger Proxy für die Evaluierung der Tokenizer-Performance, was den Rechenaufwand für Hyperparameter-Tuning und Modellselektion reduziert.
3. Interpretierbarkeit: Die Hervorhebung, dass Smirk es Forschern ermöglicht, den informationsreichen Inhalt von Atomen in Klammern direkt zu manipulieren, was den Interpretierbarkeitsvorteil der Atom-weise Tokenisierung erweitert und gleichzeitig das Risiko von Out-of-Vocabulary-Fehlern eliminiert.

Die Autoren kommen zu dem Schluss, dass, obwohl aktuelle Benchmarks (wie MoleculeNet) die Defizite von Tokenizern mit begrenzter Abdeckung aufgrund eines Mangels an Diversität in Elementen und Stereochemie möglicherweise nicht voll ausschöpfen, der Übergang zu Tokenizern, die in der Lage sind, den gesamten chemischen Raum zu kodieren, für zuverlässige molekulare Fundamentmodelle notwendig ist. Sie ermutigen die Community, die Reichweite von Benchmarks rigoros zu prüfen und Datensätze zu erweitern, um vielfältigere chemische Merkmale einzubeziehen.