StateX：通过后训练状态扩展增强 RNN 召回能力
计算与语言
2026-04-27 v3 人工智能
机器学习
摘要
循环神经网络（RNN），如线性注意力和状态空间模型，因其在处理长上下文时具有恒定的每 token 复杂度而受到欢迎。然而，这些循环模型在需要从长上下文中准确召回上下文信息的任务中表现不佳，因为所有上下文信息都被压缩到固定大小的循环状态中。先前的研究表明，召回能力与循环状态大小正相关，但直接使用大循环状态训练 RNN 会导致高昂的训练成本。在本文中，我们引入了 StateX，一个高效扩展预训练 RNN 状态的后训练框架。对于两类流行的 RNN——线性注意力和状态空间模型，我们在 StateX 中设计了后训练架构修改，以在模型参数不增加或可忽略增加的情况下扩大状态大小。在具有高达 1.3B 参数的模型上的实验表明，StateX 高效地增强了 RNN 的召回和上下文学习性能，且未带来高昂的后训练成本或损害其他能力。
引用
@article{arxiv.2509.22630,
title = {StateX: Enhancing RNN Recall via Post-training State Expansion},
author = {Xingyu Shen and Yingfa Chen and Zhen Leng Thai and Xu Han and Zhiyuan Liu and Maosong Sun},
journal= {arXiv preprint arXiv:2509.22630},
year = {2026}
}