🤖 AI

E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

E2Former-V2 est une architecture de réseau de neurones sur graphes équivariante et évolutive qui surmonte les goulots d'étranglement computationnels des modèles dominants en introduisant une sparsification équivariante alignée sur les axes et un mécanisme d'attention équivariante à la volée, atteignant une amélioration de 20 fois le TFLOPS tout en maintenant la performance prédictive sur les systèmes atomistiques 3D.

Auteurs originaux : Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Publié 2026-06-08

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayiez de prédire comment une foule immense de personnes (des atomes) va se déplacer et interagir dans une pièce. Dans le monde de la chimie et de la physique, on appelle cela la « modélisation atomistique ». Pendant longtemps, les scientifiques ont utilisé un type spécial d'IA appelé Réseau de Neurones sur Graphes Équivariant (EGNN). Ces réseaux sont excellents car ils comprennent que si l'on fait pivoter toute la pièce, la physique à l'intérieur ne doit pas changer.

Cependant, les versions actuelles de ces réseaux ont un problème majeur : elles sont maladroites et lentes quand la foule devient grande.

Voici une décomposition simple de la solution du papier, E2Former-V2, en utilisant des analogies de la vie quotidienne.

Le Problème : L'embouteillage « Centré sur les Arêtes »

Imaginez une fête où tout le monde veut discuter avec ses voisins.

L'ancienne méthode (EGNN traditionnels) : Le système crée une liste physique géante de chaque conversation se déroulant à la fête. S'il y a 1 000 personnes et que chacune parle à 50 voisins, le système doit noter 50 000 conversations distinctes sur un immense tableau blanc (mémoire) avant de pouvoir faire le moindre calcul.
Le goulot d'étranglement : À mesure que la fête grandit (plus d'atomes), ce tableau blanc devient si énorme que l'ordinateur manque d'espace. Il passe tout son temps à simplement déplacer ces notes plutôt qu'à réellement réfléchir. C'est pourquoi les anciens modèles plantent ou ralentissent considérablement lorsqu'ils simulent de grosses molécules comme les protéines.

La Solution : E2Former-V2

Les auteurs ont construit un nouveau système qui arrête d'écrire chaque note de conversation. À la place, il utilise deux astuces ingénieuses pour rendre le processus ultra-rapide.

Astuce 1 : La « Boussole Magique » (Équivariance par Sparsification Alignée sur les Axes)

Dans l'ancien système, calculer comment les atomes interagissent revient à essayer de résoudre un puzzle où chaque pièce peut être pivotée dans n'importe quelle direction. C'est désordonné et cela nécessite beaucoup de calculs.

L'analogie : Imaginez que vous avez un puzzle 3D complexe. Au lieu d'essayer de le résoudre pendant qu'il tourne, vous utilisez une « Boussole Magique » pour faire pivoter tout le puzzle de sorte qu'une pièce spécifique pointe toujours vers le haut.
Le résultat : Une fois le puzzle aligné, la plupart des pièces deviennent vides ou nulles. Vous n'avez plus besoin de faire de calculs sur les parties vides. Le papier appelle cela l'EAAS. Cela transforme un calcul dense et désordonné en un simple jeu de « ré-indexation » (déplacer simplement les pièces), ce qui est beaucoup plus rapide.

Astuce 2 : Le « Chef à la Carte » (Attention Équivariante « On-the-Fly »)

C'est la plus grande percée.

L'ancienne méthode (Le Buffet) : L'ancien système prépare tout le buffet (les 50 000 notes de conversation) et le pose sur une table avant que quiconque puisse manger. Si la table est trop petite, la nourriture déborde et la fête s'arrête.
La nouvelle méthode (Le Chef à la Carte) : Le nouveau système, E2Former-V2, agit comme un chef qui cuisine une assiette à la fois, juste devant le client, et la sert immédiatement.
- Il ne note pas les conversations.
- Il ne stocke pas la liste entière.
- Il calcule l'interaction entre deux personnes, ajoute le résultat au score final, puis oublie immédiatement les détails pour faire de la place pour la paire suivante.
Le résultat : C'est ce qu'on appelle le traitement « On-the-Fly » (à la volée). Comme il ne construit jamais le géant tableau blanc, il utilise presque aucune mémoire. Il fait tenir tout le processus dans la mémoire cache ultra-rapide de l'ordinateur (SRAM), ce qui le rend 20 fois plus rapide que les anciennes méthodes.

Ce qu'ils ont réellement accompli

Les auteurs ont testé ce nouveau système sur deux éléments principaux :

Petites Molécules (SPICE) : Ils l'ont testé sur des molécules de type médicamenteux. Le nouveau système était aussi précis que les meilleurs modèles existants, mais beaucoup plus rapide.
Grosses Molécules (OMol25) : Ils l'ont testé sur des structures chimiques massives.
- La limite : Les anciens modèles plantaient (manque de mémoire) lorsqu'ils essayaient de simuler des systèmes de plus de 10 000 atomes.
- La victoire : E2Former-V2 a réussi à simuler des systèmes de 100 000 atomes sans planter.

