Tensorflow中可扩展,高效的分层Softmax？

Question

我有兴趣实现一个可以处理大型词汇表的分层softmax模型,比如大约10M类.这样做的最佳方法是扩展到大班级和高效吗？例如,至少有一篇论文表明,当使用每个节点sqrt(N)类别的2级树时,HS可以为大型词汇实现~25倍的加速.我也对具有任意分支因子的任意深度树的更通用版本感兴趣.

我在这里看到一些选项:

1)tf.gather为每个批次运行,我们收集索引和拆分.这会产生大批量和胖树的问题,现在系数重复很多,导致OOM错误.

2)与#1类似,我们可以使用tf.embedding_lookup哪个可以帮助解决OOM错误,但是现在可以将所有内容保存在CPU上,并且可以减慢速度.

3)使用tf.map_fn与parallel_iterations=1分别处理每个样本,并返回到使用收集.这更具可扩展性,但由于序列化,它并没有真正接近25倍的加速.

有没有更好的方法来实施HS？深层和窄层与短树和宽树有不同的方式吗？

Answer 1

你提到你想要GPU级性能:

但是现在将所有内容都保存在CPU上,并将速度降低了很多

并希望使用300个单元的隐藏大小和10M字词典.

这意味着(假设float32),您需要4*300*10M*2字节= 24 GB才能存储输出层的参数和梯度.

分层Softmax(HSM)不会降低内存需求 - 它只会加快培训速度.

实际上,你需要更多的GPU内存,因为你还需要存储:

因此,如果您想在GPU上进行所有计算,您将别无选择,只能将此层分布在多个高内存GPU中.

但是,您现在有另一个问题:

为了使这个具体,我们假设您有一个2级HSM,3K类,每个类3K字(总共9M字).您在8个GPU中分发3K类,因此每个主机有384个类.

如果批处理中的所有目标词来自相同的384类,即它们属于同一GPU,该怎么办？一个GPU将完成所有工作,而另外7个等待它.

问题是,即使在一个批次中的目标词属于不同的GPU,你仍然有相同的性能,在最坏的情况下,如果你想要做这个计算中TensorFlow(这是因为TensorFlow是一个"指定并运行"框架 - 对于最佳情况和最坏情况,计算图是相同的)

这样做的最佳方法是扩展到大班级和高效吗？

上述模型并行性的低效率(每个GPU必须处理整批)表明应该尝试将所有内容保存在一个地方.

让我们假设您要么在主机上实现所有内容,要么在1个庞大的GPU上实现.

如果您没有对序列建模,或者如果您没有,但整个序列只有一个输出,那么复制参数的内存开销与上述内存要求相比可忽略不计:

400 ==批量大小<<类数== 3K

在这种情况下,您可以简单地使用gather或embedding_lookup(虽然复制效率低)
但是,如果您对每个时间步长的输出建模长度(例如100),那么参数复制就成了一个大问题.

在这种情况下,我认为您需要下载到C++/CUDA C并将整个层及其渐变实现为自定义操作.