LoLA: Low-Rank Linear Attention With Sparse Caching

TL;DR

LoLA enhances linear attention in transformers by integrating a multi-system memory augmentation, significantly improving long-term recall and performance on various tasks without increasing training complexity.

Contribution

LoLA introduces a training-free memory augmentation for linear attention, boosting associative recall and efficiency in long-context scenarios.

Findings

Achieves 97.4% accuracy on pass-key retrieval tasks.

Uses 4.6x smaller cache than Llama-3.1 8B.

Outperforms other models on zero-shot reasoning.

Abstract

The per-token cost of transformer inference scales with context length, preventing its application to lifelong in-context learning. Linear attention is an efficient alternative that maintains a constant memory footprint, even on infinite context lengths. While this is a potential candidate for lifelong learning, it falls short in memory capacity. In this paper, we propose LoLA, a training-free augmentation to linear attention that boosts associative recall. LoLA distributes past key-value pairs from context into three memory systems: (i) recent pairs in a local sliding window cache; (ii) difficult-to-memorize pairs in a sparse, global cache; and (iii) generic pairs in the recurrent hidden state of linear attention. We show through ablations that our self-recall error metric is crucial to efficiently manage long-term associative memories. On pass-key retrieval tasks, LoLA improves the base model's performance from 0.6% to 97.4% accuracy. This is achieved with a 4.6x smaller cache than Llama-3.1 8B on 4K context length. LoLA also outperforms other 1B and 8B parameter subquadratic models on zero-shot commonsense reasoning tasks.