UltraFastBERT与Fast-FeedForward-Network

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

然而参数量的减少是否会降低模型性能呢？对此，来自瑞士联邦理工学院苏黎世分校的研究人员在最新研究中给出了证明。他们设计了一种名为 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 BASE中：

词表的大小是(word list)：30522
Encoder层个数是(layer)：12
词向量的大小(vocab dim)：768
文本最大长度(seq length)：512
头个数(multi head attention)：12
Feed Forward的两层全链接层神经元个数分别是：3702, 768
BERT中Encoder包括: Embedding层，Multi-Head Attention 层，Feed-Forward Network层，LayerNorm层参数。

1. Embedding层
   该层包括三种Embedding，具体是Token Embedding, Segment Embedding, Position Embedding
   Token Embedding 层参数: 30522 * 768
   Segment Embedding层参数：2 * 768
   Position Embedding层参数：512 * 768
   因此总的参数量为：（30522 + 512 +2）* 768 = 23835648 = 22.7 M
2. Multi-Head Attention层
   该层主要是由Q、K、V三个矩阵运算组成，BERT模型中是Multi-head多头的Self-attention机制。先通过Q和K矩阵运算并通过softmax变换得到对应的权重矩阵，然后将权重矩阵与 V矩阵相乘，最后将12个头得到的结果进行concat，得到最终的Self-attention层输出。
   又因为BERT模型中包含12个Transformer Encoder层，因此改层的参数总量为：[768 * (768/12) * 3 * 12 + 768 * 768 ] * 12 = 28311552 = 27M
3. LayerNorm层
   LayerNorm层主要有weight和bias两个参数。而LN层在Embedding层、Self-attention层、Feed-Forward Network层三个层都有用到，因此LN层的参数总量为：768 * 2 + （768 * 2）* 12 + （768 * 2）* 12 = 38400 = 37.5KB
4. Feed-Forward Network层
   前馈网络FFN主要由两个全连接层组成，且W1和W2的形状分别是(768,3072)，(3072,768)，层数为12，因此该层的参数量为：
   （768 * 3072 + 3072 * 768）* 12 = 56623104 = 54M

将上面的计算结果加起来，那么BERT模型的参数总量为：23835648 + 28311552 + 56623104 + 38400 = 108808704 ≈ 104M。

Embedding层约占参数总量的20%，Transformer层约占参数总量的80%。

大语言模型的前馈层占据了它们参数的大部分。然而，并非所有神经元都需要在每次输入时被用于计算前馈层的输出。
本文介绍了 UltraFastBERT，它是一种 BERT 变体，在推理过程中使用了 0.3% 的神经元，而性能却与类似的 BERT 模型相当。UltraFastBERT 每层推理只选择性地使用 4095 个神经元中的 12 个。这是通过用快速前馈网络（FFF）取代前馈网络实现的。UltraFastBERT 的中间层在设计上比BERT快了指数级：给定前馈（FF）和快速前馈（FFF）网络，每个网络都有 n 个神经元，通过 FFF 的前向传递的时间复杂度为 $\mathcal{O}(lo{g}_{2}n)$，而FF是$\mathcal{O}(n)$这是因为 FFF 将其神经元组织成一个平衡二叉树，并且根据输入条件选择性地执行二叉树的某一个分支。从而大大减少了运算量。
在 FFF 上执行推理等同于进行条件矩阵乘法（CMM），其中输入的行与神经权重的列依次进行点乘，而选择进行的权重列取决于之前点乘操作的输出。通过这种方式，所有神经元仅被一些输入使用，并且没有任何输入需要网络处理超过少数几个神经元。这与传统前馈网络中的密集矩阵乘法（DMM）形成对比，DMM 计算所有行与所有列的点积。

模型

架构

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

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

• 更大的模型大小可以让 Crammed BERT 学习到更复杂的关系和模式；
• 更多的参数可以让 Crammed BERT 更好地捕捉文本中的细微差别。
• Crammed BERT 在各种自然语言处理任务上都取得了显著的改进。

快速前馈网络（Fast Feedforward Network，FFF）是由两部分组成的：节点网络集合$\mathcal{N}$和叶子网络集合$\mathcal{L}$。

• 节点网络集合$\mathcal{N}$ 包含了一组节点网络，每个节点网络都是一个 $⟨{\dim }_I,n,1⟩$前馈网络，并在输出上增加了一个 sigmoid 激活函数。这些节点网络按照平衡的可微分二叉树的形式排列，其中 $N_{m+1,2n}$和 $N_{m+1,2n+1}$是$N_{m,n}$ 的子节点。
• 叶子网络集合$\mathcal{L}$包含了一组叶子网络，每个叶子网络都是一个$⟨{\dim }_I,\ell ,{\dim }_O⟩$-前馈网络。叶子网络没有子节点，它们的输出直接作为 FFF 的输出。
  
  

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

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

2. 将叶大小固定为 1；

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

