Tokenization for Molecular Foundation Models

Technische Samenvatting: Tokenisatie voor Moleculaire Fundamentele Modellen

Probleemstelling
Nauwkeurige voorspelling van chemische eigenschappen is cruciaal voor sectoren variërend van energieopslag tot farmaceutische ontdekking. Hoewel transformer-architecturen de natuurlijke taalverwerking (NLP) hebben gerevolutioneerd, staat hun toepassing op moleculaire fundamentele modellen voor een fundamentele bottleneck: tokenisatie. Huidige moleculaire modellen vertrouwen voornamelijk op "atomaire" tokenisatie, waarbij Simplified Molecular Input Line Entry System (SMILES) strings worden opgesplitst in tokens op atomaire basis met behulp van vaste vocabularia.

De primaire beperking van deze aanpak is het onvermogen om de volledige OpenSMILES-specificatie te dekken. Atomaire tokenizers behandelen tussen haakjes geplaatste atomen (die isotopen, chirale centra, ladingen en expliciete waterstofaantallen coderen) als enkele, ondeelbare tokens. Om elke mogelijke permutatie van deze kenmerken te dekken, zou een vocabulaire meer dan 28 biljoen tokens vereisen. Bij consequentie maken bestaande modellen gebruik van vocabularia van minder dan 3.000 tokens, wat resulteert in aanzienlijke dekkingsverschillen. Wanneer ze met nieuwe tussen haakjes geplaatste atomen worden geconfronteerd, grijpen deze closed-vocabulary tokenizers terug op een generiek onbekend token [UNK], wat potentieel cruciale chemische informatie zoals chiraliteit of specifieke isotopische samenstelling kan verhullen. Bovendien lijden bestaande open-vocabulary alternatieven (zoals BPE-gebaseerde schema's) vaak aan ambiguïteit, waarbij verschillende chemische entiteiten (bijv. een zwavel-koolstofverbinding versus het element Scandium) in hetzelfde token worden samengevoegd.

Methodologie
De auteurs stellen een nieuw tokenisatie-framework voor, Smirk, en een gecomprimeerde variant, Smirk-GPE, ontworpen om een volledige dekking van de OpenSMILES-specificatie te bereiken terwijl de computationele efficiëntie behouden blijft.

1. Smirk Tokenisatie: Dit schema maakt gebruik van een twee-fasen, karakter-niveau decompositie van SMILES-strings gebaseerd op de glyphs gedefinieerd door de OpenSMILES-specificatie.

   • Fase 1: Decompositie in atomen (bijv. OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]).
   • Fase 2: Decompositie van tussen haakjes geplaatste atomen in constituerende glyphs (bijv. [C@@H] $\rightarrow$ [ C @ @ H ]).
   • Deze aanpak onderscheidt ambigue sequenties (bijv. Sc als een binding versus [Sc] als Scandium) door de haakjes en interne symbolen als afzonderlijke tokens te behandelen. Het resulterende vocabulaire is vastgesteld op 165 tokens, vereist geen training en garandeert dat elke OpenSMILES-gecodeerde molecuul getokeniseerd kan worden zonder een [UNK] token te gebruiken.

2. Smirk-GPE (Glyph Pair Encoding): Om de toename in sequentielengte (fertility) aan te pakken die wordt veroorzaakt door de volledige decompositie van tussen haakjes geplaatste atomen, hebben de auteurs Smirk-GPE geïmplementeerd. Deze variant past een Byte-Pair Encoding (BPE)-achtige compressiestrategie toe specifiek op de glyph-tokens. In tegen tegenstelling tot standaard BPE, dat strings samenvoegt, leert Smirk-GPE merge-regels op token-ID's, waardoor wordt gewaarborgd dat chemisch betekenisvolle samenvoegingen (bijv. het combineren van een zwavel- en koolstofglyph) geen ambiguïteit creëren met atomaire symbolen (bijv. Scandium).

