🤖 machine learning

Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning

Ce papier présente une méthode d'élagage structurel qui préserve l'équivariance SO(3) dans les modèles de base atomistiques en supprimant des blocs de représentations irréductibles, permettant ainsi de réduire considérablement le nombre de paramètres et les coûts de calcul tout en surpassant les modèles plus petits entraînés à partir de zéro tant en précision qu'en efficacité du fine-tuning en aval.

Auteurs originaux : Chen Wang, Siyu Hu, Guangming Tan, Weile Jia

Publié 2026-05-12

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Chen Wang, Siyu Hu, Guangming Tan, Weile Jia

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 avez un chef étoilé capable de préparer n'importe quel plat au monde avec une saveur parfaite. Ce chef est incroyablement talentueux, mais il est aussi immense : il porte un sac à dos rempli de milliers d'épices, de gigantesques marmites et d'outils complexes. Bien qu'il puisse préparer un repas parfait, il met beaucoup de temps à cuisiner, consomme énormément d'énergie et a besoin d'une cuisine gigantesque pour fonctionner.

Maintenant, imaginez que vous souhaitiez ouvrir un petit stand de restauration rapide. Vous n'avez pas besoin de l'ensemble massif de l'équipement du chef étoilé ; vous avez simplement besoin des compétences essentielles pour préparer rapidement d'excellents hamburgers et frites.

