SoftMax 是现代机器学习算法中无处不在的组成部分。它将输入向量映射到概率单纯形，并通过将概率质量集中在较大的条目上，来重新加权输入。然而，作为 Argmax 函数的平滑近似，SoftMax 将大量的概率质量分配给其他剩余的条目，导致可解释性差和噪声。虽然稀疏性可以通过一系列 SoftMax 变体来实现，但它们通常需要替代的损失函数，并且不保留多模态。
本文将探讨 SoftMax 在稀疏性和多模态之间面临的困境，并介绍一种名为 MultiMax 的新方法，它可以有效地解决这一问题。
SoftMax 的局限性：温度的尴尬
SoftMax 的一个重要参数是温度（temperature）。温度控制着输出分布的熵，较高的温度会导致过度平滑，降低优化效率，而较低的温度会导致多模态坍塌，使训练不稳定。
例如，在注意力机制中，较低的温度会导致除了峰值之外的所有相关位置都被忽略，而较高的温度会导致注意力在无关键上“浪费”大量的注意力。因此，在注意力层中，温度通常默认设置为 1。然而，这种折衷方案会导致最近在视觉和语言 Transformer 中观察到的过度平滑问题。
寻找突破：稀疏 SoftMax 的不足
为了克服 SoftMax 的问题，研究人员提出了稀疏 SoftMax 替代方案，这些方案允许完全忽略低于阈值的较小条目。这些稀疏 SoftMax 变体已经在不同的背景下得到研究，例如生成模型、多类分类器的输出激活和/或注意力机制。
然而，这些方法通常存在梯度信号差的问题，这会导致训练过程中的不稳定性。此外，非稀疏维度的数量通常被视为经验选择的超参数。
MultiMax：兼顾稀疏性和多模态
与稀疏性相比，多模态在之前的研究中讨论得较少。由于注意力在大多数情况下不应该是排他的，因此传统的 SoftMax 作为 Argmax 的近似值，并不容易与多模态相容。稀疏 SoftMax 替代方案甚至更倾向于不保留分布的多模态。
为了解决 SoftMax 在稀疏性和多模态之间难以兼顾的问题，本文提出了 MultiMax。MultiMax 允许学习何时强调稀疏性，何时强调多模态，从而在两者之间提供灵活的权衡。同时，它保持分段可微，以便允许稳定的基于梯度的优化。
具体来说，MultiMax 通过一个先前的参数化函数扩展了传统的 SoftMax，该函数允许分别为特定输入值范围学习不同的温度值。例如，在自注意力机制中，这有助于学习特别低的温度，这些温度会为低输入值范围诱导稀疏性，即无关的标记可以被忽略，同时为更高的输入值范围学习较高的温度，即多个相关标记可以以多模态的方式共享注意力。
MultiMax 的优势：理论分析和实验结果
通过理论分析和实验评估，本文验证了 MultiMax 在以下方面的优势：
- 提高了帕累托效率： MultiMax 在稀疏性和多模态之间取得了更好的帕累托效率，即在保持较高稀疏性的同时，也提高了多模态性。
- 有效地防止过度平滑