BRo-JEPA: Learning Modular Arithmetic in Latent Space

Résumé Technique : BRo-JEPA – Apprentissage de l'arithmétique modulaire dans l'espace latent

Énoncé du problème
L'article examine une question fondamentale en apprentissage profond : les réseaux de neurones apprennent-ils réellement des règles algébriques abstraites, ou se contentent-ils de mémoriser des motifs d'entraînement ? Si les modèles actuels excellent dans les tâches mathématiques au sein des plages observées, ils échouent souvent à extrapoler vers des plages numériques ou des opérations non vues. Les auteurs postulent que les modèles supervisés standards et les modèles JEPA (Joint Embedding Predictive Architecture) existants reposent sur des corrélations statistiques plutôt que sur une logique arithmétique sous-jacente. Plus précisément, lorsqu'on leur confie une tâche d'arithmétique modulaire (par exemple, $y = (x + k) \mod 10$) utilisant des chiffres MNIST comme états, les modèles standards s'adaptent aux opérations vues mais échouent à généraliser aux opérations non vues (généralisation zero-shot stricte).

Méthodologie
Les auteurs utilisent le cadre JEPA, où un modèle prédit la représentation latente d'un chiffre cible étant donné un chiffre de contexte (image) et une opération arithmétique (action), plutôt que de prédire des sorties de pixels bruts. L'étude compare plusieurs variations architecturales :

1. Baselines standards : Modèles supervisés (MLP et ResNet-18) utilisant une perte d'entropie croisée et des plongements d'actions compositionnels additifs. Ces plongements représentent les opérations comme des multiples scalaires d'un vecteur de base (par exemple, $op(k) = k \cdot v$) afin d'imposer une structure arithmétique dans l'espace de plongement.
2. JEPA avec plongements additifs : Un modèle JEPA utilisant les mêmes plongements additifs, augmenté d'une « Composition Consistency Loss » (perte de cohérence de composition). Cette perte auxiliaire impose que l'application d'une opération composée (par exemple, $+2$) en une seule étape doit produire le même résultat latent que l'application séquentielle d'opérations primitives (par exemple, $+1$ puis $+1$).
3. BRo-JEPA (Block-Rotation JEPA) : La contribution principale est un nouveau prédicteur qui remplace le MLP non contraint standard par un prédicteur de rotation de blocs conditionné par l'opération.
   • Mécanisme : Inspiré des Rotary Position Embeddings (RoPE), ce prédicteur traite la dynamique de l'espace latent comme une série de rotations planes en 2D. Pour un vecteur latent $z$ et une opération $k$, la prédiction est $z_p = R_{block}(k \cdot \theta) \cdot z$.
   • Biais inductif : Cette architecture impose explicitement la structure circulaire de l'arithmétique modulo-10 sur l'espace latent. Les angles de rotation $\theta$ peuvent être fixes (fixés à $2\pi/10$) ou apprenables.
   • Variantes : Les auteurs testent la Single-Frequency Rotation (SFR), où tous les blocs latents partagent un seul angle, et la Multi-Frequency Rotation (MFR), où chaque bloc possède un angle indépendant.
   • Régularisation : Pour éviter l'effondrement latent (où différents intrants sont mappés vers des représentations indiscernables), le modèle utilise la régularisation Variance-Invariance-Covariance (VICReg) et un encodeur cible par moyenne mobile exponentielle (EMA).

Contributions clés

1. Définition du banc d'essai : L'introduction de l'arithmétique modulaire de MNIST comme un banc d'essai contrôlé pour étudier l'apprentissage de règles dans les modèles de monde latents.
2. Analyse d'échec : Démonstration que les baselines supervisées et les modèles JEPA avec plongements additifs, bien qu'ajustant les opérations vues, échouent à l'extrapolation zero-shot stricte. Ils capturent les signes des opérations mais échouent à encoder la magnitude ou la géométrie compositionnelle.
3. Proposition de BRo-JEPA : Le développement d'un modèle de dynamique latente de type JEPA présentant un prédicteur de rotation de blocs conditionné par l'opération qui s'aligne sur la structure cyclique de l'arithmétique modulaire.
4. Analyse géométrique : Analyse montrant que BRo-JEPA organise les représentations latentes en une variété géométrique structurée (une variété circulaire), permettant au modèle d'internaliser la topologie de l'arithmétique modulaire.

Résultats
Les expériences ont été menées avec un entraînement sur les opérations $\{-1, +1\}$ et une évaluation sur les opérations non vues $\{\pm 2, \dots, \pm 9\}$.

• Performance des baselines : Les baselines supervisées ont atteint une précision élevée sur les opérations vues (~97-98 %) mais une précision quasi nulle sur les opérations non vues (2,42 % pour ResNet-18). Les visualisations PCA ont révélé que ces modèles compressaient arbitrairement les opérations non vues dans l'espace de représentation des opérations vues.
• JEPA avec plongements additifs : Bien que ces modèles affichent une haute précision de rollout (appliquer $+1$ de manière répétée pour simuler $+k$), ils échouent à la prédiction directe zero-shot d'opérations non vues (précision ~10 %). La perte de cohérence de composition a amélioré cela à ~99 %, mais est considérée comme un « zero-shot faible » car le modèle a été indirectement exposé à des opérations composées pendant l'entraînement.
• Performance de BRo-JEPA :
  • Précision Zero-Shot : Le BRo-JEPA basé sur ResNet avec Multi-Frequency Rotation (MFR) a atteint 99,46 % de précision zero-shot sur les opérations non vues.
  • Précision de Rollout : Le modèle a également atteint 99,46 % de précision dans les rollouts multi-étapes, confirmant la stabilité des dynamiques latentes apprises.
  • Structure Latente : Les visualisations PCA ont confirmé que BRo-JEPA a réussi à contraindre l'espace latent en une variété elliptique/circulaire continue où les classes de chiffres sont disposées séquentiellement. La variante MFR a produit des prototypes de classes mieux séparés (similitude cosinus hors diagonale plus faible) par rapport à la SFR, conduisant à une généralisation supérieure.
  • Angles apprenables : Lorsque les angles de rotation étaient apprenables, la MFR a convergé rapidement vers une généralisation efficace, tandis que la SFR présentait des retards importants et une instabilité, convergeant souvent vers des solutions harmoniques sous-optimales.

Signification et revendications
L'article affirme que les modèles de monde latents peuvent apprendre des règles de transformation symboliques, mais seulement lorsque le biais inductif de l'architecture correspond à la structure du problème. Les résultats suggèrent que :

• Les prédicteurs neuronaux standards manquent des contraintes nécessaires pour préserver la géométrie compositionnelle pour les opérations non vues.
• Modéliser explicitement la nature cyclique de l'arithmétique modulaire via des rotations de blocs dans l'espace latent permet une forte généralisation zero-shot sans nécessiter l'exposition aux opérations spécifiques pendant l'entraînement.
• Cette approche offre une direction prometteuse pour l'apprentissage de représentations au-delà de l'arithmétique, suggérant que l'incorporation de structures de transformation appropriées dans les modèles de transition latents est cruciale pour l'apprentissage de règles abstraites.

Les auteurs restent modestes, notant que leurs résultats sont spécifiques à ce cadre hautement contrôlé de l'arithmétique modulaire de MNIST avec une structure de groupe cyclique connue. Ils ne prétendent pas que l'architecture généralisera automatiquement à des tâches de raisonnement symbolique arbitraires ou à des groupes de transformation inconnus sans des priors structurels similaires.