QAT与BitNet：1-bit及1.58-bit量化训练解析

量化感知训练（QAT）让模型在训练时就「知道」自己将被量化，从而找到对舍入误差更鲁棒的权重。本文深入解析QAT的核心机制，并延伸至BitNet的1-bit与1.58-bit三值权重方案，揭示极低比特量化如何在不牺牲过多质量的前提下大幅压缩模型。

QAT：把量化放进训练循环

从PTQ到QAT

上一篇文章讨论了训练后量化（PTQ）：权重已经固定，我们只需找到最佳的整数网格来压缩它们。但如果我们让模型在训练时就意识到量化呢？这就是量化感知训练（QAT）。QAT的核心思想很简单：在前向传播中对权重进行量化，用量化后的值计算损失，然后优化器调整权重，使模型在「尽管有舍入」的情况下依然表现良好。

为什么QAT有效：损失曲面的几何学

标准训练找到一个损失最低的点，然后就不再关心该点周围损失曲面的形状。但量化打破了这个假设：每个权重现在会落在训练结束时位置附近的一个小误差范围δ内。真正重要的不仅是损失在最小值处有多低，还有损失随着远离最小值而上升的速度。

在最小值附近斜率（一阶导数）为零，损失增加量近似为 ΔL ≈ ½ · H · δ²，其中H是曲率（二阶导数）。两个损失相同的最小值可能行为迥异：

尖锐最小值（大H）：小的舍入误差会导致大幅损失上升。
平坦最小值（小H）：同样的舍入误差几乎无关紧要。

因此QAT同时优化两件事：低损失，以及围绕低损失点的平缓曲率，使量化误差保持在可控范围内。原文提供了交互式页面，可以拖动量化误差滑块，观察尖锐最小值上损失上升速度是平坦最小值的8倍。

曲率与泛化：损失曲面的曲率（Hessian矩阵的特征值）不仅决定量化鲁棒性，也与模型的泛化能力密切相关。尖锐最小值通常意味着模型对输入扰动和参数扰动都更敏感，而平坦最小值则对应更宽的最优解区域。QAT通过显式引入量化噪声，实际上是在训练过程中进行一种正则化，迫使模型收敛到更平坦的区域。这与Sharpness-Aware Minimization（SAM）等方法的思路异曲同工，只不过QAT的扰动是结构化的、由量化网格决定的，而非随机噪声。

梯度如何流动：伪量化与STE

舍入操作几乎处处梯度为零，反向传播将学不到任何东西。解决方案是直通估计器（STE）：前向传播使用量化值，反向传播则假装舍入是恒等函数，将梯度直接传过去。

import torch

def fake_quant(w, bits=4):
    qmax = 2 ** (bits - 1) - 1
    scale = w.abs().max() / qmax
    w_q = torch.round(w / scale).clamp(-qmax, qmax) * scale
    # forward: w_q. backward: gradient of w (identity)
    return w + (w_q - w).detach()

模型用舍入后的权重计算，但优化器更新的是底层的全精度权重。这个思路正是BitNet的基石。实践中QAT通常优于PTQ，在4-bit等低比特宽度下尤为明显。

STE的数学本质与变体：直通估计器（STE）本质上是用恒等函数的梯度替代舍入函数的梯度，即∂L/∂w ≈ ∂L/∂w_q。这种近似在数学上并不精确，但实践中效果显著。研究者提出了多种改进变体，例如在|w|接近量化边界时衰减梯度、使用软舍入（soft rounding）或可学习的阶梯函数来缓解梯度失配。STE的成功表明，深度学习优化对梯度的精确性有一定容忍度，只要梯度方向大致正确即可。

BitNet与1-bit权重

线性层在Transformer中的位置

线性层是神经网络最基本的构建块：输入向量乘以权重矩阵得到输出向量。模型中几乎全部参数都是这样的权重矩阵元素。在Transformer中，每个块包含两部分：

注意力机制有四个线性层：Q、K、V投影以及输出投影。中间的softmax(QKᵀ)V是激活值相乘，没有存储权重，因此不是线性层。
前馈网络有两个线性层：上投影和下投影，中间夹着GELU非线性。

BitNet正是瞄准这些线性层：压缩它们，就压缩了整个模型。

BitLinear改变了什么？

BitNet用BitLinear替换大部分Linear。任务相同（output = input × weights），但输入的数字小得多：

	标准Linear	BitLinear
权重	FP16（每权重16位）	±1（每权重1位）
激活值	FP16	INT8
矩阵乘法	FP16乘加	INT8输入对±1权重

乘以+1或-1只是保留数字或翻转符号，矩阵乘法坍缩为加减法。乘法是芯片最昂贵的操作之一，因此这既节省时间与能耗，也节省内存。

