Facebook 人工智能研究（FAIR）设计出一种新式的 softmax 函数逼近，专用于 GPU，帮助其在语言模型的基础上通过巨量词汇来有效训练神经网络。
由于在语音识别、机器翻译或语言建模等领域的优异表现，用于序列预测的神经网络最近重新获得关注。然而这些模型都需要巨量的计算，这反而限制了它们的应用。
在语言建模领域，最近的研究进展用到了海量的大型模型，这些大型模型只能在大型 GPU 集群上训练，并且一次需要几周时间。这些处理密集型工作很棒，也有利于探索大型计算基础设备，但这些计算设备对于学术界来说通常十分昂贵，投入生产也不实际，以至于限制了研究的速度、再生产能力和结果的可用性。
意识到这种计算上的瓶颈，Facebook 人工智能研究（FAIR）设计出一种新式的 softmax 函数逼近，专用于 GPUs，帮助其在语言模型的基础上通过巨量词汇来有效训练神经网络。我们的方法叫做自适应 softmax（adaptive softmax），利用不平衡词分布形成簇（cluster），这种簇能明确地减少对计算复杂度的期望，从而规避对词汇量的线性依赖。这种方法通过利用流行架构的特殊性和矩阵-矩阵向量运算（matrix-matrix vector operations）进一步减少了训练和测试时的计算成本。这使得它特别适合于 GPU，而以往的方法，如分层 softmax，NCE 和重要性采样，都是为标准的 CPU 设计的。
FAIR 也开发并正在开源一个名为 torch-rnnlib 的库，这个库允许研究者设计出新的循环模型并在 GPU 上以最小的代价测试这些原型（prototypes）。它也允许通过绑定 torch.cudnn 无缝对接快速基线。几个标准循环网络，如 RNN、LSTM 和 GRU 都已经被部署进来，下面我们将展示如何利用这个库来设计一个新的循环网络。
这些工具和技术后来被一起用来应对标准基准，如 Euro Parl 和 One Billion word，这些都是因需要使用巨大的词汇量而复杂的训练环境，让我们无法在 GPU 上拟合一个大模型和 full softmax。结果显示我们在单一的 GPU 上每秒能处理 12500 个单词，通过标准逼近，大大提升了效率，减少了从试验到结果的时间，同时得到的精确度接近于 full softmax 的精确度。这就能让学界和业界的工程师与研究者都能在短时间内在资源有限的情况下训练出最好的模型。
循环模型的定义有很多，我遵循的是这一条：循环网络随着离散时间对变量序列建模。它遵循马尔可夫的属性，其未来的状态只取决于它的现状。最简单的循环模型是 Elman 的循环模型。根据当下的输入变量 x[t] 以及之前的状态，在每个时间步骤 t，就会输出一个 y[t]。更确切的说，Elman 的循环模型可以通过下面的等式来定义：
h[t] = f(R * h[t-1] + A * x[t]),
y[t] = B * h[t]
其中，h 代表网络（隐藏的）内在的状态，f 是 sigmoid 函数。Elman 之后就有人提出了更复杂的循环模型，如 LSTM、GRU 或者 SCRNN。
语言建模的目的是在一个给定词典中的一个词序列上学习一个概率分布。该联合分布被定义为给定状态下的词的条件分布的一个乘积。确切地说，一个 T 词序列 w[1],...,w[T] 的概率被给定为的
P(w[1],..., w[T])) = P(w[T]|w[T-1],..., w[1])...P(w[1]).这个问题通常由基于计数统计的非参数模型解决（详见 Goodman, 2001）。最近，基于循环神经网络的参数模型在语言建模上才流行起来（例如，Jozefowicz 等人, 2016，obtained state-of-the-art performance on the 1B word dataset）。
我们给出了用于建构带有循环连接的三个不同的 API：
1.nn.{RNN, LSTM, GRU} 接口可以用来建构在所有层都带有相同数量隐藏单元的循环网络。
-- 'Import the library'
local rnnlib = require 'rnnlib'
-- 'Construct the LSTM network'
local lstm = nn.LSTM{inputsize = 256,
hidsize = 512,
nlayer = 2,}
2.rnnlib.recurrentnetwork 的接口能用于构建任意形状的循环网络。上一个和这个接口都为你考虑到了节省隐藏状态。
-- 'Import the library'
local rnnlib = require 'rnnlib'
-- 'Construct the LSTM network'
local lstm = rnnlib.makeRecurrent{
cellfn = rnnlib.cell.LSTM,
inputsize = 256,
hids = {512, 512},}
3.nn. SequenceTable 接口能用来将计算像一个『scan』一样链接起来。nn. RecurrentTable 构建器（constructor）仅仅是一个轻量级的封装，能随着时间为你克隆循环模块。然而，要注意的是这是最低层级的接口，你还需要 rnnlib.setupRecurrent(model, initializationfunctions) 来设定循环隐藏状态行为。
local rnnlib = require 'rnnlib'
--[['The table of cells is fed to each level of
the recurrent network to construct each layer.
The table of initialization functions helps
with the construction of the hidden inputs.'--]]
local cells, initfun = {}, {}
cells[1], initfun[1] = rnnlib.cell.LSTM(256, 512)
cells[2], initfun[2] = rnnlib.cell.LSTM(512, 512)
local lstm = nn.SequenceTable{
dim = 1,
modules = {
nn.RecurrentTable{dim = 2,
module = rnnlib.cell.gModule(cells[1])},
nn.RecurrentTable{dim = 2,
module = rnnlib.cell.gModule(cells[2])},
}
rnnlib.setupRecurrent(lstm, initfun)
你也可以通过定义一个行为像 cell 那样的函数来创建你自己模型，以及一个类似于这个 cell 状态的初始化函数。在 rnnlib.cell 中有预定义的 cell，如 LSTM、RNN、和 GRN。下面来看看如何一步步建立一个 RNN：
local rnncell = function(nin, nhid)
--[['The _make function will later be fed
into rnnlib.cell.gModule to turn it into an
nn.Module which can then be fed into an
nn.SequenceTable or nn.RecurrentTable.
It performs the cell computation.'--]]
local _make = function(prevh, input)
local i2h = nn.Linear(nin, nhid, false)(input)
local h2h = nn.Linear(nhid, nhid, false)(prevh)
local nexth = nn.Sigmoid()(nn.CAddTable(){i2h, h2h})
return nexth, nn.Identity()(nexth)
end
--[[ 'The _init function initializes the
state stored in the cell depending
on the batch size (bsz).'--]]
local _init = function(bsz)
return torch.Tensor(bsz, nhid):fill(0)
end
return _make, _init
end
在 GPU 上训练它
既然 torch-rnnlib 连着 nn 模块接口，仅需要在模型上调取：cuda（）就能将它拉进 GPU。rnnlib 的目的是允许用户自由创建新 cell，或者使用快速基线。这样就 OK 了，如果你在上一节中使用第一或第二个 API 来构建循环网络，就可以轻松使用 cudnn 来大大加速你的网络。对于 nn.{RNN, LSTM, GRU} 接口，仅需要用 usecudnn=true 就能调取这个建构器（constructor）：
nn.LSTM{ inputsize = 256,
hidsize = 512,
nlayer = 2,
usecudnn = true}
对于第二个 API，仅需要将 rnnlib.Recurrent 替换成 rnnlib.makeCudnnRecurrent 并将 cell 函数改为 cudnnAPI 中已有的一个 cell 串。例如：
rnnlib.makeCudnnRecurrent{
cellstring = 'LSTM',
inputsize = 256,
hids = {512, 512},
}
最终的结果通常是，模型的循环部分至少能提速两倍。要注意的是，不是整个模型提速两倍，尤其是如果大部分计算不是在循环部分中时。例如，如果你的模型中有一个 softmax，比循环部分需要更多的计算，那最终速度可能只提升 1.3 倍。
当处理大输出空间（如语言模型）时，分类器（classifier）可能是模型的计算瓶颈。过去已经提出了许多解决方案（分层 softmax（hierarchical softmax），噪声对比估计（noise contrastive estimation），差分 softmax（differentiated softmax）），但是它们通常被设计用于标准的 CPU，并且很少充分利用 GPU 的特性。
我们研究出了一种新的近似方法，叫做自适应 softmax（adaptive softmax）：一个使其计算载荷量（computational budget）适应数据分布的 softmax。它通过更快地访问最常用的类（class），为它们提供更多的资源，来达到优良的近似效果和快速运行时间之间的平衡。更准确地说，它学习了一个 k-通道（k-way）的层次 softmax（hierarchical softmax），它考虑了 GPU 的架构以有效地分配计算。这种计算资源的分配可以使用一个简单的动态规划算法来精确求解。几个技巧可以进一步处理分类器的计算负载问题：我们使用分枝少的树（shallow tree）来避免顺序计算（sequential computation），并且我们为每个 GPU 集群（cluster）确定类（class）的最小值，以避免浪费 GPU 的并行计算能力。
正如表格 1 中显示的，自适应 softmax 几乎可以与完整 softmax（full softmax）的结果相媲美，并且自适应 softmax 速度更快。它也优于 GPU 运行的其他近似模型。
除自适应 softmax（adatpive softmax）外，在技术细节方面，我们使用相对标准的设置：对于小模型，我们使用层数为 1、2048 个神经元单位的的 LSTM；对于大模型，我们使用层数为 2、2048 个神经元单位的 LSTM。我们使用 L2 正则化（regularization）的权重和 Adagrad 来训练模型。我们使用大小为 128 的批处理（batch），同时设置反向传播（back-propagation）窗口的大小为 20。
关于自适应 softmax，我们使用 10 亿单词的训练数据集分布的最佳设置，即 4 个 GPU 集群（cluster），每用一个集群，数据的维数减少 4 倍（更多细节详见论文）。
正如表格 2 中显示的，我们的小模型在几天内就达到了 43.9 的 perplexity 值。我们的大模型在 6 天内达到了 39.8 的 perplexity 值。目前最好的 perplexity 值（越小越好）是 30.0，由 Jozefowicz 等人在 2016 年达到。这个结果是他们用 3 周的时间，使用了 32 个 GPU 达到的。他们也声称使用 18 个 GPU 训练的更小模型，达到了数值为 44 的 perplexity 值。我们的小模型的速度为 180 毫秒/批，并在一个迭代（epoch）后（迭代时间大约为 14 小时）达到数值为 50 的 perplexity 值。不使用 cuDNN 加速库，小模型的速度为 230 毫秒/批，这比之前只慢了 30%。
参考文献
Chelba, C., Mikolov, T., Schuster, M., Ge, Q., Brants, T., Koehn, P., and Robinson, T. (2013). One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3005.
Chen, W., Grangier, D., and Auli, M. (2015). Strategies for Training Large Vocabulary Neural Language Models. arXiv preprint arXiv:1512.04906.
Elman, J. L. (1990). Finding structure in time. Cognitive science, 14(2), 179-211.
Ji, S., Vishwanathan, S.V.N., Satish, N., Anderson, M.J. and Dubey, P. (2015). BlackOut: Speeding up Recurrent Neural Network Language Models With Very Large Vocabularies. arXiv preprint arXiv:1511.06909.
Goodman, J. T. (2001). A bit of progress in language modeling. Computer Speech and Language, 15(4), 403-434.
Grave, É., Joulin, A., Cissé, M., Grangier, D., and Jégou, H. (2016). Efficient softmax approximation for GPUs. arXiv preprint arXiv:1609.04309.
Jozefowicz, R., Vinyals, O., Schuster, M., Shazeer, N., and Wu, Y. (2016). Exploring the limits of language modeling. arXiv preprint arXiv:1602.02410.
Shazeer, N. M., Pelemans, J., and Chelba, C. (2015). Sparse non-negative matrix language modeling for skip-grams.
