Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement

Cette étude introduit la décomposition de sonde compositionnelle (CPD) pour démontrer que l'alignement de la tâche d'entraînement, plutôt que l'architecture du modèle, détermine la séparation linéaire des informations géométriques et compositionnelles dans les modèles de fondation atomistiques, tout en révélant un routage de l'information par symétrie au sein des représentations MACE.

L'essentiel

Imaginez que vous essayez de comprendre comment un chef cuisinier (l'intelligence artificielle) prépare un plat complexe. Vous voulez savoir si le goût final dépend de ce qu'il y a dans l'assiette (les ingrédients : composition) ou de la façon dont ils sont disposés (la forme, la géométrie).

Le problème, c'est que dans la chimie moléculaire, ces deux choses sont souvent mélangées. Si vous avez un gâteau au chocolat, le goût dépend à la fois du chocolat (ingrédient) et de la façon dont il est cuit (forme).

Cette recherche, menée par Joshua Steier, a inventé une nouvelle méthode pour séparer ces deux saveurs dans les cerveaux numériques des modèles d'IA. Voici l'explication simple, avec des analogies :

1. Le Problème : Le "Brouillard" des Ingrédients

Les modèles d'IA modernes (comme MACE, PaiNN, SchNet) sont excellants pour prédire les propriétés des molécules. Mais quand on regarde à l'intérieur de leur "cerveau" (leurs représentations), on ne sait pas toujours si ils ont appris la forme des molécules ou s'ils ont juste mémorisé la liste des ingrédients.

C'est comme si un élève apprenait l'histoire en mémorisant les dates (ingrédients) sans comprendre les causes et les effets (la forme de l'événement). Si on lui demande de prédire un événement futur, il risque de se tromper s'il n'a pas compris la logique.

2. La Solution : La "Machine à Séparer le Sel" (CPD)

L'auteur a créé un outil appelé CPD (Décomposition par Sonde Compositionnelle).
Imaginez que vous avez un verre d'eau salée (la représentation de la molécule). Vous voulez savoir combien de "forme" (géométrie) reste après avoir retiré tout le "sel" (les ingrédients).

• L'ancienne méthode (fausse) : On utilisait des outils très puissants (des arbres de décision) pour essayer de deviner ce qui restait. Résultat : ils étaient trop intelligents ! Ils arrivaient à reconstituer le sel qu'on avait retiré en utilisant des astuces mathématiques complexes. C'était comme essayer de goûter l'eau douce alors que le détecteur de sel trouvait encore du sel parce qu'il "imaginait" le goût.
• La nouvelle méthode (CPD) : On utilise une méthode très simple et stricte (une régression linéaire) pour retirer le sel. Ensuite, on demande à un élève très simple (un détecteur linéaire) de deviner la forme restante. Si l'élève devine bien, c'est que la forme est bien rangée et accessible.

3. La Grande Découverte : Ce n'est pas l'architecture, c'est l'entraînement !

L'étude a testé 10 modèles différents. Ils s'attendaient à ce que les modèles les plus "avancés" (ceux qui respectent les lois de la physique, appelés équivariants) soient les meilleurs.

La surprise ? Ce n'est pas la complexité du modèle qui compte, c'est ce qu'on lui a appris à faire.

• L'analogie du Coach :
  • Si vous entraînez un athlète à courir en ligne droite (prédire l'énergie), il deviendra excellent pour courir en ligne droite, mais il sera nul pour faire des figures de gymnastique (prédire la forme précise).
  • Si vous l'entraînez à faire de la gymnastique (prédire les propriétés électroniques comme le "trou HOMO-LUMO"), il développera une compréhension fine de la forme.
• Le résultat : Les modèles entraînés spécifiquement sur la "forme" des molécules étaient 6,6 fois meilleurs pour révéler cette information que ceux entraînés uniquement sur l'énergie, même s'ils utilisaient la même architecture de base.

4. Le Tri des Informations : Les "Couloirs" de l'IA

Dans certains modèles très avancés (comme MACE), l'information voyage dans des "couloirs" différents selon sa nature, un peu comme des rails de train :

• Les informations qui sont des chiffres simples (comme l'énergie) voyagent dans des couloirs "scalaires".
• Les informations qui ont une direction (comme un aimant ou un dipôle) voyagent dans des couloirs "vectoriels".

C'est comme si le modèle avait organisé son bureau : les papiers importants sont dans le tiroir du haut, les outils dans le tiroir du bas. Cette organisation intelligente n'apparaît que si le modèle a été entraîné avec le bon objectif.

5. Pourquoi c'est important pour vous ?

• Pour les chercheurs : Si vous voulez utiliser une IA pour découvrir de nouveaux médicaments ou matériaux, ne choisissez pas le modèle le plus complexe. Choisissez celui qui a été entraîné sur le type de problème qui vous intéresse.
• Pour la science : Cela nous apprend que l'IA ne "pense" pas comme nous. Elle organise l'information différemment selon ce qu'on lui demande. Si on veut qu'elle soit intelligente sur la forme, il faut lui demander de jouer avec la forme, pas juste avec les ingrédients.

En résumé :
Cette recherche nous dit que pour avoir une IA qui comprend vraiment la forme des molécules, il ne suffit pas de lui donner un cerveau complexe. Il faut lui donner le bon "devoir" à faire. C'est l'entraînement, pas la structure, qui détermine si l'IA verra la forme ou seulement les ingrédients. Et attention : ne faites pas confiance aux outils trop complexes pour vérifier ce qu'elle a appris, car ils risquent de vous mentir en reconstituant ce que vous aviez caché !

Résumé technique

Voici un résumé technique détaillé de l'article de recherche intitulé "Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement".

1. Problématique

Les modèles de base atomistiques (comme MACE, SchNet, PaiNN, ViSNet) prédisent avec une grande précision les propriétés moléculaires. Cependant, une question fondamentale reste ouverte : comment ces modèles organisent-ils leurs représentations internes ?

Plus spécifiquement, les représentations séparent-elles proprement la composition (quels éléments sont présents et en quelles proportions) de la géométrie (comment ces atomes sont arrangés dans l'espace) ?

• Le défi : Les propriétés moléculaires dépendent souvent des deux facteurs, qui sont fortement corrélés. Une sonde (probe) standard ne peut pas distinguer si un modèle a appris la géométrie ou s'il a simplement mémorisé la composition.
• L'erreur méthodologique précédente : Les auteurs montrent que l'utilisation de sondes non linéaires (comme les arbres de décision ou les GBT) sur des représentations "résiduelles" (après soustraction de la composition) conduit à des résultats biaisés. Ces modèles non linéaires sont capables de reconstruire le signal de composition éliminé via des interactions de caractéristiques complexes, faussant ainsi les mesures d'accessibilité géométrique.

2. Méthodologie : Décomposition par Sonde Compositionnelle (CPD)

Pour répondre à ce problème, les auteurs introduisent la Compositional Probe Decomposition (CPD), une méthode rigoureuse pour isoler l'information géométrique.

• Principe de base :

  1. Extraction : On extrait les représentations $X$ d'un modèle figé (frozen) pour un ensemble de molécules.
  2. Projection OLS : On projette linéairement la représentation pour retirer la composante liée à la composition ($Z$). Cela se fait via une régression des moindres carrés ordinaires (OLS) : $\hat{\beta} = (Z^\top Z)^{-1}Z^\top X$.
  3. Résidu Géométrique : Le résidu $X_{geom} = X - Z\hat{\beta}$ contient l'information linéairement orthogonale à la composition (topologie, géométrie, connexions).
  4. Sondage Linéaire : On applique une régression Ridge (linéaire) sur ce résidu pour prédire une propriété cible. La performance ($R^2_{geom}$) mesure l'information géométrique linéairement accessible.

• Validation rigoureuse :

  • Benchmark d'isomères structuraux : Utilisation de paires d'isomères (même composition, géométrie différente). La composante compositionnelle doit donner un score au hasard (50%), tandis que le résidu géométrique doit permettre une classification précise.
  • Contrôle des sondes non linéaires : Démonstration que les GBTs récupèrent artificiellement un $R^2$ élevé (0.68–0.95) sur des cibles purement compositionnelles à partir des résidus, prouvant leur inadéquation pour cette tâche.
  • Robustesse : Douze vérifications (changement de features de composition, projection globale vs fold-wise, élimination de concept LEACE) confirment que le classement des modèles est invariant.

3. Contributions Clés

1. Méthodologie CPD : Une nouvelle méthode validée pour mesurer la désentanglement linéaire entre composition et géométrie, incluant la preuve que les sondes non linéaires surévaluent systématiquement l'information dans les résidus.
2. Gradient d'accessibilité linéaire : Identification d'un gradient de performance sur 10 modèles (5 familles d'architectures) variant d'un facteur 6,6x ($R^2_{geom}$ allant de 0,081 à 0,533 pour le gap HOMO-LUMO).
3. Hiérarchie des facteurs : Démonstration que l'alignement de la tâche (training objective) est le facteur dominant, surpassant l'équivariance et la diversité des données.
4. Routage d'information par symétrie : Mise en évidence que dans l'architecture MACE, les canaux scalaires ($L=0$) et vectoriels ($L=1$) encodent préférentiellement des propriétés correspondant à leur type de symétrie (ex: moment dipolaire dans les canaux vectoriels).

4. Résultats Principaux

A. Le Gradient d'Accessibilité Linéaire

Les modèles varient considérablement dans la quantité d'information géométrique accessible après suppression de la composition :

• Meilleurs modèles : PaiNN, ViSNet et MACE-HL (entraînés sur le gap HOMO-LUMO) atteignent des scores élevés ($R^2_{geom} \approx 0,44 - 0,53$).
• Pires modèles : MACE entraîné uniquement sur QM9 (énergie) atteint un score très bas ($R^2_{geom} \approx 0,08$).

B. Les Trois Facteurs Explicatifs

1. Alignement de la Tâche (Dominant) : C'est le facteur le plus important. Les modèles entraînés sur une propriété sensible à la géométrie (comme le gap HOMO-LUMO) développent des représentations où l'information géométrique est linéairement accessible.
   • Preuve causale : En réentraînant PaiNN et MACE sur l'énergie au lieu du gap HOMO-LUMO, le $R^2_{geom}$ chute drastiquement (ex: MACE passe de 0,44 à 0,08). L'architecture compte moins que l'objectif d'apprentissage.
2. Équivariance (Conditionnelle) : L'équivariance seule ne garantit pas une bonne accessibilité. Un modèle équivariant mal entraîné (MACE sur énergie) performe moins bien qu'un modèle invariant bien entraîné (SchNet sur énergie). L'équivariance amplifie le signal seulement si la tâche est alignée.
3. Diversité des Données (Compensation partielle) : Un pré-entraînement sur des données diversifiées (MPTraj pour MACE) améliore l'accessibilité géométrique ($R^2_{geom} = 0,36$) par rapport à un entraînement sur QM9 uniquement, mais ne comble pas totalement l'écart avec les modèles alignés sur la tâche.

C. Routage par Représentation Irréductible (MACE)

Dans MACE, l'information est routée selon le type de symétrie :

• Les propriétés scalaires (Gap HOMO-LUMO) sont encodées principalement dans les canaux scalaires ($L=0$).
• Les propriétés vectorielles (Moment dipolaire) sont encodées principalement dans les canaux vectoriels ($L=1$).
• Ce phénomène n'est pas observé dans ViSNet, suggérant que le routage structuré est spécifique aux architectures utilisant des produits tensoriels explicites de harmoniques sphériques.

D. Efficacité Échantillonnaire

Les représentations bien désentanglées (comme PaiNN) nécessitent beaucoup moins d'exemples pour extraire le signal géométrique. PaiNN atteint un score supérieur à SchNet avec seulement 50 molécules, là où SchNet en nécessite 2000.

5. Signification et Impact

• Pour la sélection de modèles : Pour les tâches en aval sensibles à la géométrie (propriétés électroniques, réactivité), il est plus crucial de choisir un modèle entraîné sur un objectif aligné que de privilégier une architecture équivariante spécifique.
• Pour l'interprétabilité : Les modèles comme MACE offrent une décomposition structurelle naturelle (canaux scalaires vs vectoriels) qui facilite l'analyse des mécanismes de prédiction.
• Pour la méthodologie : L'article met en garde contre l'utilisation de sondes non linéaires sur des représentations résiduelles, un problème qui pourrait affecter d'autres domaines comme le NLP ou la vision par ordinateur.
• Généralisation : Les résultats se maintiennent sur des cristaux périodiques (Materials Project), suggérant que la séparation composition-géométrie est une propriété fondamentale des représentations atomistiques, au-delà des petites molécules organiques.

En résumé, cet article démontre que l'objectif d'entraînement est le levier principal pour organiser les représentations internes des modèles atomistiques, permettant un accès linéaire et efficace à l'information géométrique, tandis que l'architecture joue un rôle secondaire et conditionnel.