Tiny Recursive Reasoning with Mamba-2 Attention Hybrid

Cette étude démontre que le remplacement des blocs Transformer par des opérateurs hybrides Mamba-2 dans le cadre de raisonnement récursif TRM améliore la couverture des solutions correctes sur le benchmark ARC-AGI-1 tout en préservant les performances de sélection initiale, validant ainsi l'efficacité des opérateurs SSM pour le raisonnement récursif.

L'essentiel

🧠 L'Idée de Base : Faire "Penser" un petit cerveau

Imaginez que vous essayez de résoudre un casse-tête très difficile, comme un Sudoku géant ou une énigme logique.

• Les gros modèles d'IA actuels sont comme des géants qui parlent fort. Pour résoudre le problème, ils écrivent chaque étape de leur réflexion à voix haute (sur du papier), mot par mot. C'est long et ils gaspillent parfois de l'énergie à écrire des phrases inutiles juste pour faire joli.
• Les "petits modèles" (comme TRM) sont des nains très intelligents. Au lieu de parler, ils pensent en silence. Ils bouillonnent dans leur tête, affinant leur idée étape par étape sans rien écrire. C'est ce qu'on appelle la "récursion latente" : une boucle de pensée interne.

Jusqu'à présent, ces petits cerveaux utilisaient un type de "moteur" très courant appelé Transformer (basé sur l'attention). C'est un moteur puissant, mais un peu gourmand.

🚀 Le Nouveau Moteur : Le "Mamba"

Les auteurs de cette étude se sont demandé : "Et si on remplaçait ce moteur classique par un nouveau type de moteur, appelé Mamba-2, qui est connu pour être très rapide et efficace ?"

Le Mamba-2 est comme un tapis roulant intelligent. Au lieu de regarder tout le passé en même temps (comme le Transformer), il glisse le long de l'information, gardant en mémoire ce qui est important et oubliant ce qui ne l'est pas, très rapidement.

Le défi : Est-ce que ce nouveau moteur, qui est très différent, peut aussi bien "penser" dans la boucle de réflexion du petit modèle ? Ou va-t-il perdre sa capacité de raisonnement ?

🏁 Le Résultat : Une Course de Vérité

Les chercheurs ont créé deux versions du même petit modèle (avec exactement le même nombre de "neurones" pour être équitables) :

1. L'ancien modèle (avec le moteur Transformer).
2. Le nouveau modèle (avec le moteur Mamba-2 hybride).

Ils les ont mis à l'épreuve sur ARC-AGI, un ensemble de puzzles visuels très difficiles (un peu comme un examen de QI pour les robots).

Ce qu'ils ont découvert :

1. Le nouveau modèle est plus "créatif" :
   Imaginez que vous devez deviner la solution à un problème.

   • L'ancien modèle (Transformer) est comme un champion de tir à l'arc. Il vise très précisément une cible. S'il rate, il rate. Mais s'il vise, il touche souvent.
   • Le nouveau modèle (Mamba) est comme un jardinier. Il fait pousser beaucoup de plantes différentes. Il ne vise pas une seule cible, mais il explore un grand champ de possibilités.

2. La statistique clé (Pass@K) :

   • Pass@1 (La meilleure chance) : Les deux modèles sont à égalité. Si vous ne pouvez faire qu'un seul essai, ils ont la même chance de gagner.
   • Pass@2 et plus (La sécurité) : C'est là que le Mamba brille. Si on lui donne le droit de faire 2, 10 ou 100 tentatives, le modèle Mamba trouve la bonne réponse beaucoup plus souvent.
   • Pourquoi ? Parce que le Mamba explore plus de chemins différents. Il génère une plus grande diversité de solutions. Même si sa "meilleure" réponse n'est pas toujours la première, il a plus de chances d'avoir la bonne réponse quelque part dans son panier de solutions.

🎨 L'Analogie du Chef Cuisinier

Pour résumer avec une image :

• Le modèle Transformer est un chef qui prépare un seul plat parfait. Il est très concentré. S'il se trompe d'assaisonnement, le plat est raté.
• Le modèle Mamba est un chef qui prépare 100 petits plats différents en même temps. Il ne sait pas toujours lequel est le meilleur au premier coup d'œil, mais il a une chance bien plus élevée d'avoir le plat parfait parmi ses 100 créations.

