Adaptive softmax模型：Efficient softmax approximation for GPUs

参考链接

• 论文链接：https://arxiv.org/pdf/1609.04309v3.pdf
• 项目链接：https://github.com/facebookresearch/adaptive-softmax

Introduction：

• 本文提出是的方法是adaptive softmax, 该算法目的是为了提高softmax函数的运算效率，适用于一些具有非常大词汇量的神经网络。
• 在NLP的大部分任务中，都会用到softmax,但是对于词汇量非常大的任务，每次进行完全的softmax会有非常大的计算量，很耗时。所以paper中提出adaptive softmax来提升softmax的运算效率。
• adaptive softmax的特点
  • 该算法的提出利用到了单词分布不均衡的特点(unbalanced word distribution)来形成将单词分成不同的类, 这样在计算softmax时可以避免对词汇量大小的线性依赖关系，降低时间复杂度；
  • 另外通过结合现代架构和矩阵乘积操作的特点，通过使其更适合GPU单元的方式进一步加速计算。

Computation time model of matrix-multiplication：

• 大部分语言模型的瓶颈在于：隐藏状态的矩阵$H∈R^{B×d}$H∈RB×d(B=batch size,d为隐藏层状态维度)与词汇表矩阵$E∈R^{d×k}$E∈Rd×k(k=词汇表的大小)之间的矩阵相乘。
• 假设将$d$d和$B$B固定，下图为两个GPU模型K40和M40矩阵乘法时间随$k$k的变化图：
  
  我们将上图的曲线记为函数$g(k)$g(k),我们可以看到，大概在$k<k_0≈50$k<k0​≈50时，时间$g(k)$g(k)是一个常量。直到$k>k_0$k>k0​时, 时间$g(k)$g(k)变成成了一个线性函数。
• 因此假设g(k)的表达式如下：$g(k)=max\{c+\lambda k_0,c+\lambda k\}=c_m+max\{0,\lambda (k+k_0)\}$g(k)=max{c+λk0​,c+λk}=cm​+max{0,λ(k+k0​)}在K40上$c_m=0.4$cm​=0.4,在M40上$c_m=0.22$cm​=0.22.
• 同样的，假设将$d$d和$K$K固定，模型K40和M40矩阵乘法时间随$B$B的变化也呈现上面这样的关系。
• 这就表明矩阵相乘时，当其中一个维度很小时，矩阵乘法是低效的。也就是说计算大量小矩阵是不划算的，很容易理解，数据传输、函数调用开销等也需要消耗时间。也说明了，按层次来组织words时，如果中间节点有较少的子节点时 (比如huffman 编码的tree)是次优的。
• 所以paper中综合考虑了$k$k与$B$B，提出了下面的计算时间模型：$g(k; B) = max(c + λk_0 B_0; c + λkB)$g(k;B)=max(c+λk0​B0​;c+λkB)

two-clusters case：

• 这部分讲解将词汇表分成两个类的Adaptive softmax模型。
• 在自然语言中，单词的分布遵循Zipf定律。大部分的概率被词典中一小部分词所覆盖，例如，在Penn TreeBank中87%的文档只被20%的词汇所覆盖。
• 首先将将词汇表$V$V分成两部分$V_h$Vh​和$V_t$Vt​,$V_h$Vh​包含在词汇表头部高频词，$V_t$Vt​中包含在词汇表尾部的低频词。高频词会经常被访问所有它应该被高效的计算。这使得我们使用不平衡基数(unbalanced cardinalities)：$|V_h |≪|V_t |$∣Vh​∣≪∣Vt​∣和不平衡概率(unbalanced probabilities):$P(V_h)≫P(V_t)$P(Vh​)≫P(Vt​)(其中$P(A)=∑_{m∈A}P_i$P(A)=∑m∈A​Pi​表示一个词在$A$A中的概率),来进行词典划分。
• 两个词类(clusters)有两种组织方式：
  • 方法一：$V_h$Vh​和$V_t$Vt​都放在树的第二层的叶节点中。
  • 方法二：$V_h$Vh​保存在根节点中而$V_t$Vt​保存在第二层的叶节点中。
