StateX : Amélioration de la capacité de rappel des RNN par expansion d'état post-entraînement

StateX: Enhancing RNN Recall via Post-training State Expansion

September 26, 2025

papers.authors: Xingyu Shen, Yingfa Chen, Zhen Leng Thai, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

papers.abstract

Bien que les modèles basés sur les Transformers aient démontré des performances remarquables en modélisation du langage, leur complexité élevée entraîne des coûts importants lors du traitement de contextes longs. En revanche, les réseaux de neurones récurrents (RNNs) tels que l'attention linéaire et les modèles d'espace d'état ont gagné en popularité grâce à leur complexité constante par token. Cependant, ces modèles récurrents peinent à accomplir des tâches nécessitant un rappel précis des informations contextuelles provenant de contextes longs, car toutes les informations contextuelles sont compressées dans un état récurrent de taille constante. Les travaux précédents ont montré que la capacité de rappel est positivement corrélée à la taille de l'état récurrent, mais l'entraînement direct des RNNs avec des états récurrents plus grands entraîne des coûts d'entraînement élevés. Dans cet article, nous présentons StateX, un pipeline d'entraînement pour étendre efficacement les états des RNNs pré-entraînés via un post-entraînement. Pour deux classes populaires de RNNs, l'attention linéaire et les modèles d'espace d'état, nous concevons des modifications architecturales de post-entraînement pour augmenter la taille de l'état sans augmentation ou avec une augmentation négligeable des paramètres du modèle. Les expériences sur des modèles allant jusqu'à 1,3 milliard de paramètres démontrent que StateX améliore efficacement la capacité de rappel et d'apprentissage en contexte des RNNs sans engendrer de coûts élevés de post-entraînement ni compromettre d'autres capacités.

English

While Transformer-based models have demonstrated remarkable language modeling performance, their high complexities result in high costs when processing long contexts. In contrast, recurrent neural networks (RNNs) such as linear attention and state space models have gained popularity due to their constant per-token complexities. However, these recurrent models struggle with tasks that require accurate recall of contextual information from long contexts, because all contextual information is compressed into a constant-size recurrent state. Previous works have shown that recall ability is positively correlated with the recurrent state size, yet directly training RNNs with larger recurrent states results in high training costs. In this paper, we introduce StateX, a training pipeline for efficiently expanding the states of pre-trained RNNs through post-training. For two popular classes of RNNs, linear attention and state space models, we design post-training architectural modifications to scale up the state size with no or negligible increase in model parameters. Experiments on models up to 1.3B parameters demonstrate that StateX efficiently enhances the recall and in-context learning ability of RNNs without incurring high post-training costs or compromising other capabilities.