💡 Pourquoi c'est important ?

Cette étude prouve deux choses essentielles :

1. On peut changer le moteur : On n'est pas obligé d'utiliser les vieux moteurs "Transformer" pour faire de l'intelligence artificielle qui réfléchit. On peut utiliser des moteurs plus rapides et plus efficaces comme le Mamba.
2. La diversité est la clé : Pour résoudre des problèmes complexes, il ne suffit pas d'avoir la "meilleure" réponse immédiate. Il faut avoir la capacité de générer beaucoup d'idées différentes (diversité) pour s'assurer que la bonne réponse est incluse.

En bref, les chercheurs ont réussi à remplacer le cœur de leur petit robot penseur par une technologie plus moderne, et non seulement il n'a pas perdu ses capacités, mais il est devenu plus robuste et plus capable de trouver la solution quand on lui laisse plusieurs chances. C'est une première étape vers des IA plus intelligentes et moins gourmandes en énergie.

Résumé technique

1. Problématique et Contexte

Les modèles de langage actuels excellent dans les tâches de raisonnement grâce à des techniques comme le Chain-of-Thought (CoT) ou l'itération explicite. Cependant, ces approches génèrent des tokens intermédiaires visibles, ce qui peut être inefficace (tokens grammaticaux non pertinents) et poser des problèmes de fidélité (le raisonnement affiché ne reflète pas toujours le calcul interne).

Une alternative émergente est la récursion latente (ou latent recursion), où le modèle affine ses représentations internes sans émettre de tokens intermédiaires. Le modèle TRM (Tiny Recursive Model) a démontré qu'un réseau minuscule (7M de paramètres) pouvait surpasser des modèles commerciaux sur le benchmark ARC-AGI grâce à cette méthode.

