二叉树遇见语言模型：浅尝二叉树在前馈神经网络上的应用

随着大模型的发展，模型参数量暴涨，以Transformer的为组成成分的隐藏神经元数量增长的越来越多。因此，降低前馈层的推理成本逐渐进入视野。

前段时间看到本文介绍的相关工作还是MNIST数据集上的实验，现在这个工作推进到BERT上面来了，再次引起兴趣记录一下。

该工作将前馈神经基于二叉树结构进行改装，加速前向传播的速度，称为：快速前馈网络（FFF），然后应用FFF，取代BERT中的前馈网络（FF），实现12个神经元加速推理。

快速前馈网络算法概述

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

• 节点网络集合  包含了一组节点网络，每个节点网络都是一个 -前馈网络，并在输出上增加了一个 sigmoid 激活函数。这些节点网络按照平衡的可微分二叉树的形式排列，其中  和  是  的子节点。

• 叶子网络集合  包含了一组叶子网络，每个叶子网络都是一个 -前馈网络。叶子网络没有子节点，它们的输出直接作为 FFF 的输出。

前向传播过程由下面算法控制。

算法的输入包括一个输入样本  和根节点 ，输出为该样本在 FFF 中的输出。

算法定义了两个函数： 和 。其中， 函数用于计算节点的输出，而  函数用于计算节点的指示值（indicator value）。

• 在  函数中，如果当前节点是叶子节点，则直接调用该节点的前馈传播函数  来计算输出。否则，首先计算当前节点的输出 ，然后递归地调用  函数来计算当前节点的两个子节点的输出，并将它们加权相加作为当前节点的输出。
• 在  函数中，如果当前节点是叶子节点，则直接调用该节点的前馈传播函数  来计算输出。否则，首先计算当前节点的输出 ，然后根据输出值的大小决定选择哪个子节点进行递归计算。

传统前馈神经网络快速前馈神经网络

与传统的前馈神经网络算法相比，该算法的主要区别在于引入了一个计算节点的指示值。指示值表示了当前节点的输出是否大于等于阈值（这里的阈值为0.5），根据指示值的大小来确定选择哪个子节点进行计算。这种方式可以大大减少计算量，提高前向传播的效率。同时，FFF 是一种具有平衡二叉树结构的前馈神经网络，其中节点网络和叶子网络分别用于处理中间层和输出层的计算。通过利用二叉树结构和递归计算，FFF 可以实现快速的前向传播。

UltraFastBERT

UltraFastBERT，一种BERT变体，在推理过程中使用0.3％的神经元，同时表现 与类似的BERT模型相当。UltraFastBERT选择性地使用4095个神经元中的12个（有选择的执行矩阵乘法（CMM））进行每层推理。这是通过用快速前馈网络(FFFs)取代前馈网络来实现的。

FFF_BMM代码

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)对结果进行处理，并最终得到输出。

结果

在推理过程中仅使用0.3％的神经元，同时表现与类似的BERT模型相当（下游任务没有降很多点）；实现78倍CPU加速，实现40倍PyTorch加速。

总结

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

文章及公开的代码还介绍了条件矩阵乘法的详细细节，因此感兴趣可以深入研究一下。

参考文献

【1】paper：Exponentially Faster Language Modelling，https://arxiv.org/abs/2311.10770 

【2】code：https://github.com/pbelcak/fastfeedforward

【3】paper：Fast Feedforward Networks，https://arxiv.org/abs/2308.14711

【4】code：https://github.com/pbelcak/UltraFastBERT

【5】model：https://huggingface.co/pbelcak/UltraFastBERT-1×11-long

关于我们

老刘，刘焕勇，NLP开源爱好者与践行者，主页：https://liuhuanyong.github.io。

老刘说NLP，将定期发布语言资源、工程实践、技术总结等内容，欢迎关注。

对于想加入更优质的知识图谱、事件图谱、大模型AIGC实践、相关分享的，可关注公众号，在后台菜单栏中点击会员社区->会员入群加入。

​​​