3. Evaluatiekader:

   • Intrinsieke Metrieken: De auteurs evalueerden tokenizers met behulp van fertility (gemiddelde sequentielengte), genormaliseerde entropie (compressie-efficiëntie), token-imbalans en de frequentie van het [UNK] token.
   • Low-Cost Proxy: In het besef dat het trainen van volledige transformer-modellen voor elke tokenizer computationeel duur is, gebruikten de auteurs n-gram modellen als een proxy. Ze trainden n-gram modellen op 1,6 miljard SMILES-strings en maten cross-entropy loss en informatieverlies (via KL-divergentie) om de downstream prestaties te schatten.
   • Extrinsieke Validatie: Om de n-gram proxy te valideren, hebben de auteurs 18 encoder-only RoBERTa-modellen (vanaf nul) gepre-traind met 11 verschillende tokenizers en drie moleculaire encodings. Deze modellen werden gefinetuned op zes regressie- en zeven classificatietaken van MoleculeNet en tmQM.

Belangrijkste Resultaten

• Dekking: Smirk en Smirk-GPE zijn de enige geëvalueerde tokenizers die een 100% dekking van de OpenSMILES-specificatie bereiken, waardoor het gebruik van het [UNK] token wordt geëlimineerd. In contrast hiermee emitteren bestaande chemie-specifieke tokenizers (inclusief SPE, APE en diverse BPE-varianten) het [UNK] token met een niet-verwaarloosbare frequentie (tot ~50% op de tmQM-dataset).
• Informatieverlies: Tokenizers met beperkte dekking vertonen substantieel informatieverlies, met name op datasets die rijk zijn aan overgangsmetalen en stereochemie (bijv. tmQM). Zo veroorzaakt de MoLFormer-tokenizer bijvoorbeeld een verlies van 40,3 nats/molecuul op tmQM door onbekende tokens, terwijl Smirk-varianten deze degradatie mitigeren.
• Prestatiecorrelatie: De studie vond een sterke lineaire correlatie tussen n-gram metrieken (cross-entropy en informatieverlies) en de downstream prestaties van transformer-gebaseerde modellen. Dit valideert het gebruik van n-grammen als een low-cost proxy voor het evalueren van de kwaliteit van een tokenizer.
• Downstream Impact:
  • Smirk toonde een positief effect op de pretraining kwaliteit en downstream prestaties op de tmQM-dataset.
  • Op MoleculeNet-taken (gedomineerd door kleine organische moleculen) presteerde Smirk vergelijkbaar met standaard atomaire tokenisatie.
  • Tokenizers met slechte dekking (SPE/APE) hadden een negatieve impact op zowel de pretraining als de downstream prestaties ten opzichte van de baseline.
  • De keuze van de moleculaire encoding (SMILES versus SELFIES) bleek een verwaarloosbare impact te hebben vergeleken met de keuze van de tokenizer.

Betekenis en Claims
Het artikel betoogt dat een fundamenteel model voor chemie de volledige breedte van de chemische ruimte moet coderen om te voorkomen dat kritieke kenmerken worden verborgen. De auteurs beweren dat huidige tokenizers onbedoeld informatie op atomair niveau (zoals chiraliteit in Cisplatin of specifieke isotopen) verhullen, wat kan leiden tot potentieel significante informatieverlies dat niet louter theoretisch is, maar ook klinisch en industrieel relevante moleculen beïnvloedt.

De betekenis van dit werk ligt in:

1. Robuustheid: Het aantonen dat open-vocabulary tokenizers (Smirk/Smirk-GPE) een robuuste dekking van de chemische ruimte bieden, waardoor het informatieverlies geassocieerd met onbekende tokens wordt voorkomen.
2. Efficiëntie: Het vaststellen dat n-gram modellen als een betrouwbare, low-cost proxy kunnen dienen voor het evalueren van de prestaties van een tokenizer, wat de computationele last van hyperparameter tuning en modelselectie vermindert.
3. Interpreteerbaarheid: Het benadrukken dat Smirk onderzoekers in staat stelt om de informatie-rijke inhoud van tussen haakjes geplaatste atomen direct te manipuleren, waarmee het voordeel van interpretabiliteit van atomaire tokenisatie wordt uitgebreid terwijl het risico op out-of-vocabulary fouten wordt weggenomen.

De auteurs concluderen dat hoewel huidige benchmarks (zoals MoleculeNet) de tekortkomingen van limited-coverage tokenizers mogelijk niet volledig blootleggen door een gebrek aan diversiteit in elementen en stereochemie, de overstap naar tokenizers die in staat zijn de volledige chemische ruimte te coderen noodzakelijk is voor betrouwbare moleculaire fundamentele modellen. Zij moedigen de gemeenschap aan om de scope van benchmarks rigoureus te beoordelen en datasets uit te breiden om diverse chemische kenmerken te bevatten.