AllMem: A Memory-centric Recipe for Efficient Long-context Modeling

TL;DR

AllMem introduces a hybrid architecture combining sliding window attention with memory networks, enabling efficient long-context modeling in LLMs with reduced computational costs and mitigated catastrophic forgetting.

Contribution

The paper presents a novel hybrid architecture and a memory-efficient fine-tuning method that allow pre-trained LLMs to effectively handle ultra-long contexts with minimal performance loss.

Findings

Achieves near-lossless performance on 37k LongBench with 4k window.

Outperforms full attention on 128k context in InfiniteBench.

Reduces computational and memory overhead during long-sequence inference.

Abstract

Large Language Models (LLMs) encounter significant performance bottlenecks in long-sequence tasks due to the computational complexity and memory overhead inherent in the self-attention mechanism. To address these challenges, we introduce \textsc{AllMem}, a novel and efficient hybrid architecture that integrates Sliding Window Attention (SWA) with non-linear Test-Time Training (TTT) memory networks. \textsc{AllMem} enables models to effectively scale to ultra-long contexts while mitigating catastrophic forgetting. This approach not only overcomes the representation constraints typical of linear memory models but also significantly reduces the computational and memory footprint during long-sequence inference. Furthermore, we implement a Memory-Efficient Fine-Tuning strategy to replace standard attention layers in pre-trained models with memory-augmented sliding window layers. This framework facilitates the efficient transformation of any off-the-shelf pre-trained LLM into an \textsc{AllMem}-based architecture. Empirical evaluations confirm that our 4k window model achieves near-lossless performance on 37k LongBench with a marginal 0.83 drop compared to full attention. Furthermore, on InfiniteBench at a 128k context, our 8k window variant outperforms full attention, which validates the effectiveness of our parameterized memory in mitigating noise and maintaining robust long-range modeling without the prohibitive costs of global attention.