用PyTorch 从零开始构建 BitNet 1.58bit

我们手动实现BitNet的编写，并进行的一系列小实验证实，看看1.58bit 模型是否与全精度的大型语言模型相媲美！

什么是量化以及为什么需要它？

量化是用更少的比特数表示浮点数的过程。当两个数字使用不同的比特数进行量化时，浮点运算的计算成本几乎按照减少的比特数的比例降低（理论上）。这使我们能够提高速度并减少机器学习模型的内存消耗。但这通常会导致信息丢失，从而降低准确性，我们可以通过对量化模型进行更多的微调来一定程度上恢复这种损失。

现有的量化方法与 BitNet 1.58bit 对比

大多数量化算法都需要一个全精度的预训练模型。人们通常会应用后训练量化（PTQ）和量化感知训练（QAT）等技术，以使这些算法有效运行。

PTQ 是一种量化技术，模型在训练完成后进行量化。QAT 是对 PTQ 模型的进一步微调，即在考虑量化的情况下进一步训练模型。

而BitNet 采用了一种截然不同的方法，即从头开始训练模型时就进行量化！

BitNet 的量化算法

上图中，通过取绝对值的平均值的一半（假设 n=2）来计算权重裁剪阈值 γ。然后，权重矩阵 W 被相同的值除，导致新的权重矩阵在原始权重值 ≥ γ 时的值 ≥ 1，原始权重值 ≤ -γ 时的值 ≤ -1。对于 -γ 和 γ 之间的值，它们被映射到 -0.99999... 到 0.9999...

当执行 roundclip 时，

对于原始值 ≥ γ，新值为 1.0，原始值 ≤ -γ，新值为 -1.0，原始值在 -γ 和 γ 之间的新值为 0.0。

理论上，结果值可以用信息编码理论表示为 1.58 位。由于位数不能是分数，我们可以用 2 位来表示。

量化函数在Pytorch中的实现

阈值计算:

def compute_adjustment_factor(self, input_tensor: torch.Tensor):

  absmean_weight = torch.mean(torch.abs(input_tensor))

  adjustment_factor = 1e-4 + absmean_weight * 2 # 1e-4 to avoid zero divison error

  return adjustment_factor

这里没有把绝对值减半，而是乘以了2。但是实验还是成功了!

RoundClip (1.58~= 2bit)

def compute_2bit_quantized_tensor(self, input_tensor: torch.Tensor):

  twobit_matrix = torch.clip(input=torch.round(input_tensor), min=-1, max=1)

  return twobit_matrix

def compute_1bit_quantized_tensor(self, input_tensor: torch.Tensor):

  return torch.sign(input_tensor)

def compute_quantized_tensor(self, input_tensor: torch.Tensor):

  if self.quantization_mode == QuantizationMode.two_bit:

      return self.compute_2bit_quantized_tensor(input_tensor)

  else:

      return self.compute_1bit_quantized_tensor(input_tensor)

量化步骤

weight_adjustment_factor = self.compute_adjustment_factor(self.weight)

adjusted_weight = self.weight / weight_adjustment_factor

quantized_weight = self.compute_quantized_tensor(adjusted_weight)

线性层操作

F.linear(weight_adjustment_factor * x, quantized_weight, self.bias)

将调整因子与输入相乘，并将其除以量化权重

如果在将权重传递给线性层函数之前对其进行量化，则对量化矩阵的更新不会通过量化函数(因为大多数更新将在1e-4到1e-2之间，当通过量化步骤反向传播时将变为零)。因为原始的权重矩阵永远不会更新，模型永远不会学习!!

但有一个巧妙的工程技巧可以做到这一点，完整的前向传播是这样的

def forward(self, x):

  weight_adjustment_factor = self.compute_adjustment_factor(self.weight)

  adjusted_weight = self.weight / weight_adjustment_factor

  if self.training:

      quantized_weight = (

          adjusted_weight

          + (

              self.compute_quantized_tensor(adjusted_weight) - adjusted_weight

          ).detach()

      )

  else:

      quantized_weight = self.compute_quantized_tensor(adjusted_weight)

  return F.linear(weight_adjustment_factor * x, quantized_weight, self.bias)

量化权重块的值无论 self.training 是否设置为 True 都是相同的。但是当 self.training 设置为 True 时，计算得到的梯度会被优雅地复制到调整后的权重中。这允许在训练过程中更新调整后的权重，同时也更新原始的权重矩阵。

这是从谷歌 DeepMind 的 VQ VAE PyTorch 实现中借鉴的简单却实用的技巧

自定义Pytorch实现的实验结果

下面的实验选择了一个小型模型和一个相对于小型模型来假设足够大的数据集。此外，为了创建目标模型的量化变体，我简单地使用以下代码块，将 nn.Linear 模块替换为这个自定义实现：

结果如下：

4层FFN的Mnist结果 ：

128维6层VIT版本训练Fashion MNIST的结果

128维8层VIT在 CIFAR100上的结果

我们可以看到，除了第一个实验外，2位和1位版本的模型与全精度的常规版本的模型表现得一样好。在第一个实验中，量化模型可能发生了灾难性遗忘。

这些实验并未使用大型语言模型（LLMs）进行，但足以证明论文关于这样的系统能与全精度模型竞争的说法。

我们的实验与论文的唯一一个区别是，这个实现并没有将量化权重存储在2位矩阵中，计算仍以fp32执行的，要真正看到计算速度的提升，需要为此专门的计算内核，我们目前没有能力编写，所以实现仅验证了论文的潜在的论点。

以上实验的所有代码和模块代码都可以在github repo中找到

作者:Chidhambararajan R

MORE

kaggle比赛交流和组队

喜欢就关注一下吧：

点个在看你最好看！