Memory-efficient Transformers via Top-$k$ Attention

• URL: http://arxiv.org/abs/2106.06899v1
• Date: Sun, 13 Jun 2021 02:30:23 GMT
• Title: Memory-efficient Transformers via Top-$k$ Attention
• Authors: Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, Jonathan Berant
• Abstract summary: In this work, we propose a simple yet highly accurate approximation for vanilla attention. We process the queries in chunks, and for each query, compute the top-$k$ scores with respect to the keys. We show our approach leads to accuracy that is nearly-identical to vanilla attention in multiple setups including training from scratch, fine-tuning, and zero-shot inference.
• Score: 23.672065688109395
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Following the success of dot-product attention in Transformers, numerous approximations have been recently proposed to address its quadratic complexity with respect to the input length. While these variants are memory and compute efficient, it is not possible to directly use them with popular pre-trained language models trained using vanilla attention, without an expensive corrective pre-training stage. In this work, we propose a simple yet highly accurate approximation for vanilla attention. We process the queries in chunks, and for each query, compute the top-$k$ scores with respect to the keys. Our approach offers several advantages: (a) its memory usage is linear in the input size, similar to linear attention variants, such as Performer and RFA (b) it is a drop-in replacement for vanilla attention that does not require any corrective pre-training, and (c) it can also lead to significant memory savings in the feed-forward layers after casting them into the familiar query-key-value framework. We evaluate the quality of top-$k$ approximation for multi-head attention layers on the Long Range Arena Benchmark, and for feed-forward layers of T5 and UnifiedQA on multiple QA datasets. We show our approach leads to accuracy that is nearly-identical to vanilla attention in multiple setups including training from scratch, fine-tuning, and zero-shot inference.

Related papers

• CORM: Cache Optimization with Recent Message for Large Language Model Inference [57.109354287786154]
  We introduce an innovative method for optimizing the KV cache, which considerably minimizes its memory footprint. CORM, a KV cache eviction policy, dynamically retains essential key-value pairs for inference without the need for model fine-tuning. Our validation shows that CORM reduces the inference memory usage of KV cache by up to 70% with negligible performance degradation across six tasks in LongBench.
  arXiv  Detail & Related papers  (2024-04-24T16:11:54Z)
• Towards Model-Size Agnostic, Compute-Free, Memorization-based Inference of Deep Learning [5.41530201129053]
  This paper proposes a novel memorization-based inference (MBI) that is compute free and only requires lookups. Specifically, our work capitalizes on the inference mechanism of the recurrent attention model (RAM) By leveraging the low-dimensionality of glimpse, our inference procedure stores key value pairs comprising of glimpse location, patch vector, etc. in a table. The computations are obviated during inference by utilizing the table to read out key-value pairs and performing compute-free inference by memorization.
  arXiv  Detail & Related papers  (2023-07-14T21:01:59Z)
• Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model [89.8764435351222]
  We propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
  arXiv  Detail & Related papers  (2023-05-24T15:52:08Z)
• AdaNPC: Exploring Non-Parametric Classifier for Test-Time Adaptation [64.9230895853942]
  Domain generalization can be arbitrarily hard without exploiting target domain information. Test-time adaptive (TTA) methods are proposed to address this issue. In this work, we adopt Non-Parametric to perform the test-time Adaptation (AdaNPC)
  arXiv  Detail & Related papers  (2023-04-25T04:23:13Z)
• CloudAttention: Efficient Multi-Scale Attention Scheme For 3D Point Cloud Learning [81.85951026033787]
  We set transformers in this work and incorporate them into a hierarchical framework for shape classification and part and scene segmentation. We also compute efficient and dynamic global cross attentions by leveraging sampling and grouping at each iteration. The proposed hierarchical model achieves state-of-the-art shape classification in mean accuracy and yields results on par with the previous segmentation methods.
  arXiv  Detail & Related papers  (2022-07-31T21:39:15Z)
• Transformers Can Do Bayesian Inference [56.99390658880008]
  We present Prior-Data Fitted Networks (PFNs) PFNs leverage in-context learning in large-scale machine learning techniques to approximate a large set of posteriors. We demonstrate that PFNs can near-perfectly mimic Gaussian processes and also enable efficient Bayesian inference for intractable problems.
  arXiv  Detail & Related papers  (2021-12-20T13:07:39Z)
• DA-Transformer: Distance-aware Transformer [87.20061062572391]
  DA-Transformer is a distance-aware Transformer that can exploit the real distance. In this paper, we propose DA-Transformer, which is a distance-aware Transformer that can exploit the real distance.
  arXiv  Detail & Related papers  (2020-10-14T10:09:01Z)
• Fast Transformers with Clustered Attention [14.448898156256478]
  We propose clustered attention, which instead of computing the attention for every query, groups queries into clusters and computes attention just for the centroids. This results in a model with linear complexity with respect to the sequence length for a fixed number of clusters. We evaluate our approach on two automatic speech recognition datasets and show that our model consistently outperforms vanilla transformers for a given computational budget.
  arXiv  Detail & Related papers  (2020-07-09T14:17:50Z)
• DeFormer: Decomposing Pre-trained Transformers for Faster Question Answering [22.178201429268103]
  Transformer-based QA models use input-wide self-attention across both the question and the input passage. We introduce DeFormer, which substitutes the full self-attention with question-wide and passage-wide self-attentions in the lower layers. We show DeFormer versions of BERT and XLNet can be used to speed up QA by over 4.3x and with simple distillation-based losses they incur only a 1% drop in accuracy.
  arXiv  Detail & Related papers  (2020-05-02T04:28:22Z)

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.