Tokenization for Molecular Foundation Models

Résumé technique : La tokenisation pour les modèles de fondation moléculaires

Énoncé du problème
La prédiction précise des propriétés chimiques est cruciale pour des secteurs allant du stockage d'énergie à la découverte pharmaceutique. Bien que les architectures Transformer aient révolutionné le traitement du langage naturel (NLP), leur application aux modèles de fondation moléculaires se heurte à un goulot d'étranglement fondamental : la tokenisation. Les modèles moléculaires actuels reposent principalement sur une tokenisation « par atome » (Atom-wise), où les chaînes de caractères SMILES (Simplified Molecular Input Line Entry System) sont découpées en jetons au niveau de l'atome à l'aide de vocabulaires fixes.

La limitation principale de cette approche est son incapacité à couvrir pleinement la spécification OpenSMILES. Les tokeniseurs par atome traitent les atomes entre crochets (qui encodent les isotopes, les centres chiraux, les charges et les comptes d'hydrogène explicites) comme des jetons uniques et indivisibles. Pour couvrir toutes les permutations possibles de ces caractéristiques, un vocabulaire nécessiterait plus de 28 billions de jetons. Par conséquent, les modèles existants utilisent des vocabulaires de moins de 3 000 jetons, ce qui entraîne des lacunes de couverture significatives. Lorsqu'ils rencontrent de nouveaux atomes entre crochets, ces tokeniseurs à vocabulaire fermé ont recours à un jeton générique inconnu [UNK], ce qui peut potentiellement occulter des informations chimiques critiques telles que la chiralité ou une composition isotopique spécifique. De plus, les alternatives à vocabulaire ouvert existantes (comme les schémas basés sur le BPE) souffrent souvent d'ambiguïté, où des entités chimiques distinctes (par exemple, une liaison soufre-carbone versus l'élément Scandium) sont confondues en un même jeton.

Méthodologie
Les auteurs proposent un nouveau cadre de tokenisation, Smirk, et une variante compressée, Smirk-GPE, conçus pour atteindre une couverture complète de la spécification OpenSMILES tout en maintenant l'efficacité computationnelle.

1. Tokenisation Smirk : Ce schéma utilise une décomposition en deux étapes, au niveau du caractère, des chaînes SMILES basée sur les glyphes définis par la spécification OpenSMILES.

   • Étape 1 : Décomposition en atomes (ex. : OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]).
   • Étape 2 : Décomposition des atomes entre crochets en glyphes constituants (ex. : [C@@H] $\rightarrow$ [ C @ @ H ]).
     Ce processus distingue les séquences ambiguës (ex. : Sc comme une liaison vs [Sc] comme Scandium) en traitant les crochets et les symboles internes comme des jetons distincts. Le vocabulaire résultant est fixé à 165 jetons, ne nécessite aucun entraînement et garantit que n'importe quelle molécule encodée en OpenSMILES peut être tokenisée sans utiliser de jeton [UNK].

2. Smirk-GPE (Glyph Pair Encoding) : Pour remédier à l'augmentation de la longueur de séquence (fertilité) causée par la décomposition complète des atomes entre crochets, les auteurs ont implémenté Smirk-GPE. Cette variante applique une stratégie de compression de type BPE (Byte-Pair Encoding) spécifiquement aux jetons de glyphes. Contrairement au BPE standard qui fusionne des chaînes de caractères, Smirk-GPE apprend des règles de fusion sur les identifiants de jetons (token IDs), garantissant que les fusions chimiquement significatives (par exemple, combiner un glyphe de soufre et de carbone) ne créent pas d'ambiguïté avec les symboles atomiques (par exemple, le Scandium).

