FBI-LLM: Scaling Up Fully Binarized LLMs from Scratch via Autoregressive Distillation

📄 arXiv: 2407.07093v1 📥 PDF

作者: Liqun Ma, Mingjie Sun, Zhiqiang Shen

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-07-09

备注: Github at https://github.com/LiqunMa/FBI-LLM

🔗 代码/项目: GITHUB | HUGGINGFACE

💡 一句话要点

FBI-LLM：通过自回归蒸馏从头训练全二值化大语言模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 二值化神经网络 大语言模型 自回归蒸馏 模型压缩 低精度计算 从头训练 知识迁移

📋 核心要点

  1. 现有大语言模型计算和存储成本高昂，全二值化模型是降低成本的有效途径，但从头训练全二值化LLM极具挑战。
  2. FBI-LLM通过自回归蒸馏损失，在保持模型维度和训练数据量不变的情况下，实现了从头训练全二值化LLM。
  3. 实验结果表明，FBI-LLM在困惑度和特定任务上可与全精度模型媲美，且无需预训练权重即可实现。

📝 摘要（中文）

本文提出了一种全二值化大语言模型（FBI-LLM），首次展示了如何从头训练一个大规模的二值语言模型（并非像BitNet b1.58那样的部分二值或三值LLM），使其性能与基于Transformer的LLM的全精度模型（例如FP16或BF16）相匹配。该方法采用自回归蒸馏（AD）损失，并保持与常规LLM预训练相同的模型维度（130M、1.3B、7B）和训练数据量，同时在困惑度和特定任务有效性方面取得了具有竞争力的结果。有趣的是，通过分析训练轨迹，我们发现从头训练二值化LLM并不需要预训练权重。这项研究鼓励了一种新的计算框架，并可能促进未来为全1比特LLM量身定制的专用硬件的设计。我们完全公开所有模型、代码和训练数据集，以支持进一步的研究。

🔬 方法详解

问题定义：现有的大语言模型通常采用FP16或BF16等高精度格式，导致计算和存储成本巨大。全二值化模型（权重和激活均为1比特）可以显著降低这些成本，但从头训练全二值化LLM非常困难，性能往往远低于全精度模型。现有方法或者采用部分二值化，或者依赖于全精度模型的预训练权重。

核心思路：FBI-LLM的核心思路是利用自回归蒸馏（Autoregressive Distillation, AD）损失，将全精度教师模型的知识迁移到二值化学生模型。通过最小化学生模型和教师模型在自回归生成过程中的输出差异，引导学生模型学习到有效的二值化表示。

技术框架：FBI-LLM的整体框架与标准的Transformer LLM预训练流程类似，包括数据预处理、模型训练和评估。关键在于训练阶段，使用AD损失函数来训练二值化模型。具体流程如下：1) 准备全精度教师模型；2) 初始化二值化学生模型；3) 使用AD损失函数训练学生模型，使其逼近教师模型的输出分布；4) 评估学生模型的性能。

关键创新：FBI-LLM最重要的创新在于成功地从头训练了全二值化LLM，而无需依赖预训练权重。这打破了以往的认知，即二值化模型的训练需要一个良好的初始化点。此外，AD损失函数的引入使得知识迁移更加有效，避免了二值化带来的信息损失。

关键设计：FBI-LLM的关键设计包括：1) 使用与全精度模型相同的模型维度和训练数据量，保证了模型容量和训练充分性；2) 采用自回归蒸馏损失，鼓励学生模型学习教师模型的输出分布；3) 探索了不同的二值化函数和优化器，以提高训练稳定性和收敛速度。AD损失函数的具体形式为交叉熵损失，用于衡量学生模型和教师模型输出概率分布的差异。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

FBI-LLM在不同规模的模型（130M, 1.3B, 7B）上都取得了显著成果，在困惑度和下游任务上与全精度模型（FP16/BF16）的性能差距显著缩小，甚至在某些情况下可以媲美。更重要的是，该研究证明了从头训练全二值化LLM的可行性，无需依赖预训练权重，为二值化LLM的研究开辟了新的方向。

🎯 应用场景

FBI-LLM的潜在应用领域包括资源受限的边缘设备、移动设备和嵌入式系统。全二值化模型可以显著降低计算和存储成本，使得大语言模型能够在这些设备上部署和运行。此外，该研究为未来设计专门针对二值化LLM的硬件加速器提供了新的思路，有望进一步提高计算效率和降低功耗。

📄 摘要（原文）

This work presents a Fully BInarized Large Language Model (FBI-LLM), demonstrating for the first time how to train a large-scale binary language model from scratch (not the partial binary or ternary LLM like BitNet b1.58) to match the performance of its full-precision counterparts (e.g., FP16 or BF16) in transformer-based LLMs. It achieves this by employing an autoregressive distillation (AD) loss with maintaining equivalent model dimensions (130M, 1.3B, 7B) and training data volume as regular LLM pretraining, while delivering competitive results in terms of perplexity and task-specific effectiveness. Intriguingly, by analyzing the training trajectory, we find that the pretrained weight is not necessary for training binarized LLMs from scratch. This research encourages a new computational framework and may facilitate the future design of specialized hardware tailored for fully 1-bit LLMs. We make all models, code, and training dataset fully accessible and transparent to support further research (Code: https://github.com/LiqunMa/FBI-LLM. Model: https://huggingface.co/LiqunMa/).