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

442次阅读
没有评论

随着大模型的发展,模型参数量暴涨,以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实践、相关分享的,可关注公众号,在后台菜单栏中点击会员社区->会员入群加入。

​​​

 

Read More 

正文完
可以使用微信扫码关注公众号(ID:xzluomor)
post-qrcode
 
评论(没有评论)
Generated by Feedzy