下游性能

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

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

推理

从工程学角度来看，其实现的核心是条件矩阵乘法运算，其伪代码如下。
前向传播过程由下面算法控制。

算法的输入包括一个输入样本 $\iota$ 和根节点$N_{0,0}$，输出为该样本在 FFF 中的输出。
算法定义了两个函数：$Forwar{d}_T$和${Forward}_I$。其中，${Forward}_T$函数用于计算节点的输出，而${Forward}_I$函数用于计算节点的指示值（indicator value）。

• 在 ${Forward}_T$函数中，如果当前节点是叶子节点，则直接调用该节点的前馈传播函数$N_{m,n}(\iota )$来计算输出。否则，首先计算当前节点的输出 $c_{m,n}=N_{m,n}(\iota )$，然后递归地调用${Forward}_T$函数来计算当前节点的两个子节点的输出，并将它们加权相加作为当前节点的输出。
  与传统的前馈神经网络算法相比，该算法的主要区别在于引入了一个计算节点的指示值。指示值表示了当前节点的输出是否大于等于阈值（这里的阈值为0.5），根据指示值的大小来确定选择哪个子节点进行计算。这种方式可以大大减少计算量，提高前向传播的效率。同时，FFF 是一种具有平衡二叉树结构的前馈神经网络，其中节点网络和叶子网络分别用于处理中间层和输出层的计算。通过利用二叉树结构和递归计算，FFF 可以实现快速的前向传播。
  

import torch
from torch import nn
import math

class FFF(nn.Module):
	def __init__(self, input_width: int, depth: int, output_width: int, *args, **kwargs):
		super().__init__(*args, **kwargs)

		self.input_width = input_width
		self.depth = depth
		self.output_width = output_width

		self.n_nodes = 2 ** (depth + 1) - 1
		self.initialise_weights()

	def initialise_weights(self):
		init_factor_l1 = 1.0 / math.sqrt(self.input_width)
		init_factor_l2 = 1.0 / math.sqrt(self.depth + 1)
		self.w1s = nn.Parameter(torch.empty(self.n_nodes, self.input_width).uniform_(-init_factor_l1, +init_factor_l1), requires_grad=True)
		self.w2s = nn.Parameter(torch.empty(self.n_nodes, self.output_width).uniform_(-init_factor_l2, +init_factor_l2), requires_grad=True)

	def forward(self, x):
		# the shape of x is (batch_size, input_width)
		# retrieve the indices of the relevant elements
		batch_size = x.shape[0]
		current_nodes = torch.zeros((batch_size,), dtype=torch.long, device=x.device)
		all_nodes = torch.zeros(batch_size, self.depth+1, dtype=torch.long, device=x.device)
		all_logits = torch.empty((batch_size, self.depth+1), dtype=torch.float, device=x.device)

		for i in range(self.depth+1):
			all_nodes[:, i] = current_nodes
			plane_coeffs = self.w1s.index_select(dim=0, index=current_nodes)			# (batch_size, input_width)
			plane_coeff_score = torch.bmm(x.unsqueeze(1), plane_coeffs.unsqueeze(-1))	# (batch_size, 1, 1)
			plane_score = plane_coeff_score.squeeze(-1).squeeze(-1) 					# (batch_size,)
			all_logits[:, i] = plane_score
			plane_choices = (plane_score >= 0).long()									# (batch_size,)

			current_nodes = current_nodes * 2 + plane_choices + 1						# (batch_size,)

		# get the weights
		selected_w2s = self.w2s.index_select(0, index=all_nodes.flatten()) \
			.view(batch_size, self.depth+1, self.output_width)	# (batch_size, depth+1, output_width)

		# forward pass
		mlp1 = torch.nn.functional.gelu(all_logits)				# (batch_size, depth+1)
		mlp2 = torch.bmm(mlp1.unsqueeze(1), selected_w2s) 		# (batch_size, output_width)
		
		# done
		return mlp2
	

从代码可以看出，与传统的批矩阵乘法(BMM)不同的是，在forward中，基于二叉树的结构，通过迭代计算节点的索引和权重，使用激活函数(GeLU)对结果进行处理，并最终得到输出。

推理性能

文章比较了几种可用的 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® Core™ 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 实现的实现速度提升。

表中缺少的两个条目是当前不可用的 BLAS level 3 和 native fused FFF 实现。

结论

将传统前馈神经网络定义成一棵二叉树，其叶子是小型神经网络，在每个非叶子节点处都有一个微小的神经网络（单个神经元也可以工作）来决定走哪条路径取决于在输入上。在训练期间，它们对所选路径进行加权平均值，从而得出树的所有叶子（在输入上评估为神经网络）的总加权平均值，但在推理过程中，它们可以只遵循投票最高的分支，从而得出建议的结果指数加速。并且，基于FFF的思想，将工作推到BERT这种语言模型上，初步证明了大模型的前馈层的神经元并不是都需要参与推理。