自回归训练方式已经成为了大语言模型(LLMs)训练的标准模式, 今天介绍一篇来自阿联酋世界第一所人工智能大学MBZUAI的VILA实验室和CMU计算机系合作的论文,题为《FBI-LLM: Scaling Up Fully Binarized LLMs from Scratch via Autoregressive Distillation》,该论文首次提出了采用自回归蒸馏的优化方式从头训练二值化的大语言模型,性能可以匹配或者接近FP16或者BF16训练的LLMs,同时效果远超之前所有二值化大语言模型将近十个点。目前该工作的训练代码,数据和模型权重已全部开源。
论文标题:
FBI-LLM: Scaling Up Fully Binarized LLMs from Scratch via Autoregressive Distillation
论文链接:
https://arxiv.org/abs/2407.07093
代码链接:
https://github.com/LiqunMa/FBI-LLM
一、核心结论和贡献
相比之前的二值化大语言模型,这是第一个从头开始训练,不使用任何预训练参数的二值化大语言模型。
训练过程仅仅使用自回归蒸馏损失,没有加入其他损失函数。
该工作是一个全量二值化模型,而不是之前一些方法采用的局部二值化或者三值化大模型。
二、背景介绍
最近几年受益于巨大的参数规模和海量的训练语料,基于Transformer的大型语言模型(LLMs),如ChatGPT和LLaMA系列,在特定领域知识的生成和复杂推理任务中都表现出色的性能。
此外,LLMs的能力随着参数规模的扩大而继续增强,给人们在通往AGI的道路上以无限遐想。然而,巨大的参数规模导致了模型需要巨大的存储和计算需求,这大大限制了LLMs的广泛应用和发展。量化技术通过将32位参数映射到更小的位数,有效地缓解了这些限制,该技术可以显著减少存储需求,并在推理过程中提升了计算速度和能源效率。
作为量化的极端情况,模型二值化仅用{-1, 1}来表示每个参数。它最大限度地实现了压缩和推理效率,但代价是牺牲一定程度的准确性。以往维持二值化LLMs性能的研究包括如何保留模型中重要参数或使用接近一位(部分二值化或者三值化)的表达方式来表示每个参数。
之前某些方法借助预训练的全精度模型参数来减少训练计算量和优化步骤, 然而二值化过程会极大地压缩原始模型的参数空间,损害全精度模型中存储的知识,因此依然需要足够的训练数据来让二值化模型重新学习这些知识并适应二值化参数的模式; 从现有预训练模型中衍生二值化模型 (使用预训练权重) 的方案不允许选择不同的参数规模或词汇表大小,从而限制了模型结构的灵活性和实际应用。
三、模型结构
此外,LLM 的另外两个核心模块(embedding和Layer Norm)中的参数也需要保持全精确。这是因为embedding模块包含有关所有标记的语义信息,并且作为模型输入的第一层,需要用来确定文本的初始表示形式。另一方面,Layer Norm 直接缩放激活值,二值化其参数将显著降低每层激活值的语义表达能力,之前其他关于LLM二值化的工作和研究也选择采用类似的设置和做法。
四、模型训练:自回归蒸馏(Autoregressive Distillation)
学生模型与教师模型输出之间的交叉熵被计算为每一步预测下一个token时的最终损失函数。它可以表示为:
五、训练数据
六、实验结果
七、模型分析和可视化
7.1 二值化大模型是 从头开始训练还是从预训练的 LLM 接着训练?
直观地说,从预训练的 LLM 继续训练可以让二值化模型从全精度原始模型中继承知识,从而可能比从头开始训练获得更好的结果。为了论证这一假设,本文进行了全面的消融和分析实验,以记录和比较模型在两种不同训练模式下的行为。
从下图 (a) 中可以观察到,在训练初期,两种训练方式的 FF ratio 趋势基本保持一致。在整个训练过程中,两种方法的 FF ratio 都处于相似的大小上,并且数值相对较小。该结果表明,两种不同的参数初始化方法对二值化优化过程的影响没有显著差异。下图 (b) 展示了两种训练模式下的训练损失变化。在训练的初始阶段,两种方法的训练损失基本相同,表明模型的训练损失不会显著受初始化方法的影响。
虽然从头开始训练的损失在中间阶段比继续训练的损失略高,但过一段时间后,从头开始训练的损失再次与接着训练的损失相当,甚至变得比继续训练的损失更为稳定。值得注意的是,在大约第 1000 步时,如图 (a) 所示,当从预训练的 LLM 继续训练时,FF ratio 开始出现明显的波动。同样,在图 (b) 所示的第 1700 步左右,训练损失也遇到了类似的问题。
这些发现挑战了之前相关工作的一些假设,即从预训练的 LLM权重开始训练二值化大模型将赋予二值化LLM继承之前预训练的知识,从而提高性能。然而,本文结果和分析暗示了二值化LLM对参数初始化的方式不敏感,或者说随机参数初始化反而更好。对于原因,本文推测二值化和全精度LLM采用不同的参数组合和配置来编码语义,这导致其参数空间模式存在实质性差异。
二值化和全精度LLM训练在之前一些工作中都被发现表现出不稳定的训练行为。FBI-LLM 也有类似的问题,具体表现为在训练 1.3B 和 7B FBI-LLM 时训练损失突然激增,有时在此之后继续训练模型也无法收敛。本文采用类似于 PaLM的解决方案:如果损失值不再趋于收敛,模型将恢复到前一个检查点,并跳过触发不稳定损失的数据块以继续训练。使用此方法,模型在相同的训练步骤中不再遇到问题。
根据该论文的观察,从头开始训练 7B FBI 模型大约有 6% 的概率导致出现损失峰值。对于 1.3B 模型,由于其模型能力较低,训练更加不稳定,损失峰值的出现概率约为 15%。这与在全精度 LLM 中看到的预训练行为相似,而尖峰的概率明显更高,这可能与二值化参数的有限表达能力有关。为了解决这个问题,FBI跳过了发生损失峰值且没法通过后续训练恢复的数据块。
7.3 存储效率分析
7.4 生成结果示例
如上图所示,尽管 FBI-LLM 的生成质量无法完全匹配全精度 LLM模型,但 FBI-LLM 仍然可以生成流畅且有意义的内容。与具有更高参数位宽的BitNet b1.58模型相比,FBI-LLM对提示的理解更好,并且在一些生成的示例中包含了更多的知识。这表明FBI-LLMs具有更强的生成能力,并包含足够的知识。此外,FBI-LLM展示了进一步扩大模型规模从而达到更高智力水平的二值化模型的潜力,这种类型的大模型 (二值化大模型) 对部署的硬件需求也更加友好。
7.5 模型配置和训练细节
具体的模型配置和训练细节如下所示:
更多细节欢迎阅读文论原文。
来源:公众号【机器之心】
llustration From IconScout By IconScout Store
“AI技术流”原创投稿计划
TechBeat是由将门创投建立的AI学习社区(www.techbeat.net)。社区上线500+期talk视频,3000+篇技术干货文章,方向覆盖CV/NLP/ML/Robotis等;每月定期举办顶会及其他线上交流活动,不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台,希望为AI人才打造更专业的服务和体验,加速并陪伴其成长。
投稿内容
// 最新技术解读/系统性知识分享 //
// 前沿资讯解说/心得经历讲述 //
投稿须知
稿件需要为原创文章,并标明作者信息。
我们会选择部分在深度技术解析及科研心得方向,对用户启发更大的文章,做原创性内容奖励
投稿方式
发送邮件到
chenhongyuan@thejiangmen.com
或添加工作人员微信(chemn493)投稿,沟通投稿详情;还可以关注“将门创投”公众号,后台回复“投稿”二字,获得投稿说明。