3. Cadre d'évaluation :

   • Métriques intrinsèques : Les auteurs ont évalué les tokeniseurs en utilisant la fertilité (longueur moyenne de séquence), l'entropie normalisée (efficacité de compression), le déséquilibre des jetons et la fréquence du jeton [UNK].
   • Proxy à faible coût : Reconnaissant que l'entraînement de modèles Transformer complets pour chaque tokeniseur est coûteux en calcul, les auteurs ont utilisé des modèles n-grammes comme proxy. Ils ont entraîné des modèles n-grammes sur 1,6 milliard de chaînes SMILES et ont mesuré la perte de cross-entropie et la perte d'information (via la divergence KL) pour estimer la performance en aval.
   • Validation extrinsèque : Pour valider le proxy n-gramme, les auteurs ont pré-entraîné 18 modèles RoBERTa de type encodeur uniquement (à partir de zéro) en utilisant 11 tokeniseurs différents et trois encodages moléculaires. Ces modèles ont été affinés sur six tâches de régression et sept de classification issues de MoleculeNet et tmQM.

Résultats clés

• Couverture : Smirk et Smirk-GPE sont les seuls tokeniseurs évalués qui atteignent une couverture de 100 % de la spécification OpenSMILES, éliminant ainsi l'utilisation du jeton [UNK]. En revanche, les tokeniseurs spécialisés en chimie (incluant SPE, APE et diverses variantes de BPE) émettent le jeton [UNK] avec une fréquence non négligeable (jusqu'à ~50 % sur le jeu de données tmQM).
• Perte d'information : Les tokeniseurs à couverture limitée présentent une perte d'information substantielle, particulièrement sur les jeux de données riches en métaux de transition et en stéréochimie (ex. : tmQM). Par exemple, le tokeniseur de MoLFormer subit une perte de 40,3 nats/molécule sur tmQM en raison des jetons inconnus, alors que les variantes de Smirk atténuent cette dégradation.
• Corrélation de performance : L'étude a trouvé une forte corrélation linéaire entre les métriques n-gramme (cross-entropie et perte d'information) et la performance en aval des modèles basés sur les Transformers. Cela valide l'utilisation des n-grammes comme un proxy à faible coût pour évaluer la qualité d'un tokeniseur.
• Impact en aval :
  • Smirk montre un effet positif sur la qualité du pré-entraînement et la performance en aval sur le jeu de données tmQM.
  • Sur les tâches de MoleculeNet (dominées par de petites molécules organiques), Smirk affiche des performances similaires à une tokenisation classique par atome.
  • Les tokeniseurs à faible couverture (SPE/APE) impactent négativement tant le pré-entraînement que la performance en aval par rapport à la référence.
  • Il a été constaté que le choix de l'encodage moléculaire (SMILES vs SELFIES) a un impact négligeable par rapport au choix du tokeniseur.

Signification et affirmations
L'article soutient qu'un modèle de fondation pour la chimie doit encoder toute l'étendue de l'espace chimique pour éviter d'occulter des caractéristiques critiques. Les auteurs affirment que les tokeniseurs actuels occultent par inadvertance des informations au niveau de l'atome (telles que la chiralité dans le Cisplatine ou des isotopes spécifiques), déclenchant une perte d'information potentiellement significative qui n'est pas seulement théorique, mais qui impacte des molécules cliniquement ou industriellement pertinentes.

La signification de ce travail réside dans :

1. Robustesse : Démontrer que les tokeniseurs à vocabulaire ouvert (Smirk/Smirk-GPE) offrent une couverture robuste de l'espace chimique, empêchant la perte d'information associée aux jetons inconnus.
2. Efficacité : Établir que les modèles n-gramme peuvent servir de proxy fiable et à faible coût pour évaluer la performance des tokeniseurs, réduisant ainsi la charge de calcul pour l'ajustement des hyperparamètres et la sélection de modèles.
3. Interprétabilité : Souligner que Smirk permet aux chercheurs de manipuler directement le contenu riche en informations des atomes entre crochets, étendant les bénéfices d'interprétabilité de la tokenisation par atome tout en supprimant le risque d'erreurs de vocabulaire inconnu.

Les auteurs concluent que, bien que les benchmarks actuels (comme MoleculeNet) puissent ne pas exposer pleinement les déficiences des tokeniseurs à couverture limitée en raison d'un manque de diversité des éléments et de la stéréochimie, la transition vers des tokeniseurs capables d'encoder l'intégralité de l'espace chimique est nécessaire pour des modèles de fondation moléculaires fiables. Ils encouragent la communauté à évaluer rigoureusement la portée des benchmarks et à élargir les jeux de données pour inclure des caractéristiques chimiques diversifiées.