• 方法一比方法二性能上要低$5\%-10\%$5%−10%。当预测$c$c类中词$w$w时：
  • 方法一：$P(c│H)*P(w│c,h)$P(c│H)∗P(w│c,h)
  • 方法二：$P(w|h)$P(w∣h)，其更为简单直接，所以效果更好。
• 因此，除非方法一比方法二的计算时间快很多，否则选择方法二的组织方式。

Minimizing the computation time：

• 令词汇表的大小为$k$k，$V_h$Vh​中词的个数为$k_h, V_h$kh​,Vh​对应的概率为$p_h$ph​.则：$k_t=k-k_h；p_t=1-p_t$kt​=k−kh​；pt​=1−pt​.
• 所以，根部计算的时间为：$g(k_h+1;B); V_t$g(kh​+1;B);Vt​叶节点的计算时间为：$g(k_t; p_t B)$g(kt​;pt​B);所以整体的计算时间为：$C = g(k_h+1;B) + g(k_t; p_t B)$C=g(kh​+1;B)+g(kt​;pt​B)
• 这样我们通过最小化$C$C来求的合适的$k_h$kh​，$C$C与$k_h$kh​的关系图如下
  

Adapting the classifier capacity for each cluster：

• 每个类(cluster)是独立的被访问，因此它们不需要有相同的能力。对应频繁的词它们使用跟高的能力，而稀少的词他们可以不被学的非常好，因为它们很少被查看。
• 在论文中，对不同的类共享隐层，通过加映射矩阵的方式降低分类器的输入大小。一般的，tail部分(低频词部分)的映射矩阵将其大小从$d$d维缩小到$d_t=d/4$dt​=d/4 维。

General case

• 现在考虑更一般的情况，将词汇表分成J分：$V = V_h ∪V_1…∪ V_J$V=Vh​∪V1​…∪VJ​，并且$V_i∩V_j=∅(i≠j)$Vi​∩Vj​=∅(i​=j),分成结构如下所示：

  

• C：总计算时间；词$w$w属于类$V_i$Vi​的概率为$P(w∈V_i )=p_i ∑_{w∈{v_i}}p(w)$P(w∈Vi​)=pi​∑w∈vi​​p(w),$k_i=|V_i|$ki​=∣Vi​∣;固定$B$B(batch size)和$d$d(隐藏层维度)，仅考虑$C$C与$k_i$ki​和概率的关系.

• 将$C$C进行如下分解:$C =C_h+∑_iC_i =g(J+k_h; B)+g(k_i; p_i B)$C=Ch​+i∑​Ci​=g(J+kh​;B)+g(ki​;pi​B)

• 回顾公式:$g(k; B)= max(c + λk_0 B_0; c + λkB)$g(k;B)=max(c+λk0​B0​;c+λkB);其中$+ λk_0 B_0$+λk0​B0​为常数部分, $c + λk$c+λkB是$kB$kB的线性函数。有上面部分分析可知，在常数部分GPU的算力并未得到充分利用，因此应尽量避免落入该部分;所以设置一个约束条件：$kB≥k_0 B_0$kB≥k0​B0​;合并上面式子得:$C=c+\lambda (J+k_h)B+\sum_i(c+\lambda k_ip_iB)=(J+1)c+\lambda B[J+k_h+\sum_ip_ik_i]$C=c+λ(J+kh​)B+i∑​(c+λki​pi​B)=(J+1)c+λB[J+kh​+i∑​pi​ki​]

• 讨论$∑_ip_i k_i$∑i​pi​ki​项:

  • 令$p_{i+j}=p_i+p_j$pi+j​=pi​+pj​,则$p_j k_j=(p_(i+j)-p_i ) k_j$pj​kj​=(p(​i+j)−pi​)kj​则可以得到：$p_i k_i+p_j k_j=p_i (k_i-k_j)+p_i+p_(i+j) k_j$pi​ki​+pj​kj​=pi​(ki​−kj​)+pi​+p(​i+j)kj​
  • 假设$k_i>k_j, p_(i+j)、k_j$ki​>kj​,p(​i+j)、kj​都固定，只有$p_i$pi​是可变的；那么当$p_i$pi​越小式上式越小，即让容量更大的第$i$i类包含更多概率小的词。

• 讨论$k_i$ki​项: 论文中说使用了动态规划。

• 讨论$J$J项：最优的计算机时间关于J的函数图形如下图所示：