arXiv 2406.065672024-12-07cs.LGcs.AIcs.CL

DHA：通过自适应头融合从Transformer检查点学习解耦头注意力

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

Institute of Information Engineering, Chinese Academy of Sciences（中国科学院信息工程研究所）；School of Cyber Security, University of Chinese Academy of Sciences（中国科学院大学网络空间安全学院）；Baidu Inc.（百度公司）

Yilong Chen, Linhao Zhang, Junyuan Shang, Zhenyu Zhang, Tingwen Liu, Shuohuan Wang, Yu Sun

备注：: Accepted at NeurIPS 2024 10 pages, 9 figures, 3 tables

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

针对大语言模型中多头注意力推理成本高的问题，提出解耦头注意力DHA机制，通过自适应头分组与相似头线性融合将MHA检查点转换为DHA模型，以极低预训练预算实现性能与效率的更优平衡，训练速度与性能表现优于GQA。

摘要：

拥有数十亿参数的大语言模型（LLM）展现出了令人瞩目的性能。然而，LLM中广泛使用的多头注意力（MHA）在推理过程中会产生大量的计算和内存成本。虽然已有一些研究通过剪枝注意力头或在头之间共享参数来优化注意力机制，但这些方法往往会导致性能下降，或是需要高昂的持续预训练成本才能恢复性能。基于对注意力冗余性的分析，我们设计了一种解耦头注意力（Decoupled-Head Attention，DHA）机制。DHA针对不同层的键（key）头和值（value）头自适应地配置分组共享策略，在性能与效率之间实现了更好的平衡。受相似头聚类现象的启发，我们提出通过逐步对相似头的参数进行线性融合，将MHA检查点渐进式转换为DHA模型，从而保留MHA检查点中的参数知识。我们在给定目标头数量预算的前提下，通过转换不同规模的MHA检查点构建了DHA模型。实验结果表明，DHA仅需原始模型0.25%的预训练预算，就能达到97.6%的性能，同时节省75%的KV缓存。与分组查询注意力（Group-Query Attention，GQA）相比，DHA实现了5倍的训练加速；在0.01%的预训练预算下，性能最大提升13.93%；在0.05%的预训练预算下，性能相对提升4%。

英文摘要：

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.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。