StateX: Улучшение запоминания в рекуррентных нейронных сетях за счет расширения состояний после обучения

StateX: Enhancing RNN Recall via Post-training State Expansion

September 26, 2025

Авторы: Xingyu Shen, Yingfa Chen, Zhen Leng Thai, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

Аннотация

Хотя модели на основе архитектуры Transformer продемонстрировали выдающуюся производительность в задачах языкового моделирования, их высокая сложность приводит к значительным затратам при обработке длинных контекстов. В то же время рекуррентные нейронные сети (RNN), такие как модели с линейным вниманием и модели пространства состояний, приобрели популярность благодаря их постоянной сложности на один токен. Однако эти рекуррентные модели сталкиваются с трудностями в задачах, требующих точного воспроизведения контекстной информации из длинных контекстов, поскольку вся контекстная информация сжимается в состояние постоянного размера. Предыдущие исследования показали, что способность к воспроизведению положительно коррелирует с размером рекуррентного состояния, однако прямое обучение RNN с увеличенными состояниями приводит к высоким затратам на обучение. В данной работе мы представляем StateX — конвейер обучения для эффективного расширения состояний предварительно обученных RNN посредством пост-обучения. Для двух популярных классов RNN — моделей с линейным вниманием и моделей пространства состояний — мы разработали архитектурные модификации для пост-обучения, позволяющие увеличить размер состояния без или с незначительным увеличением числа параметров модели. Эксперименты на моделях с числом параметров до 1,3 миллиарда демонстрируют, что StateX эффективно улучшает способность RNN к воспроизведению и обучению в контексте без высоких затрат на пост-обучение и без ущерба для других возможностей.

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.