Simple linear attention language models balance the recall-throughput tradeoff

• URL: http://arxiv.org/abs/2402.18668v1
• Date: Wed, 28 Feb 2024 19:28:27 GMT
• Title: Simple linear attention language models balance the recall-throughput tradeoff
• Authors: Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, Christopher R\'e
• Abstract summary: We propose BASED, a simple architecture combining linear and sliding window attention. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points.
• Score: 40.08746299497935
• License: http://creativecommons.org/publicdomain/zero/1.0/
• Abstract: 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.

Related papers

• Layer-Condensed KV Cache for Efficient Inference of Large Language Models [44.24593677113768]
  We propose a novel method that only computes and caches the KVs of a small number of layers. Our method achieves up to 26$times$ higher throughput than standard transformers.
  arXiv  Detail & Related papers  (2024-05-17T08:59:46Z)
• CodeArt: Better Code Models by Attention Regularization When Symbols Are Lacking [12.458135956476639]
  Transformer based code models have impressive performance in many software engineering tasks. However, their effectiveness degrades when symbols are missing or not informative. We propose a new method to pre-train general code models when symbols are lacking.
  arXiv  Detail & Related papers  (2024-02-19T05:13:22Z)
• Zoology: Measuring and Improving Recall in Efficient Language Models [42.159338928861864]
  We train a suite of 17 attention and "gated-convolution" language models. We find that gated-convolution architectures still underperform attention by up to 2.1 perplexity points on the Pile. We develop a new formalization of the task called multi-query associative recall (MQAR) that better reflects actual language.
  arXiv  Detail & Related papers  (2023-12-08T09:44:25Z)
• Efficient Streaming Language Models with Attention Sinks [72.20260088848987]
  StreamingLLM is an efficient framework that enables Large Language Models to generalize to infinite sequence lengths without any fine-tuning. We show that StreamingLLM can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more.
  arXiv  Detail & Related papers  (2023-09-29T17:59:56Z)
• Constant Memory Attention Block [74.38724530521277]
  Constant Memory Attention Block (CMAB) is a novel general-purpose attention block that computes its output in constant memory and performs updates in constant computation. We show our proposed methods achieve results competitive with state-of-the-art while being significantly more memory efficient.
  arXiv  Detail & Related papers  (2023-06-21T22:41:58Z)
• LaMemo: Language Modeling with Look-Ahead Memory [50.6248714811912]
  We propose Look-Ahead Memory (LaMemo) that enhances the recurrence memory by incrementally attending to the right-side tokens. LaMemo embraces bi-directional attention and segment recurrence with an additional overhead only linearly proportional to the memory length. Experiments on widely used language modeling benchmarks demonstrate its superiority over the baselines equipped with different types of memory.
  arXiv  Detail & Related papers  (2022-04-15T06:11:25Z)
• ABC: Attention with Bounded-memory Control [67.40631793251997]
  We show that bounded-memory control (ABC) can be subsumed into one abstraction, attention with bounded-memory control (ABC) ABC reveals new, unexplored possibilities. First, it connects several efficient attention variants that would otherwise seem apart. Last, we present a new instance of ABC, which draws inspiration from existing ABC approaches, but replaces their memory-organizing functions with a learned, contextualized one.
  arXiv  Detail & Related papers  (2021-10-06T03:53:25Z)
• Unlocking Pixels for Reinforcement Learning via Implicit Attention [61.666538764049854]
  We make use of new efficient attention algorithms, recently shown to be highly effective for Transformers. This allows our attention-based controllers to scale to larger visual inputs, and facilitate the use of smaller patches. In addition, we propose a new efficient algorithm approximating softmax attention with what we call hybrid random features.
  arXiv  Detail & Related papers  (2021-02-08T17:00:26Z)
• Data-Informed Global Sparseness in Attention Mechanisms for Deep Neural Networks [33.07113523598028]
  We propose Attention Pruning (AP), a framework that observes attention patterns in a fixed dataset and generates a global sparseness mask. AP saves 90% of attention computation for language modeling and about 50% for machine translation and GLUE tasks, maintaining result quality.
  arXiv  Detail & Related papers  (2020-11-20T13:58:21Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.