[论文解读] Improving Transformers with Probabilistic Attention Keys

Tam Nguyen, Tan M. Nguyen+6 more|arXiv (Cornell University)|Oct 16, 2021

Topic Modeling被引用 6

一句话总结

本文提出了一种新型的Transformer架构——Transformer-MGK，通过使用高斯分布混合模型替代标准注意力机制中的键（keys），以减少注意力头之间的冗余。通过概率化建模键，该方法在加速训练与推理、减少参数量和浮点运算次数（FLOPs）的同时，实现了与标准Transformer相当或更优的性能，且使用更少的注意力头。该方法在Wikitext-103和Long Range Arena基准上得到了验证。

ABSTRACT

Multi-head attention is a driving force behind state-of-the-art transformers, which achieve remarkable performance across a variety of natural language processing (NLP) and computer vision tasks. It has been observed that for many applications, those attention heads learn redundant embedding, and most of them can be removed without degrading the performance of the model. Inspired by this observation, we propose Transformer with a Mixture of Gaussian Keys (Transformer-MGK), a novel transformer architecture that replaces redundant heads in transformers with a mixture of keys at each head. These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts of the input sequence efficiently. Compared to its conventional transformer counterpart, Transformer-MGK accelerates training and inference, has fewer parameters, and requires fewer FLOPs to compute while achieving comparable or better accuracy across tasks. Transformer-MGK can also be easily extended to use with linear attention. We empirically demonstrate the advantage of Transformer-MGK in a range of practical applications, including language modeling and tasks that involve very long sequences. On the Wikitext-103 and Long Range Arena benchmark, Transformer-MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.

研究动机与目标

• 解决多头注意力机制中多个头学习到相似注意力模式所导致的冗余问题。
• 在不牺牲模型性能的前提下，降低Transformer模型的计算与参数开销。
• 构建一种概率化框架，用于建模注意力键，从而提升每个头的表征能力。
• 通过用单一、更具表达力的高斯混合键机制替代多个冗余头，实现高效推理与训练。
• 将该框架扩展至线性注意力机制，以支持长序列建模。

提出的方法

• 通过将注意力键建模为高斯混合模型（GMM），为自注意力机制提供一种概率化解释，将注意力得分与后验概率关联起来。
• 将多头注意力头中的每个标准键向量替换为高斯分量的混合，使单个头能够同时关注输入序列中多个不同且多样的位置。
• 将注意力计算表示为高斯分量的加权和，其中注意力权重对应于GMM下的后验概率。
• 引入一种可微分的训练过程，以优化键混合中高斯分量的均值、方差和混合权重。
• 将高斯混合键集成到标准Transformer模块中，保持与现有架构及优化技术的兼容性。
• 通过使用概率化键表示，将该框架扩展至线性注意力机制，从而实现对长序列的高效建模。

实验结果

研究问题

• RQ1 能否通过注意力键的概率化建模，减少Transformer中多头注意力头之间的冗余？
• RQ2 用高斯混合模型替代标准键，是否能在降低模型复杂度的同时提升每个注意力头的表征能力？
• RQ3 与标准Transformer相比，Transformer-MGK在使用更少注意力头的情况下，能在多大程度上保持或提升性能？
• RQ4 该高斯混合键机制在长序列任务和低资源设置下的表现如何？
• RQ5 所提出的方法能否高效地扩展至线性注意力机制，以实现可扩展的长上下文建模？

主要发现

• 在Wikitext-103语言建模基准上，使用4个头的Transformer-MGK实现了与使用8个头的标准Transformer相当或更优的性能。
• 在Long Range Arena基准上，使用4个头的Transformer-MGK在长上下文任务中表现与使用8个头的标准Transformer相当或更优。
• 由于减少了头冗余，该模型在降低FLOPs和参数量的同时，显著加速了训练与推理过程。
• 高斯混合键机制使每个注意力头能够同时关注输入序列中多个不同且多样的部分，从而增强了表征多样性。
• 该方法与线性注意力兼容，可实现对超长序列的高效处理。
• 实证结果表明，概率化键建模可使注意力模式在多种自然语言处理任务中更加稳定且泛化能力更强。