利用动态可组合多头注意力改进 Transformer
机器学习
2024-06-05 v2 计算与语言
摘要
多头注意力 (MHA) 是 Transformer 的关键组件。在 MHA 中，注意力头独立工作，导致注意力分数矩阵低秩瓶颈和注意力头冗余等问题。我们提出了动态可组合多头注意力 (DCMHA)，这是一种参数与计算高效的注意力架构，通过动态组合注意力头来克服 MHA 的缺陷并提升模型的表达能力。DCMHA 的核心是一个 函数，它以依赖于输入的方式变换注意力分数和权重矩阵。DCMHA 可作为 MHA 的直接替换模块用于任何 Transformer 架构，从而获得相应的 DCFormer。在语言建模任务中，DCFormer 在不同架构和模型规模上均显著优于 Transformer，其性能匹配计算量约 1.7 倍至 2.0 倍的模型。例如，DCPythia-6.9B 在预训练困惑度和下游任务评估上均优于开源的 Pythia-12B。代码和模型可在 https://github.com/Caiyun-AI/DCFormer 获取。
关键词
引用
@article{arxiv.2405.08553,
title = {Improving Transformers with Dynamically Composable Multi-Head Attention},
author = {Da Xiao and Qingye Meng and Shengping Li and Xingyuan Yuan},
journal= {arXiv preprint arXiv:2405.08553},
year = {2024}
}
备注
Accepted to the 41st International Conference on Machine Learning (ICML'24 oral)