但两侧量化会丢失尺度信息：±1权重不再说明原始大小。因此BitLinear保留两个缩放因子，分别用于权重和激活值，最后将结果拉伸回正确尺寸。

BitLinear的配方

量化输入：从FP16到INT8，使用absmax缩放（在LayerNorm之后，激活值有稳定尺度）。
二值化权重：取权重均值α，高于α变为+1，低于变为-1（先中心化使+1和-1大致平衡）。
矩阵乘法：整数输入与±1权重相乘，主要是加法，远比乘法便宜。
反量化输出：使用激活缩放因子和权重缩放因子β（权重绝对值的均值），将结果恢复到正确量级。

这是两个量化器的叠加：一个作用于激活值，一个作用于权重。

为什么不能只存储1比特？

如果训练只是在学习哪些权重「开」哪些「关」，为什么还要保留比1比特更大的东西？

因为梯度。二值化函数是sign()，其梯度处处为零，在零点处未定义。如果权重真的只是比特，梯度下降将没有东西可更新。

因此模型保留高精度的潜在权重（论文用FP16/BF16训练）。训练循环如下：

保持全精度潜在权重。
前向传播：用sign二值化，使用±1版本。
反向传播：使用STE，让梯度落到潜在权重上。
更新潜在权重。

潜在权重不能用INT8：学习率很小，INT8上的微小梯度步会舍入回同一个整数，更新直接消失。潜在权重需要足够的分辨率来累积许多微小推动，直到某个时刻翻转符号。推理时只需要±1版本，潜在权重是训练时的脚手架。

潜在权重的分辨率需求：潜在权重需要保留足够精度，核心原因是梯度更新的幅度通常远小于量化步长。以学习率1e-4、权重尺度0.01为例，单步更新量约为1e-6，而FP16的最小正规格数约为6e-8，尚能累积；若用INT8（步长约4e-3），则每次更新都被舍入抹零。因此BitNet论文选择用BF16/FP16维护潜在权重，确保微小更新能逐步累积，最终跨越符号翻转的阈值。这也解释了为什么QAT训练成本与全精度训练相当——计算图仍保留高精度权重。

「1-bit是骗局吗？」

有人怀疑BitNet随着模型增大越来越接近FP16质量的说法：这不意味着你只需要更大的模型，节省就消失了吗？

并非如此。节省是按参数计算的。1-bit模型可能需要更多参数来匹配质量，但每个参数比FP16小最多16倍。只有当需要超过16倍的参数时才会失败。论文的趋势是质量差距随模型规模缩小。

1.58比特：三值权重的崛起

后续论文《The Era of 1-bit LLMs (BitNet b1.58)》发现纯1-bit需要很大规模才能匹配全精度质量。作者增加了一个状态：零。权重变为-1、0或+1。

为什么零有帮助

零让权重可以说「这个连接不重要」，而不是被迫推高或拉低输出。权重矩阵在几乎不增加成本的情况下变得更具表达力。

三值权重的配方

计算γ，权重绝对值的均值。
每个权重除以γ，四舍五入到-1、0、+1中最接近的值。

等价地：|w| <= γ/2变为0；w > γ/2变为+1；w < -γ/2变为-1。这就是「absmean」量化。激活值切换为对称absmax量化（围绕零缩放），而非之前的0到最大值范围。训练设置（包括潜在权重和STE）保持不变。

def ternary_weights(w, eps=1e-5):
    gamma = w.abs().mean()
    return (w / (gamma + eps)).round().clamp(-1, 1), gamma

五权重示例

取w = [0.9, -0.3, 0.05, -1.2, 0.4]。1-bit版本（β=0.57）将微小权重0.05强制变为+0.57，误差0.52；三值版本（γ=0.57）将其变为0，误差仅0.05。其余四个权重不变。

为什么叫「1.58」比特？

三种状态需要log2(3) ≈ 1.58比特的信息。但硬件无法存储1.58比特，实际实现会打包数值。最简单的打包每个权重用2比特；更紧凑的方案将5个三值值放入一个字节（3⁵=243 ≤ 256），即每权重1.6比特。因此1.58是信息论数字，实际存储略高于此。相比FP16，每参数大约小10倍。

核心要点

PTQ对完成模型进行舍入；QAT在训练循环中包含舍入，让模型找到能承受舍入的权重。
其原理是几何学：偏好平坦最小值（低曲率），使量化误差代价小。
舍入没有梯度，因此使用STE并保留高精度潜在权重来累积微小更新。
BitNet用sign()实现这一点；b1.58增加零状态，使用absmean缩放，获得更好的质量/大小权衡。