By 马里奥·迈克尔·克雷尔博士, Graphcore 首席机器学习负责人 & 马杰·科塞克, Graphcore 人工智能应用专家

图片由作者提供。

 

通过使用新的打包算法，我们在训练 BERT-Large 时将自然语言处理速度提高了 2 倍以上。 我们的新打包技术去除了填充，显着提高了计算效率。

我们怀疑这也可以应用于基因组学和蛋白质折叠模型以及其他具有倾斜长度分布的模型，以在不同的行业和应用中产生更广泛的影响。

我们在一篇新论文 [1] 中介绍了 Graphcore 的高效非负最小二乘直方图打包算法（或 NNLSHP）以及我们应用于打包序列的 BERT 算法。

由于序列填充导致 NLP 中的计算浪费

 
 
我们在最近的工作中开始研究优化 BERT 训练的新方法 向 MLPerf™ 提交的基准测试. 目标是开发可以在实际应用中轻松采用的有用优化。 BERT 是这些优化的重点模型之一，这是一个自然的选择，因为它被广泛用于工业和我们的许多客户。

当我们得知在我们自己使用 Wikipedia 数据集的 BERT-Large 训练应用程序中，数据集中 50% 的标记是填充的——这导致了大量的计算浪费，这真的让我们感到惊讶。

填充序列以将它们全部对齐到相等的长度是 GPU 常用的方法，但我们认为值得尝试不同的方法。

序列的长度变化很大，原因有两个：

1. 基础维基百科数据显示文档长度变化很大
2. BERT 预处理本身随机地减少了组合以生成训练序列的提取文档的大小

将长度填充到最大长度 512 会导致所有标记的 50% 是填充标记。 用真实数据替换 50% 的填充可以导致以相同的计算工作量处理的数据多 50%，因此在最佳条件下速度提高 2 倍。

图 1：维基百科数据集分布。 作者图片。

 

这是维基百科特有的吗？ 不。

那么，它是特定于语言的吗？ 不。

事实上，倾斜的长度分布随处可见：在语言、基因组学和蛋白质折叠中。 图 2 和图 3 显示了 SQuAD 1.1 数据集和 GLUE 数据集的分布。

图 2：SQuAD 1.1 BERT 预训练数据集序列长度直方图，最大序列长度为 384。图片由作者提供。

 


图 3：最大序列长度为 128 的 GLUE 数据集序列长度直方图。图片由作者提供。

 

我们如何在避免计算浪费的同时处理不同的长度？

当前的方法需要不同长度的不同计算内核，或者工程师以编程方式删除填充，然后为每个注意力块和损失计算重复添加它。 通过炸毁代码并使其更复杂来节省计算并没有吸引力，因此我们寻找更好的方法。 我们不能将多个序列放在一个最大长度的包中并一起处理吗？ 事实证明，我们可以！

这种方法需要三个关键要素：

1. 一种有效的算法来决定将哪些样本放在一起以使剩余的填充尽可能少
2. 调整 BERT 模型以处理包而不是序列
3. 并调整超参数

包装

 
 
