TensorNet : Représentations par tenseurs cartésiens pour un apprentissage efficace des potentiels moléculaires

Le développement de modèles d’apprentissage automatique efficaces pour la représentation des systèmes moléculaires devient crucial dans la recherche scientifique. Nous introduisons TensorNet, une architecture innovante de réseau neuronal à passage de messages équivariante sous O(3), qui exploite des représentations tensorielles cartésiennes. Grâce à des embeddings atomiques tensoriels cartésiens, le mélange de caractéristiques est simplifié par des opérations de produit matriciel. De plus, la décomposition à faible coût de ces tenseurs en représentations irréductibles du groupe de rotation permet un traitement séparé des scalaires, vecteurs et tenseurs lorsque cela est nécessaire. Par rapport aux modèles de tenseurs sphériques de rang supérieur, TensorNet atteint des performances de pointe avec un nombre significativement réduit de paramètres. Pour l’énergie potentielle des petites molécules, cette performance peut être obtenue même avec une seule couche d’interaction. En raison de toutes ces propriétés, le coût computationnel du modèle est considérablement réduit. En outre, la prédiction précise de grandeurs moléculaires vectorielles et tensorielles, au-delà de l’énergie potentielle et des forces, devient possible. En résumé, le cadre de TensorNet ouvre une nouvelle voie pour la conception de modèles équivariants de pointe.