随着NLP方向的快速发展,用于支撑训练的大型语料库也如雨后春笋般层出不穷,其中有个One Billion Word的数据集词表更是达到了800K。在实践过程中,这种大词表往往会导致训练的速度极慢。许多研究者往往一周也训练不出指标,不能及时更新周报,苦不堪言。那么如何解决这类问题呢?
刚好笔者最近有个任务的词表也是特别大,所以调研了几种处理大词表的方法,顺便记录下来,欢迎各位看官评论指正。
本文主要介绍的是 Efficient softmax approximation for GPUs [1] 这篇paper中提出的adaptive softmax的策略,之所以选这篇文章,是因为这个模型能用GPU来加速,且实际工程上的效果也不错,代码的例子:
https://github.com/peterzhang2029/adaptivesoftmax_torchgithub.com1. Introduction
在一开始的文章介绍中,作者们就总结了,解决大词表一般有两种方案:
- 考虑原始分布:近似原始概率分布或近似原始概率分布的子集。比如 NCE;
- 构造近似模型,但是产生准确的概率分布。比如:Hierarchical softmax;
这篇文章主要采用的是第二种套路,并且结合了词频的信息,也可以说和第一种也有关系吧。
主要的贡献(在笔者看来)
- 能够用GPU(NCE和Hierarchical softmax GPU版本很难实现)
- 实验证明相比softmax精度损失很少
2. Approach
根据Zipf定律,大多数文章是由少数的单词组成的,比如Penn TreeBank中 20%的词,可以覆盖一般文档中87%的出现过的词。
Adaptive Softmax利用了这一定律,将不同词频区间的词分为不同的clusters,按照词频高的cluster优先访问的原则,对cluster中的每个词进行softmax来预测,所以也要求词表需要按照频率从大到小进行排列。
至于为什么想到用这种方法来优化,文章中评估了softmax计算耗时关于单词维度的变化,同时探究了Batch size和softmax计算耗时的关系,发现在矩阵乘法运算中,其中一个维度很小时计算并不高效。于是萌生想法进行改进,此处具体可以看paper。
公式之类的太麻烦,由于这篇文章主要也是实验性的文章,所以我模型就用图来展示一下吧。
先看adaptivesoftmax和softmax之间的区别:
2.1 Softmax vs AdaptiveSoftmax
- Softmax
假设softmax的输入shape为(B, H),而softmax中的weight的shape为(H, V)
那么矩阵计算的复杂度为:B * H * V
- AdaptiveSoftmax
以论文中的two-level为例:
- 第一层是对高频词进行softmax预测,同时最后的三个元素产出对应低频词所属的cluster的概率
- 第二层对应低频词的3个cluster,在第一层预测之后,在进入对应的cluster对低频词用softmax预测
权重的shape关系为:
第一层权重需要预测高频词和低频词对应的cluster(需要3位元素),所以词表大小为
这里
那么AdaptiveSoftmax矩阵计算的复杂度为:
其中
其中
后面我们就结合代码看看模型是如何计算的,代码位置:
https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/adaptive.pygithub.com2.2 AdaptiveSoftmax类定义
class
初始化的输入有5个:
- in_features 就是上文的
- n_classes 等于上文的
- cutoffs 表示低频词clusters的分界值list
- div_value 还是上文的div_value
- head_bias 表示第一层的softmax是否需要bias
这里需要注意的是 cutoffs list中的值需要大于0并小于n_classes
举个例子:假设 cutoffs = [10, 20, 30], n_classes = 100
则第一层高频词
第二层
然后定义了第一层softmax里的线性转换: self.head, 以及第二层softmax的线性转换:self.tail
这里对于第二层的线性转换都是先对输入进行了一个降维来加速计算:Linear(self.in_features, hsz, bias=False) 。上文提到,由于这里是低频词的预测,所以降维造成的效果损失应该可以容忍。
2.3 训练
def
具体的计算过程我都写在注释里了,总结起来大概是以下过程:
- 循环分界点,获取对应cluster的样本
- 对于高频词,这里只用来记录对应的target,在后续统一预测
- 对于低频词,获取低频cluster对应的target的相对位置,经过线性转化并计算当前cluster内的log_prob,记录当前cluster在第一层对应的target位置,在output中记录对应target的log_prob
- 最后,对于整个batch,进行第一层的线性转化,将之前记录的第二层的log_prob加上第一层得到的log_prob,这就得到了样本对应在target位置的 log_prob
最后一点用图画一下:
2.4 预测
def
步骤:
- 第一层的线性转化,获取样本对应第一层target的位置
- 分情况讨论:
- 如果预测的结果都为高频词,则直接返回第一层的log_prob结果
- 如果预测的结果都为低频词,需要计算低频词对应cluster中target对应的log_prob,不要忘了还要加上第一层对应cluster的log_prob
- 如果预测的结果既有高频词也有低频词,只对低频词进行对应cluser的预测
3.实验
光说不练假把式,我在wikitext-2上使用Adaptivesoftmax进行language model的实验:
得到了以下结果:
结果看起来还可以,同样的LSTM model Adaptivesoftmax能减少约40%的训练时间,精度差别在3%左右。
代码地址:https://github.com/peterzhang2029/adaptivesoftmax_torch
同时找到了一个tensorflow的实现,感兴趣的小伙伴可以试试:
https://github.com/yangsaiyong/tf-adaptive-softmax-lstm-lm
如果本文对你有帮助的话,麻烦点个赞支持一下,同时欢迎评论或者私信你想要了解的NLP知识。
参考
- ^Efficient softmax approximation for GPUs https://arxiv.org/pdf/1609.04309.pdf