0.3%参数推理，实现78倍加速！ETH团队提出UltraFastBERT，构筑语言模型巨人

2023-11-24 04:11

©PaperWeekly 原创 · 作者 | 岳廷

研究方向 | 大语言模型

对于大语言模型，如果能在推理过程中只选择性地使用模型中的一小部分参数，将大幅减少计算量，使语言模型的响应速度得到质的提升。

然而参数量的减少是否会降低模型性能呢？对此，来自瑞士联邦理工学院苏黎世分校的研究人员在最新研究中给出了证明。他们设计了一种名为 UltraFastBERT 的 BERT 模型变体，其参数量与原 BERT-BASE 持平，但推理时只使用了0.3%的参数。

实验结果表明，在多项下游语言理解任务上的表现与 BERT-BASE 可相比拟。这说明语言模型的参数在推理时存在巨大的冗余。同时实验结果显示，相对于基线获得 78 倍加速！

论文题目：

Exponentially Faster Language Modeling

论文链接：

https://arxiv.org/pdf/2311.10770

代码链接：

https://github.com/pbelcak/UltraFastBERT

方法

大语言模型的前馈层占据了它们参数的大部分。然而，并非所有神经元都需要在每次输入时被用于计算前馈层的输出。

文章提出了一种新的 BERT 变体模型 UltraFastBERT，其核心是使用快速前馈网络(FFF)来替代 BERT 的前馈层（FF）。UltraFastBERT 的中间层在设计上快得多：对于具有 n 个神经元的FF网络和FFF网络，通过 FFF 的前向传播的时间复杂度为，而 FF 的时间复杂度为。这是因为 FFF 将其神经元组织成一个平衡二叉树，并且根据输入条件选择性地执行二叉树的某一个分支。从而大大减少了运算量。

在 FFF 上执行推理等同于进行条件矩阵乘法（CMM），其中输入的行与神经权重的列依次进行点乘，而选择进行的权重列取决于之前点乘操作的输出。通过这种方式，所有神经元仅被一些输入使用，并且没有任何输入需要网络处理超过少数几个神经元。这与传统前馈网络中的密集矩阵乘法（DMM）形成对比，DMM 计算所有行与所有列的点积。

对于 CMM，除了 high-level 模拟，目前尚无高效的原生 CMM 库，也没有流行的深度学习框架有现成的接口可供实现。因此，文章提供了一套基于 BLAS 库的指针批量矩阵乘法程序的 CPU 实现。

模型

2.1 架构

文章以 Crammed BERT 架构为基准。将前馈网络替换为快速前馈网络。Crammed BERT 是一种对 BERT 进行改进的模型，它通过增加模型的大小和容量来提高性能。Crammed BERT 的模型大小是原始 BERT 模型的两倍，并且具有更多的参数。

Crammed BERT 的改进主要体现在以下两个方面：

更大的模型大小可以让 Crammed BERT 学习到更复杂的关系和模式；
更多的参数可以让 Crammed BERT 更好地捕捉文本中的细微差别。

Crammed BERT 在各种自然语言处理任务上都取得了显著的改进。

另外，文章对架构做以下简化更改：

1. 删除叶节点和非叶节点之间的所有差异。特别是在所有节点上使用相同的(GELU)激活函数，在所有节点上都配备输出权重，并删除所有输出偏差；

2. 将叶大小固定为 1；

3. 允许多个 FFF 树并行。通过对各个树的输出求和并将和呈现为中间层输出，以允许多个 FFF 树共同计算中间层输出。

2.2 下游性能

遵循与 Crammed BERT 一样的训练过程。调优所有 FastBERT 模型以用于 GLUE 基准任务，并报告评估分数。调优结果如表 1。

可以看到，在单个 A6000 GPU 上训练 1 天的所有 FastBERT 变量至少保留了原始 BERT-base 模型下游预测性能的 96.0%。

推理

从工程学角度来看，其实现的核心是条件矩阵乘法运算，其伪代码如下。

3.1 算法

在算法 1 中，B 表示批大小，H 表示层输入宽度（transformer 隐藏维度），是神经元的数量，表示 M 的第 k 列，表示M的第l行。CMM 中的 >- 比较的结果假定为整数∈{0, 1}。

另外，总体上 CMM 与现有计算流程和硬件是兼容的，但实现起来更有难度。

3.2 推理性能

文章比较了几种可用的 FF/FFF 推理实现的速度。

实现方法：对于 CPU 推理，使用 Intel oneAPI 中提供的数学内核库。

Level 1 实现仅使用 BLAS 级 1 例程和类 BLAS 级 1 扩展构建，即向量-向量点积和标量-向量积。

Level 2 实现使用批处理 BLAS 2 级例程和类 BLAS 1 级扩展，即批处理矩阵-向量乘法和批处理标量-向量积。

Level 3 实现使用（非批处理）BLAS 3 级矩阵-矩阵乘法。这是 FF 的最快 CPU 实现，但由于向量级稀疏性不被库支持，目前无法为 FFF 提供此类实现。

对于 GPU 实现，使用 PyTorch 内核或自定义 CUDA 内核：

native fused 实现使用 native fused 前馈层内核。请注意，对于 FF 图层这是最快的 GPU 实现，但由于 CMM 的性质，当前不存在这样用于 FFF 的内核；
BMM 实现对 FF 和 FFF 都使用批处理矩阵乘法和激活内核。在 FFF 的情况下，在每一步的树下降中大量使用向量复制来模拟条件性；
native CUDA 实现是自定义 CUDA 内核代码，既用于 FF 又用于 FFF。

测试方法：对于 CPU 推理，在 Intel(R) Core(TM) i7-6700HQ CPU 上执行 250 次前向传递，在 Intel MKL v2023.2.0 下，使用所有例程的 64 位变体。报告单次推理花费的平均时间，注意标准偏差的值始终远低于平均值的 2%。对于 GPU 实现，在 NVIDIA RTX A6000 GPU 上执行 1000 次前向传递，在 CUDA v11.7 和 PyTorch 2.0.1 下。测量 GPU 时间并报告平均时间，标准偏差同样低于平均值的 2%。

结果：表 2 列出了 BERT-base 和 FastBERT-1x11 中前馈层和快速前馈层的性能比较。表的每一列列出了在使用相同线性代数例程原语时 FFF 相对于 FF 实现的实现速度提升。