DHA：通过自适应头融合从Transformer检查点学习解耦头注意力
机器学习
2024-12-10 v2 人工智能
计算与语言
摘要
拥有数十亿参数的大型语言模型(LLM)显示出惊人的性能。然而，LLM中广泛使用的多头注意力(MHA)在推理期间造成了巨大的计算和内存开销。虽然一些努力通过修剪头或在头之间共享参数来优化注意力机制，但这些方法常常导致性能下降，或需要 substantial 继续预训练成本来恢复性能。基于注意力冗余性分析，我们设计了解耦头注意力(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%的相对性能提升。
引用
@article{arxiv.2406.06567,
title = {DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion},
author = {Yilong Chen and Linhao Zhang and Junyuan Shang and Zhenyu Zhang and Tingwen Liu and Shuohuan Wang and Yu Sun},
journal= {arXiv preprint arXiv:2406.06567},
year = {2024}
}
备注
Accepted at NeurIPS 2024 10 pages, 9 figures, 3 tables