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

• URL: http://arxiv.org/abs/2406.06567v2
• Date: Sat, 07 Dec 2024 13:23:39 GMT
• Title: DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion
• Authors: Yilong Chen, Linhao Zhang, Junyuan Shang, Zhenyu Zhang, Tingwen Liu, Shuohuan Wang, Yu Sun,
• Abstract summary: We propose a Decoupled-Head Attention (DHA) mechanism for large language models (LLMs)<n>DHA adaptively configures group sharing for key heads and value heads across various layers, achieving a better balance between performance and efficiency.<n>Our experiments show that DHA remarkably requires a mere 0.25% of the original model's pre-training budgets to achieve 97.6% of performance.
• Score: 29.531814426276885
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: Large language models (LLMs) with billions of parameters demonstrate impressive performance. However, the widely used Multi-Head Attention (MHA) in LLMs incurs substantial computational and memory costs during inference. While some efforts have optimized attention mechanisms by pruning heads or sharing parameters among heads, these methods often lead to performance degradation or necessitate substantial continued pre-training costs to restore performance. Based on the analysis of attention redundancy, we design a Decoupled-Head Attention (DHA) mechanism. DHA adaptively configures group sharing for key heads and value heads across various layers, achieving a better balance between performance and efficiency. Inspired by the observation of clustering similar heads, we propose to progressively transform the MHA checkpoint into the DHA model through linear fusion of similar head parameters step by step, retaining the parametric knowledge of the MHA checkpoint. We construct DHA models by transforming various scales of MHA checkpoints given target head budgets. Our experiments show that DHA remarkably requires a mere 0.25\% of the original model's pre-training budgets to achieve 97.6\% of performance while saving 75\% of KV cache. Compared to Group-Query Attention (GQA), DHA achieves a 5$\times$ training acceleration, a maximum of 13.93\% performance improvement under 0.01\% pre-training budget, and 4\% relative improvement under 0.05\% pre-training budget.

Related papers

• IAM: Efficient Inference through Attention Mapping between Different-scale LLMs [74.81417160018856]
  IAM framework achieves dual benefits of accelerated attention computation and reduced KV cache usage.<n>We show that IAM can accelerate prefill by 15% and reduce KV cache usage by 22.1% without appreciably sacrificing performance.
  arXiv  Detail & Related papers  (2025-07-16T06:39:11Z)
• SAS: Simulated Attention Score [75.1409882298863]
  We introduce Simulated Attention Score (SAS), which maintains a compact model size while simulating a larger number of attention heads and hidden feature dimension per head.<n>Comprehensive experiments on a variety of datasets and tasks demonstrate the effectiveness of the proposed SAS method.
  arXiv  Detail & Related papers  (2025-07-10T12:16:16Z)
• Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs [74.74225314708225]
  Multi-head Latent Attention (MLA) is an innovative architecture designed to ensure efficient and economical inference. This paper proposes the first data-efficient fine-tuning method for transitioning from Multi-Head Attention to MLA.
  arXiv  Detail & Related papers  (2025-02-20T18:50:42Z)
• Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA [8.305827430948654]
  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.
  arXiv  Detail & Related papers  (2024-12-30T03:05:45Z)
• On the Robustness of Distributed Machine Learning against Transfer Attacks [1.0787328610467801]
  No prior work has examined the combined robustness stemming from distributing both the learning and the inference process. We show that properly distributed ML instantiations achieve across-the-board improvements in accuracy-robustness tradeoffs against state-of-the-art transfer-based attacks.
  arXiv  Detail & Related papers  (2024-12-18T17:27:17Z)
• Hymba: A Hybrid-head Architecture for Small Language Models [65.94140459055244]
  Hymba is a family of small language models featuring a hybrid-head parallel architecture. We introduce learnable meta tokens that are prepended to prompts, storing critical information. This model is further optimized by incorporating cross-layer key-value sharing and partial sliding window attention.
  arXiv  Detail & Related papers  (2024-11-20T19:51:25Z)
