GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón +1 more
5/22/2023
cs.CLcs.LG

No official code found. Best implementation: fkodom/grouped-query-attention-pytorch (194 stars, MIT, last commit 2024-05-09). 1 more below.

Abstract

Multi-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation, and moreover it may not be desirable to train a separate model just for faster inference. We (1) propose a recipe for uptraining existing multi-head language model checkpoints into models with MQA using 5% of original pre-training compute, and (2) introduce grouped-query attention (GQA), a generalization of multi-query attention which uses an intermediate (more than one, less than number of query heads) number of key-value heads. We show that uptrained GQA achieves quality close to multi-head attention with comparable speed to MQA.

View on arXivView PDF

Code Implementations(2)

fkodom/grouped-query-attention-pytorchCommunity64%
MIT

(Unofficial) PyTorch implementation of grouped-query attention (GQA) from "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" (https://arxiv.org/pdf/2305.13245.pdf)

19414PythonJul 19, 20232 years agoMIT
AbelHutten/deformable-grouped-query-attentionCommunity62%
MIT

Combining Grouped-Query Attention (https://arxiv.org/abs/2305.13245) with Deformable Attention (https://arxiv.org/abs/2201.00520) in PyTorch.

10PythonMar 8, 20258 months agoMIT

Similar Papers

Weighted Grouped Query Attention in Transformers

91% match2024

Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

88% match1 repo172024

Sparse Query Attention (SQA): A Computationally Efficient Attention Mechanism with Query Heads Reduction

87% match2025

DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion

87% match2026

Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA

87% match1 repo32024

Beyond Uniform Query Distribution: Key-Driven Grouped Query Attention

87% match1 repo2024

Cite this paper

@article{ainslie2023gqa,
  title  = {GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints},
  author = {Joshua Ainslie and James Lee-Thorp and Michiel de Jong and Yury Zemlyanskiy and Federico Lebrón and Sumit Sanghai},
  year   = {2023},
  eprint = {2305.13245},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2305.13245v3}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2305.13245.svg)](https://codeofpaper.com/paper/2305.13245)

Discussion