Passage de messages directionnel pour les graphes moléculaires

Les réseaux de neurones graphiques ont récemment connu un grand succès dans la prédiction des propriétés mécaniques quantiques des molécules. Ces modèles représentent une molécule sous forme de graphe en ne tenant compte que des distances entre les atomes (nœuds). Ils ne prennent toutefois pas en compte la direction spatiale d’un atome vers un autre, bien que cette information directionnelle joue un rôle central dans les potentiels empiriques des molécules, par exemple dans les potentiels angulaires. Pour pallier cette limitation, nous proposons une transmission de messages directionnels, dans laquelle les messages échangés entre atomes sont eux-mêmes encodés, plutôt que les atomes eux-mêmes. Chaque message est associé à une direction dans l’espace cartésien. Ces encodages de messages directionnels sont équivariants par rotation, car les directions associées tournent avec la molécule. Nous proposons un schéma de transmission de messages analogue à la propagation de croyances, qui exploite l’information directionnelle en transformant les messages en fonction de l’angle entre eux. En outre, nous utilisons des fonctions de Bessel sphériques et des harmoniques sphériques pour construire des représentations théoriquement fondées et orthogonales, qui surpassent les représentations radiales gaussiennes actuellement prévalentes en termes de performance tout en utilisant moins d’un quart du nombre de paramètres. En exploitant ces innovations, nous avons conçu le réseau de neurones à transmission de messages directionnels, DimeNet. DimeNet dépasse en moyenne les GNN précédents de 76 % sur MD17 et de 31 % sur QM9. Notre implémentation est disponible en ligne.