起初，您似乎不太可能非常有效地打包像维基百科这样的大型数据集。 这个问题通常被称为装箱。 即使打包限制为三个或更少的序列，由此产生的问题仍然是强 NP 完全的，缺乏有效的算法解决方案。 现有的启发式打包算法并不乐观，因为它们的复杂度至少为 O(日志(n））， 在哪里 n 是序列的数量（维基百科约为 16M）。 我们对可以很好地扩展到数百万个序列的方法感兴趣。

两个技巧帮助我们大大降低了复杂性：

1. 将包中的序列数限制为三个（对于我们的第一种解决方案）
2. 仅对序列长度的直方图进行操作，每个出现的长度都有一个 bin

我们的最大序列长度是 512。因此，移动到直方图将维度和复杂性从 16 万个序列减少到 512 个长度计数。 一个包中最多允许三个序列将允许的长度组合数量减少到 22K。 这已经包含了要求序列按包中的长度排序的技巧。 那么为什么不尝试 4 个序列呢？ 这将组合数量从 22K 增加到 940K，这对于我们的第一种建模方法来说太多了。 此外，深度 3 已经实现了非常高的填充效率。

最初，我们认为在一个包中使用三个以上的序列会增加计算开销并影响训练过程中的收敛行为。 但是，为了支持推理等需要更快、实时打包的应用程序，我们开发了高效的非负最小二乘直方图打包 (NNLSHP) 算法。

非负最小二乘直方图包装 (NNLSHP)

 
 
装箱通常被表述为一个数学优化问题。 但是，对于 16 万个序列（或更多），这是不切实际的。 单独的问题变量会超过大多数机器的内存。 基于直方图的方法的数学程序非常简洁。 为简单起见，我们决定采用最小二乘法（轴=b) 与直方图向量 b. 我们通过请求策略向量来扩展它 x 为非负值并添加权重以允许较小的填充。

棘手的部分是策略矩阵。 每列的最大总和为 512，并对哪些序列打包在一起以完全匹配所需的总长度进行编码； 在我们的例子中是 XNUMX。 行对每个潜在组合进行编码以达到总长度。 策略向量 x 是我们正在寻找的，它描述了我们选择 20k 组合中的任何一个的频率。 有趣的是，最后只选择了大约 600 个组合。 为了获得准确的解决方案，该策略很重要 x 必须是正整数，但我们意识到只有非负的近似舍入解 x 就足够了。 对于近似解，可以使用简单的开箱即用求解器在 30 秒内获得结果。

图 4：序列长度 8 和打包深度 3 的策略矩阵示例。行代表长度为 1-8 的序列，这些序列被打包在一起，列代表一个包中所有可能的长度组合，没有特定的排序。 作者图片。

 

最后，我们不得不修复一些没有分配策略但最小的样本。 我们还开发了一个变体求解器，它强制每个序列都被打包，可能带有填充，并且具有依赖于填充的权重。 花了更长的时间，解决方案也好不到哪里去。

最短包装优先直方图包装

 
 
NNLSHP 为我们提供了充分的包装方法。 然而，我们想知道我们是否可以在理论上获得更快的在线能力并消除仅将 3 个序列放在一起的限制。

因此，我们从现有的打包算法中获得了一些灵感，但仍然专注于直方图。

我们的第一个算法，最短包装优先直方图包装 (SPFHP) 有四个要素：

1. 对从最长序列到最短序列的直方图计数进行操作
2. 如果当前序列长度不适合任何包，则开始一组新的包
3. 如果有多个拟合，取序列长度总和最短的包，分别修改计数
4. 再次检查剩余计数的拟合

这种方法最容易实现，只需要 0.02 秒。

一种变体是采用序列长度的最大总和而不是最短和分裂计数来获得更完美的拟合。 总的来说，这并没有太大改变效率，但增加了很多代码复杂度。

最短包装优先直方图包装的工作原理。 作者动画。

 

维基百科，SQuAD 1.1，GLUE 打包结果

 
 
表 1、2 和 3 显示了我们提出的两种算法的打包结果。 包装深度 描述了打包序列的最大数量。 打包深度 1 是基线 BERT 实现。 最大出现的填充深度，没有设置限制，用附加的“max”表示。 这 包数 描述新打包数据集的长度。 效率 是打包数据集中真实标记的百分比。 这 包装系数 描述了与填充深度 1 相比产生的潜在加速。

我们有四个主要观察结果：

1. 分布越偏，打包的好处就越大。
2. 所有数据集都受益于打包。 有些甚至超过 2 倍。
3. 当填充深度不受限制时，SPFHP 的效率会更高。
4. 对于最多 3 个打包序列，NNLSHP 越复杂，效率越高（99.75 对 89.44）。

表 1：维基百科上提出的打包算法（SPFHP 和 NNLSHP）的关键性能结果。 作者图片。

 


表 2：为 SQUaD 1.1 BERT 预训练提出的打包算法的性能结果。 作者图片。

 


表 3：针对 GLUE 数据集提出的打包算法的性能结果。 仅显示基线和 SPFHP 保压结果，不限制保压深度。 作者图片。

 

BERT 处理调整

 
 
BERT 架构的有趣之处在于，大多数处理发生在令牌级别，这意味着它不会干扰我们的打包。 只有四个组件需要调整：注意力掩码、MLM 损失、NSP 损失和准确性。

处理不同数量序列的所有四种方法的关键是矢量化和使用可以连接的最大数量的序列。 为了引起注意，我们已经有了一个掩码来解决填充问题。 将其扩展到多个序列很简单，如下面的 TensorFlow 伪代码所示。 这个概念是我们确保注意力仅限于单独的序列，不能超出这个范围。

注意掩码代码示例。

 


图 5：示例零一掩码

 

对于损失计算，原则上我们解包序列并计算单独的损失，最终获得序列（而不是包）的损失平均值。

对于传销损失，代码如下所示：

损失计算代码示例。

 

对于NSP loss和accuracy，原理是一样的。 在我们的公共示例中，您可以在我们的内部找到相应的代码 PopART 框架.

维基百科开销和加速估计

 
 
通过我们对 BERT 的修改，我们有两个问题：

1. 它带来了多少开销？
2. 开销有多少取决于放在一个包中的最大序列数？

由于 BERT 中的数据准备工作可能很麻烦，因此我们使用了快捷方式并针对多个不同的打包深度编译了代码，并比较了各自的（测量的）周期。 结果如表 4 所示。 开销，我们表示由于模型更改以启用打包（例如用于注意力的屏蔽方案和更改的损失计算）而导致吞吐量下降的百分比。 这 实现加速 是打包带来的加速的组合（ 包装系数) 和吞吐量下降 开销.

