Probing the Limits of Compressive Memory: A Study of Infini-Attention in Small-Scale Pretraining

TL;DR

This paper explores Infini-attention, a memory-augmented attention mechanism, in small-scale language models, demonstrating its potential to improve long-context understanding and retrieval despite some performance trade-offs.

Contribution

It provides an empirical evaluation of Infini-attention in 300M-parameter models, highlighting its benefits and limitations for long-context processing in small language models.

Findings

Infini-attention improves long-context retrieval accuracy.

Performance drops with repeated memory compressions.

Achieves up to 31% higher accuracy at 16,384 tokens.

Abstract

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.