Cet article porte sur une nouvelle méthode permettant de prendre ce « chef étoilé » (un modèle d'intelligence artificielle géant et très précis pour prédire le comportement des atomes) et de le tailler chirurgicalement en un « chef de restauration rapide » (un modèle plus petit et plus rapide) sans perdre sa capacité à préparer des plats délicieux.

Voici le détail de leur méthode, en utilisant des analogies simples :

1. Le Problème : Le « Sac à Dos Lourd »

Les meilleurs modèles d'IA actuels pour la chimie (appelés modèles équivariants SO(3)) sont comme des chefs qui comprennent parfaitement la physique de la rotation. Si vous faites tourner une molécule, le modèle sait exactement comment les atomes se déplacent les uns par rapport aux autres. Cela les rend incroyablement précis.

Cependant, pour ce faire, ils portent un « sac à dos lourd » de mathématiques complexes (tenseurs d'ordre supérieur).

Le compromis : Plus le sac à dos est gros, plus la cuisine est précise, mais plus l'exécution est lente et coûteuse.
L'ancienne méthode : Si vous vouliez un chef plus petit, vous deviez généralement former un nouveau chef, tout petit, à partir de zéro. Ce nouveau chef commence sans expérience, met beaucoup de temps à apprendre et finit souvent par préparer de moins bons plats que le chef étoilé.

2. La Solution : « L'Élagage Structurel » (La Taille Chirurgicale)

Au lieu de former un nouveau chef à partir de zéro, les auteurs ont pris le chef étoilé existant et ont effectué un « élagage structurel ». Imaginez cela comme une coupe de cheveux très soignée ou une taille du sac à dos.

Le défi : Vous ne pouvez pas simplement retirer au hasard un bocal d'épices ou un outil. Dans ces modèles, les « outils » (les parties mathématiques) sont étroitement liés. Si vous coupez une pièce d'un outil, tout l'outil se brise, et le modèle perd sa capacité à comprendre la rotation (il cesse d'être « équivariant »).
L'innovation : Les auteurs ont trouvé comment retirer des blocs entiers d'outils à la fois. Ils traitent un ensemble spécifique d'outils liés comme une « unité atomique » unique. S'ils décident de retirer une unité, ils retirent tout le bloc, garantissant que les outils restants fonctionnent toujours parfaitement ensemble.

3. Comment Ils Ont Décidé Quoi Couper (Le « Test de Sensibilité »)

Comment savoir quels outils garder et lesquels jeter ? Vous ne pouvez pas simplement deviner.

Les auteurs ont utilisé un test astucieux basé sur l'énergie et la force :

Imaginez que le modèle tient une sculpture en verre délicate (la molécule).
Ils se sont demandé : « Si je retire cet outil spécifique, la sculpture se fissure-t-elle ou tombe-t-elle ? »
Ils ont mesuré dans quelle mesure la prédiction du modèle concernant « l'énergie » et la « force » changeait lorsqu'un outil était retiré.
La règle : Si retirer un outil modifie à peine le résultat, c'est un « outil paresseux » et il est coupé. Si le retirer fait trébucher le modèle, c'est un « outil critique » et il est conservé.

4. Le Processus : Une Recette en Quatre Étapes

L'article décrit un processus en quatre étapes pour transformer le modèle géant en un modèle compact :

Calibration : Exécuter quelques plats tests avec le chef étoilé pour voir quels outils sont réellement utilisés et quelle est leur importance.
Élagage : Sur la base du test, retirer chirurgicalement les blocs d'outils « paresseux ».
Reformation : Le modèle est maintenant plus petit et présente un vide là où se trouvaient les outils. On lui donne un rapide « cours de remise à niveau » sur une petite quantité de données pour l'aider à s'adapter à son nouveau corps plus petit.
Ajustement fin : Enfin, ils enseignent à ce nouveau modèle compact une tâche spécifique (comme prédire le comportement d'une molécule de médicament spécifique).

5. Les Résultats : Plus Rapide, Moins Cher, Toujours Délicieux

Les résultats étaient impressionnants :

Mieux que repartir de zéro : Le modèle élagué était plus précis qu'un nouveau petit modèle formé à partir de zéro, même s'ils avaient la même taille.
Économies massives :
- Entraînement : Il a fallu 2,5 à 4 fois moins de temps et d'argent informatiques pour créer le modèle élagué par rapport à l'entraînement d'un petit modèle à partir de zéro.
- Vitesse : Lors de l'utilisation du modèle pour faire des prédictions, il était 2,7 fois plus rapide.
- Mémoire : Il nécessisait considérablement moins de mémoire informatique (jusqu'à 5,7 fois moins).
Polyvalence : Cette méthode a fonctionné non seulement sur un type de modèle (MACE), mais aussi sur d'autres (SevenNet, eSCN), prouvant qu'il s'agit d'une recette générale pour ce type de chefs d'IA.

6. La Combinaison « Bonus »

L'article note également que cette méthode d'élagage peut être combinée avec d'autres astuces d'efficacité :

Quantification : Comme passer de la vidéo haute définition à la définition standard pour économiser de l'espace.
Distillation de connaissances : Comme demander au chef étoilé d'enseigner des astuces spécifiques au petit chef.
Lorsqu'elles sont combinées, ces méthodes rendent le modèle encore plus rapide et plus petit sans perdre en qualité.

Résumé

En bref, les auteurs ont trouvé un moyen de réduire la taille des modèles d'IA géants et coûteux utilisés pour la chimie en retirant soigneusement les parties inutiles tout en maintenant intact le « cerveau physique » central. Le résultat est un modèle plus petit, plus rapide et moins cher, qui reste plus intelligent que n'importe quel petit modèle que vous pourriez construire à partir de zéro. Cela rend la découverte de nouveaux matériaux avancés et la conception de médicaments accessibles à davantage de chercheurs qui ne disposent pas de supercalculateurs.

Résumé technique : Modèles de fondation atomistiques compacts équivariants SO(3) par élagage structurel

Énoncé du problème

Les réseaux de neurones graphiques équivariants SO(3) (GNN) se sont imposés comme le paradigme dominant pour les modèles de fondation atomistiques, atteignant une haute précision et une efficacité des données en codant explicitement la symétrie rotationnelle. Cependant, ces modèles reposent sur des opérations de produit tensoriel d'ordre supérieur pour coupler les caractéristiques géométriques, entraînant des coûts de calcul qui peuvent atteindre $O(L_{max}^6)$ , où $L_{max}$ est l'ordre maximal des représentations irréductibles. Cela crée un compromis significatif entre la précision du modèle et l'efficacité de l'inférence.

Les stratégies existantes de compression de modèles — spécifiquement la distillation de connaissances (KD) et la quantification — présentent des limites dans ce domaine. La KD nécessite souvent l'entraînement d'un élève à partir de zéro et ne réutilise pas directement les poids de l'enseignant, tandis que la quantification réduit la précision des bits mais laisse inchangé le nombre de paramètres et les chemins de produit tensoriel. De plus, les méthodes d'élagage structurel standard développées pour les grands modèles de langage (LLM) ne peuvent pas être directement appliquées aux modèles équivariants SO(3). Une suppression naïve des dimensions de caractéristiques brise la symétrie rotationnelle du chemin de produit tensoriel, violant la propriété d'équivariance fondamentale requise pour la cohérence physique. Les tentatives précédentes d'élagage de modèles atomistiques se sont limitées à une suppression de blocs grossiers (par exemple, des couches d'interaction entières), ce qui réduit le champ récepteur du modèle et est particulièrement préjudiciable aux architectures peu profondes.

Méthodologie

Les auteurs proposent un cadre d'élagage structurel spécifiquement conçu pour les modèles de fondation atomistiques équivariants SO(3). La méthode opère sur des points de contrôle pré-entraînés pour produire des modèles plus petits et denses qui conservent l'équivariance SO(3) et une haute précision.

1. Granularité de l'élagage équivariant

L'innovation centrale réside dans la définition de l'unité d'élagage. Dans les modèles équivariants SO(3), les caractéristiques sont décomposées en représentations irréductibles d'ordre $l \in [0, L_{max}]$ . Un vecteur de caractéristiques pour un canal spécifique $k$ et un ordre $l$ se compose de $2l+1$ composantes indexées par $m \in [-l, l]$ .

Contrainte : L'opération de rotation dans SO(3) couple toutes les composantes $m$ au sein d'un ordre $l$ spécifique. La suppression de composantes individuelles (élagage élément par élément) brise l'équivariance.
Solution : La méthode traite le sous-espace engendré par les $2l+1$ composantes d'un ordre $l$ et d'un canal $k$ spécifiques comme une unité atomique. L'élagage est effectué sur le tuple $(k, l)$ , représenté par un masque binaire $z^{(t)}_{k,l}$ . Cela garantit que si un bloc de caractéristiques est conservé, toutes ses composantes rotationnelles sont préservées, maintenant strictement l'équivariance SO(3).

2. Critère d'importance informé par la physique

Pour déterminer quels blocs $(k, l)$ conserver, les auteurs proposent un score d'importance basé sur la sensibilité de la surface d'énergie potentielle.

Objectif dual : Les MLIP doivent prédire à la fois l'énergie potentielle ( $E$ ) et les forces ( $F_i = -\nabla_{r_i} E$ ).
Métrique : Le score d'importance $I^{(t)}_{k,l}$ est dérivé d'un développement de Taylor du premier ordre du changement d'énergie lors de la suppression d'un bloc de caractéristiques. Il agrège le produit de la magnitude des caractéristiques et du gradient de l'énergie par rapport à cette caractéristique sur tous les atomes et les dimensions de représentation :
$I^{(t)}_{k,l} = \mathbb{E}_{R,Z \sim \mathcal{D}} \left| \sum_{j,m} \frac{\partial E}{\partial h^{(t)}_{j,klm}} \cdot h^{(t)}_{j,klm} \right|$
Avantage : Ce critère est invariant SO(3) et engendre des frais généraux négligeables car les gradients $\frac{\partial E}{\partial h}$ sont déjà calculés lors de l'inférence standard des forces.

3. Alignement structurel et pipeline

Le cadre suit un pipeline en quatre étapes :

Calibration : Un petit sous-ensemble de données de pré-entraînement est transmis au grand modèle pour calculer les scores d'importance.
Élagage piloté par la cible : Une architecture cible est définie. Les blocs $(k, l)$ les mieux classés sont conservés en fonction des scores d'importance.
Alignement structurel : La structure du modèle est physiquement modifiée pour correspondre aux dimensions réduites. Cela implique :
- Reconfiguration du produit tensoriel : Tranchage des poids radiaux et des chemins de Clebsch-Gordan pour s'aligner sur les canaux conservés.
- Tranchage des poids : Suppression physique des lignes et des colonnes des poids de transformation linéaire pour créer un tenseur de poids plus petit et dense.
Re-entraînement et affinement : Le modèle élagué subit un re-entraînement léger sur un sous-ensemble de données de pré-entraînement pour récupérer des erreurs d'approximation, suivi d'un affinement sur des tâches en aval spécifiques.

Contributions clés

Méthode d'élagage équivariant : Introduction d'une méthode d'élagage structurel pour les tenseurs équivariants SO(3) d'ordre supérieur à travers les dimensions de canal et d'ordre, qui préserve strictement l'équivariance SO(3) en traitant les blocs $(k, l)$ comme des unités atomiques.
Cadre en quatre étapes : Un pipeline complet (Calibration, Élagage, Re-entraînement, Affinement) pour la compression des modèles de fondation.
Généralisabilité : Démonstration que la méthode s'applique à plusieurs architectures, notamment MACE, SevenNet et eSCN.
Combinabilité : Validation que l'élagage structurel peut être combiné avec la quantification et la distillation de connaissances pour des gains d'efficacité cumulatifs.

Résultats expérimentaux

La méthode a été évaluée principalement sur MACE-MP (matériaux inorganiques) et MACE-OFF (molécules organiques), avec des tests de généralisabilité sur SevenNet et eSCN.

Précision vs entraînement à partir de zéro

Supériorité : Les modèles élagués surpassent systématiquement les modèles de même taille entraînés à partir de zéro.
- Sur MACE-MP, les modèles élagués ont obtenu un MAE de force plus faible et de meilleures performances hors distribution sur le jeu de données sAlex par rapport aux bases de référence entraînées à partir de zéro.
- Sur Matbench Discovery, un modèle MACE-MP élagué (n'utilisant que 25 % des données de pré-entraînement) a surpassé le petit modèle officiel entraîné à partir de zéro sur 7 des 9 métriques.
- Sur MACE-OFF, les modèles élagués ont égalé ou dépassé les bases de référence entraînées à partir de zéro sur divers sous-ensembles moléculaires, même en n'utilisant que 50 % des données de pré-entraînement.
Performances en aval : L'affinement des modèles élagués sur huit jeux de données inorganiques en aval a réduit les erreurs d'énergie et de force de 70,1 % et 34,4 %, respectivement, par rapport à l'entraînement de modèles spécifiques à la tâche à partir de zéro.

Gains d'efficacité

Réduction des paramètres : Les modèles compressés contiennent 1,5 à 4 fois moins de paramètres.
Coût d'entraînement : Le pipeline complet nécessite 2,5 à 4 fois moins de calcul de pré-entraînement que l'entraînement d'un petit modèle à partir de zéro (par exemple, 181 heures-GPU contre 750 heures-GPU pour une configuration spécifique).
Vitesse d'inférence : Les modèles élagués atteignent une accélération de l'inférence de 2,7 à 3,0 fois et jusqu'à une réduction de 5,7 fois de l'utilisation maximale de la mémoire.

Combinaison avec d'autres méthodes

Quantification : La combinaison de l'élagage avec la quantification FP32 a produit une amélioration du débit de 9,6 fois (par rapport au modèle de base FP64).
Distillation de connaissances : La combinaison de l'élagage avec une KD basée sur l'hessienne a affiné davantage les performances du modèle, démontrant que les deux méthodes sont complémentaires.

Importance et revendications

L'article revendique que l'élagage structurel offre une voie naturelle et efficace pour combler le fossé précision-efficacité dans les modèles de fondation atomistiques. En héritant directement des poids de points de contrôle pré-entraînés de grande taille et en les compressant tout en préservant les symétries physiques, la méthode produit des modèles compacts qui sont :

Plus précis que les modèles de même taille entraînés à partir de zéro.
Significativement moins chers à entraîner et à déployer.
Physiquement cohérents, maintenant la stricte équivariance SO(3) requise pour des simulations de dynamique moléculaire fiables (validé via des simulations MD d'eau liquide montrant une dérive d'énergie stable et des fonctions de distribution radiale précises).

Les auteurs soulignent que le principe fondamental — évaluer l'importance des représentations irréductibles couplées — est largement applicable à toute architecture équivariante représentant des caractéristiques comme des sommes de représentations irréductibles SO(3). Le travail suggère que l'élagage permet aux praticiens de déployer des modèles de fondation de haute fidélité dans des environnements à ressources contraintes sans sacrifier la transférabilité apprise lors du pré-entraînement à grande échelle.