StateX: Miglioramento del Richiamo nelle RNN tramite Espansione degli Stati Post-addestramento

StateX: Enhancing RNN Recall via Post-training State Expansion

September 26, 2025

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

cs.AI

Abstract

Sebbene i modelli basati su Transformer abbiano dimostrato prestazioni eccezionali nella modellazione del linguaggio, le loro elevate complessità comportano costi significativi quando si elaborano contesti lunghi. Al contrario, le reti neurali ricorrenti (RNN) come l'attenzione lineare e i modelli a spazio di stato hanno guadagnato popolarità grazie alle loro complessità costanti per token. Tuttavia, questi modelli ricorrenti faticano nei compiti che richiedono un richiamo accurato delle informazioni contestuali da contesti lunghi, poiché tutte le informazioni contestuali vengono compresse in uno stato ricorrente di dimensioni costanti. Studi precedenti hanno dimostrato che la capacità di richiamo è positivamente correlata alla dimensione dello stato ricorrente, ma addestrare direttamente RNN con stati ricorrenti più grandi comporta costi di addestramento elevati. In questo articolo, introduciamo StateX, una pipeline di addestramento per espandere in modo efficiente gli stati di RNN pre-addestrati attraverso un post-addestramento. Per due classi popolari di RNN, l'attenzione lineare e i modelli a spazio di stato, progettiamo modifiche architetturali post-addestramento per aumentare la dimensione dello stato senza o con un aumento trascurabile dei parametri del modello. Esperimenti su modelli fino a 1,3 miliardi di parametri dimostrano che StateX migliora in modo efficiente la capacità di richiamo e di apprendimento contestuale delle RNN senza comportare costi elevati di post-addestramento o compromettere altre capacità.

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.