Simple linear attention language models balance the recall-throughput
  tradeoff
Recent work has shown that attention-based language models excel at recall,
the ability to ground generations in tokens previously seen in context.
However, the efficiency of attention-based models is bottle-necked during
inference by the KV-cache's aggressive memory consumption. In this work, we
explore whether we can improve language model efficiency (e.g. by reducing
memory consumption) without compromising on recall. By applying experiments and
theory to a broad set of architectures, we identify a key tradeoff between a
model's state size and recall ability. We show that efficient alternatives to
attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but
struggle at recall. We propose BASED a simple architecture combining linear and
sliding window attention. By varying BASED window size and linear attention
feature dimension, we can dial the state size and traverse the pareto frontier
of the recall-memory tradeoff curve, recovering the full quality of attention
on one end and the small state size of attention-alternatives on the other. We
train language models up to 1.3b parameters and show that BASED matches the
strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them
on real-world recall-intensive tasks by 6.22 accuracy points. Implementations
of linear attention are often less efficient than optimized standard attention
implementations. To make BASED competitive, we develop IO-aware algorithms that
enable 24x higher throughput on language generation than FlashAttention-2, when
generating 1024 tokens using 1.3b parameter models. Code for this work is
provided at: https://github.com/HazyResearch/based.
65e157690890daf03e88f25a
Simran Arora
claimed_verified
2024-10-21T07:54:44.122Z
62703f4bbd9c82ff64c2f99f
/avatars/8ed16c6b38a06fd009c39d26f279a6a9.svg
Simran
user
simarora
65e157690890daf03e88f25b
Sabri Eyuboglu
admin_assigned
2024-03-01T10:46:21.661Z
6337537b267cee4d068f604d
/avatars/15267f0759a6570c98ee6a150558fcc0.svg
Sabri Eyuboglu
user
sabrieyuboglu
65e157690890daf03e88f25c
Michael Zhang
65e157690890daf03e88f25d
Aman Timalsina
claimed_verified
2024-03-01T12:44:44.616Z
65e1bb842b435c2dcb5e2199
/avatars/5d4c91401fb989a8005ce43bc3fc2779.svg
Aman Timalsina
user
amantimalsina
65e157690890daf03e88f25e
Silas Alberti
admin_assigned
2024-03-01T10:47:25.387Z
63378a6cba7895577bcd53c3
/avatars/b504bb2088e996325c93f2b56de37c81.svg
Silas Alberti
user
alberti
65e157690890daf03e88f25f
Dylan Zinsley
65e157690890daf03e88f260
James Zou
admin_assigned
2024-03-01T10:48:00.307Z
648a769003fc4a3938bb7943
/avatars/7647f99abdcca4251fcac7783b6fcc8d.svg
zou
user
jameszou707
65e157690890daf03e88f261
Atri Rudra
65e157690890daf03e88f262
Christopher Ré
Simple linear attention language models balance the recall-throughput
  tradeoff
Recent work has shown that attention-based language models excel at recall,
the ability to ground generations in tokens previously seen in context.
However, the efficiency of attention-based models is bottle-necked during
inference by the KV-cache's aggressive memory consumption. In this work, we
explore whether we can improve language model efficiency (e.g. by reducing
memory consumption) without compromising on recall. By applying experiments and
theory to a broad set of architectures, we identify a key tradeoff between a
model's state size and recall ability. We show that efficient alternatives to
attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but
struggle at recall. We propose BASED a simple architecture combining linear and
sliding window attention. By varying BASED window size and linear attention
feature dimension, we can dial the state size and traverse the pareto frontier
of the recall-memory tradeoff curve, recovering the full quality of attention
on one end and the small state size of attention-alternatives on the other. We
train language models up to 1.3b parameters and show that BASED matches the
strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them
on real-world recall-intensive tasks by 6.22 accuracy points. Implementations
of linear attention are often less efficient than optimized standard attention
implementations. To make BASED competitive, we develop IO-aware algorithms that
enable 24x higher throughput on language generation than FlashAttention-2, when
generating 1024 tokens using 1.3b parameter models. Code for this work is
provided at: https://github.com/HazyResearch/based.
BASED, a hybrid architecture combining linear and sliding window attention, improves language model efficiency without sacrificing recall quality.
attention-based language models KV-cache recall H3 Mamba RWKV BASED linear attention sliding window attention state size perplexity real-world recall-intensive tasks IO-aware algorithms FlashAttention-2