自适应抽样Softmax及反向多索引：方法、理论与应用

Softmax函数是多类分类的基石，构成广泛应用的机器学习领域，从大规模检索和排序模型到先进的大型语言模型。然而，其计算成本随类别数线性增长，在拥有数百万甚至数十亿类别的场景中变得不可接受。基于自归一化重要抽样的抽样Softmax已成为强大的替代方案，显著降低计算复杂度。然而，其估计器仅在抽样分布匹配真实Softmax分布时才无偏。为提高近似精度和抽样效率，我们提出了MIDX采样器，一种基于反向多索引方法的自适应抽样策略。具体而言，我们将Softmax概率分解为多个多项分布概率，每种概率与特定的译码字集合相关联，最后一个译码字集合则与查询的残余得分相关联，从而将时间复杂度降低到译码字数量而非类别数量。为进一步提升效率，我们将查询特定的残余概率替换为简单均匀分布，在简化计算的同时保持高性能。我们的方法基于严格的理论分析，解决了抽样偏差、梯度偏差、收敛速度及泛化误差界等关键问题。结果表明，与理想Softmax分布的偏差越小，收敛速度越快，泛化能力越强。大量实验在大规模语言模型、顺序推荐系统和极端多类分类任务上表明，MIDX采样器在有效性和效率方面优于现有方法。