[论文解读] Breaking the Softmax Bottleneck via Learnable Monotonic Pointwise Non-linearities

Octavian-Eugen Ganea, Sylvain Gelly+2 more|arXiv (Cornell University)|Feb 21, 2019

Model Reduction and Neural Networks参考文献 30被引用 6

一句话总结

本文提出线性单调Softmax（LMS），一种可学习的、单调的逐点非线性变换，应用于Softmax前的logits，以突破大规模词汇语言模型中的秩缺陷瓶颈。该方法在计算开销可忽略的前提下，显著提升了交叉熵损失和模式匹配性能，优于标准Softmax和混合Softmax模型，在PennTreeBank和WikiText-2数据集上实现了最先进水平的困惑度，且效率远超MoS。

ABSTRACT

The Softmax function on top of a final linear layer is the de facto method to output probability distributions in neural networks. In many applications such as language models or text generation, this model has to produce distributions over large output vocabularies. Recently, this has been shown to have limited representational capacity due to its connection with the rank bottleneck in matrix factorization. However, little is known about the limitations of Linear-Softmax for quantities of practical interest such as cross entropy or mode estimation, a direction that we explore here. As an efficient and effective solution to alleviate this issue, we propose to learn parametric monotonic functions on top of the logits. We theoretically investigate the rank increasing capabilities of such monotonic functions. Empirically, our method improves in two different quality metrics over the traditional Linear-Softmax layer in synthetic and real language model experiments, adding little time or memory overhead, while being comparable to the more computationally expensive mixture of Softmaxes.

研究动机与目标

• 为解决Linear-Softmax层在建模大规模输出词汇上的复杂概率分布时所面临的表征限制。
• 探究可学习的、单调的逐点非线性函数是否能有效提升矩阵秩并缓解Softmax瓶颈。
• 开发一种方法，在合成与真实语言建模任务中均能提升交叉熵最小化与模式匹配性能。
• 为单调非线性函数在Softmax上下文中的秩提升能力提供理论保证。
• 设计一种高效替代昂贵的混合Softmax模型的方法，在保持高性能的同时实现低内存与时间开销。

提出的方法

• 提出一种可学习的、连续的、单调递增的逐点函数（PLIF），应用于最终Softmax层之前的logits，构成LMS架构。
• 通过约束非线性函数为单调性，以保持秩的性质并确保优化过程稳定。
• 该方法通过学习参数化的非线性变换，而非使用固定函数，推广了Sigsoftmax。
• 采用可微分的、参数化的函数（如基于ReLU或样条的函数）来建模logits的非线性失真。
• 在AWD-LSTM模型上使用标准训练流程（SGD配合学习率调度），用PLIF层替换最终的Softmax层。
• 在合成数据和真实语言建模范例（PennTreeBank、WikiText-2）上评估性能，并对函数斜率统计量进行消融分析。

实验结果

研究问题

• RQ1 可学习的、单调的逐点非线性函数是否能有效提升logits到概率矩阵的秩，并突破Softmax瓶颈？
• RQ2 在合成分布上，LMS模型在交叉熵和模式匹配方面与Linear-Softmax及混合Softmax相比表现如何？
• RQ3 LMS能否在真实语言建模任务中实现具有竞争力的困惑度，同时保持低计算与内存开销？
• RQ4 所学习的非线性函数的函数形式是什么？其在实际中是否表现出显著的非线性行为？
• RQ5 最大熵原理在具有线性约束条件下的理论与LMS目标之间是否存在理论关联？

主要发现

• 在合成实验中，LMS在最小化交叉熵和匹配真实模式方面显著优于Linear-Softmax和混合Softmax，尤其在低维嵌入（低D）和大规模词汇（大M）设置下表现突出。
• 在PennTreeBank数据集上，LMS-PLIF实现测试困惑度107.5，优于基线AWD-LSTM（110.2），并达到远为昂贵的MoS模型的性能水平。
• 在WikiText-2上，LMS-PLIF实现测试困惑度101.7，优于Linear-Softmax，并与最先进的MoS模型性能相当，但计算成本低得多。
• 所学习的PLIF函数表现出强烈的非线性行为，其斜率统计量显示均值为1.10，标准差为0.62，最大斜率为5.16，表明对logits存在显著的非线性失真。
• LMS的计算开销可忽略不计：训练时间与GPU内存使用量与标准Linear-Softmax几乎完全一致，而MoS则昂贵数个数量级。
• 将MoS与PLIF层结合（MoS + PLIF）在PennTreeBank上取得最佳结果（困惑度106.8），优于所有其他基线模型，证明了LMS组件的模块化与高效性。