StateX: Enhancing RNN Recall via Post-training State Expansion
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866908990444666880 |
|---|---|
| author | Shen, Xingyu Chen, Yingfa Thai, Zhen Leng Han, Xu Liu, Zhiyuan Sun, Maosong |
| author_facet | Shen, Xingyu Chen, Yingfa Thai, Zhen Leng Han, Xu Liu, Zhiyuan Sun, Maosong |
| contents | Recurrent neural networks (RNNs), such as linear attention and state-space models, have gained popularity due to their constant per-token complexity when processing long contexts. 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 fixed-size recurrent state. Previous studies have shown that recall ability is positively correlated with the recurrent state size, yet directly training RNNs with large recurrent states results in high training costs. In this paper, we introduce StateX, a post-training framework that efficiently expands the states of pre-trained RNNs. For two popular classes of RNNs, linear attention and state-space models, we design post-training architectural modifications in StateX, to scale up the state size with no or negligible increase in model parameters. Experiments on models with up to 1.3B parameters demonstrate that StateX efficiently enhances the recall and in-context learning performance of RNNs without incurring high post-training costs or compromising other capabilities. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2509_22630 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | StateX: Enhancing RNN Recall via Post-training State Expansion Shen, Xingyu Chen, Yingfa Thai, Zhen Leng Han, Xu Liu, Zhiyuan Sun, Maosong Computation and Language Artificial Intelligence Machine Learning Recurrent neural networks (RNNs), such as linear attention and state-space models, have gained popularity due to their constant per-token complexity when processing long contexts. 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 fixed-size recurrent state. Previous studies have shown that recall ability is positively correlated with the recurrent state size, yet directly training RNNs with large recurrent states results in high training costs. In this paper, we introduce StateX, a post-training framework that efficiently expands the states of pre-trained RNNs. For two popular classes of RNNs, linear attention and state-space models, we design post-training architectural modifications in StateX, to scale up the state size with no or negligible increase in model parameters. Experiments on models with up to 1.3B parameters demonstrate that StateX efficiently enhances the recall and in-context learning performance of RNNs without incurring high post-training costs or compromising other capabilities. |
| title | StateX: Enhancing RNN Recall via Post-training State Expansion |
| topic | Computation and Language Artificial Intelligence Machine Learning |
| url | https://arxiv.org/abs/2509.22630 |