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

Question

我对实现一个可以处理大型词汇表的分层 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 分别处理每个样本并返回使用gather。这更具可扩展性，但由于序列化，并没有真正接近 25 倍的加速。

有没有更好的方法来实现 HS？深而窄的树与短而宽的树有不同的方法吗？

Answer 1

你提到你想要 GPU 级的性能：

但现在将所有内容都保留在 CPU 上并减慢速度

并希望使用 300 个单位的隐藏大小和 1000 万字的字典。

这意味着（假设 float32），您需要 4 * 300 * 10M * 2 字节 = 24 GB 来存储输出层的参数和梯度。

分层 Softmax (HSM) 不会降低内存需求 - 它只是加快了训练速度。

实际上，您将需要更多的 GPU 内存，因为您还需要存储：

• 其他参数及其梯度

• 优化器数据，例如动量训练中的速度

• 激活和反向传播的临时数据

• 特定于框架的开销

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

但是，你现在有另一个问题：

为了具体说明，假设您有一个包含 3K 类的 2 级 HSM，每个类有 3K 字（总共 9M 字）。您将 3K 类分布在 8 个 GPU 上，因此每个 GPU 托管 384 个类。

如果一个批次中的所有目标词都来自相同的 384 个类，即它们属于同一个 GPU，该怎么办？一个 GPU 将完成所有工作，而其他 7 个等待它。

问题在于，即使一个批次中的目标词属于不同的 GPU，如果您想在 TensorFlow 中进行此计算（这是因为 TensorFlow是一个“指定并运行”的框架——计算图对于最好的情况和最坏的情况是一样的）

什么是做到这一点的最佳方法，既可扩展到大类数量又高效？

上述模型并行性的低效率（每个 GPU 必须处理整个批次）表明应该尝试将所有内容放在一个位置。

让我们假设您要么在主机上实现所有东西，要么在 1 个巨大的 GPU 上实现。

1. 如果您不是对序列进行建模，或者如果您是，但整个序列只有一个输出，那么与上述内存要求相比，复制您提到的参数的内存开销可以忽略不计：

   400 == 批量大小

   在这种情况下，您可以简单地使用gather 或embedding_lookup（虽然复制效率低下）

2. 但是，如果您对长度为 100 的模型序列进行建模，并且在每个时间步都有输出，那么参数复制就会成为一个大问题。

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