1.58bit 量化,内存仅需 1/10,但表现不输 FP16?
微软最新推出的蒸馏框架 BitNet Distillation(简称 BitDistill),实现了几乎无性能损失的模型量化。
该框架在 4B 及以下的 Qwen、Gemma 上已被证实有效,理论上可用于其他 Transformer 模型。
BitDistill 包含三个依次衔接的阶段,分别是模型结构优化(Modeling Refinement)、继续预训练(Continue Pre-training)和蒸馏式微调(Distillation-based Fine-tuning)。
建模结构优化的主要目标是为 1.58-bit 模型训练提供结构层面的支持,缓解低精度训练中常见的优化不稳定问题。
在传统的全精度 Transformer 模型中,隐藏状态的方差通常在预训练时已被良好控制。然而,当模型被压缩到极低位宽(如 1.58-bit)后,激活值在经过量化前的分布可能会出现方差膨胀等问题,从而导致训练过程震荡甚至失败。
为了应对这一问题,BitDistill 在每一个 Transformer 层中引入了一个名为SubLN(Sub-layer LayerNorm)的归一化模块。
具体来说,SubLN 的插入位置有两个,一是在多头自注意力模块的输出投影之前,二是在前馈网络的输出投影之前。
这样的插入方式,不改变主干计算路径,仅在关键位置对信号做规范化调整,使得量化后模型具备更好的收敛性。
这种设计使得量化前的表示能够在进入下一计算阶段前被重新归一化,有效抑制激活尺度的发散,提升训练稳定性。
也就是说,随着模型参数增大,1.58-bit 模型与其全精度版本之间的性能差距反而进一步扩大。
为了缓解这一问题,BitDistill 设计了一个轻量级的继续预训练阶段。在此阶段中,模型会在少量通用语料上进行自回归语言建模训练,训练目标为最大化条件概率。
这一过程并不涉及特定任务数据,也不需精调标签,仅是让模型权重从全精度空间缓慢迁移到适合 1.58-bit 表示的分布上。
换句话说,这个阶段的本质是一种预适配训练,让模型 " 学会如何被量化 ",避免在微调阶段才仓促适应低位宽带来的信息丢失。
为了弥补量化后模型在表达能力上的损失,BitDistill 采用了一种双重蒸馏机制—— Logits 蒸馏与多头注意力蒸馏。
这一阶段的目的是从原始的全精度模型中提取关键行为模式,并引导低位宽模型在具体任务上学习这些模式,从而恢复性能。
Logits 蒸馏是将全精度模型输出的类概率分布作为 " 软标签 ",引导量化模型在预测分布上向其靠拢。具体做法是使用 Kullback – Leibler 散度(KL 散度)来最小化两者输出分布之间的差异。
具体而言,对于选定的某一层(通常是模型后部的一层),分别从教师模型与学生模型中提取 Q、K、V 三组张量,并计算它们之间的点积相关性,形成关系分布矩阵。
然后通过 KL 散度使两者对齐,训练学生模型还原出与教师模型相似的结构依赖。
BitDistill 展示出在多个下游任务中几乎等同于全精度模型的表现,同时显著降低了内存开销并提升了推理速度。作者在两个典型任务类型上进行了全面实验,分别是文本分类与文本摘要。
以 Qwen3 为基础模型,测试中的分类任务包括 MNLI、QNLI 与 SST-2,摘要任务则采用 CNN/DailyMail 数据集作为标准。
分类任务中,BitDistill 的 1.58-bit 模型在准确率与生成质量指标上与全精度微调模型(FP16-SFT)几乎一致,而显著优于直接对量化模型进行微调的 BitNet-SFT。
以 ROUGE 和 BLEU 等标准指标衡量,在 CNN/DailyMail 上,BitDistill 所生成文本的 BLEU 为 14.41,ROUGE-L 为 27.49,与 FP16 模型的 13.98 和 27.72 几乎等同,甚至在 BLEU 上略有超出。
相比之下,直接量化后的模型在 BLEU 与 ROUGE 上普遍下降 2 至 3 个百分点。
进一步的实验表明,BitDistill 在不同量化策略下也具备良好的兼容性。作者将其与常见的 Block-Quant、GPTQ、AWQ 等量化方法结合,在分类任务上依然能够稳定地恢复原始性能,证明该方法可作为一个独立于量化算法的上层蒸馏方案,适用于多种后量化优化场景。
One More Thing
BitStill 的作者全部来自微软研究院,而且均为华人。
通讯作者为微软亚洲研究院副总裁、武汉大学校友韦福如博士。
他读博期间就曾在 MSRA 实习,毕业后到 IBM 工作,又于 2010 年重回微软,工作至今。
一键三连「点赞」「转发」「小心心」
欢迎在评论区留下你的想法!
— 完 —
年度科技风向标「2025 人工智能年度榜单」评选报名开启啦!我们正在寻找 AI+ 时代领航者 点击了解详情
企业、产品、人物 3 大维度,共设立了 5 类奖项,欢迎企业报名参与
一键关注 点亮星标
科技前沿进展每日见