DREAMSTATE: Diffusing States and Parameters for Recurrent Large Language Models

Résumé Technique : DREAMSTATE

Énoncé du Problème

Les réseaux de neurones récurrents (RNN) modernes, particulièrement les architectures comme RWKV, offrent des avantages significatifs par rapport aux Transformers standards concernant l'efficacité de l'inférence et la représentation d'état à taille fixe. Cependant, l'article identifie deux limitations critiques dans les conceptions actuelles de RNN :

1. Potentiel de Représentation d'État Inexploité : Bien que les états des RNN soient puissants pour la modélisation à courte portée, ils sont rarement traités comme des représentations de connaissances éditables. Il existe un manque de recherche sur la modélisation de la distribution de l'état interne pour permettre une génération ou une manipulation directe.
2. Bruit Structurel dans la Récurrence Statique : Similaire au « bruit d'attention » dans les Transformers où les modèles allouent trop d'attention à un contexte non pertinent, les RNN souffrent de « bruit structurel ». Le mécanisme de récurrence central dans RWKV repose sur des paramètres WKV ($\theta_{WKV}$) statiques et fixes après l'entraînement. Cette nature statique agit comme un biais de pondération, sous-optimal pour des contextes divers, empêchant le modèle d'ajuster dynamiquement ses « lois de fonctionnement » pour se concentre sur les informations saillantes.

Les auteurs postulent que si l'état interne est une représentation structurée du contexte, il devrait être possible de le modéliser et de le manipuler directement. De plus, l'hypothèse de la statique des paramètres de récurrence doit être remise en question pour atténuer le bruit structurel.

Méthodologie

L'article propose DREAMSTATE, un cadre fondé sur la modélisation générative (spécifiquement les modèles de diffusion) pour répondre à ces problèmes via une approche en deux étapes :

1. Synthèse Générative d'État (Cadre DREAMSTATE)

Les auteurs traitent d'abord l'état interne de RWKV comme un point de donnée échantillonné à partir d'une variété apprenable ($\mathcal{M}$).

• Modèle : Un Transformer de Diffusion Conditionnel (DiT) est entraîné pour modéliser la distribution de probabilité $p(S|c)$, où $S$ est l'état et $c$ est un plongement (embedding) de texte.
• Processus : Le tenseur d'état de haute dimension (états multi-têtes aplatis) est traité comme une séquence de patchs. Le DiT apprend à prédire le bruit ajouté à un état, conditionné par le contexte d'entrée.
• Objectif : Valider que la variété de l'état est bien définie et permettre la génération contrôlée d'états pour des modes opérationnels spécifiques (ex: une personnalité spécifique).

2. Synthèse de Paramètres Dynamiques (Architecture Hybride)

Pour traiter le bruit structurel, les auteurs proposent une nouvelle architecture hybride où les paramètres de récurrence statiques sont générés dynamiquement.

• Mécanisme : Un « Parameter DiT » suit les caractéristiques globales d'un contexte d'entrée de longueur variable (implémenté initialement comme l'embedding du premier jet/token, $x[:,0,:]$) pour générer un vecteur de paramètres spécifique au contexte $\theta_{gen}$.
• Paramètres Cibles : Le DiT génère les matrices de projection linéaire aplaties ($W_r, W_k, W_v$) qui constituent les paramètres WKV.
• Fusion : Les paramètres générés sont fusionnés avec les paramètres de base statiques originaux ($\theta_{static}$) via interpolation :
  $$\theta_{WKV-final} = \alpha \cdot \theta_{static} + (1 - \alpha) \cdot \theta_{gen}$$
  où $\alpha$ est un hyperparamètre apprenable. Cela permet au modèle de conserver une connaissance stable tout en introduisant une modulation spécifique au contexte.
• Entraînement : Le système est entraîné de bout en bout en utilisant une fonction de perte multi-objectif :
  $$L_{total} = \lambda_1 L_{LLM}(\theta_{WKV-final}) + \lambda_2 L_{param\_diff}(\psi)$$
  Cela équilibre la prédiction du jet suivant (cross-entropie standard) avec la perte de diffusion pour la génération de paramètres.

Contributions Clés

1. Cadre DREAMSTATE : Une méthode pour modéliser générativement l'état dynamique d'un modèle RWKV. Les auteurs valident que l'état interne est une variété structurée où différents concepts (ex: « programmeur » vs « poète ») forment des clusters distincts, permettant une inférence contrôlée via la manipulation d'état.
2. Architecture de Paramètres Dynamiques : Un nouveau design hybride où un modèle de diffusion génère des paramètres de récurrence basés sur un suivi de contexte global. Cela est proposé comme une méthode rigoureuse pour contrer le « bruit structurel » en rendant les lois de fonctionnement de l'RNN dynamiques plutôt que statiques.
3. Paradigme d'Entraînement Multi-Objectif : La définition et la validation empirique réussie d'une stratégie d'entraînement qui optimise conjointement la génération de paramètres et la modélisation du langage, prouvant la faisabilité de cette approche hybride.
4. Autoregressivité par Concept : Ce travail fournit une preuve de concept que la combinaison de l'autoregressivité traditionnelle par jet (token-wise) et de l'« autoregressivité par concept » (via la génération d'état/paramètre) est une direction viable pour la conception de modèles futurs.

Résultats Expérimentaux

• Visualisation de la Variété d'État : En utilisant t-SNE, les auteurs ont démontré que les états réels de RWKV-7 forment des clusters distincts correspondant à des personas d'experts spécifiques (ex: terminal Linux, développeur Ethereum, conteur), confirmant que l'état est un objet structuré et bien défini.
• Inférence Contrôlée : En initialisant le modèle RWKV-7 avec des états générés par DREAMSTATE (ou des états interpolés), le modèle a produit des textes nettement plus cohérents thématiquement et créatifs par rapport aux bases de référence. Par exemple, l'interpolation entre un prompt de « persévérance » et un prompt de « changement climatique dans le désert » a produit une histoire mêlant métaphoriquement les deux concepts.
• Stabilité de l'Entraînement : L'architecture hybride a été entraînée avec succès sur un sous-ensemble de The Pile. Les courbes de perte d'entraînement ont montré une diminution stable et constante tant pour la perte de modélisation du langage que pour la perte de diffusion des paramètres, confirmant la stabilité du processus d'optimisation conjointe.

Signification et Revendications

L'article affirme ouvrir une nouvelle direction de recherche pour la représentation d'état des RNN, dépassant l'état en tant que simple variable transitoire pour en faire un objet probabiliste manipulable.

• Contrôlabilité : Ce travail démontre que les états internes peuvent être directement générés et manipulés pour préparer un modèle à des modes opérationnels spécifiques, offrant un contrôle sans précédent sur l'inférence.
• Adaptabilité : En traitant les paramètres de récurrence comme des variables dynamiques générées par un suiveur de contexte global, l'architecture vise à construire des modèles de monde plus flexibles et expressifs.
• Référence Architecturale : L'article positionne cet hybride de synthèse de paramètres par diffusion et d'inférence récurrente comme une référence architecturale concrète pour les futurs modèles de langage de grande taille (LLM), suggérant que le « comment » de la modélisation de l'état est la clé de l'avenir des LLM.

Les auteurs soulignent que leur travail sert de preuve de concept, validant la faisabilité de la conception de la synthèse dynamique de paramètres pour atténuer le bruit structurel, sans prétendre à une supériorité immédiate sur tous les modèles existants dans tous les benchmarks, mais plutôt en établissant la viabilité du paradigme proposé.