AIxiv专栏是AI在线发布学术、技术内容的栏目。过去数年，AI在线AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]
作者介绍：论文一作Liqun Ma目前是MBZUAI机器学习系的博士生，导师为Zhiqiang Shen助理教授，同时也是该论文的最后作者，其在加入MBZUAI之前为CMU博士后，研究领域主要为机器学习，基础大模型等等。Liqun的主要研究方向为高效的大模型预训练和微调，他本科毕业于天津大学。论文二作Mingjie Sun目前为CMU计算机系在读博士，导师为Zico Kolter教授。
自回归训练方式已经成为了大语言模型（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是完全可行的，其性能与全精度模型相比只有很小的差距。
实验结果
如下图所示，首先是对现有的二值化 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-LLaMA-7B 的中间层输出分布与从全精度模型继续训练二值化模型的中间层输出分布进行了比较。图中展示了第 8 层、16 层、24 层、32 层的特征分布对比。研究发现，尽管模型初始化方式不同，但两个二值化LLMs的特征输出分布在所有层上都非常接近，且从头训练的模型收敛效果更为稳定。
结论
本文工作展示了二值化大语言模型从零开始训练的可行性，提出的自回归蒸馏损失函数显著提升了二值化LLMs的训练稳定性和性能。研究结果表明，从头训练的二值化大语言模型不仅具有优异的性能，还在存储和计算效率上有显著优势。这为未来大规模、高效的二值化语言模型的发展提供了新的方向。
更多细节请阅读论文原文。