[论文解读] Normalized Attention Without Probability Cage

Oliver Richter, Roger Wattenhofer|arXiv (Cornell University)|May 19, 2020

Topic Modeling参考文献 37被引用 10

一句话总结

本文提出归一化注意力（NAP），一种无需Softmax的注意力机制，通过使用L2归一化替代概率单纯形约束，避免了Softmax注意力的凸包限制。该方法提升了模型对超参数和数据偏差的鲁棒性，在序列任务中实现最先进性能——在25,000次实验中优于Transformer、最大池化和求和池化，尤其在长上下文和大词汇量设置下表现突出。

ABSTRACT

Attention architectures are widely used; they recently gained renewed popularity with Transformers yielding a streak of state of the art results. Yet, the geometrical implications of softmax-attention remain largely unexplored. In this work we highlight the limitations of constraining attention weights to the probability simplex and the resulting convex hull of value vectors. We show that Transformers are sequence length dependent biased towards token isolation at initialization and contrast Transformers to simple max- and sum-pooling - two strong baselines rarely reported. We propose to replace the softmax in self-attention with normalization, yielding a hyperparameter and data-bias robust, generally applicable architecture. We support our insights with empirical results from more than 25,000 trained models. All results and implementations are made available.

研究动机与目标

• 探究通过Softmax约束注意力权重至概率单纯形所带来的理论与实际限制。
• 挑战注意力必须表示为概率分布的假设，尤其关注其对模型表达能力与初始化偏差的影响。
• 提出一种通用、鲁棒的注意力机制，避免Softmax注意力的凸包约束。
• 在多种任务与超参数设置下，对NAP与Transformer、最大池化、求和池化及其他注意力变体进行实证比较。
• 证明移除概率单纯形约束可增强模型泛化能力，尤其在长上下文与大词汇量场景下。

提出的方法

• 用L2归一化替代自注意力中的Softmax操作，将注意力权重投影至单位球面而非概率单纯形。
• 将注意力机制定义为归一化加权和：o_i^m = Σ_j (a_i,j^m / ||a_i^m||_2) · v_j^m，其中 a_i,j^m = <q_i^m, k_j^m> / √d_h。
• 在注意力logits中引入可学习偏置项b，使模型在归一化失效时能表现出类似求和池化的残差行为。
• 使用标准反向传播训练模型，并在多个任务上评估性能：案例区分、模式查找与序列分类。
• 开展大规模消融实验，训练超过25,000个模型，涵盖不同序列长度、模型深度、学习率与词汇量。
• 使用RGB热图可视化超参数空间中的模型性能，突出其鲁棒性与泛化能力。

实验结果

研究问题

• RQ1 将注意力权重约束于概率单纯形如何限制自注意力机制的表达能力？
• RQ2 Softmax注意力在初始化层面引入了哪些偏差，特别是关于标记隔离与凸包约束？
• RQ3 基于归一化的注意力机制是否能在鲁棒性与准确性方面超越Softmax注意力与传统池化方法（求和/最大）？
• RQ4 所提出的归一化注意力机制在分布偏移（如更长序列长度或更大词汇量）下的表现如何？
• RQ5 新注意力机制在多大程度上降低了对学习率与模型深度等超参数的敏感性？

主要发现

• 在序列长度为128的案例区分任务中，NAP实现93.8%的平均验证准确率，显著优于BERT（71.5%）与MTE（57.7%）。
• 在词汇量为256的模式查找任务中，NAP实现84.6%的平均验证准确率，优于BERT（74.4%）、MTE（64.9%）与最大池化（3.1%）。
• NAP在超参数选择上表现出卓越鲁棒性：在广泛的学习率与模型维度范围内保持高性能，尤其在长上下文设置下。
• 最大池化在案例区分任务中表现优异（准确率90.9%），但在大词汇量模式查找任务中失败，表明其在分布偏移下泛化能力差。
• 求和池化在案例区分任务中表现欠佳（准确率29.1%），凸显了缺乏注意力机制时建模复杂依赖关系的局限性。
• 在模式查找任务中，NAP架构在更长序列（N=256）下泛化良好，验证准确率达64.3%，而最大池化下降至2.1%。