Tests en conditions réelles mentionnés

Les auteurs n'ont pas seulement fait des calculs ; ils ont simulé de la vraie physique :

Eau liquide : Ils ont simulé une boîte de molécules d'eau. Le nouveau modèle a prédit comment les molécules d'eau s'organisent (la fonction de distribution radiale) de manière presque parfaite, correspondant mieux aux expériences réelles que les modèles précédents.
Protéines : Ils ont simulé une grosse protéine (environ 30 000 atomes) flottant dans l'eau. Le modèle a maintenu la stabilité de la protéine et l'a aidée à se replier selon une forme correspondant aux données expérimentales réelles, prouvant qu'il peut gérer des structures biologiques complexes.

Résumé

E2Former-V2 est comme passer d'un système qui écrit chaque conversation sur un immense et lent tableau blanc à un système où un chef intelligent cuisine et sert les conversations une par une, instantanément. Cela permet aux scientifiques de simuler des systèmes moléculaires massifs (comme de grosses protéines) sur des ordinateurs standards sans manquer de mémoire, tout en conservant la précision de la physique.

Résumé Technique : E2Former-V2 : Attention Équivariante « On-the-Fly » avec Mémoire d'Activation Linéaire

1. Énoncé du Problème

Les réseaux de neurones sur graphes équivariants (EGNN) sont devenus la norme pour la modélisation des systèmes atomistiques 3D en raison de leur capacité à préserver les symétries de rotation et de translation. Cependant, les architectures dominantes (par exemple, MACE, EquiformerV2, eSCN) souffrent de goulots d'étranglement critiques en matière d'efficacité, ancrés dans leur conception centrée sur les arêtes (edge-centric). Ces modèles construisent explicitement des caractéristiques géométriques ou effectuent des produits tensoriels denses sur chaque arête, ce qui entraîne :

Complexité Mémoire : Une complexité de mémoire d'activation de $O(|E|)$ , où $|E| \approx kN$ (avec $k$ étant le nombre moyen de voisins). Pour les systèmes de grande taille, cela se traduit par des empreintes mémoire massives qui dépassent la capacité des GPU.
Inefficacité Matérielle : La matérialisation explicite des tenseurs d'arêtes force le mouvement de données entre la mémoire à haute bande passante (HBM) et la SRAM sur puce, créant un mur de mémoire qui limite le débit.
Limites de Scalabilité : À mesure que la taille du système ( $N$ ) augmente, la latence croît mal, rendant les simulations de dynamique moléculaire (MD) à grande échelle prohibitive sur du matériel grand public.

Bien que les avancées récentes dans les Transformers (par exemple, FlashAttention) aient résolu des goulots d'étranglement similaires pour les données séquentielles via une exécution par flux et par blocs (tile-based), ces principes sont restés inexplorés dans les architectures SO(3)-équivariantes.

2. Méthodologie

E2Former-V2 répond à ces problèmes en intégrant une sparsification algébrique à un modèle d'exécution centré sur les nœuds (node-centric) conscient du matériel. L'architecture repose sur deux innovations centrales :

2.1. Sparsification Axis-Aligned Équivariante (EAAS)

Les contractions de tenseurs SO(3) standard impliquent des sommes denses sur les coefficients de Clebsch–Gordan, résultant en une complexité de $O(L^6)$ et des schémas d'accès mémoire irréguliers. E2Former-V2 introduit l'EAAS pour transformer ces opérations denses en opérations de sparsification basées sur des permutations :

