面向大词表神经网络快速Softmax推理的筛选学习
机器学习
2018-10-31 v1 计算与语言
机器学习
摘要
神经语言模型已被广泛用于各类NLP任务，包括机器翻译、下一词预测与对话代理。然而，由于预测速度慢，这些模型难以部署在移动设备上，其瓶颈在于计算softmax层中的top候选。本文中，我们利用上下文向量的聚类结构，引入一种新颖的softmax层近似算法。我们的算法使用轻量级筛选模型基于给定上下文预测小得多的候选词集合，然后仅在该子集内执行精确softmax。端到端训练此类过程具有挑战性，因为传统聚类方法离散且不可微，因而无法在训练过程中与反向传播结合使用。利用Gumbel softmax，我们能够在训练集上端到端训练筛选模型以利用数据分布。该算法在机器翻译束搜索或下一词预测等任务中预测top-词时，比原始softmax层实现数量级更快的推理。例如，在约25K词表的德英机器翻译数据集上，我们可实现20.4倍加速，precision@1达98.9%、precision@5达99.3%(与原始softmax层预测相比)，而state-of-the-art ~\citep{MSRprediction}在同一任务下仅实现6.7倍加速，precision@1为98.7%、precision@5为98.1%。
引用
@article{arxiv.1810.12406,
title = {Learning to Screen for Fast Softmax Inference on Large Vocabulary Neural Networks},
author = {Patrick H. Chen and Si Si and Sanjiv Kumar and Yang Li and Cho-Jui Hsieh},
journal= {arXiv preprint arXiv:1810.12406},
year = {2018}
}