表 4：维基百科上提议的打包算法（SPFHP 和 NNLSHP）的估计加速比较。 作者图片。

 

多亏了矢量化技术，开销小得惊人，而且将多个序列打包在一起也没有任何缺点。

超参数调整

 
 
通过包装，我们将有效批量（平均）增加了一倍。 这意味着我们需要调整训练超参数。 一个简单的技巧是将梯度累积计数减半，以保持与训练前相同的有效平均批量大小。 通过使用带有预训练检查点的基准设置，我们可以看到准确度曲线完全匹配。

图 6：打包和解包处理的学习曲线比较 减少批量大小 对于打包方法。 作者的图片。

 

准确率匹配：MLM 训练损失在开始时可能略有不同，但很快就赶上了。 这种初始差异可能来自于注意力层的轻微调整，这些调整可能在之前的训练中偏向于短序列。

为避免速度变慢，有时有助于保持原始批次大小相同，并将超参数调整为增加的有效批次大小（加倍）。 要考虑的主要超参数是 beta 参数和学习率。 一种常见的方法是将批量大小加倍，在我们的案例中这会降低性能。 查看 LAMB 优化器的统计数据，我们可以证明将 beta 参数提高到包装因子的幂对应于连续训练多个批次以保持动量和速度的可比性。

图 7：打包和解包处理的学习曲线比较 启发式 应用。 作者的图片。

 

我们的实验表明，将 beta 设为 XNUMX 的幂是一种很好的启发式方法。 在这种情况下，预计曲线不会匹配，因为增加批次大小通常会降低样本/时期意义上的收敛速度，直到达到目标精度。

现在的问题是，在实际场景中，我们是否真的获得了预期的加速？

图 8：打包和未打包处理的学习曲线比较 优化设置. 作者的图片。

 

是的，我们愿意！ 我们获得了额外的加速，因为我们压缩了数据传输。

结论

 
 
将句子打包在一起可以节省计算工作量和环境。 这种技术可以在任何框架中实现，包括 PyTorch 和 TensorFlow。 我们获得了明显的 2 倍加速，并且在此过程中，我们扩展了打包算法的最新技术。

我们感兴趣的其他应用是基因组学和蛋白质折叠，可以观察到类似的数据分布。 Vision Transformers 也可能是一个有趣的领域，可以应用不同大小的打包图像。 您认为哪些应用程序运行良好？ 我们很想听到您的声音！

阅读论文

在 GitHub 上访问代码

谢谢

 
 
感谢 Graphcore 应用工程团队的同事 Sheng Fu 和 Mrinal Iyer 对这项工作的贡献，并感谢 Graphcore 研究团队的 Douglas Orr 提供的宝贵反馈。

参考资料

 
 
[1] M. Kosec, S. Fu, MM Krell, 打包：迈向 2 倍 NLP BERT 加速 （2021），arXiv

 
马里奥·迈克尔·克雷尔博士 是 Graphcore 的首席机器学习负责人。 马里奥一直在研究和开发机器学习算法超过 12 年，为机器人、汽车、电信和医疗保健等不同行业开发软件。 在 Graphcore，他为我们令人印象深刻的 MLPerf 提交 并且热衷于加速新的非标准模型，例如用于统计 COVID-19 数据分析的近似贝叶斯计算。

马杰·科塞克 是 Palo Alto Graphcore 的 AI 应用专家。 他曾在圣何塞的 NIO 担任自动驾驶人工智能科学家，并拥有斯坦福大学航空航天硕士学位。

原版。 经许可重新发布。

相关新闻：

• 使用 BERT 构建求职知识图谱
• 如何从头开始训练 BERT 模型
• 用拥抱脸理解BERT