Scaling Laws and Symmetry, Evidence from Neural Force Fields

Ce papier démontre que l'intégration des symétries de tâche via des architectures équivariantes, en particulier celles dotées de représentations d'ordre supérieur, améliore considérablement les lois d'échelle pour les potentiels interatomiques, suggérant que les biais inductifs fondamentaux devraient être explicitement conçus plutôt que laissés aux modèles pour les découvrir à mesure qu'ils évoluent.

L'essentiel

Imaginez que vous essayiez d'enseigner à un robot comment prédire le mouvement et les interactions des atomes dans une molécule. C'est un peu comme enseigner à un enfant à comprendre comment une structure complexe en Lego tient ensemble. Vous pouvez donner au robot deux types différents de manuels d'instructions :

1. Le manuel « Aveugle » : Vous montrez simplement au robot des millions d'images de structures en Lego et vous lui dites : « Déduis les règles toi-même. » Le robot doit tout apprendre à partir de zéro, y compris le fait que si vous faites pivoter l'ensemble de la structure, la physique ne change pas.
2. Le manuel « Symétrie » : Vous donnez au robot un manuel qui indique explicitement : « Hé, souviens-toi, si tu fais tourner cette structure, c'est toujours la même structure. Si tu la retournes, les règles restent les mêmes. » Vous intégrez directement les lois de la physique (la symétrie) dans le cerveau du robot.

Pendant longtemps, de nombreux chercheurs ont cru en l'approche « Aveugle ». Ils pensaient que si l'on donnait simplement au robot suffisamment de données et de puissance de calcul (un « cerveau plus gros »), il finirait par déduire les règles de symétrie par lui-même. Ils croyaient qu'enseigner explicitement les règles était inutile et qu'un modèle simple et flexible finirait par rattraper son retard.

Ce papier dit : « En fait, non. Le manuel « Symétrie » est bien meilleur, et l'écart s'élargit à mesure que vous grandissez. »

Voici une analyse détaillée de leurs découvertes utilisant des analogies simples :

1. La course : Vitesse contre Efficacité

Les chercheurs ont organisé une course entre différents types de cerveaux de robots (architectures) pour voir à quelle vitesse ils pouvaient apprendre à prédire les forces atomiques.

• Les robots « Aveugles » (Non contraints) : Ils sont flexibles mais inefficaces. Ils doivent « réapprendre » le fait qu'une molécule tournée est la même molécule à chaque fois qu'ils la voient.
• Les robots « Symétrie » (Équivariants) : Ils ont les règles de rotation et de translation intégrées. Ils ne gaspillent pas d'énergie à réapprendre les bases de la physique.

La découverte : Lorsque les robots étaient petits, la différence n'était pas énorme. Mais à mesure que les chercheurs rendaient les robots massifs (en augmentant les données et la puissance de calcul), les robots « Symétrie » ne se contentaient pas de rester en tête ; ils prenaient une avance spectaculaire. Les robots « Aveugles » butaient contre un mur où l'ajout de plus de données ne les aidait guère, tandis que les robots « Symétrie » continuaient de devenir de plus en plus intelligents.

2. Le « degré » de symétrie compte

Tous les robots « Symétrie » ne se valent pas. Certains ne comprennent que des rotations simples (comme une pièce de monnaie plate), tandis que d'autres comprennent des rotations 3D complexes (comme un globe qui tourne).

• Symétrie d'ordre faible : Comprend des règles de base.
• Symétrie d'ordre élevé : Comprend des règles très complexes et détaillées sur la façon dont les formes interagissent dans l'espace 3D.

La découverte : Plus les règles de symétrie intégrées au robot sont complexes, plus il apprend vite. Un robot avec une symétrie « d'ordre élevé » apprend tellement plus vite que l'écart entre lui et le robot « Aveugle » devient un canyon. C'est comme comparer un élève qui connaît l'alphabet à un élève qui connaît déjà la grammaire et le vocabulaire de la langue ; à mesure que le livre s'épaissit, le second élève laisse le premier dans la poussière.

3. La « leçon amère » contre la réalité

Il existe une idée célèbre en intelligence artificielle appelée la « leçon amère », qui suggère que nous devrions arrêter d'essayer de coder à la main des connaissances humaines (comme la symétrie) dans l'IA et simplement laisser l'IA les apprendre à partir de données brutes, car c'est moins cher et cela s'adapte mieux à l'échelle.

• Ce papier soutient : Dans le monde des atomes et des molécules, la « leçon amère » est fausse. Si vous essayez de laisser un modèle découvrir la symétrie par lui-même, c'est comme demander à un élève de redécouvrir la gravité. C'est possible, mais incroyablement inefficace. Au moment où l'élève finit par comprendre, l'élève à qui la gravité a été enseignée vole déjà.

4. L'équilibre « Juste »

