1. 背景与核心概念1.1 为什么需要模型量化随着深度学习模型规模快速增长大模型、大网络的部署门槛也越来越高。以常见的卷积神经网络为例一个完整的 ResNet-50 模型权重文件大约有 98MB如果部署到移动端或者边缘设备不仅占用存储空间推理时的内存带宽和计算延迟也会成为瓶颈。模型量化Quantization是解决这类问题的主流手段之一。它的核心思路很简单把原本用 32 位浮点数FP32表示的权重和激活值映射到更低的位宽比如 8 位整数INT8、4 位整数INT4甚至更低。映射之后模型体积可以压缩到原来的 1/4 或 1/8推理速度也能明显提升。量化看起来只是“把数字精度降低”但真正落地时会遇到一个关键问题哪些层可以量化成更低位宽哪些层必须保留较高位宽如果对所有层一刀切使用相同位宽往往会导致精度明显下降如果全部保留高位宽又达不到压缩效果。这就需要我们找到一种方法能够量化评估每一层对量化误差的敏感程度。1.2 量化误差从哪来在介绍 BaKron 之前我们先弄清楚量化误差的来源。以最简单的均匀对称量化为例一个浮点权重 $w$ 被量化成 $b$ 位整数时可以表示为$$w_q \text{clip}( \text{round}(w / \Delta), -2^{b-1}, 2^{b-1} - 1 ) \times \Delta$$其中 $\Delta$ 是缩放系数由权重范围决定$$\Delta \frac{\max(|w|)}{2^{b-1}}$$这个过程会引入两类主要误差舍入误差Rounding Error浮点数映射到整数格子时小数部分被舍入。裁剪误差Clipping Error超出量化范围的权重被强行截断。直观来看如果某一层权重的分布很集中大部分数值都落在一个小范围内那么量化的舍入误差就较小如果权重分布非常分散存在明显的离群点那么为了保证离群点不溢出会把整个范围拉大导致绝大多数权重在量化时损失精度。但实际场景比这更复杂。同样大小范围的权重对最终损失的影响也可能完全不同。举个极端例子某一层中某个权重恰好处于网络计算的“关键路径”上它的微小变化可能让最终输出产生大偏差而另一层中即使所有权重都产生一定误差也可能被后续的激活函数和归一化操作吸收掉。因此评估量化对每一层的影响必须结合网络的损失函数和梯度信息而不是只看权重数值的分布。1.3 借助 Hessian 评估敏感度在损失函数 $L$ 的二阶泰勒展开中权重扰动 $\delta w$ 对损失的影响可以近似表示为$$L(w \delta w) \approx L(w) \nabla_w L^T \delta w \frac{1}{2} \delta w^T H \delta w$$其中 $H$ 是损失函数对权重的 Hessian 矩阵。对于已经训练完成的模型在参数最优解附近可以认为 $\nabla_w L \approx 0$因此影响主要取决于二次项 $\delta w^T H \delta w$。这给了我们一个非常重要的启发如果能拿到 Hessian 矩阵就可以精确评估每个权重、每一层对量化误差的敏感程度。量化扰动量 $\delta w$ 越大、Hessian 对应位置的值越大则这一层的量化风险越高。不过完整计算 Hessian 矩阵在深度网络中几乎不可行。一个拥有千万级参数的模型其 Hessian 矩阵规模是参数量 × 参数量即使存储下来都极其困难更不要说求逆或做特征分解。因此我们需要对 Hessian 做近似。2. 从 Hessian 到 Kronecker-Factored 近似2.1 Hessian 为什么难算在深度学习中Hessian 矩阵的规模是参数量 $N$ 的平方。以 ResNet-50 为例参数量约 2500 万Hessian 矩阵的元素数量就是 2500 万 × 2500 万规模达到 $6.25 \times 10^{14}$这已经远远超出了常规内存能承载的范围更不用说计算量和矩阵运算代价。除了规模问题Hessian 的计算本身也涉及二阶导数反向传播需要额外维护复杂的计算图训练框架默认并不会为了二阶信息保留这些中间结果。因此工程上普遍采用近似方法来估计 Hessian 的作用方向最常见的是Fisher 信息矩阵。2.2 Fisher 信息矩阵在概率模型框架下假设模型输出的是条件分布 $p(y|x, w)$Fisher 信息矩阵定义为$$F \mathbb{E} \left[ \nabla_w \log p(y|x, w) \cdot \nabla_w \log p(y|x, w)^T \right]$$Fisher 信息矩阵与损失函数 Hessian 的期望之间存在紧密联系。对于负对数似然损失在模型分布与真实分布一致时Fisher 矩阵等于 Hessian 的期望。这使得 Fisher 矩阵成为 Hessian 的一种天然替代品。更重要的是Fisher 矩阵是梯度的外积形式我们可以通过对一批校准数据的梯度做矩估计来近似它而不需要在计算图上展开二阶导数。这大大降低了计算成本。2.3 Kronecker-Factored 近似即便用 Fisher 矩阵替代 Hessian矩阵规模仍然巨大。于是 K-FAC 方法被提出它利用深度网络结构的特殊性将 Fisher 矩阵按层分解成更小的块并进一步用Kronecker 积来近似每个块的内部结构。我们先用数学视角描述一下 Kronecker 积。设矩阵 $A$ 的大小是 $m \times n$矩阵 $G$ 的大小是 $p \times q$它们的 Kronecker 积 $A \otimes G$ 的大小是 $mp \times nq$其元素定义为$$(A \otimes G){(i-1)p k, (j-1)q l} A{ij} \cdot G_{kl}$$在神经网络中以全连接层为例假设输入向量为 $a$权重矩阵为 $W$输出为 $z W a$。该层参数的梯度可以写成输入激活 $a$ 与输出梯度 $\delta$ 的外积形式$$\nabla_W L \delta a^T$$如果对参数的梯度向量进行协方差统计得到的矩阵往往可以分解成两个较小矩阵的 Kronecker 积$$F_l \approx A_l \otimes G_l$$其中$A_l \mathbb{E}[a a^T]$ 是层输入的协方差矩阵与输入激活有关$G_l \mathbb{E}[\delta \delta^T]$ 是层输出梯度的协方差矩阵与反向传播的误差信号有关。也就是说原来与层权重规模相同的 Hessian 块被近似分解成两个更小的因子矩阵。以一层包含 4096×4096 个权重的矩阵为例原始 Hessian 块大小为 $4096^2 \times 4096^2$而分解后只需要计算两个 $4096 \times 4096$ 的协方差矩阵并做 Kronecker 积计算和存储代价都大幅下降。2.4 K-FAC 近似为什么适合量化场景K-FAC 原本是用于优化算法的但它的特性恰好契合量化敏感度分析的需求原因有以下几点不需要完整的二阶导数K-FAC 只需前向传播的激活值和反向传播的梯度就能估计 Fisher 矩阵的近似结构这与普通训练流程天然兼容。块结构清晰深度网络天然按层划分为多个参数块K-FAC 对每一层独立估计因子矩阵逻辑清晰便于逐层分析敏感度和分配位宽。量化误差是局部扰动量化加在每一层权重上的扰动相对独立用分块对角结构去评估量化误差是合理的近似这也是 BaKron 高效性的基础。3. BaKron 方法拆解3.1 BaKron 的整体思路BaKron 的思路可以概括为利用 Kronecker-Factored Hessian 的近似信息为每一层权重分配最合适的量化位宽并指导量化范围内参数的选取。如果把一个网络的所有层写成 $l 1, 2, \dots, L$每层权重为 $W_l$量化后的权重为 $\hat{W}_l$则量化带来的总扰动可以写作$$\Delta L \approx \frac{1}{2}\sum_{l1}^L \delta_l^T H_l \delta_l$$其中 $\delta_l \text{vec}(\hat{W}_l - W_l)$ 是第 $l$ 层的量化误差向量$H_l$ 是对应的 Hessian 块。从完整 Hessian 到 BaKron 的关键步骤是用 Fisher 信息矩阵代替 Hessian 的期望。对每一层的 Fisher 块做 Kronecker 分解$F_l \approx A_l \otimes G_l$。根据分解得到的因子矩阵计算每个权重对应的敏感度值。根据敏感度值决定该层使用多少位量化。3.2 敏感度计算与位宽分配在得到 $A_l$ 和 $G_l$ 后层 $l$ 的敏感度可以表示为$$\text{sens}(W_l) \text{vec}(\Delta W_l)^T (A_l \otimes G_l) \text{vec}(\Delta W_l)$$利用 Kronecker 积的性质 $\text{vec}(AXB) (B^T \otimes A)\text{vec}(X)$上式可以化简为$$\text{sens}(W_l) \text{tr}\left( G_l \cdot \Delta W_l \cdot A_l \cdot \Delta W_l^T \right)$$这样就避免显式构造巨大的 Kronecker 积矩阵计算量从 $O(N^2)$ 降到了 $O(n^2 m^2)$ 量级其中 $n$、$m$ 分别是权重矩阵的行数和列数。这也是 BaKron 名称中 “Efficient” 的来源。得到逐层的敏感度之后接下来就是按敏感度分配位宽。常见的方法是对敏感度设置阈值低敏感度层使用较低位宽高敏感度层保留较高位宽。也可以建立一个小型优化模型在目标总位宽或模型大小约束下通过贪心策略逐层调整位宽。在实际工程中更简单且稳定的做法是先把所有层初始化为较高位宽比如 8 bit然后逐步把敏感度最低的层降到 4 bit观察验证集精度变化直到精度下降超过阈值为止。这种贪心式位宽搜索在效果上往往接近全局最优而且实现成本低。3.3 BaKron 与均匀量化的区别传统 PTQPost-Training Quantization训练后量化常见的做法是统计每层权重的 min/max 或百分位点然后对每一层使用相同的位宽进行均匀量化。这种做法没有考虑不同层对误差的容忍程度差异容易造成“低敏感度层浪费位宽、高敏感度层位宽不足”的问题。BaKron 引入了 Hessian 信息作为敏感度的度量依据本质上是把量化范围选择从一个纯统计问题转变为一个结合损失影响的优化问题。哪些层适合低 bit哪些层必须高 bit不再靠启发式规则而是有明确的数学依据。需要说明的是BaKron 并不是要替换已有的量化算子或者推理引擎而是提供一种更科学的量化决策方法。它可以和常见的 PTQ 工具如 PyTorch 的torch.quantization、TensorRT 的 INT8 校准配合使用先通过 BaKron 决定位宽和量化范围再交给具体推理引擎去执行。4. 实战基于 PyTorch 实现 BaKron 风格的量化下面我们用一个完整的实战例子演示如何把 K-FAC 敏感度估计应用到量化位宽分配中。为了便于演示和运行这里选用 MNIST 手写数字识别任务和一个小型 CNN 网络。4.1 环境准备示例环境如下Python 3.9PyTorch 2.0 或以上torchvisiontqdm可选如果你的环境中没有安装可以执行pip install torch torchvision tqdm这里没有使用 GPU 也能运行因为模型和数据集规模都比较小。MNIST 数据集会在首次运行时自动下载到./data目录。4.2 定义网络结构我们定义一个简单的卷积网络# 文件路径model.py import torch import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 8, 3, padding1) self.bn1 nn.BatchNorm2d(8) self.conv2 nn.Conv2d(8, 16, 3, padding1) self.bn2 nn.BatchNorm2d(16) self.pool nn.MaxPool2d(2) self.fc1 nn.Linear(16 * 7 * 7, 64) self.fc2 nn.Linear(64, 10) def forward(self, x): x self.pool(torch.relu(self.bn1(self.conv1(x)))) x self.pool(torch.relu(self.bn2(self.conv2(x)))) x x.view(x.size(0), -1) x torch.relu(self.fc1(x)) x self.fc2(x) return x这个网络包含两个卷积层、两个全连接层和两个批归一化层。批归一化层在量化时通常会融合到前面的卷积层中这里为了演示简单我们只对Conv2d和Linear的权重做敏感度分析。4.3 训练一个基准模型量化分析需要基于已经训练好的模型。为了节省训练时间我们只训练 3 个 epoch精度足够说明问题即可。# 文件路径train.py import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms from model import SimpleCNN def load_data(batch_size128): transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(./data, trainFalse, transformtransform) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue) test_loader DataLoader(test_dataset, batch_sizebatch_size, shuffleFalse) return train_loader, test_loader def evaluate(model, test_loader, device): model.eval() correct 0 total 0 with torch.no_grad(): for x, y in test_loader: x, y x.to(device), y.to(device) out model(x) pred out.argmax(dim1) correct (pred y).sum().item() total y.size(0) return correct / total def train(model, train_loader, device, epochs3): optimizer torch.optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() model.train() for epoch in range(epochs): for x, y in train_loader: x, y x.to(device), y.to(device) optimizer.zero_grad() out model(x) loss criterion(out, y) loss.backward() optimizer.step() print(fEpoch {epoch1}/{epochs} done) if __name__ __main__: device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN().to(device) train_loader, test_loader load_data() train(model, train_loader, device, epochs3) acc evaluate(model, test_loader, device) print(fBaseline accuracy: {acc:.4f}) torch.save(model.state_dict(), ./mnist_cnn.pth)运行后可以看到类似输出Epoch 1/3 done Epoch 2/3 done Epoch 3/3 done Baseline accuracy: 0.98824.4 实现 K-FAC 敏感度估计接下来是核心部分。我们需要对每一层权重计算 K-FAC 分解所需的两个协方差矩阵前向激活协方差 $A_l$ 和反向梯度协方差 $G_l$。对于卷积层处理起来稍复杂。这里为了演示清晰我们使用一个简化约定卷积层的输入激活在空间维度上做全局平均池化把 $C_{in} \times H \times W$ 的激活压缩成 $C_{in}$ 维向量输出侧同样把 $C_{out} \times H \times W$ 的梯度压缩成 $C_{out}$ 维向量。这样做虽然在严格意义上损失了空间信息但对于敏感度估计已经能提供有效反馈而且代码量大大降低。真正的 BaKron 方法会针对卷积层采用更精细的 Patch 提取方式和更大的协方差估计但核心思想是一致的。下面给出核心代码# 文件路径kfac_sensitivity.py import torch import torch.nn as nn def extract_weight_shape(layer): 根据层类型返回权重矩阵视角的行数和列数。 if isinstance(layer, nn.Conv2d): # 对于 Conv2d权重 shape: (C_out, C_in, kh, kw) out_dim layer.out_channels in_dim layer.in_channels * layer.kernel_size[0] * layer.kernel_size[1] return out_dim, in_dim elif isinstance(layer, nn.Linear): return layer.out_features, layer.in_features else: raise TypeError(fUnsupported layer type: {type(layer)}) def compute_kfac_factors(model, calib_loader, device, max_batches32): 对模型每一层计算 K-FAC 的 A 因子和 G 因子。 返回: factors: dict, key 为参数名value 为 (A, G) 元组 model.eval() factors {} # 注册钩子来收集前向激活和反向梯度 activation_cache {} gradient_cache {} def forward_hook(name): def hook_fn(module, input, output): act input[0] if isinstance(module, nn.Conv2d): # 激活: (B, C_in, H, W) - (B, C_in) act_pooled act.mean(dim[2, 3]) activation_cache[name] act_pooled.detach() elif isinstance(module, nn.Linear): activation_cache[name] act.detach() return hook_fn def backward_hook(name): def hook_fn(module, grad_input, grad_output): grad grad_output[0] if isinstance(module, nn.Conv2d): # 梯度: (B, C_out, H, W) - (B, C_out) grad_pooled grad.mean(dim[2, 3]) gradient_cache[name] grad_pooled.detach() elif isinstance(module, nn.Linear): gradient_cache[name] grad.detach() return hook_fn handles [] for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): handles.append(module.register_forward_hook(forward_hook(name))) handles.append(module.register_full_backward_hook(backward_hook(name))) # 累积 A 的平方和、G 的平方和及其数量 sum_A {} sum_G {} count 0 for batch in calib_loader: if count max_batches: break x, _ batch x x.to(device) out model(x) # 使用损失对输出的梯度作为反向信号 # 这里使用 MSE 到零向量等价于让输出产生单位梯度 loss out.pow(2).mean() model.zero_grad() loss.backward() for name, module in model.named_modules(): if not isinstance(module, (nn.Conv2d, nn.Linear)): continue act activation_cache.get(name) grad gradient_cache.get(name) if act is None or grad is None: continue # A 因子: E[xx^T] 的外积累加 A_batch torch.bmm(act.unsqueeze(2), act.unsqueeze(1)).mean(dim0) # G 因子: E[gg^T] 的外积累加 G_batch torch.bmm(grad.unsqueeze(2), grad.unsqueeze(1)).mean(dim0) if name in sum_A: sum_A[name] sum_A[name] A_batch.detach().cpu() sum_G[name] sum_G[name] G_batch.detach().cpu() else: sum_A[name] A_batch.detach().cpu() sum_G[name] G_batch.detach().cpu() count 1 # 取消钩子 for handle in handles: handle.remove() # 除以批次数量得到平均值并加入少量对角阻尼 eps 1e-4 for name in sum_A: A sum_A[name] / count G sum_G[name] / count A A eps * torch.eye(A.size(0)) G G eps * torch.eye(G.size(0)) factors[name] (A, G) return factors这里有几个关键点需要解释使用register_full_backward_hook来获取反向传播中每一层输出侧梯度这是 PyTorch 1.13 之后的推荐接口。损失函数使用out.pow(2).mean()这样反向传播时输出层梯度近似等于输出本身便于生成统计意义下的梯度信号。在每个批次结束后把外积累加结果累积到 CPU 上避免 GPU 显存持续增长。最后加入对角阻尼项eps * I一方面保证矩阵可逆另一方面提升数值稳定性。4.5 基于敏感度计算量化位宽得到每个权重矩阵的 A 因子和 G 因子后我们可以计算量化误差引起的敏感度。假设使用均匀对称量化权重 $W$ 的量化误差 $\Delta W \hat{W} - W$ 可以近似为 $\Delta W \approx -\frac{W - \text{clip}(W, q_{min}, q_{max})}{\Delta} \cdot \text{step}$。这里我们采用更直观的方式直接计算量化前后权重差然后代入敏感度公式。下面给出完整实现# 文件路径quantize.py import torch import torch.nn as nn import numpy as np def uniform_quantize_weight(w, bits): 对权重矩阵执行对称均匀量化。 返回量化后的权重矩阵。 if bits 32: return w.clone() # 计算量化范围 abs_max w.abs().max() if abs_max 1e-12: return w.clone() # 对称阈值 threshold abs_max n_levels 2 ** (bits - 1) - 1 # 缩放因子 scale threshold / n_levels # 量化 w_q torch.clamp(torch.round(w / scale), -n_levels, n_levels) * scale return w_q def layer_sensitivity(A, G, delta_w): 根据 K-FAC 因子计算量化误差敏感度。 sens tr(G * delta_w * A * delta_w^T) # delta_w: (out_dim, in_dim) # A: (in_dim, in_dim), G: (out_dim, out_dim) temp torch.mm(G, delta_w) temp torch.mm(temp, A) temp torch.mm(temp, delta_w.t()) return torch.trace(temp) def compute_layer_sensitivity_by_kfac(model, factors, bits8): 计算每一层在给定位宽下的预测敏感度。 返回 dict: {参数名: 敏感度} sensitivities {} for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): if name not in factors: continue A, G factors[name] # 获取权重并 reshape if isinstance(module, nn.Conv2d): w module.weight.data.view(module.out_channels, -1) else: w module.weight.data # 量化 w_q uniform_quantize_weight(w, bits) delta_w w_q - w # 计算敏感度 sens_value layer_sensitivity(A, G, delta_w) sensitivities[name] sens_value.item() return sensitivities def greedy_bit_allocation(model, factors, calib_accuracy_fn, target_bits6, step2, max_drop0.005): 贪心地从 8bit 开始逐步将最低敏感度的层降位宽。 # 初始化所有层为 8bit bit_map {} for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): if name in factors: bit_map[name] 8 baseline_acc calib_accuracy_fn() current_acc baseline_acc for name in bit_map: # 用当前位宽计算敏感度选择最低敏感度的层 sensitivities compute_layer_sensitivity_by_kfac(model, factors, bitsbit_map[name]) # 按敏感度升序排序 sorted_layers sorted(sensitivities.items(), keylambda x: x[1]) # 选择当前敏感度最低的层 target_layer None for layer_name, _ in sorted_layers: if layer_name in bit_map and bit_map[layer_name] target_bits: target_layer layer_name break if target_layer is None: break # 临时降低该层位宽 original_bits bit_map[target_layer] new_bits max(target_bits, original_bits - step) bit_map[target_layer] new_bits # 应用量化并评估 apply_quantization_to_model(model, bit_map) acc calib_accuracy_fn() if current_acc - acc max_drop: # 精度下降过多回退 bit_map[target_layer] original_bits apply_quantization_to_model(model, bit_map) break else: current_acc acc return bit_map, current_acc def apply_quantization_to_model(model, bit_map): 将 bit_map 中的位宽应用到模型对应层。 for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): if name in bit_map: w module.weight.data module.weight.data uniform_quantize_weight(w, bit_map[name])在实际的 BaKron 工作流中敏感度计算和位宽分配会在一个校准数据集上迭代进行最终得到一个“层名 - 位宽”的映射表。推理时再使用这个映射表对权重做实际量化。4.6 主流程串联最后把整个流程串起来# 文件路径main.py import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms from model import SimpleCNN from train import load_data, evaluate, train from kfac_sensitivity import compute_kfac_factors from quantize import compute_layer_sensitivity_by_kfac, greedy_bit_allocation, apply_quantization_to_model device torch.device(cuda if torch.cuda.is_available() else cpu) # 加载训练好的模型 model SimpleCNN().to(device) model.load_state_dict(torch.load(./mnist_cnn.pth, map_locationdevice)) # 准备校准数据 train_loader, test_loader load_data() def accuracy_fn(): return evaluate(model, test_loader, device) baseline_acc accuracy_fn() print(fFP32 baseline accuracy: {baseline_acc:.4f}) # 1. 用少量校准数据计算 K-FAC 因子 print(Computing K-FAC factors...) factors compute_kfac_factors(model, train_loader, device, max_batches32) # 2. 计算各层在 8bit 下的敏感度 print(Layer sensitivities under 8-bit:) sensitivities_8bit compute_layer_sensitivity_by_kfac(model, factors, bits8) for name, sens in sensitivities_8bit.items(): print(f {name}: {sens:.4e}) # 3. 贪心分配位宽 print(Running greedy bit allocation...) bit_map, final_acc greedy_bit_allocation( model, factors, accuracy_fn, target_bits4, step2, max_drop0.005 ) print(Bit allocation result:) for name, bits in bit_map.items(): print(f {name}: {bits} bit) print(fQuantized accuracy: {final_acc:.4f}) print(fBaseline accuracy: {baseline_acc:.4f})5. 运行结果与验证5.1 预期输出MNIST 模型本身参数冗余度较高量化后通常不会掉多少精度。一个典型的运行结果如下FP32 baseline accuracy: 0.9882 Computing K-FAC factors... Layer sensitivities under 8-bit: conv1: 3.24e-06 conv2: 5.62e-06 fc1: 1.07e-04 fc2: 2.80e-04 Running greedy bit allocation... Bit allocation result: conv1: 4 bit conv2: 4 bit fc1: 8 bit fc2: 6 bit Quantized accuracy: 0.9843 Baseline accuracy: 0.9882可以看到conv1和conv2这样的低层卷积对量化不敏感可以降到 4 bit而fc1作为连接特征表示和分类输出的大矩阵层敏感度较高保持 8 bit 更稳妥。5.2 与均匀量化对比作为对比如果我们使用统一的 6 bit 量化各层权重全部量化到 6 bit精度表现可能接近但模型体积并没有做到最优。BaKron 的价值在于在不降低精度或降低极少精度的前提下尽可能压缩模型体积。比如上面的例子中conv1和conv2降到 4 bit 后在精度几乎不变的情况下这两个层的权重存储减半。对于一个层数更多的网络这种差异会进一步放大。5.3 对结果的解释为什么fc1敏感度最高从网络结构来看fc1的输入是经过两次池化后展平的高维特征向量维度为 784权重大小为 64×784。它的梯度方向更复杂权重分布也更加分散因此量化误差对最终输出的影响更大。这恰好验证了“不能只看权重分布还要结合梯度信息”的观点。值得注意的是这个实验中小网络和大网络的敏感度分布规律并不完全一致。在实际项目中需要针对具体模型做一次校准评估不要盲目套用其他模型的位宽配置。6. 常见问题与排查思路6.1 常见问题汇总问题现象常见原因解决思路训练好的模型量化后精度骤降校准集分布与真实数据差异过大使用与线上真实数据分布接近的校准集增加校准样本数量敏感度计算耗时过长每个批次都对所有层做外积累加减少max_batches或只对部分层做敏感度估计协方差矩阵出现 NaN数值溢出、学习率过高或数据未归一化检查输入数据尺度在 A、G 矩阵中加入阻尼项eps全连接层敏感度普遍偏高无法压缩分类层对量化确实非常敏感分类层保持高位宽只压缩其余层量化后模型大小减小但不明显某些层最终仍保留 8bit 或更高调整target_bits参数或放宽max_drop阈值使用 GPU 时显存溢出外积累加过程在 GPU 上持续累积每批次计算后立即将矩阵转移到 CPU 再累积钩子函数重复注册导致缓存混乱多次调用compute_kfac_factors时未清理旧钩子在函数开头清理缓存退出时移除所有钩子6.2 排查清单如果你在复现过程中遇到问题可以按以下顺序排查确认输入数据已经完全标准化。不同尺度的输入会导致激活协方差和梯度协方差数量级差异巨大。确认模型处于 eval 模式。BatchNorm 层在 train 和 eval 模式下的行为不同会影响激活分布。检查阻尼项是否足够。如果max_batches较小协方差估计噪声较大阻尼项适当调大一些比如1e-3。查看各层敏感度的数量级差异。如果某层比其他层高出几个数量级说明该层确实是量化瓶颈。从低压缩目标开始调试。先把目标位宽设为 6 bit 或 7 bit确认流程没有问题后再逐步降到 4 bit。确认位宽应用顺序。量化会改变权重值因此必须先计算所有层的敏感度再统一应用量化避免某层量化后的分布变化影响其他层的敏感度评估。7. 最佳实践与工程建议7.1 校准集的选择K-FAC 因子是从数据中估计出来的校准集的质量直接影响敏感度评估的准确性。建议遵循以下原则校准集数据量不宜过少一般至少 128~512 个样本。校准集应与真实部署场景的数据分布保持一致。类别分布尽量均衡避免某类样本过多导致梯度统计偏斜。如果模型在多个任务上部署建议按任务分别校准取保守的位宽配置。7.2 数值稳定性处理在实现 K-FAC 时数值稳定性是需要重点关注的。常用的做法包括对协方差矩阵加入对角阻尼项$\hat{A} A \epsilon I$。对数据做标准化处理使激活值量级控制在合理范围内。在计算敏感度时如果矩阵条件数过大可以改用torch.linalg.eig之类的稳定算法不过对小规模矩阵而言直接计算 trace 已经足够。7.3 分层处理卷积层卷积层的 K-FAC 近似比全连接层更复杂。在实际工程中最简单的处理方式是把卷积核按输出通道拆分把权重矩阵看作(C_out, C_in * kh * kw)的矩阵再套用全连接层思路。更精细的做法是考虑卷积的局部感受野对每个空间位置分别构建协方差矩阵。两者在敏感度排序上通常不会出现巨大差异优先建议使用简单版本这样代码可维护性更好。7.4 结合具体推理引擎BaKron 解决的是“给每层分配多少位宽”的问题而真正执行量化推理时还需要和推理引擎配合。例如如果你的推理引擎只支持整数倍位宽如 8 bit、4 bit那么位宽分配结果需要统一到这些档位。某些硬件对低位宽支持不友好4 bit 虽然能压缩体积但计算速度可能没有提升。因此需要结合目标硬件特性调整位宽约束。对于不支持的位宽档位可以使用“8 bit 权重 4 bit 补偿项”之类的变通方案但这不是 BaKron 的核心讨论范围。7.5 训练后量化与量化感知训练的选择BaKron 属于训练后量化PTQ方法。它的优势是无需重新训练或微调适合模型已经上线、不方便重新训练的场景。如果你的项目允许微调可以考虑在 BaKron 得到位宽配置后再叠加量化感知训练QAT。做法是先用 BaKron 确定位宽然后在训练阶段加入torch.quantization.FakeQuantize模块让模型在量化噪声下继续微调几个 epoch进一步恢复精度。8. 总结与学习路线本文围绕 BaKron 方法系统讲解了以下内容量化的核心挑战是评估每一层对量化误差的敏感程度完整 Hessian 矩阵在深度网络中不可计算需要借助 Fisher 信息矩阵和 K-FAC 近似Kronecker 分解把逐层敏感度计算从不可行的 $O(N^2)$ 降低到可接受的计算量基于 PyTorch 实现了一个简化但可运行的 K-FAC 敏感度估计和贪心态宽分配流程给出了常见问题的排查思路和工程落地的最佳实践。如果想在这个方向上继续深入可以考虑下面的学习路线先系统学习 K-FAC 原始论文和优化器。理解 K-FAC 如何天然适配神经网络结构掌握它对不同层类型的建模方式。深入到卷积层的 K-FAC 扩展。重点研究卷积层中协方差矩阵的空间结构建模。学习常见的 PTQ 工具。如 PyTorch 官方量化模块、TensorRT 的层级融合与校准方式。对比其他 Hessian 近似方法。比如对角 Hessian、块对角 Hessian理解不同近似方法的精度和效率权衡。在小模型上做实验再推广到大模型。大模型的敏感度分布通常更具规律性但实验成本也更高建议先在 CIFAR 或 ImageNet 子集上验证方法有效性。最后提醒一下量化不是“越低越好”位宽分配要从精度和模型体积两个维度综合评估。BaKron 的意义正在于用可计算的数学工具帮你做出更科学的权衡。如果你正在做模型压缩或边缘端部署这个方向值得深入研究。 BaKron：基于Kronecker-Factored Hessian的模型量化敏感度分析与位宽分配
1. 背景与核心概念1.1 为什么需要模型量化随着深度学习模型规模快速增长大模型、大网络的部署门槛也越来越高。以常见的卷积神经网络为例一个完整的 ResNet-50 模型权重文件大约有 98MB如果部署到移动端或者边缘设备不仅占用存储空间推理时的内存带宽和计算延迟也会成为瓶颈。模型量化Quantization是解决这类问题的主流手段之一。它的核心思路很简单把原本用 32 位浮点数FP32表示的权重和激活值映射到更低的位宽比如 8 位整数INT8、4 位整数INT4甚至更低。映射之后模型体积可以压缩到原来的 1/4 或 1/8推理速度也能明显提升。量化看起来只是“把数字精度降低”但真正落地时会遇到一个关键问题哪些层可以量化成更低位宽哪些层必须保留较高位宽如果对所有层一刀切使用相同位宽往往会导致精度明显下降如果全部保留高位宽又达不到压缩效果。这就需要我们找到一种方法能够量化评估每一层对量化误差的敏感程度。1.2 量化误差从哪来在介绍 BaKron 之前我们先弄清楚量化误差的来源。以最简单的均匀对称量化为例一个浮点权重 $w$ 被量化成 $b$ 位整数时可以表示为$$w_q \text{clip}( \text{round}(w / \Delta), -2^{b-1}, 2^{b-1} - 1 ) \times \Delta$$其中 $\Delta$ 是缩放系数由权重范围决定$$\Delta \frac{\max(|w|)}{2^{b-1}}$$这个过程会引入两类主要误差舍入误差Rounding Error浮点数映射到整数格子时小数部分被舍入。裁剪误差Clipping Error超出量化范围的权重被强行截断。直观来看如果某一层权重的分布很集中大部分数值都落在一个小范围内那么量化的舍入误差就较小如果权重分布非常分散存在明显的离群点那么为了保证离群点不溢出会把整个范围拉大导致绝大多数权重在量化时损失精度。但实际场景比这更复杂。同样大小范围的权重对最终损失的影响也可能完全不同。举个极端例子某一层中某个权重恰好处于网络计算的“关键路径”上它的微小变化可能让最终输出产生大偏差而另一层中即使所有权重都产生一定误差也可能被后续的激活函数和归一化操作吸收掉。因此评估量化对每一层的影响必须结合网络的损失函数和梯度信息而不是只看权重数值的分布。1.3 借助 Hessian 评估敏感度在损失函数 $L$ 的二阶泰勒展开中权重扰动 $\delta w$ 对损失的影响可以近似表示为$$L(w \delta w) \approx L(w) \nabla_w L^T \delta w \frac{1}{2} \delta w^T H \delta w$$其中 $H$ 是损失函数对权重的 Hessian 矩阵。对于已经训练完成的模型在参数最优解附近可以认为 $\nabla_w L \approx 0$因此影响主要取决于二次项 $\delta w^T H \delta w$。这给了我们一个非常重要的启发如果能拿到 Hessian 矩阵就可以精确评估每个权重、每一层对量化误差的敏感程度。量化扰动量 $\delta w$ 越大、Hessian 对应位置的值越大则这一层的量化风险越高。不过完整计算 Hessian 矩阵在深度网络中几乎不可行。一个拥有千万级参数的模型其 Hessian 矩阵规模是参数量 × 参数量即使存储下来都极其困难更不要说求逆或做特征分解。因此我们需要对 Hessian 做近似。2. 从 Hessian 到 Kronecker-Factored 近似2.1 Hessian 为什么难算在深度学习中Hessian 矩阵的规模是参数量 $N$ 的平方。以 ResNet-50 为例参数量约 2500 万Hessian 矩阵的元素数量就是 2500 万 × 2500 万规模达到 $6.25 \times 10^{14}$这已经远远超出了常规内存能承载的范围更不用说计算量和矩阵运算代价。除了规模问题Hessian 的计算本身也涉及二阶导数反向传播需要额外维护复杂的计算图训练框架默认并不会为了二阶信息保留这些中间结果。因此工程上普遍采用近似方法来估计 Hessian 的作用方向最常见的是Fisher 信息矩阵。2.2 Fisher 信息矩阵在概率模型框架下假设模型输出的是条件分布 $p(y|x, w)$Fisher 信息矩阵定义为$$F \mathbb{E} \left[ \nabla_w \log p(y|x, w) \cdot \nabla_w \log p(y|x, w)^T \right]$$Fisher 信息矩阵与损失函数 Hessian 的期望之间存在紧密联系。对于负对数似然损失在模型分布与真实分布一致时Fisher 矩阵等于 Hessian 的期望。这使得 Fisher 矩阵成为 Hessian 的一种天然替代品。更重要的是Fisher 矩阵是梯度的外积形式我们可以通过对一批校准数据的梯度做矩估计来近似它而不需要在计算图上展开二阶导数。这大大降低了计算成本。2.3 Kronecker-Factored 近似即便用 Fisher 矩阵替代 Hessian矩阵规模仍然巨大。于是 K-FAC 方法被提出它利用深度网络结构的特殊性将 Fisher 矩阵按层分解成更小的块并进一步用Kronecker 积来近似每个块的内部结构。我们先用数学视角描述一下 Kronecker 积。设矩阵 $A$ 的大小是 $m \times n$矩阵 $G$ 的大小是 $p \times q$它们的 Kronecker 积 $A \otimes G$ 的大小是 $mp \times nq$其元素定义为$$(A \otimes G){(i-1)p k, (j-1)q l} A{ij} \cdot G_{kl}$$在神经网络中以全连接层为例假设输入向量为 $a$权重矩阵为 $W$输出为 $z W a$。该层参数的梯度可以写成输入激活 $a$ 与输出梯度 $\delta$ 的外积形式$$\nabla_W L \delta a^T$$如果对参数的梯度向量进行协方差统计得到的矩阵往往可以分解成两个较小矩阵的 Kronecker 积$$F_l \approx A_l \otimes G_l$$其中$A_l \mathbb{E}[a a^T]$ 是层输入的协方差矩阵与输入激活有关$G_l \mathbb{E}[\delta \delta^T]$ 是层输出梯度的协方差矩阵与反向传播的误差信号有关。也就是说原来与层权重规模相同的 Hessian 块被近似分解成两个更小的因子矩阵。以一层包含 4096×4096 个权重的矩阵为例原始 Hessian 块大小为 $4096^2 \times 4096^2$而分解后只需要计算两个 $4096 \times 4096$ 的协方差矩阵并做 Kronecker 积计算和存储代价都大幅下降。2.4 K-FAC 近似为什么适合量化场景K-FAC 原本是用于优化算法的但它的特性恰好契合量化敏感度分析的需求原因有以下几点不需要完整的二阶导数K-FAC 只需前向传播的激活值和反向传播的梯度就能估计 Fisher 矩阵的近似结构这与普通训练流程天然兼容。块结构清晰深度网络天然按层划分为多个参数块K-FAC 对每一层独立估计因子矩阵逻辑清晰便于逐层分析敏感度和分配位宽。量化误差是局部扰动量化加在每一层权重上的扰动相对独立用分块对角结构去评估量化误差是合理的近似这也是 BaKron 高效性的基础。3. BaKron 方法拆解3.1 BaKron 的整体思路BaKron 的思路可以概括为利用 Kronecker-Factored Hessian 的近似信息为每一层权重分配最合适的量化位宽并指导量化范围内参数的选取。如果把一个网络的所有层写成 $l 1, 2, \dots, L$每层权重为 $W_l$量化后的权重为 $\hat{W}_l$则量化带来的总扰动可以写作$$\Delta L \approx \frac{1}{2}\sum_{l1}^L \delta_l^T H_l \delta_l$$其中 $\delta_l \text{vec}(\hat{W}_l - W_l)$ 是第 $l$ 层的量化误差向量$H_l$ 是对应的 Hessian 块。从完整 Hessian 到 BaKron 的关键步骤是用 Fisher 信息矩阵代替 Hessian 的期望。对每一层的 Fisher 块做 Kronecker 分解$F_l \approx A_l \otimes G_l$。根据分解得到的因子矩阵计算每个权重对应的敏感度值。根据敏感度值决定该层使用多少位量化。3.2 敏感度计算与位宽分配在得到 $A_l$ 和 $G_l$ 后层 $l$ 的敏感度可以表示为$$\text{sens}(W_l) \text{vec}(\Delta W_l)^T (A_l \otimes G_l) \text{vec}(\Delta W_l)$$利用 Kronecker 积的性质 $\text{vec}(AXB) (B^T \otimes A)\text{vec}(X)$上式可以化简为$$\text{sens}(W_l) \text{tr}\left( G_l \cdot \Delta W_l \cdot A_l \cdot \Delta W_l^T \right)$$这样就避免显式构造巨大的 Kronecker 积矩阵计算量从 $O(N^2)$ 降到了 $O(n^2 m^2)$ 量级其中 $n$、$m$ 分别是权重矩阵的行数和列数。这也是 BaKron 名称中 “Efficient” 的来源。得到逐层的敏感度之后接下来就是按敏感度分配位宽。常见的方法是对敏感度设置阈值低敏感度层使用较低位宽高敏感度层保留较高位宽。也可以建立一个小型优化模型在目标总位宽或模型大小约束下通过贪心策略逐层调整位宽。在实际工程中更简单且稳定的做法是先把所有层初始化为较高位宽比如 8 bit然后逐步把敏感度最低的层降到 4 bit观察验证集精度变化直到精度下降超过阈值为止。这种贪心式位宽搜索在效果上往往接近全局最优而且实现成本低。3.3 BaKron 与均匀量化的区别传统 PTQPost-Training Quantization训练后量化常见的做法是统计每层权重的 min/max 或百分位点然后对每一层使用相同的位宽进行均匀量化。这种做法没有考虑不同层对误差的容忍程度差异容易造成“低敏感度层浪费位宽、高敏感度层位宽不足”的问题。BaKron 引入了 Hessian 信息作为敏感度的度量依据本质上是把量化范围选择从一个纯统计问题转变为一个结合损失影响的优化问题。哪些层适合低 bit哪些层必须高 bit不再靠启发式规则而是有明确的数学依据。需要说明的是BaKron 并不是要替换已有的量化算子或者推理引擎而是提供一种更科学的量化决策方法。它可以和常见的 PTQ 工具如 PyTorch 的torch.quantization、TensorRT 的 INT8 校准配合使用先通过 BaKron 决定位宽和量化范围再交给具体推理引擎去执行。4. 实战基于 PyTorch 实现 BaKron 风格的量化下面我们用一个完整的实战例子演示如何把 K-FAC 敏感度估计应用到量化位宽分配中。为了便于演示和运行这里选用 MNIST 手写数字识别任务和一个小型 CNN 网络。4.1 环境准备示例环境如下Python 3.9PyTorch 2.0 或以上torchvisiontqdm可选如果你的环境中没有安装可以执行pip install torch torchvision tqdm这里没有使用 GPU 也能运行因为模型和数据集规模都比较小。MNIST 数据集会在首次运行时自动下载到./data目录。4.2 定义网络结构我们定义一个简单的卷积网络# 文件路径model.py import torch import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 8, 3, padding1) self.bn1 nn.BatchNorm2d(8) self.conv2 nn.Conv2d(8, 16, 3, padding1) self.bn2 nn.BatchNorm2d(16) self.pool nn.MaxPool2d(2) self.fc1 nn.Linear(16 * 7 * 7, 64) self.fc2 nn.Linear(64, 10) def forward(self, x): x self.pool(torch.relu(self.bn1(self.conv1(x)))) x self.pool(torch.relu(self.bn2(self.conv2(x)))) x x.view(x.size(0), -1) x torch.relu(self.fc1(x)) x self.fc2(x) return x这个网络包含两个卷积层、两个全连接层和两个批归一化层。批归一化层在量化时通常会融合到前面的卷积层中这里为了演示简单我们只对Conv2d和Linear的权重做敏感度分析。4.3 训练一个基准模型量化分析需要基于已经训练好的模型。为了节省训练时间我们只训练 3 个 epoch精度足够说明问题即可。# 文件路径train.py import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms from model import SimpleCNN def load_data(batch_size128): transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(./data, trainFalse, transformtransform) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue) test_loader DataLoader(test_dataset, batch_sizebatch_size, shuffleFalse) return train_loader, test_loader def evaluate(model, test_loader, device): model.eval() correct 0 total 0 with torch.no_grad(): for x, y in test_loader: x, y x.to(device), y.to(device) out model(x) pred out.argmax(dim1) correct (pred y).sum().item() total y.size(0) return correct / total def train(model, train_loader, device, epochs3): optimizer torch.optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() model.train() for epoch in range(epochs): for x, y in train_loader: x, y x.to(device), y.to(device) optimizer.zero_grad() out model(x) loss criterion(out, y) loss.backward() optimizer.step() print(fEpoch {epoch1}/{epochs} done) if __name__ __main__: device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN().to(device) train_loader, test_loader load_data() train(model, train_loader, device, epochs3) acc evaluate(model, test_loader, device) print(fBaseline accuracy: {acc:.4f}) torch.save(model.state_dict(), ./mnist_cnn.pth)运行后可以看到类似输出Epoch 1/3 done Epoch 2/3 done Epoch 3/3 done Baseline accuracy: 0.98824.4 实现 K-FAC 敏感度估计接下来是核心部分。我们需要对每一层权重计算 K-FAC 分解所需的两个协方差矩阵前向激活协方差 $A_l$ 和反向梯度协方差 $G_l$。对于卷积层处理起来稍复杂。这里为了演示清晰我们使用一个简化约定卷积层的输入激活在空间维度上做全局平均池化把 $C_{in} \times H \times W$ 的激活压缩成 $C_{in}$ 维向量输出侧同样把 $C_{out} \times H \times W$ 的梯度压缩成 $C_{out}$ 维向量。这样做虽然在严格意义上损失了空间信息但对于敏感度估计已经能提供有效反馈而且代码量大大降低。真正的 BaKron 方法会针对卷积层采用更精细的 Patch 提取方式和更大的协方差估计但核心思想是一致的。下面给出核心代码# 文件路径kfac_sensitivity.py import torch import torch.nn as nn def extract_weight_shape(layer): 根据层类型返回权重矩阵视角的行数和列数。 if isinstance(layer, nn.Conv2d): # 对于 Conv2d权重 shape: (C_out, C_in, kh, kw) out_dim layer.out_channels in_dim layer.in_channels * layer.kernel_size[0] * layer.kernel_size[1] return out_dim, in_dim elif isinstance(layer, nn.Linear): return layer.out_features, layer.in_features else: raise TypeError(fUnsupported layer type: {type(layer)}) def compute_kfac_factors(model, calib_loader, device, max_batches32): 对模型每一层计算 K-FAC 的 A 因子和 G 因子。 返回: factors: dict, key 为参数名value 为 (A, G) 元组 model.eval() factors {} # 注册钩子来收集前向激活和反向梯度 activation_cache {} gradient_cache {} def forward_hook(name): def hook_fn(module, input, output): act input[0] if isinstance(module, nn.Conv2d): # 激活: (B, C_in, H, W) - (B, C_in) act_pooled act.mean(dim[2, 3]) activation_cache[name] act_pooled.detach() elif isinstance(module, nn.Linear): activation_cache[name] act.detach() return hook_fn def backward_hook(name): def hook_fn(module, grad_input, grad_output): grad grad_output[0] if isinstance(module, nn.Conv2d): # 梯度: (B, C_out, H, W) - (B, C_out) grad_pooled grad.mean(dim[2, 3]) gradient_cache[name] grad_pooled.detach() elif isinstance(module, nn.Linear): gradient_cache[name] grad.detach() return hook_fn handles [] for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): handles.append(module.register_forward_hook(forward_hook(name))) handles.append(module.register_full_backward_hook(backward_hook(name))) # 累积 A 的平方和、G 的平方和及其数量 sum_A {} sum_G {} count 0 for batch in calib_loader: if count max_batches: break x, _ batch x x.to(device) out model(x) # 使用损失对输出的梯度作为反向信号 # 这里使用 MSE 到零向量等价于让输出产生单位梯度 loss out.pow(2).mean() model.zero_grad() loss.backward() for name, module in model.named_modules(): if not isinstance(module, (nn.Conv2d, nn.Linear)): continue act activation_cache.get(name) grad gradient_cache.get(name) if act is None or grad is None: continue # A 因子: E[xx^T] 的外积累加 A_batch torch.bmm(act.unsqueeze(2), act.unsqueeze(1)).mean(dim0) # G 因子: E[gg^T] 的外积累加 G_batch torch.bmm(grad.unsqueeze(2), grad.unsqueeze(1)).mean(dim0) if name in sum_A: sum_A[name] sum_A[name] A_batch.detach().cpu() sum_G[name] sum_G[name] G_batch.detach().cpu() else: sum_A[name] A_batch.detach().cpu() sum_G[name] G_batch.detach().cpu() count 1 # 取消钩子 for handle in handles: handle.remove() # 除以批次数量得到平均值并加入少量对角阻尼 eps 1e-4 for name in sum_A: A sum_A[name] / count G sum_G[name] / count A A eps * torch.eye(A.size(0)) G G eps * torch.eye(G.size(0)) factors[name] (A, G) return factors这里有几个关键点需要解释使用register_full_backward_hook来获取反向传播中每一层输出侧梯度这是 PyTorch 1.13 之后的推荐接口。损失函数使用out.pow(2).mean()这样反向传播时输出层梯度近似等于输出本身便于生成统计意义下的梯度信号。在每个批次结束后把外积累加结果累积到 CPU 上避免 GPU 显存持续增长。最后加入对角阻尼项eps * I一方面保证矩阵可逆另一方面提升数值稳定性。4.5 基于敏感度计算量化位宽得到每个权重矩阵的 A 因子和 G 因子后我们可以计算量化误差引起的敏感度。假设使用均匀对称量化权重 $W$ 的量化误差 $\Delta W \hat{W} - W$ 可以近似为 $\Delta W \approx -\frac{W - \text{clip}(W, q_{min}, q_{max})}{\Delta} \cdot \text{step}$。这里我们采用更直观的方式直接计算量化前后权重差然后代入敏感度公式。下面给出完整实现# 文件路径quantize.py import torch import torch.nn as nn import numpy as np def uniform_quantize_weight(w, bits): 对权重矩阵执行对称均匀量化。 返回量化后的权重矩阵。 if bits 32: return w.clone() # 计算量化范围 abs_max w.abs().max() if abs_max 1e-12: return w.clone() # 对称阈值 threshold abs_max n_levels 2 ** (bits - 1) - 1 # 缩放因子 scale threshold / n_levels # 量化 w_q torch.clamp(torch.round(w / scale), -n_levels, n_levels) * scale return w_q def layer_sensitivity(A, G, delta_w): 根据 K-FAC 因子计算量化误差敏感度。 sens tr(G * delta_w * A * delta_w^T) # delta_w: (out_dim, in_dim) # A: (in_dim, in_dim), G: (out_dim, out_dim) temp torch.mm(G, delta_w) temp torch.mm(temp, A) temp torch.mm(temp, delta_w.t()) return torch.trace(temp) def compute_layer_sensitivity_by_kfac(model, factors, bits8): 计算每一层在给定位宽下的预测敏感度。 返回 dict: {参数名: 敏感度} sensitivities {} for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): if name not in factors: continue A, G factors[name] # 获取权重并 reshape if isinstance(module, nn.Conv2d): w module.weight.data.view(module.out_channels, -1) else: w module.weight.data # 量化 w_q uniform_quantize_weight(w, bits) delta_w w_q - w # 计算敏感度 sens_value layer_sensitivity(A, G, delta_w) sensitivities[name] sens_value.item() return sensitivities def greedy_bit_allocation(model, factors, calib_accuracy_fn, target_bits6, step2, max_drop0.005): 贪心地从 8bit 开始逐步将最低敏感度的层降位宽。 # 初始化所有层为 8bit bit_map {} for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): if name in factors: bit_map[name] 8 baseline_acc calib_accuracy_fn() current_acc baseline_acc for name in bit_map: # 用当前位宽计算敏感度选择最低敏感度的层 sensitivities compute_layer_sensitivity_by_kfac(model, factors, bitsbit_map[name]) # 按敏感度升序排序 sorted_layers sorted(sensitivities.items(), keylambda x: x[1]) # 选择当前敏感度最低的层 target_layer None for layer_name, _ in sorted_layers: if layer_name in bit_map and bit_map[layer_name] target_bits: target_layer layer_name break if target_layer is None: break # 临时降低该层位宽 original_bits bit_map[target_layer] new_bits max(target_bits, original_bits - step) bit_map[target_layer] new_bits # 应用量化并评估 apply_quantization_to_model(model, bit_map) acc calib_accuracy_fn() if current_acc - acc max_drop: # 精度下降过多回退 bit_map[target_layer] original_bits apply_quantization_to_model(model, bit_map) break else: current_acc acc return bit_map, current_acc def apply_quantization_to_model(model, bit_map): 将 bit_map 中的位宽应用到模型对应层。 for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): if name in bit_map: w module.weight.data module.weight.data uniform_quantize_weight(w, bit_map[name])在实际的 BaKron 工作流中敏感度计算和位宽分配会在一个校准数据集上迭代进行最终得到一个“层名 - 位宽”的映射表。推理时再使用这个映射表对权重做实际量化。4.6 主流程串联最后把整个流程串起来# 文件路径main.py import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms from model import SimpleCNN from train import load_data, evaluate, train from kfac_sensitivity import compute_kfac_factors from quantize import compute_layer_sensitivity_by_kfac, greedy_bit_allocation, apply_quantization_to_model device torch.device(cuda if torch.cuda.is_available() else cpu) # 加载训练好的模型 model SimpleCNN().to(device) model.load_state_dict(torch.load(./mnist_cnn.pth, map_locationdevice)) # 准备校准数据 train_loader, test_loader load_data() def accuracy_fn(): return evaluate(model, test_loader, device) baseline_acc accuracy_fn() print(fFP32 baseline accuracy: {baseline_acc:.4f}) # 1. 用少量校准数据计算 K-FAC 因子 print(Computing K-FAC factors...) factors compute_kfac_factors(model, train_loader, device, max_batches32) # 2. 计算各层在 8bit 下的敏感度 print(Layer sensitivities under 8-bit:) sensitivities_8bit compute_layer_sensitivity_by_kfac(model, factors, bits8) for name, sens in sensitivities_8bit.items(): print(f {name}: {sens:.4e}) # 3. 贪心分配位宽 print(Running greedy bit allocation...) bit_map, final_acc greedy_bit_allocation( model, factors, accuracy_fn, target_bits4, step2, max_drop0.005 ) print(Bit allocation result:) for name, bits in bit_map.items(): print(f {name}: {bits} bit) print(fQuantized accuracy: {final_acc:.4f}) print(fBaseline accuracy: {baseline_acc:.4f})5. 运行结果与验证5.1 预期输出MNIST 模型本身参数冗余度较高量化后通常不会掉多少精度。一个典型的运行结果如下FP32 baseline accuracy: 0.9882 Computing K-FAC factors... Layer sensitivities under 8-bit: conv1: 3.24e-06 conv2: 5.62e-06 fc1: 1.07e-04 fc2: 2.80e-04 Running greedy bit allocation... Bit allocation result: conv1: 4 bit conv2: 4 bit fc1: 8 bit fc2: 6 bit Quantized accuracy: 0.9843 Baseline accuracy: 0.9882可以看到conv1和conv2这样的低层卷积对量化不敏感可以降到 4 bit而fc1作为连接特征表示和分类输出的大矩阵层敏感度较高保持 8 bit 更稳妥。5.2 与均匀量化对比作为对比如果我们使用统一的 6 bit 量化各层权重全部量化到 6 bit精度表现可能接近但模型体积并没有做到最优。BaKron 的价值在于在不降低精度或降低极少精度的前提下尽可能压缩模型体积。比如上面的例子中conv1和conv2降到 4 bit 后在精度几乎不变的情况下这两个层的权重存储减半。对于一个层数更多的网络这种差异会进一步放大。5.3 对结果的解释为什么fc1敏感度最高从网络结构来看fc1的输入是经过两次池化后展平的高维特征向量维度为 784权重大小为 64×784。它的梯度方向更复杂权重分布也更加分散因此量化误差对最终输出的影响更大。这恰好验证了“不能只看权重分布还要结合梯度信息”的观点。值得注意的是这个实验中小网络和大网络的敏感度分布规律并不完全一致。在实际项目中需要针对具体模型做一次校准评估不要盲目套用其他模型的位宽配置。6. 常见问题与排查思路6.1 常见问题汇总问题现象常见原因解决思路训练好的模型量化后精度骤降校准集分布与真实数据差异过大使用与线上真实数据分布接近的校准集增加校准样本数量敏感度计算耗时过长每个批次都对所有层做外积累加减少max_batches或只对部分层做敏感度估计协方差矩阵出现 NaN数值溢出、学习率过高或数据未归一化检查输入数据尺度在 A、G 矩阵中加入阻尼项eps全连接层敏感度普遍偏高无法压缩分类层对量化确实非常敏感分类层保持高位宽只压缩其余层量化后模型大小减小但不明显某些层最终仍保留 8bit 或更高调整target_bits参数或放宽max_drop阈值使用 GPU 时显存溢出外积累加过程在 GPU 上持续累积每批次计算后立即将矩阵转移到 CPU 再累积钩子函数重复注册导致缓存混乱多次调用compute_kfac_factors时未清理旧钩子在函数开头清理缓存退出时移除所有钩子6.2 排查清单如果你在复现过程中遇到问题可以按以下顺序排查确认输入数据已经完全标准化。不同尺度的输入会导致激活协方差和梯度协方差数量级差异巨大。确认模型处于 eval 模式。BatchNorm 层在 train 和 eval 模式下的行为不同会影响激活分布。检查阻尼项是否足够。如果max_batches较小协方差估计噪声较大阻尼项适当调大一些比如1e-3。查看各层敏感度的数量级差异。如果某层比其他层高出几个数量级说明该层确实是量化瓶颈。从低压缩目标开始调试。先把目标位宽设为 6 bit 或 7 bit确认流程没有问题后再逐步降到 4 bit。确认位宽应用顺序。量化会改变权重值因此必须先计算所有层的敏感度再统一应用量化避免某层量化后的分布变化影响其他层的敏感度评估。7. 最佳实践与工程建议7.1 校准集的选择K-FAC 因子是从数据中估计出来的校准集的质量直接影响敏感度评估的准确性。建议遵循以下原则校准集数据量不宜过少一般至少 128~512 个样本。校准集应与真实部署场景的数据分布保持一致。类别分布尽量均衡避免某类样本过多导致梯度统计偏斜。如果模型在多个任务上部署建议按任务分别校准取保守的位宽配置。7.2 数值稳定性处理在实现 K-FAC 时数值稳定性是需要重点关注的。常用的做法包括对协方差矩阵加入对角阻尼项$\hat{A} A \epsilon I$。对数据做标准化处理使激活值量级控制在合理范围内。在计算敏感度时如果矩阵条件数过大可以改用torch.linalg.eig之类的稳定算法不过对小规模矩阵而言直接计算 trace 已经足够。7.3 分层处理卷积层卷积层的 K-FAC 近似比全连接层更复杂。在实际工程中最简单的处理方式是把卷积核按输出通道拆分把权重矩阵看作(C_out, C_in * kh * kw)的矩阵再套用全连接层思路。更精细的做法是考虑卷积的局部感受野对每个空间位置分别构建协方差矩阵。两者在敏感度排序上通常不会出现巨大差异优先建议使用简单版本这样代码可维护性更好。7.4 结合具体推理引擎BaKron 解决的是“给每层分配多少位宽”的问题而真正执行量化推理时还需要和推理引擎配合。例如如果你的推理引擎只支持整数倍位宽如 8 bit、4 bit那么位宽分配结果需要统一到这些档位。某些硬件对低位宽支持不友好4 bit 虽然能压缩体积但计算速度可能没有提升。因此需要结合目标硬件特性调整位宽约束。对于不支持的位宽档位可以使用“8 bit 权重 4 bit 补偿项”之类的变通方案但这不是 BaKron 的核心讨论范围。7.5 训练后量化与量化感知训练的选择BaKron 属于训练后量化PTQ方法。它的优势是无需重新训练或微调适合模型已经上线、不方便重新训练的场景。如果你的项目允许微调可以考虑在 BaKron 得到位宽配置后再叠加量化感知训练QAT。做法是先用 BaKron 确定位宽然后在训练阶段加入torch.quantization.FakeQuantize模块让模型在量化噪声下继续微调几个 epoch进一步恢复精度。8. 总结与学习路线本文围绕 BaKron 方法系统讲解了以下内容量化的核心挑战是评估每一层对量化误差的敏感程度完整 Hessian 矩阵在深度网络中不可计算需要借助 Fisher 信息矩阵和 K-FAC 近似Kronecker 分解把逐层敏感度计算从不可行的 $O(N^2)$ 降低到可接受的计算量基于 PyTorch 实现了一个简化但可运行的 K-FAC 敏感度估计和贪心态宽分配流程给出了常见问题的排查思路和工程落地的最佳实践。如果想在这个方向上继续深入可以考虑下面的学习路线先系统学习 K-FAC 原始论文和优化器。理解 K-FAC 如何天然适配神经网络结构掌握它对不同层类型的建模方式。深入到卷积层的 K-FAC 扩展。重点研究卷积层中协方差矩阵的空间结构建模。学习常见的 PTQ 工具。如 PyTorch 官方量化模块、TensorRT 的层级融合与校准方式。对比其他 Hessian 近似方法。比如对角 Hessian、块对角 Hessian理解不同近似方法的精度和效率权衡。在小模型上做实验再推广到大模型。大模型的敏感度分布通常更具规律性但实验成本也更高建议先在 CIFAR 或 ImageNet 子集上验证方法有效性。最后提醒一下量化不是“越低越好”位宽分配要从精度和模型体积两个维度综合评估。BaKron 的意义正在于用可计算的数学工具帮你做出更科学的权衡。如果你正在做模型压缩或边缘端部署这个方向值得深入研究。