Softpick: No Attention Sink, No Massive Activations with Rectified Softmax

我们提出了 Softpick，这是对变压器注意力机制中 softmax 的一种修正版、非归一化（即不强制概率和为一）、且可直接替代的方案，能够消除注意力汇聚问题以及过大的激活值。我们通过包含 3.4 亿参数的模型进行的实验表明，Softpick 在标准基准测试中与 softmax 表现相当，同时实现了 0% 的汇聚率。采用 Softpick 的变压器生成的隐藏状态具有显著更低的峰度（340 对比 33,510），并且生成了稀疏的注意力图（稀疏率为 46.97%）。在量化场景下，使用 Softpick 的模型性能始终优于 softmax，尤其是在低比特精度下优势尤为明显。我们的分析与讨论展示了 Softpick 如何为量化、低精度训练、稀疏性优化、剪枝以及模型可解释性等领域开辟新的可能性。我们的代码已开源，可在以下地址获取：https://github.com/zaydzuhri/softpick-attention。

该论文试图解决Transformer模型中使用softmax时出现的'attention sink'问题以及隐藏状态激活值过大的问题。这些问题会导致模型在低精度量化时性能下降，并可能影响模型的稀疏性和可解释性。这是一个值得研究的问题，但并非全新的问题，因为已有研究探讨过softmax的局限性。

论文提出了一种名为Softpick的新方法，作为softmax的替代方案。Softpick是一种非归一化（not sum-to-one）的注意力机制，能够有效消除attention sink现象，并显著降低隐藏状态的尖峰分布（kurtosis）。相比传统softmax，Softpick在保持性能的同时，改善了模型的稀疏性和量化友好性。

实验表明，Softpick在340M参数规模的模型上与softmax性能相当，同时实现了0%的sink率和较低的隐藏状态尖峰分布（从33,510降至340）。此外，Softpick生成的注意力图更加稀疏（46.97%的稀疏度），并在量化场景下表现出明显优势，特别是在低比特精度条件下。作者开源了代码，提供了进一步研究的基础。未来可以深入研究Softpick在更大规模模型中的表现及其对剪枝和可解释性的潜在影响。

最近的相关研究包括：1) 'Revisiting Softmax in Attention Mechanisms'，探讨了softmax在不同任务中的表现；2) 'Sparse Transformers'，提出了稀疏注意力机制以提高效率；3) 'Low-Rank Approximation for Efficient Transformer Attention'，通过低秩近似优化注意力计算；4) 'Quantization-Friendly Attention Mechanisms'，专注于设计适合量化的注意力机制。这些研究共同推动了Transformer架构在效率、稀疏性和量化方面的改进。