探究压缩记忆的极限:小规模预训练中 Infini-attention 的研究
机器学习
2026-01-01 v1 人工智能
计算与语言
摘要
本研究调查了小语言模型(SLMs)的小规模预训练,以实现有限数据和计算资源的高效利用,提高低资源环境下的可及性并降低成本。为了增强紧凑模型的长上下文外推能力,我们重点关注 Infini-attention,它在保留局部注意力的同时从过去的片段构建压缩记忆。在我们的工作中,我们使用以 Infini-attention 预训练的 300M 参数 LLaMA 模型进行了实证研究。该模型表现出训练稳定性,并在长上下文检索中优于基线。我们确定平衡因子是模型性能的关键部分,并发现随着长序列上重复的记忆压缩,检索准确率会下降。即便如此,Infini-attention 仍然有效地弥补了 SLM 有限的参数量。特别是,尽管在 16,384 token 的上下文下性能有所下降,Infini-attention 模型的准确率仍比基线高出 31%。我们的研究结果表明,在 SLM 中实现稳健的长上下文能力得益于像 Infini-attention 这样的架构记忆。
引用
@article{arxiv.2512.23862,
title = {Probing the Limits of Compressive Memory: A Study of Infini-Attention in Small-Scale Pretraining},
author = {Ruizhe Huang and Kexuan Zhang and Yihao Fang and Baifeng Yu},
journal= {arXiv preprint arXiv:2512.23862},
year = {2026}
}