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

L'article présente FragmentNet, un modèle de graphe vers séquence qui utilise un nouveau tokenizer adaptatif pour décomposer les molécules en fragments chimiquement valides de granularité ajustable, démontrant que le pré-entraînement à ce niveau de fragments améliore considérablement les performances de prédiction des propriétés en aval par rapport aux approches traditionnelles au niveau des atomes ou basées sur des règles rigides.

L'essentiel

Imaginez que vous essayiez d'enseigner à un ordinateur à comprendre la chimie. Traditionnellement, les scientifiques ont appris aux ordinateurs à observer les molécules de deux manières principales, toutes deux présentant des défauts :

1. L'approche « atome par atome » : C'est comme essayer de comprendre un roman en le lisant lettre par lettre. Vous voyez le « t », puis le « h », puis le « e », mais vous manquez complètement le mot « the ». En chimie, cela signifie que l'ordinateur voit les atomes individuels mais peine à comprendre comment ils se regroupent pour former des parties fonctionnelles (comme le moteur d'une voiture ou une poignée de porte).
2. L'approche « règle rigide » : C'est comme utiliser un dictionnaire qui ne contient que des mots prédéfinis et immuables. Si un nouveau type de mot apparaît, le dictionnaire ne peut pas le gérer. En chimie, cela signifie utiliser des règles fixes pour découper les molécules en morceaux. Cela fonctionne à peu près, mais c'est rigide et incapable de s'adapter à la vaste variété de formes chimiques présentes dans la nature.

Voici FragmentNet : l'approche « Lego intelligent »

L'article présente FragmentNet, une nouvelle façon d'enseigner aux ordinateurs les molécules. Au lieu de regarder des atomes individuels ou d'utiliser des règles rigides, FragmentNet utilise un tokeniseur appris et adaptatif.

Imaginez une molécule comme une structure géante et complexe construite avec des briques Lego.

• Les anciennes méthodes regardaient soit chaque tout petit pic en plastique sur les briques (atomes), soit essayaient de forcer la structure dans quelques catégories prédéfinies.
• FragmentNet observe la structure et apprend à regrouper les briques en morceaux significatifs par lui-même. Il peut décider qu'un groupe spécifique de briques forme une « roue », un autre forme un « siège » et un autre forme un « moteur ». Ces morceaux sont les « fragments ».

Comment cela fonctionne (les trois tours de magie)

1. Apprendre à regrouper (le tokeniseur adaptatif) :
   Le modèle ne devine pas simplement comment regrouper les briques. Il étudie des millions de molécules et apprend quels groupes d'atomes ont tendance à rester ensemble chimiquement. Il crée un dictionnaire personnalisé où un « token » n'est pas seulement une lettre ou un atome, mais une pièce chimiquement valide d'une molécule (comme un groupe fonctionnel entier). C'est comme enseigner à l'ordinateur à reconnaître que « ing » est un suffixe, ou que « car » est une racine, plutôt que de simplement voir « c-a-r ».

2. Conserver la carte (encodages de position spatiale) :
   Lorsque vous transformez un château Lego 3D en une liste 1D de mots (une séquence), vous perdez généralement l'information sur l'emplacement des pièces les unes par rapport aux autres. FragmentNet résout ce problème en ajoutant une « étiquette GPS » spéciale à chaque fragment. Ces étiquettes disent à l'ordinateur : « Cette pièce de moteur est connectée à cette pièce de roue, et elles sont à trois pas du siège. » Cela garantit que l'ordinateur se souvient de la forme de la molécule même lorsqu'elle est aplatie en une liste.

3. Le jeu « Complétez la phrase » (modélisation masquée de fragments) :
   Pour devenir vraiment intelligent, le modèle joue à un jeu similaire aux « Mad Libs » ou à un mot croisé.

   • L'ordinateur voit une molécule composée de fragments.
   • Il cache (masque) l'un des fragments.
   • Il doit deviner quelle est cette pièce manquante en se basant sur le contexte environnant.
   • Parce qu'il devine des morceaux entiers (fragments) plutôt que des atomes individuels, il apprend la « grammaire » de la chimie beaucoup plus vite. Il apprend que si vous voyez une « roue » et un « siège », la pièce manquante est probablement un « moteur », et pas simplement une brique en plastique aléatoire.

Ce que l'article a découvert

Les auteurs ont testé cette nouvelle méthode contre les anciennes méthodes « atome par atome » sur plusieurs tests standards de chimie (prédire des choses comme la solubilité d'un médicament dans l'eau ou sa capacité à traverser la barrière hémato-encéphalique).

