我对实现可以处理大型词汇表的分层softmax模型感兴趣，比如大约10M类。这样做的最佳方法是扩展到大班级和高效吗？例如，至少one paper表明，当使用每个节点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？深层和窄层与短树和宽树有不同的方式吗？
答案 0 :(得分：11)
你提到你想要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是一个＆＃34;指定并运行＆＃34;框架 - 计算图对于最佳情况和最坏情况是相同的）
这样做的最佳方法是扩展到大班级和高效吗？
模型并行性的上述低效率（每个GPU必须处理整批）表明应该尝试将所有内容保存在一个地方。
让我们假设您要么在主机上实现所有内容，要么在1个庞大的GPU上实现。
如果您没有对序列进行建模，或者如果您没有，但整个序列只有一个输出，那么复制参数的内存开销与内存要求相比可忽略不计如上所述：
400 ==批量大小＆lt;＆lt;班级数== 3K
在这种情况下，您只需使用gather或embedding_lookup（尽管复制效率低下）
但是，如果您对每个时间步长的输出建模长度（例如100），那么参数复制就成了一个大问题。
在这种情况下，我认为您需要下载到C ++ / CUDA C并将整个图层及其渐变实现为自定义操作。