Tokenization for Molecular Foundation Models

Sintesi Tecnica: Tokenizzazione per Modelli di Fondazione Molecolare

Problematica
La previsione accurata delle proprietà chimiche è fondamentale per settori che spaziano dall'accumulo di energia alla scoperta farmaceutica. Sebbene le architetture transformer abbiano rivoluzionato l'elaborazione del linguaggio naturale (NLP), la loro applicazione ai modelli di fondazione molecolare affronta un collo di bottiglia fondamentale: la tokenizzazione. Gli attuali modelli molecolari si affidano prevalentemente a una tokenizzazione "atom-wise" (per atomi), in cui le stringhe SMILES (Simplified Molecular Input Line Entry System) vengono suddivise in token a livello di atomo utilizzando vocabolari fissi.

Il limite principale di questo approccio è la sua incapacità di coprire completamente la specifica OpenSMILES. I tokenizer atom-wise trattano gli atomi tra parentesi quadre (che codificano isotopi, centri chirali, cariche e conteggi espliciti di idrogeno) come singoli token indivisibili. Per coprire ogni possibile permutazione di queste caratteristiche, un vocabolario richiederebbe oltre 28 trilioni di token. Di conseguenza, i modelli esistenti utilizzano vocabolari inferiori a 3.000 token, risultando in significativi gap di copertura. Quando incontrano nuovi atomi tra parentesi, questi tokenizer a vocabolario chiuso ricorrono a un generico token sconosciuto [UNK], potenzialmente oscurando informazioni chimiche critiche come la chiralità o la specifica composizione isotopica. Inoltre, le alternative open-vocabulary esistenti (come gli schemi basati su BPE) soffrono spesso di ambiguità, dove entità chimiche distinte (ad esempio, un legame zolfo-carbonio rispetto all'elemento Scandio) vengono confuse nello stesso token.

Metodologia
Gli autori propongono un nuovo framework di tokenizzazione, Smirk, e una variante compressa, Smirk-GPE, progettati per raggiungere la copertura completa della specifica OpenSMILES mantenendo l'efficienza computazionale.

1. Tokenizzazione Smirk: Questo schema impiega una decomposizione a due stadi a livello di carattere delle stringhe SMILES basata sui glifi definiti dalla specifica OpenSMILES.

   • Stadio 1: Decomposizione in atomi (ad es., OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]).
   • Stadio 2: Decomposizione degli atomi tra parentesi in glifi costituenti (ad es., [C@@H] $\rightarrow$ [ C @ @ H ]).
   • Questo approccio distingue tra sequenze ambigue (ad es., Sc come legame vs. [Sc] come Scandio) trattando le parentesi e i simboli interni come token distinti. Il vocabolario risultante è fisso a 165 token, non richiede addestramento e garantisce che qualsiasi molecola codificata in OpenSMILES possa essere tokenizzata senza utilizzare un token [UNK].

2. Smirk-GPE (Glyph Pair Encoding): Per affrontare l'aumento della lunghezza della sequenza (fertilità) causato dalla decomposizione completa degli atomi tra parentesi, gli autori hanno implementato Smirk-GPE. Questa variante applica una strategia di compressione simile al Byte-Pair Encoding (BPE) specificamente ai token dei glifi. A differenza del BPE standard che unisce le stringhe, Smirk-GPE apprende le regole di fusione sui ID dei token, assicurando che le fusioni chimicamente significative (ad es., combinare un glifo di zolfo e uno di carbonio) non creino ambiguità con i simboli atomici (ad es., Scandio).

3. Framework di Valutazione:

   • Metriche Intrinseche: Gli autori hanno valutato i tokenizer utilizzando la fertilità (lunghezza media della sequenza), l'entropia normalizzata (efficienza di compressione), lo squilibrio dei token e la frequenza del token [UNK].
   • Proxy a Basso Costo: Riconoscendo che l'addestramento di modelli transformer completi per ogni tokenizer è computazionalmente costoso, gli autori hanno utilizzato i modelli n-gram come proxy. Hanno addestrato modelli n-gram su 1,6 miliardi di stringhe SMILES e misurato la perdita di cross-entropy e la perdita di informazione (tramite divergenza KL) per stimare le prestazioni a valle.
   • Validazione Estrinseca: Per validare il proxy n-gram, gli autori hanno pre-addestrato 18 modelli RoBERTa encoder-only (da zero) utilizzando 11 diversi tokenizer e tre codifiche molecolari. Questi modelli sono stati sottoposti a fine-tuning su sei compiti di regressione e sette di classificazione da MoleculeNet e tmQM.