La question centrale de l'article : Le bloc de base du TRM est actuellement un Transformer (basé sur l'attention). Puisque le Mamba-2 (un modèle à espace d'état ou SSM) repose lui-même sur une récurrence itérative ($h_t = a_t h_{t-1} + B_t x_t$), est-il possible de remplacer les blocs Transformer par des opérateurs hybrides Mamba-2 dans le cadre de la récursion latente sans dégrader les capacités de raisonnement ? L'objectif est de déterminer si les SSMs peuvent devenir des candidats viables pour la conception d'opérateurs de raisonnement récursif.

2. Méthodologie

Les auteurs proposent une variante du TRM où les blocs Transformer sont remplacés par des opérateurs hybrides Mamba-2 + Attention, tout en maintenant une parité stricte des paramètres pour isoler l'effet de l'architecture.

Architecture Hybride

Le modèle conserve la structure récursive du TRM (cycles externes $H$ et internes $L$ avec deux états latents $z_H$ et $z_L$), mais modifie l'opérateur de mise à jour par pas :

• TR-mamba2attn (Proposition) : Remplace les blocs d'attention par une pile : Mamba-2 → Mamba-2 → Attention → MLP.
  • Raison : Mamba-2 traite les séquences de manière causale (unidirectionnelle). Pour des tâches comme ARC-AGI, Sudoku ou les labyrinthes, une communication bidirectionnelle est cruciale. L'ajout d'un bloc Attention (ou MLP transposé) permet de capturer les dépendances spatiales globales.
• TR-mamba2mlpt : Variante utilisant un MLP transposé à la place de l'Attention pour la communication "tout-à-tout".

Détails Techniques

• Parité des paramètres : Le modèle de base (TRM-attn) a 6,83M de paramètres, contre 6,86M pour le modèle hybride (TR-mamba2attn).
• Normalisation Post-Norm : L'article souligne l'importance critique de la normalisation post-norm ($h_{t+1} = \text{Norm}(h_t + F(h_t))$) par rapport au pre-norm. Dans une récursion déroulée, le pre-norm peut entraîner une explosion de la magnitude du flux résiduel, causant des instabilités (NaN), tandis que le post-norm borne l'échelle de l'état caché quelle que soit la profondeur de récursion.
• Entraînement : Utilisation de cycles de mise à jour latente (3 cycles externes, 4-6 internes) sur les tâches ARC-AGI-1, Sudoku-Extreme et Maze-30x30.

3. Contributions Clés (C1, C2, C3)

1. C1 : Premier modèle hybride Mamba pour le raisonnement récursif latent. C'est la première étude intégrant des opérateurs Mamba-2 dans le cadre du TRM.
2. C2 : Validation empirique sur ARC-AGI-1. Le modèle hybride améliore la métrique officielle pass@2 de +2,0 % (45,88 % contre 43,88 %) et maintient une performance équivalente sur pass@1.
3. C3 : Analyse du compromis Couverture vs Sélection. L'analyse des métriques pass@K révèle que le modèle hybride améliore la diversité des candidats générés (couverture) sans sacrifier la qualité de la sélection du meilleur candidat (top-1).

4. Résultats Expérimentaux

Les expériences ont été menées sur trois benchmarks :

A. ARC-AGI-1 (Raisonnement Abstrait)

• Performance : Le modèle hybride (TR-mamba2attn) surpasse le modèle de base (TRM-attn) sur l'ensemble des métriques pass@K pour $K \ge 2$.
  • pass@2 : +2,0 % d'amélioration.
  • pass@100 : +4,75 % d'amélioration.
  • pass@1 : Parité quasi-totale (-0,25 %).
• Interprétation : Le modèle hybride génère une plus grande variété de solutions correctes au sein de son pool de candidats (couverture accrue), ce qui se traduit par une meilleure probabilité que la bonne réponse soit présente dans le top-K, même si la probabilité de la placer en première position reste similaire.

B. Sudoku (Contraintes sur grille fixe 9x9)

• Les variantes basées sur MLP-t (communication dense) dominent ici (87,4 % pour TRM-mlp-t).
• Les modèles hybrides Mamba-2+Attention performant moins bien (66,5 %) que les modèles MLP-t, suggérant que pour des grilles petites et fixes, une communication dense "tout-à-tout" est plus efficace que le traitement séquentiel de Mamba couplé à l'attention.

C. Maze (Grilles 30x30)

• Les variantes MLP-t échouent complètement (0,0 %), tandis que le modèle hybride TR-mamba2attn atteint 80,6 % contre 60,8 % pour le modèle de base.
• Cela indique que pour des tâches de raisonnement spatial à grande échelle, la capacité de Mamba-2 à traiter les dépendances séquentielles est cruciale, là où le mélange dense échoue à s'adapter.

Analyse Statistique (Figure 3)

L'analyse des statistiques de vote confirme le compromis :

• Hybride : Génère plus de candidats uniques (+27 %) et une entropie de vote plus élevée (exploration diverse).
• Base (Attention) : Concentre davantage les votes sur le candidat top-1 (sélection plus décisive).
• Conclusion : Mamba-2 introduit des trajectoires de solution distinctes lors de l'augmentation des données, enrichissant le pool de candidats sans dégrader la qualité du meilleur choix.

5. Signification et Implications

Ce travail valide que les opérateurs basés sur les Modèles à Espace d'État (SSM) comme Mamba-2 peuvent intégrer l'espace de conception des opérateurs de raisonnement récursif.

• Efficacité et Scalabilité : L'intégration de Mamba-2 offre une complexité linéaire par rapport à la complexité quadratique des Transformers, promettant une inférence plus rapide pour le calcul à l'exécution (test-time compute).
• Nouvelle Stratégie de Mélange : L'article suggère que la combinaison de la récurrence interne de Mamba (itération latente) avec la récursion externe du TRM est une stratégie puissante.
• Stabilité : La démonstration de l'importance du post-norm pour la stabilité des boucles récursives est une contribution technique majeure pour les futurs modèles de raisonnement profond.

En résumé, cette étude ouvre la voie à l'utilisation de modèles hybrides Mamba-2/Attention pour le raisonnement latent, prouvant qu'ils peuvent non seulement égaler, mais surpasser les architectures purement basées sur l'attention en termes de couverture de solutions, tout en offrant des avantages potentiels en termes d'efficacité computationnelle.