Le papier a également examiné comment dépenser l'argent (la puissance de calcul) de la manière la plus efficace.

• L'ancienne méthode : Peut-être devriez-vous acheter un cerveau plus gros (plus de paramètres) ou obtenir plus de manuels (plus de données).
• La nouvelle découverte : Il s'avère que vous devez acheter les deux en même temps. Si vous doublez vos données, vous devez également doubler la taille de votre modèle. Cette « mise à l'échelle tandem » fonctionne mieux pour tous les types de robots, mais les robots « Symétrie » sont simplement beaucoup plus efficaces pour utiliser cette puissance combinée.

5. Et tricher avec les fonctions de perte ?

Certains chercheurs ont essayé de piéger les robots « Aveugles » en ajoutant un score de pénalité s'ils commettaient une erreur sur la symétrie (par exemple : « Si vous dites qu'une molécule tournée est différente, vous avez une mauvaise note »).

• La découverte : Cela n'a pas bien fonctionné. C'est comme dire à un élève : « N'oublie pas les règles », sans réellement lui enseigner les règles. Le robot devait toujours lutter pour apprendre le motif. Il était bien mieux d'intégrer directement la règle dans le cerveau du robot dès le début.

La conclusion

Si vous voulez construire une IA super-intelligente pour comprendre les molécules, ne vous contentez pas de jeter plus de données sur un modèle simple et flexible en espérant qu'il déduira les lois de la physique. Intégrez directement les lois de la physique dans la conception du modèle.

À mesure que vous passez à des tailles massives, les modèles qui respectent les symétries fondamentales de l'univers (rotation, translation) ne seront pas juste légèrement meilleurs ; ils seront exponentiellement plus puissants que ceux qui tentent d'apprendre ces règles à partir de zéro. L'approche « Symétrie » change la nature même de la courbe d'apprentissage, rendant la tâche plus facile et les résultats meilleurs.

Résumé technique

Énoncé du problème
L'article traite du comportement d'échelle des potentiels interatomiques de réseaux de neurones (NNIP), qui sont des modèles d'apprentissage profond conçus pour prédire les propriétés mécaniques quantiques (spécifiquement l'énergie potentielle et les forces atomiques) de systèmes atomistiques. Alors que la littérature récente en langage naturel et en vision suggère que les lois d'échelle (relations de puissance entre la performance et les données/paramètres/calcul) sont largement indépendantes de l'architecture — impliquant que les modèles peuvent apprendre seuls les biais inductifs nécessaires comme la symétrie à mesure qu'ils évoluent — cette vision est contestée dans les domaines géométriques. Les auteurs examinent si l'équivariance architecturale explicite (imposant les symétries de rotation et de permutation) offre un avantage distinct dans les lois d'échelle pour les NNIP, ou si des modèles plus simples, non équivariants, peuvent atteindre des performances comparables compte tenu d'une puissance de calcul suffisante.

Méthodologie
Les auteurs mènent une étude empirique complète sur l'ensemble de données de molécules neutres OpenMol (environ 34 millions d'échantillons d'entraînement, ~9,2 × 10⁸ tokens). Ils comparent quatre familles architecturales distinctes représentant divers degrés de contraintes de symétrie :

1. MPNN non contraint : Un réseau de neurones à passage de messages (Message Passing Neural Network) standard traitant des caractéristiques géométriques (positions relatives) sans contraintes de symétrie.
2. Scalars invariants (GemNet-OC) : Utilise des caractéristiques invariantes (distances, angles, dièdres) mais approxime les fonctions équivariantes via un passage de messages basé sur les arêtes ; classé comme une architecture à 4 corps, ordre de tenseur $\ell=0$.
3. Vecteurs cartésiens (EGNN) : Un GNN $E(n)$-équivariant utilisant des canaux vectoriels (ordre de tenseur $\ell=1$).
4. Tenseurs sphériques d'ordre élevé (eSEN) : Un réseau équivariant utilisant des représentations irréductibles d'ordre supérieur du groupe de rotation ($\ell \ge 2$), employant l'alignement de repère pour espacer les produits tensoriels.

L'étude emploie un régime d'entraînement sur une seule époque pour s'aligner sur la littérature théorique des lois d'échelle, utilisant des optimiseurs AdamW sans planificateur pour atténuer les artefacts liés aux calendriers de taux d'apprentissage. Les lois d'échelle sont ajustées par rapport à trois métriques :

• Calcul : À la fois les FLOPs théoriques ($C$) et le temps d'entraînement réel (heures GPU, $H$).
• Données : Nombre de tokens d'entraînement ($D$).
• Paramètres : Taille du modèle ($N$).

