Résumé Technique : Apprentissage du suivi d'état par le code via des RNN Linéaires

Énoncé du Problème

Le suivi d'état est une capacité fondamentale pour la compréhension de l'exécution de programmes, des états de jeux et des modèles de mondes. Bien que les travaux théoriques aient établi une distinction où les Transformers excellent dans le rappel associatif et les réseaux de neurones récurrents (RNN) sont performants dans le suivi d'état, un écart important subsiste entre la manière dont ces modèles sont évalués et la manière dont ils sont entraînés.

La plupart des benchmarks existants pour le suivi d'état (par exemple, la composition de permutations) utilisent une configuration de type séquence-à-séquence : mapper une séquence d'actions directement vers une séquence d'états. Cela diffère fondamentalement du paradigme de prédiction du prochain jeton utilisé pour entraîner les modèles de langage de grande taille (LLM) modernes, où le modèle doit prédire le jeton suivant dans un flux qui entrelace des actions avec des révélations partielles d'état (par exemple, des instructions print). De plus, l'exécution de code dans le monde réel implique souvent une observabilité partielle et des transitions probabilistes (par exemple, des branchements sur des états cachés ou des choix aléatoires), créant un contexte où le véritable état n'est pas entièrement observable, nécessitant le maintien d'une distribution de croyance plutôt qu'un état déterministe.

Méthodologie

1. Reformulation du suivi d'état via le code (Traces REPL)

Pour combler l'écart entre les benchmarks abstraits d'automates et les objectifs d'entraînement des LLM, les auteurs convertissent les tâches de composition de permutations en traces REPL Python.

• Mécanisme : Au lieu de mapper une séquence de permutations $\sigma_1, \dots, \sigma_m$ vers des états cumulatifs, le modèle traite une séquence de commandes Python (assignations de variables, échanges/swaps) entrelacées avec des instructions print qui révèlent des informations partielles sur l'état.
• Objectif : Le modèle est entraîné via la prédiction du prochain jeton pour prédire le jeton suivant dans la trace, qui peut être une commande ou une valeur d'état révélée. Cela force le modèle à apprendre implicitement le suivi d'état à travers la tâche de prédiction.

2. Comparaison Architecturale

L'étude compare trois classes architecturales sur ces tâches :

• Transformers : Modèles standards basés sur l'attention.
• RNN Linéaires (DeltaNet [0, 1]) : RNN linéaires restreints à des valeurs propres dans $[0, 1]$, connus pour être incapables de réaliser des suivis d'état complexes comme la parité ou la composition de permutations.
• RNN Linxaires (DeltaNet [-1, 1]) : RNN linéaires avec des plages de valeurs propres étendues $[-1, 1]$, théoriquement capables de résoudre la composition de permutations déterministe.

3. Cadre Théorique : PFSA-SR

Pour analyser la difficulté du suivi d'état dans le code réel, les auteurs introduisent l'Automate à États Finis Probabiliste avec Révélations d'État (PFSA-SR).

• Définition : Un tuple $(Q, \Sigma, \delta, \rho, q_0)$ où $\delta$ est un noyau de transition probabiliste et $\rho$ est une fonction de révélation mappant les entrées vers des sous-ensembles d'états.
• Dynamique : Le système évolue via des transitions probabilistes, mais l'observateur reçoit des contraintes "dures" (révélations) qui élaguent l'ensemble des états possibles (élagage du support) plutôt que de fournir des probabilités douces.
• Mise à jour de la croyance : Le suivi nécessite de maintenir une distribution de croyance $b_t$ sur les états. La mise à jour implique une étape de révélation (mise à zéro des états incohérents et renormalisation) suivie d'une étape de transition. Crucialement, l'étape de renormalisation $f(x) = x/\|x\|_1$ est non linéaire.

4. Analyse de Représentation

L'article examine deux représentations pour les états de croyance dans les RNN linéaires :

• Représentation Conjointe : Énumère explicitement toutes les $n!$ permutations. Les auteurs soutiennent que cela est numériquement instable pour les RNN linéaires car le terme d'injection d'entrée $B(x_t)$ ne peut pas reconstituer de manière adaptative la masse de probabilité des configurations survivantes après une révélation partielle, entraînant une décroissance exponentielle de la norme.
• Représentation Marginale : Suit la probabilité marginale de chaque position indépendamment (une matrice doublement stochastique $n \times n$). Bien que cela permette à $B(x_t)$ d'injecter directement de la certitude pour les entrées révélées, les auteurs montrent qu'elle reste susceptible de décroissance exponentielle sous des séquences de révélations adverses où des entrées spécifiques non révélées ne sont jamais directement reconstituées.

Résultats Clés