Risultati Chiave

• Copertura: Smirk e Smirk-GPE sono gli unici tokenizer valutati che raggiungono il 100% della copertura della specifica OpenSMILES, eliminando l'uso del token [UNK]. Al contrario, i tokenizer specifici per la chimica esistenti (inclusi SPE, APE e varie varianti BPE) emettono il token [UNK] con una frequenza non trascurabile (fino a circa il 50% sul dataset tmQM).
• Perdita di Informazione: I tokenizer con copertura limitata mostrano una sostanziale perdita di informazione, particolarmente nei dataset ricchi di metalli di transizione e stereochimica (ad es., tmQM). Ad esempio, il tokenizer di MoLFormer incorre in una perdita di 40,3 nats/molecola su tmQM a causa dei token sconosciuti, mentre le varianti Smirk mitigano questa degradazione.
• Correlazione delle Prestazioni: Lo studio ha trovato una forte correlazione lineare tra le metriche n-gram (cross-entropy e perdita di informazione) e le prestazioni a valle dei modelli basati su transformer. Ciò valida l'uso degli n-gram come proxy a basso costo per valutare la qualità del tokenizer.
• Impatto a Valle:
  • Smirk ha mostrato un effetto positivo sulla qualità del pre-addestramento e sulle prestazioni a valle sul dataset tmQM.
  • Su compiti MoleculeNet (dominati da piccole molecole organiche), Smirk si è comportato in modo simile alla tokenizzazione standard atom-wise.
  • I tokenizer con scarsa copertura (SPE/APE) hanno influenzato negativamente sia il pre-addestramento che le prestazioni a valle rispetto al baseline.
  • Si è scoperto che la scelta della codifica molecolare (SMILES vs. SELFIES) ha un impatto trascurabile rispetto alla scelta del tokenizer.

Significato e Rivendicazioni
L'articolo sostiene che un modello di fondazione per la chimica deve codificare l'intera ampiezza dello spazio chimico per evitare di oscurare caratteristiche critiche. Gli autori affermano che gli attuali tokenizer oscurano involontariamente l'informazione a livello di atomo (come la chiralità in Cisplatino o specifici isotopi), innescando una potenziale perdita di informazioni significativa che non è solo teorica ma impatta molecole clinicamente e industrialmente rilevanti.

Il significato di questo lavoro risiede nella:

1. Robustezza: Dimostrare che i tokenizer open-vocabulary (Smirk/Smirk-GPE) forniscono una copertura robusta dello spazio chimico, prevenendo la perdita di informazioni associata ai token sconosciuti.
2. Efficienza: Stabilire che i modelli n-gram possono servire come un proxy affidabile e a basso costo per valutare le prestazioni del tokenizer, riducendo il carico computazionale per la sintonizzazione degli iperparametri e la selezione del modello.
3. Interpretabilità: Evidenziare come Smirk permetta ai ricercatori di manipolare direttamente il contenuto ricco di informazioni degli atomi tra parentesi, espandendo i benefici di interpretabilità della tokenizzazione atom-wise pur rimuovendo il rischio di errori di fuori vocabolario.

Gli autori concludono che, sebbene gli attuali benchmark (come MoleculeNet) possano non esporre completamente le carenze dei tokenizer a copertura limitata a causa della mancanza di diversità negli elementi e nella stereochimica, la transizione verso tokenizer capaci di codificare l'intero spazio chimico è necessaria per modelli di fondazione molecolare affidabili. Incoraggiano la comunità a valutare rigorosamente l'ambito dei benchmark ed espandere i dataset per includere diverse caratteristiche chimiche.