Changement de Base SO(3) $\to$ SO(2)$ : La méthode aligne le système de coordonnées global avec le vecteur local $\vec{r}$ (l'axe reliant deux atomes). Dans ce référentiel aligné, les harmoniques sphériques solides deviennent maximalement creuses (Lemme 4.1), où seul le composant $m=0$ est non nul.
Opérateur de Ré-indexation : En faisant commuter les rotations à travers le produit tensoriel, la contraction dense se réduit en une règle de ré-indexation déterministe (Éq. 13). Cet opérateur mappe les caractéristiques d'entrée vers les caractéristiques de sortie via des applications linéaires par blocs creux, éliminant le besoin de sommation dense sur les indices magnétiques.
Résultat : Cela réduit la complexité de l'étape de convolution et permet un accélération d'environ $\sim 1,5\times$ par rapport aux implémentations standards.

2.2. Attention Équivariante « On-the-Fly »

Pour éliminer le goulot d'étranglement de mémoire $O(|E|)$ , les auteurs proposent un mécanisme d'attention entièrement fusionné et centré sur les nœuds, implémenté via un noyau (kernel) Triton personnalisé :

Factorisation Centrée sur les Nœuds : Au lieu de matérialiser les caractéristiques d'arêtes, le mécanisme d'attention découple le calcul. Les vecteurs Query et Key sont construits à l'aide de projections scalaires invariantes SO(3) (Éq. 6), tandis que l'information géométrique équivariante est propagée uniquement via le chemin de la valeur (value path).
Réduction par Flux (Streaming Reduction) : Le noyau traite les voisins de manière séquentielle (streaming). Il itère sur les voisins $j \in \mathcal{N}(i)$ , calculant les scores d'attention et accumulant les mises à jour de valeurs sur le vif au sein de la SRAM.
Élimination de la Matérialisation des Arêtes : Les tenseurs intermédiaires de forme $N \times K \times H \times D$ ne sont jamais écrits dans la HBM. Le noyau maintient des statistiques courantes (max, normalisation et accumulateurs de valeurs) sur puce, réduisant la complexité de la mémoire d'activation de $O(|E|)$ à $O(|V|)$ .
Softmax en Ligne : Pour garantir la stabilité numérique lors du streaming, une formulation de softmax en ligne est utilisée pour mettre à jour les accumulateurs de manière incrémentale.

3. Contributions Clés

Sparsification Axis-Aligned Équivariante (EAAS) : Une nouvelle technique algébrique qui unifie le recouplage Wigner-6j avec l'alignement d'axe SO(2), convertissant les contractions de tenseurs SO(3) denses en opérations de ré-indexation creuses.
Attention Équivariante « On-the-Fly » : Un noyau Triton entièrement fusionné qui évite la matérialisation des arêtes, maximisant la réutilisation de la SRAM et atteignant jusqu'à 20 $\times$ plus de TFLOPS par rapport aux implémentations standards.
Performance Scalable : Capacité démontrée à gérer des systèmes allant jusqu'à 100 000 atomes sur un seul GPU, un régime où les précédents Transformers équivariants échouent en raison d'erreurs de type Out-Of-Memory (OOM).
Fidélité Physique : Validation en dynamique moléculaire montrant une reproduction précise des structures de l'eau liquide et un raffinement stable des structures protéiques prédites par AlphaFold ( $\sim$ 30k atomes).

4. Résultats Expérimentaux

Les auteurs ont évalué E2Former-V2 sur les jeux de données SPICE (chimie médicinale) et OMol25 (échelle de modèle de fondation), ainsi qu'en simulations de MD.

Précision Prédictive :
- Sur SPICE, E2Former-V2 a obtenu les plus faibles erreurs absolues moyennes (MAE) d'énergie et de force sur la plupart des sous-ensembles, surpassant MACE-Large de 48 % en MAE d'énergie sur les dimères.
- Sur OMol25, la variante conservatrice a atteint une MAE d'énergie agrégée de 1,27 meV/atome, égalant eSEN-small et surpassant nettement UMA-S (3,62 meV/atome).
- Sur OC20 S2EF-2M, il a obtenu une précision comparable à des bases solides comme EquiformerV2 et eSCN.
Efficacité et Scalabilité :
- Débit (Throughput) : Le noyau fusionné a atteint un accélération moyenne de 20 $\times$ en TFLOPS par rapport aux implémentations PyTorch naïves.
- Mémoire : L'utilisation de la mémoire de pointe a été réduite de 16,9 $\times$ à 19,5 $\times$ par rapport aux modèles de référence.
- Scalabilité : Alors que les modèles précédents (ex. EquiformerV2, MACE-Large) rencontraient des échecs OOM au-delà de 10k–50k atomes, E2Former-V2 a scalé de manière fiable jusqu'à 100k atomes.
- Vitesse d'Inférence : À 100k atomes, E2Former-V2 (Conservative) a atteint 0,12 étape/s, soit près de 3 $\times$ plus vite que UMA-S. Dans le réglage Direct, il est un ordre de grandeur plus rapide qu'Allegro et EquiformerV2 à plus petite échelle.
Dynamique Moléculaire :
- Les simulations de l'eau liquide périodique (216 molécules) ont montré qu'E2Former-V2 correspond étroitement aux fonctions de distribution radiale (RDF) expérimentales, surpassant MACE-OFF.
- Une simulation de 2 ns d'une protéine solvatée ( $\sim$ 30k atomes) a maintenu des scores TM élevés par rapport à la structure native, démontant une dynamique stable.

5. Signification et Revendications

L'article affirme qu'E2Former-V2 démontre la faisabilité pratique des grands Transformers équivariants sur des GPU grand public. En déplaçant le paradigme d'exécution d'un mode centré sur les arêtes vers un mode centré sur les nœuds et en exploitant un streaming conscient du matériel, la méthode résout le goulot d'étranglement fondamental de la mémoire qui a limité la scalabilité des modèles équivariants.

Les auteurs soulignent qu'il ne s'agit pas simplement d'une amélioration algorithmique, mais d'une réalisation au niveau système qui comble le fossé entre la factorisation théorique (réduisant la complexité de $O(|E|)$ à $O(|V|)$ ) et l'efficacité matérielle. Ce travail suggère que les modèles de fondation moléculaires à grande échelle et les simulations de MD à longue durée temporelle, auparavant prohibitifs sur le plan computationnel, sont désormais accessibles en utilisant une infrastructure GPU standard. Le code est publié pour faciliter l'adoption ultérieure par la communauté.