Attamba: Attending To Multi-Token States

TL;DR

Attamba introduces a novel architecture combining state-space models with transformers to efficiently compress token chunks, achieving significant improvements in perplexity and computational efficiency while maintaining flexible sequence handling.

Contribution

The paper presents Attamba, a new model that integrates state-space models into transformers for improved efficiency and flexibility in sequence processing.

Findings

24% improved perplexity with similar KV-Cache and attention footprint

Approximately 4 times smaller KV-Cache and Attention FLOPs for 5% perplexity trade-off

Enables smooth transition between quadratic and linear scaling in attention complexity

Abstract

When predicting the next token in a sequence, vanilla transformers compute attention over all previous tokens, resulting in quadratic scaling of compute with sequence length. State-space models compress the entire sequence of tokens into a fixed-dimensional representation to improve efficiency, while other architectures achieve sub-quadratic complexity via low-rank projections or sparse attention patterns over the sequence. In this paper, we introduce Attamba, a novel architecture that uses state-space models to compress chunks of tokens and applies attention on these compressed key-value representations. We find that replacing key and value projections in a transformer with SSMs can improve model quality and enable flexible token chunking, resulting in 24% improved perplexity with transformer of similar KV-Cache and attention footprint, and ~4 times smaller KV-Cache and Attention FLOPs for 5% perplexity trade-off. Attamba can perform attention on chunked-sequences of variable length, enabling a smooth transition between quadratic and linear scaling, offering adaptable efficiency gains.