[论文解读] Low-Rank Bottleneck in Multi-head Attention Models

Srinadh Bhojanapalli, Chulhee Yun+3 more|arXiv (Cornell University)|Feb 17, 2020

Topic Modeling参考文献 29被引用 14

一句话总结

本文识别出多头注意力模型中因头数增加而导致头尺寸缩小所引发的低秩瓶颈，该瓶颈限制了表征能力，并迫使模型依赖大嵌入维度。作者提出将头尺寸固定为输入序列长度，与头数无关，从而提升模型表达能力，实现在更小嵌入尺寸下的更好性能，并在不损失准确率的前提下减少模型参数。

ABSTRACT

Attention based Transformer architecture has enabled significant advances in the field of natural language processing. In addition to new pre-training techniques, recent improvements crucially rely on working with a relatively larger embedding dimension for tokens. Unfortunately, this leads to models that are prohibitively large to be employed in the downstream tasks. In this paper we identify one of the important factors contributing to the large embedding size requirement. In particular, our analysis highlights that the scaling between the number of heads and the size of each head in the current architecture gives rise to a low-rank bottleneck in attention heads, causing this limitation. We further validate this in our experiments. As a solution we propose to set the head size of an attention unit to input sequence length, and independent of the number of heads, resulting in multi-head attention layers with provably more expressive power. We empirically show that this allows us to train models with a relatively smaller embedding dimension and with better performance scaling.

研究动机与目标

• 探究现代Transformer模型中高嵌入维度需求的根本原因。
• 识别出将头尺寸与头数成反比缩放会引发低秩瓶颈，从而限制表征能力。
• 提出一种新型架构设计，将头尺寸固定为输入序列长度，与头数和嵌入维度解耦。
• 证明该改变可实现更小、更高效的模型训练，且性能扩展更优。
• 验证头尺寸是关键的容量控制参数，而不仅仅是头数或嵌入维度。

提出的方法

• 提出固定头尺寸设置，使每个注意力头在大小等于输入序列长度的子空间中运行。
• 将头尺寸与头数和嵌入维度解耦，实现无需秩约束的任意头数。
• 理论证明（定理1）表明，当头尺寸小于序列长度时，会引发低秩瓶颈，降低表征能力。
• 证明增加头尺寸可单调提升性能，而仅增加头数则不然。
• 在LM1B、SQuAD和MNLI基准上，使用更小嵌入尺寸训练并评估固定头尺寸的模型。
• 与采用传统头尺寸启发式方法（与头数成比例）的标准模型进行性能对比。

实验结果

研究问题

• RQ1 为何现代Transformer模型尽管架构不断进步，仍需大嵌入维度？
• RQ2 随着头数增加而减小头尺寸，如何影响多头注意力的表征容量？
• RQ3 能否通过将头尺寸固定为输入序列长度来消除注意力头中的低秩瓶颈？
• RQ4 固定头尺寸是否能实现更优的性能扩展，同时减少参数量？
• RQ5 头尺寸是否比头数或嵌入维度更具容量控制作用？

主要发现

• 标准头尺寸启发式方法（与头数成反比）会引发低秩瓶颈，限制表征能力，该结论已在定理1中得到证明。
• 在语言建模任务中，固定头尺寸（128）且嵌入尺寸更小（256）的模型，性能优于标准模型（嵌入尺寸512）。
• 在SQuAD和MNLI上，固定头尺寸（512嵌入尺寸，8–32个头）的模型在F1和准确率上均优于标准BERT-Large（1024嵌入尺寸）。
• 当头尺寸固定时，性能随头数增加而单调提升，而标准模型在头数超过某一阈值后性能反而下降。
• 将头尺寸从32提升至128可显著改善性能，证实头尺寸是关键容量参数。
• 固定头尺寸方法使训练更小、参数更少的模型成为可能，同时保持或提升性能，从而降低推理成本。