GNN Molecular

Sept GNN pour les propriétés moléculaires

Recherche · oct. 2024 · Terminé

Comparer sept réseaux de neurones sur graphes pour prédire la solubilité des molécules.

RMSE ESOL 0,498
R² 0,926
7 GNN comparés

Construit avecPyTorchPyTorch GeometricRDKitMoleculeNet (ESOL)pytest

Le code est pour l'instant dans un dépôt privé.

Présentation

Une molécule est naturellement un graphe : les atomes sont des nœuds et les liaisons des arêtes. Je voulais voir comment différentes architectures de réseaux de neurones sur graphes exploitent cette structure pour prédire une vraie propriété chimique, la solubilité aqueuse, à partir du benchmark ESOL de MoleculeNet.

J'ai d'abord implémenté trois architectures à la main dans des notebooks : un réseau convolutif sur graphes (GCN), un réseau à passage de messages (MPNN) qui exploite aussi les caractéristiques des liaisons, et un réseau à convolution de diffusion (DCNN) qui propage l'information sur plusieurs sauts. Je les ai ensuite transformées en un package Python testé et ajouté quatre modèles standard construits à partir des couches de PyTorch Geometric : GIN, GAT (v2), GraphSAGE et AttentiveFP.

Les sept tournent sur le même split à graine fixée avec les mêmes métriques. AttentiveFP (RMSE de test 0.498, R² 0.926) et GIN (0.528, 0.916) mènent, suivis du MPNN (0.645), de GAT (0.769), de GraphSAGE (0.816), du DCNN (0.907) et du GCN (0.981). Les quatre modèles qui lisent les types de liaisons devancent les trois qui ne voient que les atomes.

La question

Comment sept architectures de réseaux de neurones sur graphes se comparent-elles pour prédire la solubilité moléculaire (ESOL) directement à partir des graphes moléculaires ?

Fonctionnalités clés

Sept GNN pour la régression au niveau du graphe, partageant une même interface : GCN, MPNN, DCNN, GIN, GAT, GraphSAGE et AttentiveFP.
Conversion SMILES vers graphe avec RDKit : 9 caractéristiques d'atome et 3 caractéristiques de liaison.
Modèles sensibles aux liaisons (MPNN, GIN, GAT, AttentiveFP) à côté de modèles atomes seuls (GCN, DCNN, GraphSAGE) pour une comparaison directe.
Deux commandes : mgnn-train entraîne n'importe quel modèle et rapporte les métriques de test, mgnn-predict évalue n'importe quelle chaîne SMILES.
75 tests pytest, dont l'invariance et l'équivariance par permutation et l'indépendance des éléments d'un batch, exécutés sur chaque modèle.
Split 80/10/10 à graine fixée avec MAE, MSE, RMSE et R² suivis à chaque époque, plus un arrêt anticipé sur la perte de validation.

Comment ça marche

Des molécules aux graphes
Les molécules ESOL sont chargées depuis MoleculeNet et représentées comme des graphes d'atomes et de liaisons, les molécules contenant des atomes isolés étant filtrées lorsque nécessaire.
Regrouper les graphes en batch
PyTorch Geometric fusionne de nombreux petits graphes moléculaires en un seul graphe non connexe par batch et garde la trace de l'appartenance de chaque atome à sa molécule.
Propager l'information
Chaque architecture met à jour les représentations des atomes à sa manière : convolution sur graphe, messages appris le long des liaisons, diffusion multi-sauts, somme puis MLP (GIN), attention sur les voisins (GAT, AttentiveFP) ou poids distincts pour soi et pour les voisins (GraphSAGE).
Agréger et prédire
Les représentations des atomes sont agrégées en un vecteur par molécule, qu'une tête de régression transforme en valeur de solubilité, entraînée avec l'erreur quadratique moyenne.

Exemples

AttentiveFP, le meilleur des sept : le checkpoint de meilleure validation (époque 222) obtient une RMSE de test de 0.498 et un R² de 0.926.

GIN avec caractéristiques de liaison : RMSE de test 0.528 et R² 0.916, proche d'AttentiveFP avec 25 fois moins de paramètres.

Métriques d'entraînement (gauche) et de validation (droite) du MPNN sur ESOL. Le checkpoint de meilleure validation (époque 195) obtient une RMSE de test de 0.645 et un R² de 0.875.

La première molécule d'ESOL (l'amygdaline), rendue avec RDKit : 32 atomes deviennent 32 nœuds de graphe avec 9 caractéristiques chacun.

Entrée

mgnn-predict --checkpoint checkpoints/mpnn/best.pt "CCO" "c1ccccc1"

Sortie

1.195 log(mol/L) CCO -1.633 log(mol/L) c1ccccc1

Sortie réelle du MPNN entraîné pour l'éthanol et le benzène. Leurs valeurs mesurées sont 1.10 et -1.64 ; les deux molécules font partie d'ESOL, ceci montre donc que l'outil fonctionne, pas sa généralisation.

Stack technique

PyTorch
Définition des modèles et boucles d'entraînement
PyTorch Geometric
Couches sur graphes (GCNConv, GINEConv, GATv2Conv, SAGEConv, AttentiveFP), passage de messages et batching
RDKit
Analyse des SMILES et inspection des molécules
MoleculeNet (ESOL)
Benchmark de solubilité moléculaire
pytest
Tests unitaires des modèles, du pipeline de données et de la boucle d'entraînement

Ce que j’ai appris

Les caractéristiques des liaisons comptent : chaque modèle qui lit les types de liaisons bat chaque modèle qui ne le fait pas
La somme bat la moyenne : l'agrégation par somme de GIN peut compter les voisins, ce qui lui permet de battre nettement GraphSAGE et GCN, fondés sur la moyenne
Les tests attrapent de vrais bugs : l'un a révélé un pooling moyen du DCNN qui incluait le padding, et une vérification en mode eval a montré que GraphSAGE avait appris à dépendre du dropout

Limites

Un seul jeu de données (ESOL) et une seule graine ; dépôt privé

Mots-clés

réseaux de neurones sur graphes · GNN · GCN · réseau convolutif sur graphes · MPNN · réseau à passage de messages · DCNN · AttentiveFP · GraphSAGE · réseau d'attention sur graphes · GAT · réseau d'isomorphisme de graphes · GIN · réseau de neurones à convolution de diffusion · PyTorch Geometric · PyTorch · RDKit · chémoinformatique · prédiction de propriétés moléculaires · graphes moléculaires · SMILES · MoleculeNet · ESOL · prédiction de solubilité aqueuse · découverte de médicaments · deep learning