1. Densité de Supervision et Architecture

• Transformers : Nécessitent une supervision d'état dense pour apprendre le suivi d'état. À mesure que l'espacement entre les instructions print (révélations) augmente, la performance du Transformer chute rapidement. Ils échouent à généraliser aux permutations complètes, même avec une supervision dense.
• RNN Linéaires (DeltaNet [-1, 1]) : Excellent dans ce contexte. Ils apprennent à suivre les états de manière fiable, même avec une supervision parcimonieuse, et démontrent une extrapolation de longueur robuste (généralisation à des séquences plus longues que l'entraînement).
• RNN Linéaires (DeltaNet [0, 1]) : Échouent à résoudre la tâche, confirmant que l'extension de la plage de valeurs propres à $[-1, 1]$ est nécessaire pour le suivi d'état complexe.

2. Interprétabilité du Suivi d'État

L'analyse du DeltaNet [-1, 1] entraîné révèle que le suivi d'état est localisé dans des composants spécifiques :

• Extension des Valeurs Propres : Au moins une tête d'attention dans une couche spécifique (par exemple, Couche 12, Tête 3) apprend à régler son paramètre de décroissance $\beta$ à 2, résultant en une valeur propre de $-1$.
• Structure de Sous-espace : Cette tête spécifique opère dans un sous-espace linéaire nettement de plus faible dimension que les autres têtes, ce qui est cohérent avec les exigences théoriques pour le suivi de permutations.
• Intervention : Réduire la plage de $\beta$ de cette tête spécifique vers $[0, 1]$ provoque un effondrement immédiat des performances, tandis que l'échelle des autres têtes n'a aucun effet, confirmant cette tête comme le mécanisme dédié au suivi d'état.

3. Limites des RNN Linéaires dans les Contextes Probabilistes

L'article démontre que si les RNN linéaires résolvent des tâches déterministes, ils font face à des barrières inhérentes dans les contextes probabilistes (PFSA-SR) :

• Décroissance Exponentielle : Sous la représentation conjointe, les révélations partielles provoquent une décroissance exponentielle de la norme de l'état non normalisé ($1/2 \to 1/4 \to 1/8 \dots$) car la récurrence linéaire ne peut pas effectuer la renormalisation non linéaire nécessaire.
• Vulnérabilité Adversaire : Même avec la représentation marginale, un adversaire peut construire une séquence de swaps probabilistes et de révélations sélectives qui fait décroître la masse de probabilité des entrées non révélées de manière exponentielle jusqu'à ce qu'elle disparaisse en précision finie.
• Conclusion : Le maintien stable de la croyance dans des contextes probabilistes avec observabilité partielle semble irréalisable pour les RNN linéaires sans soit une renormalisation non linéaire dans la récurrence, soit des réinitialisations complètes fréquentes.

Signification et Revendications

L'article affirme fournir un banc d'essai plus réaliste pour le suivi d'état qui s'aligne sur l'objectif de prédiction du prochain jeton des modèles de langage, allant au-delà des benchmarks synthétiques de séquence-à-séquence.

1. Combler l'Écart : Il démontre que les RNN linéaires capables de suivi d'état (spécifiquement DeltaNet avec des valeurs propres étendues) peuvent apprendre ces capacités via la prédiction du prochain jeton sur des traces de type code, là où les Transformers échouent.
2. Barrière Théorique : Il identifie une limitation fondamentale des RNN linéaires dans les scénarios de code réel. Le code réel implique souvent des transitions probabilistes et une observabilité partielle (ex: branchements cachés, choix aléatoires). Les auteurs montrent que le suivi d'états dans de tels environnements nécessite une renormalisation non linéaire, ce que les RNN linéaires ne peuvent pas effectuer en interne. Cela suggère que l'écart de complexité observé dans les tâches réelles peut provenir de la nature probabiliste de l'exécution du code plutôt que de la seule longueur ou complexité de la séquence.
3. Implications pour le Benchmarking : Les auteurs suggèrent que le benchmarking des architectures uniquement sur des tâches d'Automates à États Finis Déterministes (DFA) occulte le potentiel des RNN non linéaires et les limitations des RNN linéaires. Le cadre PFSA-SR proposé et ses modes de défaillance associés (décroissance exponentielle de la norme) offrent une explication théorique de la raison pour laquelle les modèles récurrents linéaires standards peinent lorsque le raisonnement sur l'incertitude est requis.

L'article conclut modestement, notant que ses expériences se concentrent sur des groupes de permutations synthétiques et que des travaux futurs sont nécessaires pour étendre ces résultats aux traces d'exécution du monde réel et pour étudier les architectures hybrides qui entrelacent la récurrence linéaire avec une normalisation non linéaire périodique.