量化感知训练 (QAT) 在训练或微调过程中模拟量化效果。这种模拟通过在模型的计算图中插入操作(常被称为“伪量化”节点)来实现。这些节点接收高精度输入(如 FP32 激活值或权重),并生成低精度输出(例如模拟 INT8 或 INT4),这些输出随后用于后续操作。然而,这给依赖梯度下降的训练算法带来了重要难题。伪量化中的核心操作是舍入函数,它将连续的输入范围映射到离散的输出值集合。从数学上看,这个函数是一个阶梯函数。考虑一个简单的舍入函数 Round(x)。它的导数在几乎所有地方(平坦的阶梯上)都为零,并在值跳变点处未定义。标准的反向传播依赖于计算梯度来更新模型权重。如果梯度为零或未定义,更新就无法通过量化操作反向传播到之前的层和权重。这会有效阻碍位于量化节点之前的参数的学习过程。如果梯度被阻碍,模型如何学习适应量化呢?这时,直通估计器 (STE) 便发挥作用了。问题所在:不可微分的量化我们来表示量化函数(包括缩放、舍入以及反量化回模拟低精度步骤的浮点数)为 $y = q(x)$。在训练的前向传播中,我们从输入 $x$ 计算 $y$,并在后续计算中使用 $y$。在反向传播过程中,我们需要计算损失 $L$ 对输入 $x$ 的梯度,表示为 $\frac{\partial L}{\partial x}$。使用链式法则,这通常会这样计算:$$ \frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x} $$问题出在项 $\frac{\partial y}{\partial x} = \frac{\partial q(x)}{\partial x}$。如前所述,这个导数有问题(大部分为 0,在跳变点未定义)。如果我们使用这个真实梯度,$\frac{\partial L}{\partial x}$ 在几乎所有地方都变为零,阻止了权重更新。解决方案:直通估计器 (STE)直通估计器 (STE) 为这个问题提供了一个实际可行的变通方法。它是在反向传播过程中专门使用的近似。核心思想很简单:前向传播: 照常计算量化操作:$y = q(x)$。量化效果(精度损失)被完全应用。反向传播: 在计算梯度 $\frac{\partial L}{\partial x}$ 时,忽略 $q(x)$ 的真实导数,转而进行近似。最常见的 STE 近似是将梯度 $\frac{\partial y}{\partial x}$ 设为 1。因此,在使用 STE 的反向传播中,梯度计算变为:$$ \frac{\partial L}{\partial x} \approx \frac{\partial L}{\partial y} \cdot 1 = \frac{\partial L}{\partial y} $$这意味着为量化节点输出计算的梯度 $\frac{\partial L}{\partial y}$ 被“直通”传递,成为用于输入的梯度 $\frac{\partial L}{\partial x}$,如同量化操作是一个恒等函数($y=x$)仅用于梯度计算。digraph STE { rankdir=LR; node [shape=box, style=filled, fontname="Arial", color="#adb5bd", fillcolor="#e9ecef"]; edge [fontname="Arial", color="#495057"]; subgraph cluster_forward { label = "前向传播"; style=dashed; color="#74c0fc"; X_f [label="输入 (x)", fillcolor="#a5d8ff"]; Q_f [label="量化\ny = q(x)", shape=ellipse, fillcolor="#ffec99"]; Y_f [label="量化输出 (y)", fillcolor="#a5d8ff"]; X_f -> Q_f; Q_f -> Y_f; } subgraph cluster_backward { label = "反向传播 (带 STE)"; style=dashed; color="#f783ac"; dLdy [label="梯度 (dL/dy)", fillcolor="#fcc2d7"]; STE [label="STE 近似\ndy/dx ≈ 1", shape=ellipse, fillcolor="#ffec99"]; dLdx [label="梯度 (dL/dx)", fillcolor="#fcc2d7"]; dLdy -> STE; STE -> dLdx; } Y_f -> dLdy [style=invis]; // Keep layout separate dLdx -> X_f [style=invis]; // Keep layout separate }此图表展示了 STE 过程。前向传播应用实际的量化 $q(x)$。反向传播使用 STE 近似($\frac{\partial y}{\partial x} \approx 1$),使传入梯度 $\frac{\partial L}{\partial y}$ 不变地通过,成为 $\frac{\partial L}{\partial x}$。STE 为何有效?在反向传播过程中使用一个忽略量化函数真实属性的近似,这看起来可能与直觉相悖。然而,STE 在实际中表现良好,原因有几点:梯度流动: 它确保梯度不被阻碍,使优化信号能够到达量化节点之前的权重。学习适应性: 尽管梯度计算是近似的,但前向传播仍然使用实际的量化值。损失函数反映了这种量化引入的误差。通过允许梯度流动,STE 使模型权重能够以最小化这种量化引起的误差的方式调整。网络学习到的参数对前向传播中应用的舍入效应更具弹性。简洁性: STE 在标准深度学习框架中易于实现。大多数提供 QAT 能力的框架都在内部集成了 STE。STE 的变体虽然恒等近似($\frac{\partial y}{\partial x} = 1$)是 STE 最常见的形式,但也存在一些变体。例如,另一种方法涉及根据量化所用的输入范围剪裁梯度。如果量化函数在量化前将输入 $x$ 有效地剪裁到范围 $[c_{min}, c_{max}]$ 内,那么 STE 可以定义为:$$ \frac{\partial y}{\partial x} = \begin{cases} 1 & \text{如果 } c_{min} \le x \le c_{max} \ 0 & \text{否则} \end{cases} $$这种变体阻止了对于已经超出量化范围的输入,梯度反向传播,这有时有助于稳定训练。然而,简单的恒等近似通常就足够了。总之,直通估计器是使量化感知训练成为可能的底层技术。通过为梯度提供一条通过不可微分量化操作的路径,它使深度学习模型能够在训练或微调期间调整其权重,与训练后量化相比,这能使量化模型获得明显更好的精度,尤其是在极低比特宽度时。