StateX: Verbetering van RNN-recall via post-training staatsexpansie

StateX: Enhancing RNN Recall via Post-training State Expansion

September 26, 2025

Auteurs: Xingyu Shen, Yingfa Chen, Zhen Leng Thai, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

Samenvatting

Hoewel Transformer-gebaseerde modellen opmerkelijke prestaties hebben getoond op het gebied van taalmodellering, leiden hun hoge complexiteiten tot hoge kosten bij het verwerken van lange contexten. Daarentegen hebben recurrente neurale netwerken (RNN's), zoals lineaire aandacht en state space-modellen, aan populariteit gewonnen vanwege hun constante complexiteit per token. Deze recurrente modellen hebben echter moeite met taken die een nauwkeurige herinnering van contextuele informatie uit lange contexten vereisen, omdat alle contextuele informatie wordt samengeperst in een constante grootte van de recurrente staat. Eerdere werken hebben aangetoond dat het herinneringsvermogen positief gecorreleerd is met de grootte van de recurrente staat, maar het direct trainen van RNN's met grotere recurrente staten resulteert in hoge trainingskosten. In dit artikel introduceren we StateX, een trainingspijplijn voor het efficiënt uitbreiden van de staten van vooraf getrainde RNN's via post-training. Voor twee populaire klassen van RNN's, lineaire aandacht en state space-modellen, ontwerpen we architectonische aanpassingen voor post-training om de staatsgrootte op te schalen zonder of met een verwaarloosbare toename van modelparameters. Experimenten met modellen tot 1,3B parameters demonstreren dat StateX efficiënt het herinnerings- en in-context leervermogen van RNN's verbetert zonder hoge post-trainingskosten te veroorzaken of andere capaciteiten in gevaar te brengen.

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.