[论文解读] Talking-Heads Attention

Noam Shazeer, Zhenzhong Lan+3 more|arXiv (Cornell University)|Mar 5, 2020

Topic Modeling参考文献 9被引用 49

一句话总结

介绍了 talking-heads 注意力，一种在 softmax 之前和之后沿注意力头维学习线性投影的多头注意力变体，在参数增量适中的情况下提升 perplexity 和下游任务表现。

ABSTRACT

We introduce "talking-heads attention" - a variation on multi-head attention which includes linearprojections across the attention-heads dimension, immediately before and after the softmax operation.While inserting only a small number of additional parameters and a moderate amount of additionalcomputation, talking-heads attention leads to better perplexities on masked language modeling tasks, aswell as better quality when transfer-learning to language comprehension and question answering tasks.

研究动机与目标

• 动机与解决传统多头注意力的局限性：当头的维度太小时，头特定投影可能降低表达能力。
• 提出一种通过在头维上学习投影来实现注意力头之间信息交换的机制。
• 在语言建模和理解基准的预训练与微调任务中展示经验收益。
• 提供与标准多头注意力相比的新机制的分析与复杂性考量。

提出的方法

• 引入两个附加的可学习线性投影（P_l 与 P_w），在 softmax 之前和之后连接注意力头，以实现跨头信息交换。
• 将 TalkingHeadsAttention 定义为对多头注意力的推广，具有三个头维度（用于查询/键的 h_k、用于 logits/weights 的 h、用于值的 h_v），并给出两种等价的表达形式（简洁和显式）。
• 推导计算代价，并表明 talking-heads 投影在标准代价上增加了一项，但在 h_k、h、h_v 选取恰当时可能更有利。
• 给出统一视图（GBMA），将标准 MHA 与 talking-heads 作为通用双线性多头注意力的特例相连接。
• 在 Text-to-Text Transfer Transformer（T5）、ALBERT 和 BERT 风格设置中进行实验，比较多头注意力与 talking-heads 在预训练困惑度和下游任务中的表现。

实验结果

研究问题

• RQ1 在预训练任务中注入跨头投影是否相较于标准多头注意力能提升模型困惑度？
• RQ2 talking-heads 是否在编码器/解码器组件以及不同头维配置下提供一致的增益？
• RQ3 在使用 talking-heads 注意力与传统多头注意力时，复杂度与硬件效率的权衡是什么？
• RQ4 在实践中，talking-heads 如何与投影粒度（logits/weights）以及动态投影与静态投影相互作用？
• RQ5 talking-heads 注意力的优势能否转移到广泛使用的架构如 BERT 和 ALBERT？

主要发现

• 与标准多头注意力在若干头维配置下相比，talking-heads 注意力在掩码语言建模任务中的困惑度降低，相较于标准多头注意力，在多种头尺寸配置下。
• 增加 softmax-heads 维度 h 通常带来显著收益，而在传统 MHA 中非常小的 key/value 维度 (d_k, d_v) 可能降低性能，但 talking-heads 可以缓解这一问题。
• 在编码器自注意力中应用 talking-heads 能带来比在解码器注意力中更大的下游收益，在若干设置中。
• 配置三个独立的头维度（h_k、h、h_v），并同时使用 logits 投影与权重投影，在许多实验中得到最强的改进（例如在 T5 和 ALBERT 上）。
• 动态投影可以降低预训练困惑度，但在他们的实验中并未稳定改善下游任务。