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

• URL: http://arxiv.org/abs/2412.20677v1
• Date: Mon, 30 Dec 2024 03:05:45 GMT
• Title: Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA
• Authors: Qingyun Jin, Xiaohui Song, Feng Zhou, Zengchang Qin,
• Abstract summary: We propose a low-cost method for pruning MHA models into GQA models with any compression ratio of key-value heads. Our strategy can compress up to 87.5% of key-value heads of the LLaMA2-7B model without too much performance degradation.
• Score: 8.305827430948654
• License:
• Abstract: Large language models have been shown to perform well on a variety of natural language processing problems. However, as the model size and the input sequence's length increase, the rapid increase of KV Cache significantly slows down inference speed. Therefore GQA model, as an alternative to MHA model, has been widely introduced into LLMs. In this work, we propose a low-cost method for pruning MHA models into GQA models with any compression ratio of key-value heads. Our method is based on $\mathit{L_0}$ masks to gradually remove redundant parameters. In addition, we apply orthogonal transformations to attention heads without changing the model to increase similarity between attention heads before pruning training, in order to further improve performance of the model. Our method can be compatible with rotary position embedding (RoPE), which means the model after training can be fully adapted to the mainstream standard GQA framework. Experiments demonstrate that our strategy can compress up to 87.5% of key-value heads of the LLaMA2-7B model without too much performance degradation, just achieved through supervised fine-tuning.

Related papers

• Nudging: Inference-time Alignment via Model Collaboration [18.530367090350605]
  We propose nudging, a plug-and-play algorithm that aligns any base model at inference time using a small aligned model. Nudging is motivated by recent findings that alignment primarily alters the model's behavior on a small subset of stylistic tokens. We evaluate the effectiveness of nudging across 3 model families and 13 tasks, covering reasoning, general knowledge, instruction following, and safety benchmarks.
  arXiv  Detail & Related papers  (2024-10-11T23:24:38Z)
• Promises and Pitfalls of Generative Masked Language Modeling: Theoretical Framework and Practical Guidelines [74.42485647685272]
  We focus on Generative Masked Language Models (GMLMs) We train a model to fit conditional probabilities of the data distribution via masking, which are subsequently used as inputs to a Markov Chain to draw samples from the model. We adapt the T5 model for iteratively-refined parallel decoding, achieving 2-3x speedup in machine translation with minimal sacrifice in quality.
  arXiv  Detail & Related papers  (2024-07-22T18:00:00Z)
• EMR-Merging: Tuning-Free High-Performance Model Merging [55.03509900949149]
  We show that Elect, Mask & Rescale-Merging (EMR-Merging) shows outstanding performance compared to existing merging methods. EMR-Merging is tuning-free, thus requiring no data availability or any additional training while showing impressive performance.
  arXiv  Detail & Related papers  (2024-05-23T05:25:45Z)
• Predictable MDP Abstraction for Unsupervised Model-Based RL [93.91375268580806]
  We propose predictable MDP abstraction (PMA) Instead of training a predictive model on the original MDP, we train a model on a transformed MDP with a learned action space. We theoretically analyze PMA and empirically demonstrate that PMA leads to significant improvements over prior unsupervised model-based RL approaches.
  arXiv  Detail & Related papers  (2023-02-08T07:37:51Z)
• Masked Autoencoding for Scalable and Generalizable Decision Making [93.84855114717062]
  MaskDP is a simple and scalable self-supervised pretraining method for reinforcement learning and behavioral cloning. We find that a MaskDP model gains the capability of zero-shot transfer to new BC tasks, such as single and multiple goal reaching.
  arXiv  Detail & Related papers  (2022-11-23T07:04:41Z)
• CAMERO: Consistency Regularized Ensemble of Perturbed Language Models with Weight Sharing [83.63107444454938]
  We propose a consistency-regularized ensemble learning approach based on perturbed models, named CAMERO. Specifically, we share the weights of bottom layers across all models and apply different perturbations to the hidden representations for different models, which can effectively promote the model diversity. Our experiments using large language models demonstrate that CAMERO significantly improves the generalization performance of the ensemble model.
  arXiv  Detail & Related papers  (2022-04-13T19:54:51Z)
• Scaling Hidden Markov Language Models [118.55908381553056]
  This work revisits the challenge of scaling HMMs to language modeling datasets. We propose methods for scaling HMMs to massive state spaces while maintaining efficient exact inference, a compact parameterization, and effective regularization.
  arXiv  Detail & Related papers  (2020-11-09T18:51:55Z)
• Stochastic Attention Head Removal: A simple and effective method for improving Transformer Based ASR Models [40.991809705930955]
  We propose to randomly remove attention heads during training and keep all attention heads at test time, thus the final model is an ensemble of models with different architectures. Our method gives consistent performance gains over strong baselines on the Wall Street Journal, AISHELL, Switchboard and AMI datasets.
  arXiv  Detail & Related papers  (2020-11-08T15:41:03Z)

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.