自回归训练方式已经成为了大语言模型（LLMs）训练的标准模式， 今天介绍一篇来自阿联酋世界第一所人工智能大学MBZUAI的VILA实验室和CMU计算机系合作的论文，题为《FBI-LLM: Scaling Up Fully Binarized LLMs from Scratch via Autoregressive Distillation》，该论文首次提出了采用自回归蒸馏的优化方式从头训练二值化的大语言模型，性能可以匹配或者接近FP16或者BF16训练的LLMs，同时效果远超之前所有二值化大语言模型将近十个点。目前该工作的训练代码，数据和模型权重已全部开源。
- 论文：https://arxiv.org/abs/2407.07093
- 代码：https://github.com/LiqunMa/FBI-LLM
核心结论和贡献
- 相比之前的二值化大语言模型，这是第一个从头开始训练，不使用任何预训练参数的二值化大语言模型。
- 训练过程仅仅使用自回归蒸馏损失，没有加入其他损失函数。
- 该工作是一个全量二值化模型，而不是之前一些方法采用的局部二值化或者三值化大模型。
背景介绍
最近几年受益于巨大的参数规模和海量的训练语料，基于Transformer的大型语言模型（LLMs），如ChatGPT和LLaMA系列，在特定领域知识的生成和复杂推理任务中都表现出色的性能。
此外，LLMs的能力随着参数规模的扩大而继续增强，给人们在通往AGI的道路上以无限遐想。然而，巨大的参数规模导致了模型需要巨大的存储和计算需求，这大大限制了LLMs的广泛应用和发展。量化技术通过将32位参数映射到更小的位数，有效地缓解了这些限制，该技术可以显著减少存储需求，并在推理过程中提升了计算速度和能源效率。
作为量化的极端情况，模型二值化仅用{-1, 1}来表示每个参数。它最大限度地实现了压缩和推理效率，但代价是牺牲一定程度的准确性。以往维持二值化LLMs性能的研究包括如何保留模型中重要参数或使用接近一位（部分二值化或者三值化）的表达方式来表示每个参数。
虽然这些方法展现出了不错的性能和潜力，但它们在存储和效率方面仍有优化的空间，并且额外的全精度参数或采用非2的幂来表示参数编码在适配特定硬件时会产生额外的开销。某些全二值化LLMs的研究基于最小化层级损失的优化目标，或使用预训练的全精度LLM在其基础上继续训练，然后用少量训练数据进行二值化模型参数纠正，这些方法面临如下几个问题：
- 之前某些方法借助预训练的全精度模型参数来减少训练计算量和优化步骤， 然而二值化过程会极大地压缩原始模型的参数空间，损害全精度模型中存储的知识，因此依然需要足够的训练数据来让二值化模型重新学习这些知识并适应二值化参数的模式；
- 从现有预训练模型中衍生二值化模型 (使用预训练权重) 的方案不允许选择不同的参数规模或词汇表大小，从而限制了模型结构的灵活性和实际应用。
本文作者提出了一种从头开始训练的全二值化LLMs（FBI-LLM）。为了实现从头开始稳定地训练二值化LLMs，文章提出了一种基于全精度教师模型的自回归蒸馏的新型训练模式。具体来说，在训练过程中，作者逐步从全精度教师模型中生成蒸馏使用的软标签，并采用基于自回归蒸馏的方案来匹配教师模型在每个token位置的预测概率。
通过这种简单直接的自回归蒸馏损失，可以成功地从随机初始化中训练二值化LLMs。由于该方法相比一般LLM训练的改动主要集中在损失函数上，FBI-LLM可以轻松地融入现有的LLM预训练过程。此外，这种方法中的二值化操作与模型训练是分离的，因此任何增强LLM训练效率或者性能的技术都可以直接应用于本文提出的FBI-LLM。
作者对框架FBI-LLM的有效性进行了详细评估，训练了从130M、1.3B到7B规模的不同模型。作者使用广泛使用的Transformer架构进行LLMs的训练，结果表明从头开始训练全二值化的LLMs是完全可行的，其性能与全精度模型相比只有很小的差距。
与其他基准线方法相比，训练过程在困惑度和多个下游任务上表现更为出色。这些结果表明自回归蒸馏是训练二值化LLMs的关键，此外，通过对预训练更加深入的研究（如权重翻转比和梯度范数）的分析表明，从全精度LLMs继承权重与从头开始训练二值化LLMs之间没有显著差异。
本文的贡献可以总结如下：首先，该论文首次证明可以成功地从头开始训练具有二值权重的LLMs；其次，本文提出了一种新的蒸馏损失函数，以稳定二值化LLMs的训练，其采用自回归蒸馏来匹配教师模型的概率分布；第三，本文进行了广泛的实验和分析，以更好地理解所提出的方法的有效性。
下面介绍一下文章具体细节。
模型结构
模型结构主要基于LLaMA的结构。如上左图，首先是LLM 模块，其中包含使用可学习的α和β组成的 FBI-Linear层。右图为自回归蒸馏和模型训练相关过程。具体而言，由于在LLM 中，大多数参数都位于线性模块中。FBI-LM 将除causal head以外的所有线性模块替换为 FBI-linear层。由于causal head直接影响每个步骤中的输出token分布，因此对其参数进行二值化将显著影响模型输出的准确性，因此本文选择保留其精度。
此外，LLM 的另外两个核心模块（embedding和Layer Norm）中的参数也需要保持全精确。这是因为embedding模块包含有关所有标记的语义信息，并且作为模型输入的第一层，需要用来确定文本的初始表示形式。另一方面，Layer Norm 直接缩放激活值，二值化其参数将显著降低每层激活值的语义表达能力，之前其他关于LLM二值化的工作和研究也选择采用类似的设置和做法。
模型训练：自回归蒸馏（Autoregressive Distillation）
给定一个训练语料，标准的自回归语言模型的目标函数是最大化如下似然函数：
其中k表示上下文窗口的大小，条件概率p通过参数为θ的神经网络建模。不同于一般的自回归语言模型，本文使用自回归蒸馏训练 FBI-LLM。在训练过程中，一个全精度预训练 LLM 被用作教师模型，二值化目标模型作为学生模型。假设每个训练数据实例由输入token序列x^1,…x^m组成，教师模型对下一个标记的预测概率可以表示为：
其中表示最后一层 transformer 模块的激活，表示用于预测下一个token概率的线性输出层的参数。
学生模型与教师模型输出之间的交叉熵被计算为每一步预测下一个token时的最终损失函数。它可以表示为：
其中n表示输入标记的数量。表示教师模型预测的第i步词汇表上的标记分布，而是学生模型的相应预测分布。
训练数据
本文使用的训练数据集跟一般的LLM训练相似，包含 Refined-Web 、StarCoder 和 RedPajama-v1的混合数据集，总共包含 1.26T tokens。
实验结果
如下图所示，首先是对现有的二值化 LLM 和 FBI-LLM 在 Wikitext2 的困惑度 (Perplexity) 比较。与其他二值化 LLM 相比，FBI-LLM 在相同规模大小的模型上获得相似或更低的困惑程度。
其次是在下游任务上的性能表现，如下表所示，由于 130M 大小的FBI-LLM没有对应的基准模型，本文将 130M 模型与之前 700M 规模的 BitNetb1.58 进行比较。尽管模型规模相差五倍，权重量化程度也存在较大的差异，但FBI的模型在 BoolQA 和 OpenbookQA 上的表现仍然优于 BitNet b1.58。
对于 1.3B 规模的二值化模型，FBI-LLM 在大多数下游任务和困惑度中都取得了最佳性能，甚至接近或超过了某些 7B 规模的二值化模型（如 BiLLM-LLaMA2-7B）的性能。与相同规模的全精度模型相比， FBI-LLM 1.3B 在下游任务中可以达到其 87% 的性能。在 7B 规模中，FBI模型依然显著优于之前几乎所有的基准线模型，具体来说，FBI-LLM相比之前最好的方法平均提升了将近十个点。
此外，由于计算资源的限制，FBI-LLM 7B当前汇报的结果并不是最终结果。作者只使用了整个 数据集的 8.6%（31 个块）。下图展示了FBI-LLM-7B训练过程中下游任务准确率和困惑度的变化。显然，从目前的训练进度来看，FBI-LLM-7B的性能将持续提高，更进一步的训练可能会得到更好的效果。
模型分析和可视化
二值化大模型是 从头开始训练还是从预训练的 LLM 接着训练？
直观地说，从预训练的 LLM 继续训练可以让二值化模型从全精度原始模型中继承知识，从而可能比从头开始训练获得更好的结果。为了论证这一假设，本文进行了全面的消融和分析实验，以记录和比较模型在两种不同训练模式下的行为。
从下图 (a) 中可以观察到，在训练初期，两种训练方式的 FF ratio 趋势基本保持一致。在整个训练过程中，两种方法的 FF ratio 都处于相似的大小上，并且数值相对较小。该结果表明，两种不同的参数初始化方法对二值化优化过程的影响没有显著差异。下图 (b) 展示了两种训练模式下的训练损失变化。在训练的初始阶段，两种方法的训练损失基本相同，表明模型的训练损失不会显著受初始化方法的影响。
虽然从头开始训练的损失在中间阶段比继续训练的损失略高，但过一段时间后，从头开始训练的损失再次与接着训练的损失相当，甚至变得比继续训练的损失更为稳定。值得注意的是，在大约第 1000 步时，如图 (a) 所示，当从预训练的 LLM 继续训练时，FF ratio 开始出现明显的波动。同样，在图 (b) 所示的第 1700 步左右，训练损失也遇到了类似的问题。
这些发现挑战了之前相关工作的一些假设，即从预训练的 LLM权重开始训练二值化大模型将赋予二值化LLM继承之前预训练的知识，从而提高性能。然而，本文结果和分析暗示了二值化LLM对参数初始化的方式不敏感，或者说随机参数初始化反而更好。对于原因，本文推测二值化和全精度LLM采用不同的参数组合和配置来编码语义，这导致其参数空间模式存在实质性差异。为了适应这种模式，通过从预训练的 LLM 继续训练来优化二值化的过程可能需要更大幅度的参数数值调整和变换，从而破坏预训练参数模式结构，继而破坏其中保存的知识。这个理论可以部分地解释为什么在训练期间与从头开始训练相比，继续预训练的权重反而让二值化LLM变得更不稳定。
训练的稳定性分析
二值化和全精度LLM训练在之前一些工作中都被发现表现出不稳定的训练行为。 FBI-LLM 也有类似的问题，具体表现为在训练 1.3B 和 7B FBI-LLM 时训练损失突然激增，有时在此之后继续训练模型也无法收敛。本文采用类似于 PaLM的解决方案：如果损失值不再趋于收敛，模型将恢复到前一个检查点，并跳过触发不稳定损失的数据块以继续训练。使用此方法，模型在相同的训练步骤中不再遇到问题。根据该论文的观察，从头开始训练 7B FBI 模型大约有 6% 的概率导致出现损失峰值。对于 1.3B 模型，由于其模型能力较低，训练更加不稳定，损失峰值的出现概率约为 15%。这与在全精度 LLM 中看到的预训练行为相似，而尖峰的概率明显更高，这可能与二值化参数的有限表达能力有关。为了解决这个问题，FBI跳过了发生损失峰值且没法通过后续训练恢复的数据块。
存储效率分析
上表显示了不同规模的FBI-LLM与相同结构的全精度LLaMA相比所需的理论存储空间，同时还详细给出了 FBI-LLM 引入的附加参数（α 和 β）的比例。表中的比较表明，FBI-LLM可以实现高压缩比，大大减轻了LLMs的存储负担。尽管 FBI-LLM 引入的用于缩放和移动的额外参数需要保留全精确，但它们的比例相比整个模型很小，因此其对存储的影响可以忽略不计。
生成结果示例
如上图所示，尽管 FBI-LLM 的生成质量无法完全匹配全精度 LLM模型，但 FBI-LLM 仍然可以生成流畅且有意义的内容。与具有更高参数位宽的BitNet b1.58模型相比，FBI-LLM对提示的理解更好，并且在一些生成的示例中包含了更多的知识。这表明FBI-LLMs具有更强的生成能力，并包含足够的知识。此外，FBI-LLM展示了进一步扩大模型规模从而达到更高智力水平的二值化模型的潜力，这种类型的大模型 (二值化大模型) 对部署的硬件需求也更加友好。
模型配置和训练细节
具体的模型配置和训练细节如下所示：
更多细节欢迎阅读文论原文。
