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

TL;DR

This work investigates adapting Infini-attention, a memory-augmented attention mechanism, to a compact 300M-parameter LLaMA model trained on limited data. By integrating cross-segment memory with local attention and a trainable balance factor, the study evaluates effects on long-context retrieval, training dynamics, and fine-tuning performance. While the short training regime and dataset bias limit cross-segment memory learning, Infini-attention demonstrates improved long-context retrieval after supervision and gains on several general reasoning benchmarks, along with notable training stability advantages. The results suggest that architectural memory mechanisms can benefit small-scale language models, though future work should leverage longer documents, larger models, and multi-task fine-tuning to realize robust long-context capabilities in resource-constrained settings.

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.

Figures (6)

• Figure 1: Training loss comparison between baseline and Infini-attention models.
• Figure 2: Gradient norm comparison between baseline and Infini-attention.
• Figure 3: Mean activated balance factor across attention heads during Infini-attention training. The values converge around 0.30, reflecting a dataset bias toward shorter sequences where local attention is more heavily weighted than memory retrieval.
• Figure 4: Distribution of balance factors across attention heads. The distribution shows significant variation, with most heads clustering around lower values (favoring local attention) while some heads specialize in memory retrieval with higher balance factors.
• Figure 5: Layer-wise memory preference rate in Infini-attention. Lower layers show strong memory utilization, while higher layers show increased reliance on local attention.