NLP 神经网络训练慎用 Adam 优化器

AllenNLP 在 EMNLP 2018 的演讲 PPT 《Writing Code for NLP Research》中提到在 NLP 的神经网络中使用 sparse 版的 Adam 优化器替代普通的 Adam 优化器。

通常自然语言处理模型的输入是非常稀疏的。对于包含几十万上百万词的词表，在训练的每个 Batch 中能出现的独立词数不超过几万个。也就是说，在每一轮梯度计算过程中，只有几万个词的 embedding 的梯度是非 0 的，其它 embedding 的梯度都是 0。

Adam 优化器可以说是目前使用最广泛、收敛速度较快且收敛过程较稳定的优化器。Adam 的计算公式如图所示。可以看到公式中梯度的计算使用了动量原理，每一轮用于梯度下降的梯度是当前计算的真实梯度与上一轮用于梯度下降的梯度的加权和。这样动量的引入可以防止训练时产生震荡。Adam 优化器的学习率对于不同参数也是不同的，由该参数历史每一轮的真实梯度的大小决定。好处是对于 NLP 这种输入极度稀疏且输入特征极度不平衡（例如整个预料库中“我”这样的词经常出现，而“拉姆塞”这样的词只出现几次）的任务，学习率是自适应的，一些在一次训练 epoch 中只更新几次的 embedding，在训练后期还是会有较大的学习率。

NLP 输入稀疏的特点与 Adam 使用动量计算梯度的特点相结合就引入了麻烦。每一轮更新参数时，只有极少数 embedding 的梯度是非 0 的，大部分 embedding 的梯度是 0 即上图公式中的 g_t 是 0。但是，计算了动量之后，这些原本梯度都应该是 0 的 embedding 有了非零梯度 m_t 用于梯度下降更新。想象一个极端的例子，“拉姆塞”这个词在一个 epoch 中只在第一个 batch 出现了，于是第一个 batch 计算了“拉姆塞”这个 embedding 的真实梯度 g₀ 用于更新参数，在以后的每个 batch 中虽然“拉姆塞”这个词没有出现过，Adam 都会计算它的动量梯度 m_t，并用于更新“拉姆塞”这个 embedding，实际上方向与 g₀ 完全相同，只是每一轮做一次 β₁ 倍的衰减。这样的做法就相当于对这些出现次数较少的低频词的 embedding，每次梯度下降的等效学习率是非常大的，容易引起类似过拟合的问题。

知道了问题的根节，解决方法就很简单了，每轮迭代只更新这个 batch 中出现过的词的 embedding 即可。TensorFlow 中可以使用 tf.contrib.opt.LazyAdamOptimizer，也可参考 https://www.zhihu.com/question/265357659/answer/580469438 的实现。