小模型用压缩记忆提升长文本处理能力,成本低效果好。
Probing the Limits of Compressive Memory: A Study of Infini-Attention in Small-Scale Pretraining
- 用压缩记忆机制保留历史信息,实现小模型长序列推理。
- 在16,384词长上下文中,检索准确率比基线高31%。
- 适合资源受限场景,对轻量级语言模型有实用价值。
本研究针对小型语言模型(SLMs)的小规模预训练,探索在有限数据与算力下提升效率的路径,旨在降低使用成本、增强低资源环境下的可访问性。为改善紧凑模型的长上下文外推能力,聚焦于Infini-attention机制——该机制通过压缩过往片段构建压缩记忆,同时保留局部注意力。本文采用300M参数的LLaMA模型进行实证研究,结果显示模型训练稳定,并在长上下文检索任务中优于基线。我们发现平衡因子是影响性能的关键,长期序列中重复压缩会导致检索准确率下降。尽管如此,Infini-attention仍能有效弥补小模型参数不足的问题。尤其在16,384词上下文长度下,其检索准确率较基线最高提升31%。研究表明,采用如Infini-attention这类结构化记忆架构,有助于小模型实现稳健的长上下文能力。
原文摘要 · Abstract (English)
This study investigates small-scale pretraining for Small Language Models (SLMs) to enable efficient use of limited data and compute, improve accessibility in low-resource settings and reduce costs. To enhance long-context extrapolation in compact models, we focus on Infini-attention, which builds a compressed memory from past segments while preserving local attention. In our work, we conduct an empirical study using 300M-parameter LLaMA models pretrained with Infini-attention. The model demonstrates training stability and outperforms the baseline in long-context retrieval. We identify the balance factor as a key part of the model performance, and we found that retrieval accuracy drops with repeated memory compressions over long sequences. Even so, Infini-attention still effectively compensates for the SLM's limited parameters. Particularly, despite performance degradation at a 16,384-token context, the Infini-attention model achieves up to 31% higher accuracy than the baseline. Our findings suggest that achieving robust long-context capability in SLMs benefits from architectural memory like Infini-attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。