摘要

本文介绍了一种全新的全二值化大型语言模型（FBI-LLM），首次展示了如何从头开始训练大规模二值化语言模型，以匹配其全精度对应模型的性能。通过采用自回归蒸馏（AD）损失，同时保持与常规LLM预训练相同的模型维度（130M、1.3B、7B）和训练数据量，FBI-LLM在困惑度和任务特定有效性方面取得了竞争性结果。此外，分析训练轨迹发现，预训练权重对于从头开始训练二值化LLM并非必要。这项研究鼓励了一种新的计算框架，并可能促进未来为全1位LLM设计的专用硬件的发展。所有模型、代码和训练数据集均完全开放和透明，以支持进一步研究。

原理

FBI-LLM的核心在于其全二值化策略和自回归蒸馏技术。全二值化意味着模型中的每个参数仅使用{-1, 1}表示，这种极端的量化方式极大地压缩了存储需求并提高了计算效率，尽管牺牲了一定的准确性。自回归蒸馏是一种训练过程，其中使用一个全精度的预训练LLM作为教师模型，二值化目标模型作为学生模型。训练过程中，学生模型通过自回归方式逐步蒸馏教师模型的预测概率，从而在每个token位置匹配教师模型的输出。这种蒸馏损失使得二值化模型能够从随机初始化开始成功训练，而不依赖于预训练权重。此外，由于修改集中在损失函数上，FBI-LLM可以轻松集成到现有的LLM预训练流程中。

流程

FBI-LLM的工作流程包括以下几个关键步骤：

1. 模型结构转换：将标准LLM中的线性模块替换为全二值化线性模块（FBI-Linear），除了因果头之外，因为因果头的参数直接影响输出token分布。
2. 自回归蒸馏训练：使用全精度预训练LLM作为教师模型，二值化模型作为学生模型。在训练过程中，学生模型通过自回归蒸馏损失逐步学习教师模型的预测概率。
3. 训练细节：模型在Amber数据集上进行训练，使用Adam优化器，初始学习率为3e-4，遵循余弦学习率调度。训练过程中使用梯度裁剪和BF16精度。
4. 评估与优化：通过在多个下游任务和困惑度上评估模型性能，调整训练策略和参数，以优化二值化模型的性能。

应用

FBI-LLM的应用前景广泛，特别是在资源受限的环境中。由于其极低的存储和计算需求，FBI-LLM非常适合部署在边缘设备或低功耗系统中。此外，全二值化模型可能促进专用硬件的发展，这些硬件可以更有效地处理二值化数据，从而进一步提高能效和速度。尽管存在性能损失和训练成本增加的限制，FBI-LLM在简化模型部署和降低能耗方面的潜力巨大，特别是在需要大规模语言模型但资源有限的应用场景中。