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

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.

[2305.13245] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints Skip to main content arXiv is now an independent nonprofit! Learn more × Search arXiv Press Enter to search · Advanced search --> Computer Science > Computation and Language arXiv:2305.13245 (cs) [Submitted on 22 May 2023 ( v1 ), last revised 23 Dec 2023 (this version, v3)] Title: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints Authors: Joshua Ainslie , James Lee-Thorp , Michiel de Jong , Yury Zemlyanskiy , Federico Lebrón , Sumit Sanghai View a PD

related reading

• A Mathematical Framework for Transformer Circuitstransformer-circuits.pub
• The Illustrated Transformer – Jay Alammar – Visualizing machine learning one concept at a time.jalammar.github.io
• Multi Query Attention (MQA) and Grouped-Query Attention (GQA)tinkerd.net
• A Gentle Introduction to Multi-Head Latent Attention (MLA) - MachineLearningMastery.commachinelearningmastery.com
• transformer_attention.pdfarxiv.org
• Transformers Inference Optimization Toolset | AstraBlogastralord.github.io
• 1706.03762arxiv.org
• Transformer Explainer: LLM Transformer Model Visually Explainedpoloclub.github.io
• Mamba: The Easy Wayjackcook.com
• Getting Caught Up to Modern LLM Research | Samarth Goeldev.samarthgoel.com
• frontier model training methodologies | Alex Wa's Blogdjdumpling.github.io
• The Annotated Transformernlp.seas.harvard.edu