通过训练后扩展状态，提升RNN长文本记忆能力。
StateX: Enhancing RNN Recall via Post-training State Expansion
- 训练后动态扩展RNN状态大小，不增加参数量。
- 1.3B参数模型实验显示，回忆和上下文学习性能显著提升。
- 适合需要长程记忆的RNN应用，如长文档处理。
循环神经网络（RNN）如线性注意力和状态空间模型因处理长上下文时每标记复杂度恒定而受到欢迎。然而，这些模型在需要准确回忆长上下文信息的任务中表现不佳，因为所有上下文信息被压缩进固定大小的循环状态。先前研究发现，回忆能力与循环状态大小正相关，但直接训练大状态的RNN会导致高训练成本。本文提出StateX，一种训练后框架，可高效扩展预训练RNN的状态。针对线性注意力和状态空间模型两类主流RNN，StateX设计了无需或极少增加参数的架构修改，以扩大状态规模。在高达1.3B参数的模型上实验表明，StateX在几乎无额外训练成本的情况下显著提升了RNN的回忆能力和上下文学习性能，且未损害其他能力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。