基于概率注意力键改进 Transformer
机器学习
2022-06-14 v2 计算与语言
机器学习
摘要
多头注意力是前沿 Transformer 背后的驱动力，后者在一系列自然语言处理（NLP）和计算机视觉任务中取得了卓越性能。据观察，对于许多应用，这些注意力头学习到了冗余的嵌入，且其中大多数可在不降低模型性能的情况下被移除。受此启发，我们提出带高斯键混合的 Transformer（Transformer-MGK），一种新颖的 Transformer 架构，它在每个注意力头中以键的混合替代 Transformer 中冗余的头。这些键的混合遵循高斯混合模型，并允许每个注意力头高效地聚焦于输入序列的不同部分。与常规 Transformer 对应模型相比，Transformer-MGK 加速了训练与推理，参数更少，且计算所需的 FLOPs 更少，同时在各项任务上达到可比或更优的精度。Transformer-MGK 也可轻松扩展以配合线性注意力使用。我们在一系列实际应用（包括语言建模及涉及极长序列的任务）中实证展示了 Transformer-MGK 的优势。在 Wikitext-103 与 Long Range Arena 基准上，具有 4 个头的 Transformer-MGK 取得了与具有 8 个头的基线 Transformer 可比或更优的性能。
引用
@article{arxiv.2110.08678,
title = {Improving Transformers with Probabilistic Attention Keys},
author = {Tam Nguyen and Tan M. Nguyen and Dung D. Le and Duy Khuong Nguyen and Viet-Anh Tran and Richard G. Baraniuk and Nhat Ho and Stanley J. Osher},
journal= {arXiv preprint arXiv:2110.08678},
year = {2022}
}
备注
27 pages, 16 figures, 10 tables