Les auteurs examinent également les effets de la régularisation par perte de symétrie (pénalisant les écarts par rapport à l'équivariance dans les modèles non équivariants), l'entraînement multi-époque avec augmentation de données, et la moyenne de groupe au moment du test.

Contributions clés

1. Exposants d'échelle dépendants de l'architecture : L'article démontre que les exposants d'échelle ne sont pas constants à travers les architectures. À mesure que le « degré » d'équivariance augmente (de non contraint à ordre faible, puis à ordre élevé), les exposants de loi de puissance pour les données ($\beta$) et les paramètres ($\alpha$) augmentent significativement.
2. Échelle supérieure des modèles équivariants : Les architectures équivariantes, en particulier celles avec des représentations tensorielles d'ordre supérieur (eSEN), présentent des courbes d'échelle plus raides. Cela implique que l'écart de performance entre les modèles équivariants et non équivariants s'élargit à mesure que le calcul et les données augmentent, contredisant l'idée que les modèles peuvent simplement « apprendre » la symétrie plus tard.
3. Allocation optimale de calcul : L'étude constate que pour un entraînement optimal en calcul, la taille du modèle ($N$) et la taille de l'ensemble de données ($D$) doivent évoluer de concert ($N \propto D$) pour toutes les architectures, reflétant les découvertes en modélisation du langage (échelle Chinchilla). Cependant, la constante de proportionnalité et la réduction de perte résultante diffèrent selon le biais de symétrie de l'architecture.
4. Inefficacité de la perte de symétrie : Imposer la symétrie via un terme de perte (régularisation) dans les modèles non contraints ne produit pas les mêmes avantages d'échelle que d'intégrer l'équivariance dans l'architecture. Bien que cela améliore légèrement l'efficacité des données, cela ne parvient pas à égaler les exposants d'échelle des modèles équivariants natifs.
5. Insights sur l'entraînement multi-époque et l'augmentation : Dans des contextes à faible volume de données et multi-époque, l'augmentation de données est requise pour les modèles non contraints afin d'éviter le surapprentissage et de retrouver une échelle de loi de puissance. Cependant, même avec l'augmentation, les modèles non contraints n'atteignent pas les exposants d'échelle des modèles équivariants.

Résultats

• Lois d'échelle : La perte de validation suit une loi de puissance $L \propto C^{-\gamma}$. L'exposant $\gamma$ augmente avec la complexité architecturale :
  • MPNN non contraint : $\gamma \approx 0,14$
  • EGNN : $\gamma \approx 0,17$
  • GemNet-OC : $\gamma \approx 0,25$
  • eSEN (Ordre élevé) : $\gamma \approx 0,40$
• Échelle des données et des paramètres :
  • Les exposants d'échelle des données ($\beta$) varient de 0,31 (non contraint) à 0,75 (eSEN).
  • Les exposants d'échelle des paramètres ($\alpha$) varient de 0,28 (non contraint) à 0,82 (eSEN).
• Perte de symétrie : L'ajout d'un terme de perte de symétrie à un modèle non contraint augmente légèrement l'exposant des données ($\beta$) mais diminue l'exposant des paramètres ($\alpha$), résultant en aucun gain net de la pente de la frontière optimale en calcul par rapport à la référence non contrainte.
• Profondeur : Pour les modèles équivariants, la profondeur optimale du réseau augmente avec l'ordre de la représentation de rotation, tandis que les modèles non contraints souffrent d'un lissage excessif à des profondeurs plus élevées.

Signification et affirmations
L'article soutient que, contrairement à l'hypothèse de la « leçon amère » (qui suggère que les modèles devraient apprendre les biais inductifs à partir des données), la symétrie architecturale explicite est cruciale pour l'échelle dans les tâches géométriques. Les auteurs affirment que la symétrie n'est pas simplement une technique de réduction de données, mais qu'elle modifie fondamentalement la difficulté inhérente de la tâche et ses lois d'échelle.

La signification principale réside dans la découverte que les représentations équivariantes d'ordre supérieur se traduisent par de meilleurs exposants d'échelle. Cela suggère que pour les NNIP à grande échelle, investir dans des architectures complexes et conscientes de la symétrie (comme eSEN) est plus efficace que de simplement augmenter l'échelle de modèles plus simples et non équivariants. Les auteurs concluent que les biais inductifs fondamentaux comme la symétrie doivent être encodés dans l'architecture plutôt que laissés à la découverte du modèle, car ils modifient la trajectoire d'échelle elle-même.

L'article reste modeste quant à son périmètre, notant des limitations telles que le focus sur l'entraînement sur une seule époque, l'ensemble de données spécifique utilisé (molécules neutres), et l'exclusion des stratégies de pré-entraînement par débruitage utilisées dans d'autres travaux récents. Il appelle à un travail théorique futur pour expliquer pourquoi la symétrie modifie les exposants d'échelle et suggère d'étendre ces études à des types moléculaires plus diversifiés et à des régimes multi-époque.