[论文解读] Learning to Screen for Fast Softmax Inference on Large Vocabulary Neural Networks

Patrick H. Chen, Si Si+3 more|arXiv (Cornell University)|Oct 29, 2018

Topic Modeling参考文献 13被引用 10

一句话总结

本文提出 L2S（Learning to Screen），一种通过训练轻量级筛选模型来预测上下文向量聚类中可能的 top-k 词语子集，从而加速大词汇量神经语言模型中的 softmax 推理的方法。通过使用 Gumbel softmax 技巧实现端到端可微训练，L2S 相较于精确 softmax 实现了高达 20.4 倍的加速，且精度几乎完全一致（precision@1 为 98.9%），显著优于先前的 FGD 和 Zhang et al. (2018) 等方法。

ABSTRACT

Neural language models have been widely used in various NLP tasks, including machine translation, next word prediction and conversational agents. However, it is challenging to deploy these models on mobile devices due to their slow prediction speed, where the bottleneck is to compute top candidates in the softmax layer. In this paper, we introduce a novel softmax layer approximation algorithm by exploiting the clustering structure of context vectors. Our algorithm uses a light-weight screening model to predict a much smaller set of candidate words based on the given context, and then conducts an exact softmax only within that subset. Training such a procedure end-to-end is challenging as traditional clustering methods are discrete and non-differentiable, and thus unable to be used with back-propagation in the training process. Using the Gumbel softmax, we are able to train the screening model end-to-end on the training set to exploit data distribution. The algorithm achieves an order of magnitude faster inference than the original softmax layer for predicting top-$k$ words in various tasks such as beam search in machine translation or next words prediction. For example, for machine translation task on German to English dataset with around 25K vocabulary, we can achieve 20.4 times speed up with 98.9\% precision@1 and 99.3\% precision@5 with the original softmax layer prediction, while state-of-the-art ~\citep{MSRprediction} only achieves 6.7x speedup with 98.7\% precision@1 and 98.1\% precision@5 for the same task.

研究动机与目标

• 解决大词汇量神经语言模型中 softmax 推理的计算瓶颈，尤其是在移动设备上的问题。
• 在不牺牲预测精度的前提下，减少 top-k 词语预测的推理时间。
• 利用上下文向量和词嵌入的聚类结构，缩小候选搜索空间。
• 开发一种可端到端训练的筛选模型，联合优化聚类分配与候选集选择。
• 在机器翻译和下一个词预测等 NLP 任务中，实现比现有基于 MIPS 和近似方法更优的速度-精度权衡。

提出的方法

• 提出两阶段推理流程：首先，轻量级筛选模型基于上下文向量预测一个候选词语的小簇；其次，仅在该子集中计算精确 softmax。
• 通过 Gumbel-Softmax 技巧实现可微聚类机制，使离散聚类分配在神经网络中支持端到端反向传播。
• 制定联合优化目标，同时学习筛选模型与每个聚类的候选词语集合。
• 在完整训练分布上训练筛选模型，以利用上下文向量中数据驱动的聚类模式。
• 将筛选模型集成到束搜索与下一个词预测流程中，将 softmax 计算限制在候选子集内。
• 在损失函数中设置时间预算约束，以平衡速度与精度，实现在实际应用中可调节的权衡。

实验结果

研究问题

• RQ1 一种可学习的筛选模型，若能利用上下文向量中的聚类结构，是否能显著减少大词汇量语言模型中 softmax 推理的时间？
• RQ2 如何使离散聚类分配在神经网络中实现可微，以支持端到端训练？
• RQ3 与传统的 MIPS 算法相比，使用学习筛选在推理加速与预测精度之间的权衡如何？
• RQ4 在真实 NLP 任务中，该方法是否能在实现数量级加速的同时，保持高 precision@k（如 98% 或以上）？
• RQ5 该方法对聚类数量和时间预算约束的变化是否具有鲁棒性？

主要发现

• 在德语到英语翻译任务中，L2S 实现了 20.4 倍的加速，precision@1 达到 98.9%，precision@5 达到 99.3%，显著优于 Zhang et al. (2018) 实现的 6.7 倍加速与 98.7% precision@1。
• 在 EN-VE 翻译任务中，L2S 实现了 20 倍加速，BLEU 分数仅下降 0.08（beam size 为 5），而 FGD 的加速不足 6 倍，且精度相近。
• 在 PTB-Small 和 PTB-Large 语言建模任务中，L2S 实现了超过 5 倍的加速，且困惑度增加不足 5%，相比完整 softmax。
• 该方法对聚类数量具有鲁棒性：在 100–200 个聚类下，保持高精度（P@1 > 0.99）与低推理时间（0.14–0.17ms）。
• 在束搜索中，L2S 保持了几乎相同的 BLEU 分数（例如 30.19 vs. 30.33），同时在 DE-EN 任务中将推理时间减少了高达 13.4 倍（beam size 为 5）。
• 与 Gumbel-Softmax 的联合优化使端到端训练更加有效，使模型能够学习数据驱动的聚类模式，从而同时提升速度与精度。