二值神经网络(BNN)把权重和激活值压缩到只有+1和-1两种状态,相比32位浮点模型,存储占用可以减少32倍,卷积运算也可以从浮点乘加变成位运算和加减法。然而这些收益的代价非常明显:同样的网络结构,二值化后准确率往往比全精度模型低很多,尤其在ImageNet这种大规模图像分类任务上,简单二值化的ResNet-18可能从约70%的top-1准确率掉到50%附近。精度低的原因并不是网络容量简单减半,而是离散化在前向传播和反向传播两个环节同时制造了障碍。解决思路也不止一种,但实践证明,梯度近似和通道级缩放是两个非常关键、并且可以稳定复现的优化方向。本文会先剖析BNN精度崩塌的根因,再分别讨论梯度近似与通道级缩放的具体做法,最后给出可运行的PyTorch融合实现。

二值化为什么会造成精度崩塌
二值化通常使用符号函数sign,当输入大于等于0时输出+1,否则输出-1。这个函数在数学上有一个致命问题:它在几乎所有的点导数为0,在0点处不可导。反向传播依赖链式法则逐层计算梯度,如果某一层的导数全是0,梯度信号就会在这里断掉,前面的层完全无法更新。这正是BNN训练的硬伤。
为了绕过符号函数不可导的问题,一种常用做法是前向仍然使用离散的sign,反向却使用一个近似函数的导数,这就是所谓的直通估计器(STE)。最简单的STE把sign的导数近似为1,通常在[-1,1]区间内,超出该区间则置为0。这样梯度可以穿过二值化节点,但代价是前向函数与反向导数并不对应,梯度存在系统偏差。网络越深,这种偏差累积越严重,最终导致权重更新方向偏离真实损失下降方向。
除了梯度问题,二值化还会抹掉权重的幅值信息。全精度卷积核中,不同输出通道的权重可能具有完全不同的L1范数或标准差,这种尺度差异决定了各通道在后续特征图中的重要性。二值化把所有权重变成+1或-1后,每个通道的卷积核幅值都变成相同的固定范数,通道之间的表达能力差异被强行拉平。即便后面跟有批归一化(BatchNorm),BN层也需要在训练早期承受更大的分布偏移,容易造成训练不稳定和精度下降。
梯度近似:让不可导节点重新获得学习信号
梯度近似要解决的核心问题是:在前向保持离散二值化,以保证推理时低比特特性;在反向给出一个有意义、噪声可控的梯度,让优化器可以继续更新全精度权重。最简单的STE实现可以看作一个自定义的autograd.Function,前向返回二值结果,反向直接透传上游梯度。
import torch
import torch.nn as nn
class STEBinaryActivation(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
# 前向使用 sign 进行硬二值化
return torch.where(x >= 0, torch.ones_like(x), -torch.ones_like(x))
@staticmethod
def backward(ctx, grad_output):
# 反向直通,直接返回上游梯度
return grad_output
class BinaryActivation(nn.Module):
def forward(self, x):
return STEBinaryActivation.apply(x)
这段代码很容易理解,但它有一个明显缺陷:对所有输入都返回相同的梯度,哪怕输入已经远离0点,也仍然会获得完整梯度。对于非常负或非常正的值,它们的二值化结果已经确定,继续给出强梯度会让优化器不断调整这些已经稳定的权重,引入额外噪声。改进方案是采用区间截断的STE,只在[-1,1]区间内透传梯度,超出区间则返回0。
class ClippedSTE(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return torch.where(x >= 0, 1.0, -1.0)
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
mask = (x >= -1) & (x <= 1)
return grad_output * mask.float()
区间截断策略能减少远离0点样本的无效更新,但它的导数形状仍然是硬边界。更平滑的做法是使用近似符号函数的导数,例如用tanh的梯度或二次函数来软化边界。有的工作进一步引入可学习的温度参数,在训练初期使用较软的二值化曲线,让梯度更丰富;随着训练推进逐步降低温度,使前向逐渐逼近硬二值函数。这类方法统称为代理梯度,它们与STE的根本区别在于,反向导数不再是一个人为指定的常数或掩码,而是来自某个连续可导函数的真实导数,从而减少梯度偏差。
实现代理梯度时,通常需要在前向保存输入变量,并在反向根据输入值计算近似导数。例如可以把hardtanh的导数作为sign的反向梯度,其本质等价于区间截断STE。如果要使用更平滑的代理,可以用tanh乘上一个温度系数,并在反向时显式计算1-tanh^2。不论选择哪种代理函数,梯度幅值都应该通过实验调节,避免梯度爆炸或消失。
通道级缩放:恢复二值权重的幅值信息
二值化卷积核的另一个问题在于,+1/-1卷积核无法表示原始浮点卷积核的幅值。例如一个浮点卷积核的权重全部接近0.01,另一个卷积核的权重集中在0.5附近,二值化后它们都变成同样幅值的+1/-1,但前者实际对输出的贡献本应远小于后者。修复这个问题的直接思路是给二值化后的卷积核乘上一个缩放因子。XNOR-Net最早使用了整体缩放,即对整个卷积核计算一个标量alpha,让它等于原始浮点权重绝对值的平均值。这样二值卷积核乘上alpha后,在最小二乘意义上最接近原始浮点卷积核。
整体缩放实现简单,但对不同输出通道一视同仁。在实际卷积层中,输出通道是独立的检测器,不同通道的特征尺度经常差异很大。如果只用同一个alpha缩放所有通道,那些原本响应较弱的通道仍然会被放大到和强响应通道相同的量级,无法恢复通道级重要性。因此更合理的做法是把缩放因子设计成每个输出通道一个,形状为out_channels x 1 x 1 x 1,对二值化后的卷积核逐通道加权。这样可以让不同输出通道保留各自的幅值先验。
下面的PyTorch示例实现了一个带通道级缩放和批归一化的二值卷积层。二值化权重时,先按每个输出通道计算均值作为该通道的二值化阈值,而不是固定使用0阈值,这样更能保留原始权重的符号分布。缩放因子初始化为1,并作为可学习参数参与反向传播。
import torch
import torch.nn as nn
import torch.nn.functional as F
class ChannelScaleBinaryConv2d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
super().__init__()
self.stride = stride
self.padding = padding
self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size) * 0.05)
self.scale = nn.Parameter(torch.ones(out_channels, 1, 1, 1))
self.bn = nn.BatchNorm2d(out_channels)
def binarize_weight(self):
# 每个输出通道使用自己的均值作为二值化阈值
w_mean = self.weight.view(self.weight.size(0), -1).mean(dim=1).view(-1, 1, 1, 1)
binary_w = torch.where(self.weight >= w_mean, 1.0, -1.0)
# 通道级缩放
return binary_w * self.scale
def forward(self, x):
bw = self.binarize_weight()
x = F.conv2d(x, bw, stride=self.stride, padding=self.padding)
x = self.bn(x)
return x
这个实现中,全精度权重self.weight仍然是需要训练的主体参数,二值化只发生在前向计算时。反向传播时,binarize_weight中的torch.where本身不会被梯度穿透,通常需要配合代理梯度或STE来处理。通道级缩放self.scale在反向时可以正常更新,它承担了恢复二值卷积核幅值的作用。训练完成后,推理阶段可以把scale与批归一化的参数融合,或者直接保留缩放后的二值卷积核,不会增加额外推理开销。
联合优化与训练调试
梯度近似和通道级缩放解决的是两个不同层面的问题,前者负责修复反向传播信号,后者负责修复前向表达容量。单独使用其中一种往往不够:如果只做梯度近似而没有通道缩放,二值化权重仍然丢失幅值信息;如果只做通道缩放而继续使用错误梯度,优化过程仍然可能震荡。把两者组合起来,才容易在训练中同时获得稳定更新和更强的表示能力。
训练BNN时有几个细节值得注意。首先是学习率,全精度模型适用的学习率对二值网络通常偏大,因为二值化会放大权重更新带来的输出变化。建议先使用全精度模型做预训练,然后切换到二值化微调,学习率可以降为原来的十分之一或更小。其次是优化器,Adam配合较小的学习率在很多BNN实验里表现稳定,动量较大的SGD则容易让二值权重频繁翻转。再次是批归一化的动量,二值网络的特征分布变化比全精度网络更剧烈,适当减小BN动量,或者增加BN层前的缩放因子,有助于训练早期保持稳定。
还有一个常被忽略的点是二值化阈值的设定。固定使用0作为阈值虽然简单,但当某个通道的权重均值偏离0较远时,固定阈值会让二值化结果严重偏向某一类符号,等效于该通道只输出局部响应。因此在通道级缩放的基础上,使用通道均值作为阈值,可以让每个通道在二值化后保留更多原始符号结构。结合代理梯度的截断区间,可以进一步减少边界附近的梯度噪声。
从实验结果来看,梯度近似与通道级缩放并不是昂贵的补偿手段。它们几乎不增加推理计算量,所有额外参数都可以在训练后融合。对于中小规模数据集,这种组合经常能把二值模型的准确率从不可用拉升到接近全精度模型的95%以上。即便在大规模数据集上,它们也是几乎所有先进BNN方法的基础组件。理解并实现这两个机制,是解决BNN精度极低问题最实用的切入点。