StateX: Aprimorando a Capacidade de Recuperação de RNNs por meio de Expansão de Estados Pós-treinamento

StateX: Enhancing RNN Recall via Post-training State Expansion

September 26, 2025

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

cs.AI

Resumo

Embora os modelos baseados em Transformers tenham demonstrado um desempenho notável na modelagem de linguagem, suas altas complexidades resultam em custos elevados ao processar contextos longos. Em contraste, redes neurais recorrentes (RNNs), como modelos de atenção linear e modelos de espaço de estados, ganharam popularidade devido às suas complexidades constantes por token. No entanto, esses modelos recorrentes enfrentam dificuldades em tarefas que exigem a recuperação precisa de informações contextuais de contextos longos, pois todas as informações contextuais são comprimidas em um estado recorrente de tamanho constante. Trabalhos anteriores mostraram que a capacidade de recuperação está positivamente correlacionada com o tamanho do estado recorrente, mas o treinamento direto de RNNs com estados recorrentes maiores resulta em custos elevados de treinamento. Neste artigo, apresentamos o StateX, um pipeline de treinamento para expandir eficientemente os estados de RNNs pré-treinados por meio de pós-treinamento. Para duas classes populares de RNNs, atenção linear e modelos de espaço de estados, projetamos modificações arquitetônicas de pós-treinamento para aumentar o tamanho do estado sem aumento ou com aumento insignificante no número de parâmetros do modelo. Experimentos em modelos com até 1,3 bilhão de parâmetros demonstram que o StateX melhora eficientemente a capacidade de recuperação e aprendizado em contexto das RNNs sem incorrer em altos custos de pós-treinamento ou comprometer outras capacidades.

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.