StateX: Verbesserung der RNN-Erinnerungsfähigkeit durch Post-Training-Zustandserweiterung

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

Während Transformer-basierte Modelle bemerkenswerte Leistungen im Sprachmodellieren gezeigt haben, führen ihre hohen Komplexitäten zu hohen Kosten bei der Verarbeitung langer Kontexte. Im Gegensatz dazu haben rekurrente neuronale Netze (RNNs) wie lineare Attention und State-Space-Modelle aufgrund ihrer konstanten Komplexität pro Token an Popularität gewonnen. Diese rekurrenten Modelle haben jedoch Schwierigkeiten mit Aufgaben, die ein präzises Abrufen von Kontextinformationen aus langen Kontexten erfordern, da alle Kontextinformationen in einen konstant großen rekurrenten Zustand komprimiert werden. Frühere Arbeiten haben gezeigt, dass die Fähigkeit zum Abrufen positiv mit der Größe des rekurrenten Zustands korreliert, doch führt das direkte Training von RNNs mit größeren rekurrenten Zuständen zu hohen Trainingskosten. In diesem Artikel stellen wir StateX vor, eine Trainingspipeline zur effizienten Erweiterung der Zustände vortrainierter RNNs durch Nachschulung. Für zwei beliebte Klassen von RNNs, lineare Attention und State-Space-Modelle, entwerfen wir architektonische Modifikationen für die Nachschulung, um die Zustandsgröße zu skalieren, ohne oder mit nur geringfügiger Erhöhung der Modellparameter. Experimente mit Modellen bis zu 1,3 Milliarden Parametern zeigen, dass StateX die Abruf- und In-Context-Lernfähigkeit von RNNs effizient verbessert, ohne hohe Nachschulungskosten zu verursachen oder andere Fähigkeiten zu beeinträchtigen.

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.