• Le résultat : L'approche « Lego intelligent » (FragmentNet) a gagné la plupart du temps.
• Pourquoi ? Parce qu'elle a appris le contexte. En s'entraînant sur des fragments entiers, l'ordinateur a compris que certains groupes d'atomes fonctionnent ensemble, conduisant à de meilleures prédictions.
• Fonctionnalité bonus : L'article montre également que, parce que le modèle comprend ces morceaux, il peut facilement remplacer un « morceau Lego » par un autre pour créer une nouvelle molécule valide. C'est comme prendre une voiture, retirer le moteur et cliquer un autre moteur à la place sans que la voiture ne s'effondre.

Le bémol (limites)

L'article est honnête sur ses limites. Ils ont mené cette expérience sur un seul ordinateur portable (un MacBook Pro) en raison de contraintes budgétaires. Ils ont utilisé un ensemble de données relativement petit (2 millions de molécules) par rapport aux milliards utilisés par les modèles d'IA massifs. Ils n'ont également testé que deux niveaux de « granularité » (très petits morceaux contre morceaux de taille moyenne).

En résumé

FragmentNet est un nouvel outil qui apprend aux ordinateurs à lire la chimie non pas en fixant des atomes individuels, mais en reconnaissant des « mots » significatifs (fragments) et en comprenant comment ces mots s'assemblent pour former une phrase. Cela rend l'ordinateur un étudiant de la chimie bien meilleur, conduisant à des prédictions plus précises sur le comportement des molécules.

Résumé technique

Résumé Technique : FragmentNet

Énoncé du Problème

L'apprentissage de représentations moléculaires a traditionnellement reposé sur la tokenisation des molécules en atomes individuels ou sur l'utilisation de décompositions de fragments rigides et basées sur des règles (par exemple, BRICS). Ces approches présentent des limitations significatives :

1. La tokenisation au niveau des atomes échoue souvent à capturer le contexte chimique plus large, conduisant à un « transfert négatif » où les modèles pré-entraînés sous-performent par rapport à des bases de référence plus simples. Le masquage d'atomes individuels peut créer des environnements chimiquement incohérents qui entravent l'apprentissage des règles de liaison et des interactions des groupes fonctionnels.
2. La fragmentation basée sur des règles manque de flexibilité et peine à généraliser à travers des espaces chimiques diversifiés.
3. Les méthodes basées sur les séquences (par exemple, la tokenisation SMILES) perdent souvent des informations topologiques critiques inhérentes aux graphes moléculaires.

Les stratégies existantes de modélisation de langage masqué (MLM) appliquées aux graphes masquent souvent des atomes, ce qui brise la cohérence chimique. À l'inverse, les méthodes qui masquent des sous-graphes (par exemple, SimSGT) ne modélisent pas explicitement les interactions entre eux, limitant la capture des dépendances à long terme.

Méthodologie

Les auteurs introduisent FragmentNet, un modèle graphe-à-séquence conçu pour combler le fossé entre la topologie du graphe et la modélisation de séquences grâce à une tokenisation adaptative et apprise.

1. Tokeniseur Adaptatif et Appris

Contrairement aux méthodes basées sur des règles, FragmentNet emploie un tokeniseur piloté par les données qui décompose les graphes moléculaires en fragments chimiquement valides d'une granularité ajustable.

• Fusion Itérative de Paires : Le tokeniseur commence par des atomes individuels et fusionne itérativement des paires connectées sur la base d'un historique de fusion appris dérivé du corpus d'entraînement.
• Contrôle de la Granularité : Le nombre d'itérations de fusion ($T$) contrôle la taille des tokens. Une molécule peut être tokenisée en utilisant les $t$ premières fusions ($t \le T$) sans réentraînement, permettant une optimisation de la granularité spécifique à la tâche.
• Gestion des Liaisons Pendantes : Les liaisons brisées sont représentées par des « atomes factices » (numéro atomique 0). Les fragments sont distingués par le nombre et le type de liaisons brisées (par exemple, un carbone avec une liaison simple brisée contre deux).
• Unicité : Pour distinguer les stéréoisomères et les tautomères, les auteurs utilisent l'algorithme de hachage de graphes Weisfeiler-Lehman (WL), garantissant que les graphes non isomorphes reçoivent des hachages distincts.

2. Encodeur Hiérarchique (VQVAE + GCN)

Le modèle intègre des caractéristiques au niveau des atomes et des fragments en utilisant un encodeur hybride :

• VQ-VAE : Encode des caractéristiques discrètes au niveau atomique dans un espace latent quantifié.
• GCN : Agrège les caractéristiques des nœuds voisins au sein des fragments discrets pour capturer les relations structurelles.
• Intégration : Les plongements d'atomes sont moyennés pour former des représentations de fragments, qui sont ensuite combinées avec les sorties du GCN pour générer des plongements de caractéristiques de fragments compressés au niveau des fragments.

3. Encodages de Position Spatiale Conscients de la Chimie (SPEs)

Pour préserver la topologie moléculaire lors de la sérialisation des graphes en séquences, FragmentNet emploie trois types d'encodages de position :

