StateX: Mejora de la Recuperación en RNN mediante Expansión de Estados Post-entrenamiento

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

Resumen

Aunque los modelos basados en Transformers han demostrado un rendimiento notable en el modelado del lenguaje, sus altas complejidades resultan en costos elevados al procesar contextos largos. En contraste, las redes neuronales recurrentes (RNN), como los modelos de atención lineal y los modelos de espacio de estados, han ganado popularidad debido a sus complejidades constantes por token. Sin embargo, estos modelos recurrentes enfrentan dificultades en tareas que requieren una recuperación precisa de información contextual de contextos largos, ya que toda la información contextual se comprime en un estado recurrente de tamaño constante. Trabajos previos han demostrado que la capacidad de recuperación está positivamente correlacionada con el tamaño del estado recurrente, pero entrenar directamente RNNs con estados recurrentes más grandes resulta en costos de entrenamiento elevados. En este artículo, presentamos StateX, una canalización de entrenamiento para expandir eficientemente los estados de RNNs preentrenados mediante post-entrenamiento. Para dos clases populares de RNNs, la atención lineal y los modelos de espacio de estados, diseñamos modificaciones arquitectónicas de post-entrenamiento para escalar el tamaño del estado sin aumentar o con un aumento insignificante en los parámetros del modelo. Experimentos en modelos de hasta 1.3 mil millones de parámetros demuestran que StateX mejora eficientemente la capacidad de recuperación y aprendizaje en contexto de las RNNs sin incurrir en altos costos de post-entrenamiento ni comprometer otras 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.