title: Simple linear attention language models balance the recall-throughput
    tradeoff
summary: 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.
paper_authors._id: 65e157690890daf03e88f25a
paper_authors.name: Simran Arora
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-21T07:54:44.122Z
paper_authors._id: 62703f4bbd9c82ff64c2f99f
paper_authors.avatarUrl: /avatars/8ed16c6b38a06fd009c39d26f279a6a9.svg
paper_authors.fullname: Simran
paper_authors.type: user
paper_authors.user: simarora
paper_authors._id: 65e157690890daf03e88f25b
paper_authors.name: Sabri Eyuboglu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-03-01T10:46:21.661Z
paper_authors._id: 6337537b267cee4d068f604d
paper_authors.avatarUrl: /avatars/15267f0759a6570c98ee6a150558fcc0.svg
paper_authors.fullname: Sabri Eyuboglu
paper_authors.type: user
paper_authors.user: sabrieyuboglu
paper_authors._id: 65e157690890daf03e88f25c
paper_authors.name: Michael Zhang
paper_authors._id: 65e157690890daf03e88f25d
paper_authors.name: Aman Timalsina
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-03-01T12:44:44.616Z
paper_authors._id: 65e1bb842b435c2dcb5e2199
paper_authors.avatarUrl: /avatars/5d4c91401fb989a8005ce43bc3fc2779.svg
paper_authors.fullname: Aman Timalsina
paper_authors.type: user
paper_authors.user: amantimalsina
paper_authors._id: 65e157690890daf03e88f25e
paper_authors.name: Silas Alberti
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-03-01T10:47:25.387Z
paper_authors._id: 63378a6cba7895577bcd53c3
paper_authors.avatarUrl: /avatars/b504bb2088e996325c93f2b56de37c81.svg
paper_authors.fullname: Silas Alberti
paper_authors.type: user
paper_authors.user: alberti
paper_authors._id: 65e157690890daf03e88f25f
paper_authors.name: Dylan Zinsley
paper_authors._id: 65e157690890daf03e88f260
paper_authors.name: James Zou
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-03-01T10:48:00.307Z
paper_authors._id: 648a769003fc4a3938bb7943
paper_authors.avatarUrl: /avatars/7647f99abdcca4251fcac7783b6fcc8d.svg
paper_authors.fullname: zou
paper_authors.type: user
paper_authors.user: jameszou707
paper_authors._id: 65e157690890daf03e88f261
paper_authors.name: Atri Rudra
paper_authors._id: 65e157690890daf03e88f262
paper_authors.name: Christopher Ré
paper_title: Simple linear attention language models balance the recall-throughput
    tradeoff
paper_summary: 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.
paper_ai_summary: BASED, a hybrid architecture combining linear and sliding window attention, improves language model efficiency without sacrificing recall quality.
paper_ai_keywords: 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