• Encodage basé sur les Sauts : Capture la connectivité relative via les distances du plus court chemin.
• Encodage de Position Absolu WL : Assigne des identifiants de rôle uniques basés sur la structure du graphe pour distinguer les isomères.
• Encodage de Matrice de Coulomb : Modélise les interactions basées sur les distances de la loi en carré inverse et les charges atomiques.
  Ces éléments sont agrégés pour fournir un contexte spatial complet au Transformer.

4. Modélisation de Fragments Masqués (MFM)

L'objectif de pré-entraînement consiste à masquer des fragments entiers chimiquement valides plutôt que des atomes individuels.

• Processus : Un fragment est remplacé par un token [MASK], et le modèle prédit le fragment original en utilisant le contexte des fragments non masqués.
• Avantage : Cela préserve des contextes chimiquement significatifs, analogue à la reconstruction de phrases multi-mots en TALN, facilitant l'apprentissage des règles de liaison et des relations fonctionnelles.
• Configuration : Les auteurs limitent le masquage à un seul token par séquence pour préserver le contexte, entraîné sur 2 millions de molécules.

5. Architecture

Les plongements de fragments sérialisés, enrichis par des SPEs et un token CLS de Descripteur Moléculaire (dérivé des descripteurs RDKit), sont traités par un encodeur Transformer. Une tête de prédiction de propriétés utilise un max-pooling sur la séquence pour les tâches en aval.

Contributions Clés

1. Nouveau Tokeniseur Adaptatif Appris : Une méthode pour décomposer les graphes moléculaires en fragments chimiquement valides tout en préservant la connectivité structurelle, permettant une granularité ajustable.
2. Encodages de Position Spatiale : Un ensemble d'encodages (Sauts, WL, Coulomb) qui capturent la topologie du graphe moléculaire dans un format compatible avec les séquences, permettant une modélisation efficace graphe-à-séquence.
3. Étude Empirique sur la Granularité : Une démonstration que la granularité de la tokenisation est un choix de conception critique. L'article montre que la tokenisation au niveau des fragments, lorsqu'elle est combinée à un pré-entraînement MFM, surpasse la tokenisation au niveau des atomes sur la majorité des tâches de prédiction de propriétés.

Résultats

Le modèle a été évalué sur les benchmarks MoleculeNet et Malaria en utilisant une division par échafaudage (80-10-10).

• Impact du Pré-entraînement : FragmentNet pré-entraîné avec MFM a constamment surpassé les modèles non pré-entraînés.
• Fragments vs Atomes : Avec un pré-entraînement MFM, la variante au niveau des fragments (100 itérations de fusion) a surpassé la variante au niveau des atomes (0 itération de fusion) sur 5 des 7 jeux de données (BBBP, Tox21, ToxCast, BACE, ESOL, Lipo, Malaria). Sans pré-entraînement, la tokenisation au niveau des atomes fonctionnait souvent mieux, suggérant que les avantages d'une tokenisation plus grossière sont débloqués spécifiquement grâce au pré-entraînement.
• Interprétabilité : Les cartes d'attention ont révélé des motifs chimiquement intuitifs, tels que des têtes d'attention se concentrant sur les groupes hydroxyles pour la solubilité (ESOL) ou sur les noyaux quinazoline pour l'activité antipaludique, s'alignant sur les pharmacophores connus.
• Échange de Fragments : Le tokeniseur appris a permis à un module d'échange de fragments de générer des analogues chimiquement valides (par exemple, la modification de l'Ibuprofène) sans correspondance de sous-structure, démontrant son utilité dans l'édition moléculaire.

Signification et Revendications

L'article postule que la granularité de la tokenisation est un levier clé pour améliorer les représentations moléculaires. En passant d'une modélisation au niveau des atomes à une modélisation au niveau des fragments, FragmentNet résout les problèmes de transfert négatif courants dans le masquage au niveau des atomes et capture des motifs structurels de plus haut niveau.

Les auteurs soulignent que leur approche est « informée par la chimie », réduisant les longueurs de séquence et abaissant les coûts de calcul par rapport aux modèles Transformer standards. Malgré un entraînement sur une configuration modeste (un seul ordinateur portable avec 2 millions de molécules et un petit vocabulaire), le modèle de fragments pré-entraîné a montré des gains substantiels par rapport aux variantes non pré-entraînées.

Ce travail établit que la tokenisation adaptative et apprise combinée à la modélisation de fragments masqués est une stratégie viable et efficace pour l'apprentissage de représentations moléculaires, offrant une amélioration des performances en aval et une interprétabilité chimique accrue. Les auteurs reconnaissent les limitations concernant l'échelle de leurs expériences (un seul ordinateur portable, petit jeu de données) et suggèrent que les travaux futurs devraient explorer la granularité optimale pour des tâches spécifiques et passer à l'échelle vers des modèles et des jeux de données plus vastes.