• MoH: Multi-Head Attention as Mixture-of-Head Attention [63.67734699877724]
  We upgrade the multi-head attention mechanism, the core of the Transformer model, to improve efficiency while maintaining or surpassing the previous accuracy level. We propose Mixture-of-Head attention (MoH), a new architecture that treats attention heads as experts in the Mixture-of-Experts mechanism. MoH has two significant advantages: First, MoH enables each token to select the appropriate attention heads, enhancing inference efficiency without compromising accuracy or increasing the number of parameters.
  arXiv  Detail & Related papers  (2024-10-15T17:59:44Z)
• Zero-Shot Embeddings Inform Learning and Forgetting with Vision-Language Encoders [6.7181844004432385]
  The Inter-Intra Modal Measure (IIMM) functions as a strong predictor of performance changes with fine-tuning. Fine-tuning on tasks with higher IIMM scores produces greater in-domain performance gains but also induces more severe out-of-domain performance degradation. With only a single forward pass of the target data, practitioners can leverage this key insight to evaluate the degree to which a model can be expected to improve following fine-tuning.
  arXiv  Detail & Related papers  (2024-07-22T15:35:09Z)
• Effectively Compress KV Heads for LLM [28.0801697946958]
  We propose a novel approach for compressing Key-Value ( KV) caches. Our method can compress half or even three-quarters of KV heads while maintaining performance comparable to the original LLMs.
  arXiv  Detail & Related papers  (2024-06-11T08:37:33Z)
• Augmenting Unsupervised Reinforcement Learning with Self-Reference [63.68018737038331]
  Humans possess the ability to draw on past experiences explicitly when learning new tasks. We propose the Self-Reference (SR) approach, an add-on module explicitly designed to leverage historical information. Our approach achieves state-of-the-art results in terms of Interquartile Mean (IQM) performance and Optimality Gap reduction on the Unsupervised Reinforcement Learning Benchmark.
  arXiv  Detail & Related papers  (2023-11-16T09:07:34Z)
• Mitigating the Alignment Tax of RLHF [76.4300447532456]
  aligning LLMs under Reinforcement Learning with Human Feedback can lead to forgetting pretrained abilities, also known as the alignment tax. We propose model averaging to maximize alignment performance while incurring minimal alignment tax. We validate HMA's performance across a range of RLHF algorithms over OpenLLaMA-3B and further extend our findings to Mistral-7B.
  arXiv  Detail & Related papers  (2023-09-12T14:16:54Z)
• Robust Learning with Progressive Data Expansion Against Spurious Correlation [65.83104529677234]
  We study the learning process of a two-layer nonlinear convolutional neural network in the presence of spurious features. Our analysis suggests that imbalanced data groups and easily learnable spurious features can lead to the dominance of spurious features during the learning process. We propose a new training algorithm called PDE that efficiently enhances the model's robustness for a better worst-group performance.
  arXiv  Detail & Related papers  (2023-06-08T05:44:06Z)
• Newer is not always better: Rethinking transferability metrics, their peculiarities, stability and performance [5.650647159993238]
  Fine-tuning of large pre-trained image and language models on small customized datasets has become increasingly popular. We show that the statistical problems with covariance estimation drive the poor performance of H-score. We propose a correction and recommend measuring correlation performance against relative accuracy in such settings.
  arXiv  Detail & Related papers  (2021-10-13T17:24:12Z)
• Cascaded Head-colliding Attention [28.293881246428377]
  Transformers have advanced the field of natural language processing (NLP) on a variety of important tasks. We present cascaded head-colliding attention (CODA) which explicitly models the interactions between attention heads through a hierarchical variational distribution.
  arXiv  Detail & Related papers  (2021-05-31T10:06:42Z)
• Cauchy-Schwarz Regularized Autoencoder [68.80569889599434]
  Variational autoencoders (VAE) are a powerful and widely-used class of generative models. We introduce a new constrained objective based on the Cauchy-Schwarz divergence, which can be computed analytically for GMMs. Our objective improves upon variational auto-encoding models in density estimation, unsupervised clustering, semi-supervised learning, and face analysis.
  arXiv  Detail & Related papers  (2021-01-06T17:36:26Z)

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.