做学习型图像压缩Learned Image Compression, LIC模型部署时大家很容易被一个问题卡住模型效果很好但参数量大、计算量大直接搬到移动端或边缘设备上跑不动。这里最常用的手段就是量化但量化一做狠了重建图像又明显变差。最近读到一篇很有意思的工作专门针对 LIC 模型做后训练量化方法名是 HAMP-LIC核心思路是“看海森矩阵来决定每一层用多少位精度”。本文会对这篇论文的方法体系做一次完整拆解先解释背景和量化为什么难做再讲 Hessian 感知的思路然后按流程解构 HAMP-LIC 的混合精度分配方式最后给出一份便于理解和复现的实现思路与踩坑清单。适合正在做模型压缩、端侧部署或者想深入了解混合精度后训练量化的同学阅读。1. 背景与问题为什么 LIC 模型部署需要量化1.1 学习型图像压缩模型的基本结构学习型图像压缩不是一个单一网络而是一套完整系统。主流方案通常包含自编码器结构编码器把图像映射成隐空间特征量化器把连续特征变成离散符号熵模型负责估计每个离散符号的概率分布进而通过算术编码器转成真正的二进制码流。解码端再通过解码器重建图像。这个结构里推理路径上的计算主要是卷积、非线性激活和熵参数估计。看上去并不复杂但实际部署时有两个非常现实的问题模型体积较大编码器、解码器、超先验网络和上下文模型叠加起来参数往往有几十 MB 甚至上百 MB。推理速度不理想隐空间特征通道数多逐像素或者逐块做熵估计时上下文模型的串行依赖会拖慢解码速度。所以在移动端、实时通信或嵌入式设备上直接加载全精度模型很难满足内存和延迟约束。模型量化因此成了 LIC 部署绕不开的一步。1.2 量化带来的精度损失问题量化可以简单理解为将 FP32 浮点权重和激活值映射到低比特整数比如 INT8、INT4 甚至 INT2。对于普通分类模型INT8 量化已经很成熟精度损失通常很小。但图像压缩模型不一样。压缩模型的目标是重建图像的感知质量和码率整个系统对量化噪声极为敏感。尤其是隐空间特征经过了严格的率失真优化特征分布本身高度非均匀稍微损失一点精度可能就让重建图像出现明显伪影。如果直接把所有层都压到同一比特宽度比如一刀切 INT8某些敏感层就会成为瓶颈图像整体质量被拖垮。这个问题促使研究者开始关注“哪些层敏感、哪些层不敏感”的细粒度量化。1.3 为什么混合精度是必然方向混合精度量化的基本思想是不同层或者不同通道对量化的敏感度不同因此应该分配不同的位宽。敏感度高的层保留较高精度敏感度低的层可以用很低的位宽甚至二值化。这个思路本身不新鲜关键是“如何快速、准确地找到每一层的敏感度”。传统做法是逐层尝试量化观察损失曲线变化但这在 LIC 模型上行不通因为压缩模型的损失函数包含重建损失和码率估计两部分层的量化误差会同时影响到像素空间和码流空间单纯试错的时间成本很高。HAMP-LIC 的关键贡献就是引入 Hessian 矩阵来刻画这种敏感度让混合精度位宽分配有一个理论支撑而不是凭经验或者穷举。2. 核心概念PTQ 与混合精度量化2.1 后训练量化 PTQ 是什么后训练量化Post-Training Quantization, PTQ指的是已经有一个训练好的全精度模型我们在不重新训练、或者只做极少训练的前提下把模型权重和激活值量化到低比特精度。它和量化感知训练QAT的主要区别如下对比项PTQQAT是否需要完整训练不需要只用少量校准数据需要重新训练或微调训练开销很低较高适用范围追求快速部署的场景对精度要求高且时间预算充足对敏感度估计要求更高因为可调空间有限低通过训练自适应恢复精度常见位宽INT8、混合 INT8/INT4INT8、INT4、更低比特LIC 模型通常网络结构复杂重新训练成本很高而且压缩模型的训练需要联合优化率失真损失QAT 改起来比较麻烦。因此 PTQ 是更实用的部署方案。但 PTQ 的难点在于不能通过大量训练迭代来恢复精度必须在一开始就把量化误差控制在可接受范围内这就对位宽分配策略提出了更高要求。2.2 均匀量化的局限最简单的 PTQ 是均匀量化即所有层都使用同一位宽例如全部 INT8。均匀量化实现简单但存在明显的资源浪费和精度瓶颈如果某层对量化极不敏感给它 INT8 就浪费了位数完全可以降到 INT4 来减小内存带宽。如果某层对量化极敏感即使给了 INT8可能仍然不够需要 INT16 或 FP16 混合使用。这个现象在 LIC 中尤为明显因为编码器的浅层特征和深层特征的能量分布差异很大熵模型中的参数又直接决定码率估计精度。用统一位宽本质上是在用最高敏感度的层决定整模型的最低精度效率很低。2.3 混合精度的求解目标混合精度量化的目标是在给定平均位宽约束或模型大小约束的前提下最小化量化后模型与全精度模型之间的输出偏差。形式化地看需要找到一个位宽分配策略使得总比特数不超过预算法则量化后模型的损失函数变化尽可能小尽量保留重建图像的主观质量。搜索这样一个最优分配是一个组合优化问题。如果模型有 100 层每层候选位宽有 4 种组合空间就是 4 的 100 次方不可能穷举。所以必须依靠敏感度指标来引导搜索。HAMP-LIC 采用的敏感度指标正是基于 Hessian 矩阵的损失变化近似。3. Hessian 矩阵量化敏感度的“度量尺”3.1 从泰勒展开理解量化误差影响假设模型损失函数为 L(θ)θ 表示模型权重。量化可以看作在权重上叠加了一个扰动 Δθ。那么量化后的损失变化可以近似用泰勒展开表示L(θ Δθ) ≈ L(θ) gᵀΔθ (1/2)ΔθᵀHΔθ其中 g 是梯度H 是损失函数对权重的二阶偏导矩阵也就是 Hessian 矩阵。对于已经收敛的模型梯度 g 通常接近零所以一阶项影响很小主导损失变化的是二阶项 (1/2)ΔθᵀHΔθ。这意味着如果某个参数对应的 Hessian 值很大那么即使给它很小的扰动也会造成较大的损失变化如果 Hessian 值很小则可以放心地把它压到很低精度。这就是 Hessian-Aware 的核心思想用 Hessian 信息来估计“这一层能承受多少量化噪声”。3.2 为什么用 Hessian 而不是直接看权重范围很多人会问判断一层是否容易量化直接看权重分布范围不就行了吗权重范围大、分布分散的层量化误差不是更大吗这个观点有道理但不全面。权重分布只反映了参数本身的数值范围没有反映这一层在损失函数中的重要性。有的层权重范围很大但它对最终输出影响很小有的层权重范围很小却是整个模型的性能命脉。Hessian 矩阵则把“参数变化”与“损失变化”连接起来相当于给每个参数标注了“重要程度”所以更适合做敏感度分析。L-BFGS、Hessian 对角线、Hutchinson 估计等都是在试图高效地求解或近似 Hessian 信息避免显式构造 N×N 的庞大矩阵。这也是 HAMP-LIC 这类方法能够落地的原因之一。3.3 层级别还是通道级别Hessian 信息的粒度决定了位宽分配细度。层级别layer-wise意味着整层统一位宽实现简单硬件友好通道级别channel-wise可以更细腻地分配精度但会引入更复杂的索引与反量化逻辑。HAMP-LIC 这类方法通常先在层或模块级别做敏感度分析因为压缩模型中的每个残差块、注意力模块、熵参数模块对最终码率和失真影响不同。先粗粒度定位“最容易受伤”的模块再决定是否继续细分是工程上更稳妥的做法。4. HAMP-LIC 方法拆解4.1 总体流程HAMP-LICHessian-Aware Mixed-Precision Post-Training Quantization for Learned Image Compression不是一个简单的量化算法而是一整套流程。从名字可以看出它的三大关键词Hessian 感知、混合精度、后训练量化。大致流程可以分成四步加载一个已经训练好的全精度 LIC 模型引入校准集计算模型各模块对量化扰动的 Hessian 敏感度根据敏感度和位宽预算搜索每一层或每个模块的最佳位宽组合在目标位宽配置下执行实际量化并通过轻量校准来修正统计参数。4.2 敏感度计算与位宽分配论文的核心在于第二步和第三步。敏感度计算不是简单取 Hessian 对角线的绝对值而是结合量化误差的二阶近似来估计“如果这一层用 k bit 量化损失会增加多少”。这一步通常需要对不同位宽候选分别估算因为位宽越低量化误差 Δθ 越大二阶项的值也越大。得到每个候选位宽下的敏感度后混合精度分配就变成了一个资源分配问题在总位宽预算受限的前提下选择每一层的位宽使整体敏感度最小。由于组合空间大实际做法常用贪心算法或动态规划搜索。这里要专门强调HAMP-LIC 的损失函数不只是重建损失还包括码率估计损失。图像压缩模型训练时用的是率失真损失因此在 Hessian 计算时也要同时考虑重建误差和码率估计误差否则敏感性分析可能失真。4.3 量化与校准环节确定位宽配置后还不能直接完成量化。LIC 模型的某些特殊层例如熵模型中的概率参数层、最后的图像重建卷积层往往需要单独处理。论文的思路是保留关键层为较高位宽其余层尽量压到 4 bit 甚至更低。同时在 PTQ 过程中通常还需要对 BatchNorm 的统计量进行重新统计或者对偏置进行修正。虽然很多代码库中量化流程会自动处理这些但理解这一步骤有助于我们在复现时排查问题。5. 与相关量化方案对比为了更好理解 HAMP-LIC 的定位我们把它和几种常见量化方案放在一起看方案类型位宽策略敏感度依据适合场景均匀 INT8 量化全模型统一 8bit无快速部署精度要求一般均匀 INT4 量化全模型统一 4bit无极低带宽场景精度下降明显手工敏感层保留高精度少数层 8bit其余 4bit专家经验参数调优成本低但依赖经验Hessian 引导混合精度各层不同位宽二阶敏感度自动搜索、位宽预算受限HAMP-LICLIC 场景下 Hessian 引导混合精度率失真损失二阶近似学习型图像压缩模型部署HAMP-LIC 与其他通用混合精度 PTQ 的区别在于它专门面向 LIC 模型设计了敏感度度量方式把重建质量和码率同时纳入 Hessian 分析。这使得量化后的模型不仅能保持视觉质量也能尽可能稳定地保持码流大小。6. 实现思路与示例流程由于论文的完整代码和细节可能在公开仓库中更新这里提供一套基于论文思想的最小实现思路方便理解 Hessian-Aware 混合精度在 LIC 模型中如何落地。下面的伪代码基于 PyTorch 风格的 API只用于表达算法流程不直接对应论文原版代码使用时需要根据实际模型结构调整。6.1 准备校准集与模型import torch # 假设已有预训练 LIC 模型 # model load_pretrained_lic_model() # 准备一小批校准图像 # calibrate_loader build_calibrate_loader(num_samples64)校准集不需要很大但必须能代表真实场景常见做法是选择 64 到 256 张包含纹理、边缘、平坦区域分布合理的图像。校准集过大增加 Hessian 估计开销过小则敏感度估计不稳定。6.2 计算每层 Hessian 敏感度from torch.autograd.functional import hessian def compute_layer_sensitivity(model, calibrate_loader, layer_names): 基于 Hessian 对角线估计每层对量化扰动的敏感度。 实际实现一般使用 Hutchinson 估计避免显式构造完整 Hessian。 sensitivities {} for name, layer in model.named_modules(): if name not in layer_names: continue # 对每一层计算损失对层权重的 Hessian 近似值 # 方法1随机投影方式估计 Hessian 对角线 # 方法2计算损失对权重的二阶梯度取范数 # 这里只展示逻辑不包含完整实现 hessian_diag estimate_hessian_diagonal(model, layer, calibrate_loader) sensitivities[name] hessian_diag return sensitivities实际工程中很少直接调用torch.autograd.functional.hessian因为模型参数量大显存压力高。常用做法是Hutchinson 方法用随机向量 v计算 vᵀHv通过多次采样近似 Hessian 的迹或对角线分组估计只对目标层参数求二阶梯度而不是对整个模型求分块迭代逐层估计敏感度释放中间计算结果。6.3 混合精度位宽搜索def assign_bitwidth(sensitivities, candidate_bitwidths, budget): 基于敏感度做贪心位宽分配。 fixed_bits 用于强制保留高精度的层例如熵参数层。 assignments {} remaining_budget budget # 先给关键层分配最高精度 for layer, bits in fixed_bits.items(): assignments[layer] bits remaining_budget - bits * layer_size[layer] # 其余层按敏感度从高到低排序敏感度越高分配越高位宽 sorted_layers sorted( sensitivities.items(), keylambda x: x[1], reverseTrue ) for layer, sens in sorted_layers: if layer in assignments: continue # 选择在预算内可行的最高位宽 best_bits min(candidate_bitwidths) for bits in candidate_bitwidths: cost bits * layer_size[layer] if cost remaining_budget: best_bits bits else: break assignments[layer] best_bits remaining_budget - best_bits * layer_size[layer] return assignments贪心算法不一定能得到全局最优解但胜在效率高。如果预算约束比较严格可以考虑动态规划或遗传算法做更精细搜索但真实场景中贪心已经足够体现 Hessian 敏感度的排序价值。6.4 模拟量化与验证def quantize_layer(weight, bits): 对权重执行对称均匀量化。 这里只做模拟量化用于验证精度趋势。 scale weight.abs().max() / (2 ** (bits - 1) - 1) weight_q torch.round(weight / scale) * scale return weight_q完成位宽分配后可以用模拟量化fake quantize评估 PSNR、MS-SSIM、bpp 等指标。重点观察重建图像是否出现块效应码率估计是否明显偏移熵模型的概率输出是否出现极端值。如果某个指标异常优先怀疑对应模块的位宽是否过低。7. 常见问题与排查思路7.1 Hessian 计算导致显存溢出问题现象常见原因解决思路CUDA out of memory对全模型直接求 Hessian改用 Hutchinson 近似或按层分块计算训练速度极慢敏感度估计迭代次数过多减少随机投影次数使用更小校准集估计结果不稳定校准集单一或样本过少增加图片数量保证纹理多样性Hessian 相关计算确实有额外开销这也是它通常只用于离线 PTQ 的原因。在线推理阶段不会计算 Hessian所以部署端没有增加延迟。7.2 混合位宽配置后重建图像出现伪影如果分配完位宽后发现重建图像有局部伪影常见原因有三个没有对最后一层重建卷积保留高位宽导致输出层误差直接叠加到图像像素上熵参数层的量化粒度过粗导致概率估计偏差间接影响解码端反量化没有对量化后模型的偏置项做校正。排查顺序建议先逐模块做消融把某一层强制设为 FP16 或 INT16观察是哪个模块引起质量下降再决定是否调整位宽候选集。7.3 量化后 bpp 明显变化bppbits per pixel是压缩模型的关键指标。量化后 bpp 变化通常意味着熵模型的概率估计不准。建议把所有熵参数相关模块单独保留高精度并且校准集中要包含不同复杂度的图像低纹理图会让概率估计的敏感度被低估。7.4 校准集数量的选择校准集不是越大越好。一般实验里128 张图像已经足够稳定。如果继续增加样本Hessian 估计可能趋向于过于“平均”反而会掩盖某些极端敏感层。建议至少尝试 32、64、128 三档校准集观察敏感度排序是否发生剧烈变化。8. 工程实践与最佳实践建议8.1 按模块分组做敏感度分析LIC 模型通常由编码器、超先验编码器、上下文模型、熵参数网络和解码器组成。推荐的做法是先按大模块分组统计敏感度在敏感度较高的模块内部再按层分析对低敏感度模块直接使用较低位宽减少搜索空间。这样既能缩减搜索时间又能保证精度关键模块不会因为一刀切而被误伤。8.2 把码率损失纳入敏感度指标只关注重建损失来算 Hessian可能得到“码率膨胀但像素变化不大”的结果。压缩模型部署后码率和质量都要符合预期所以敏感度计算必须覆盖率失真损失。这也是 HAMP-LIC 这类方法最值得学习的地方它没有把压缩模型的损失函数简化成纯 MSE而是同时考虑像素域和码流域。8.3 保留关键层的固定位宽在混合精度搜索中可以直接把某些层从优化列表里剔除强制保留为 INT16 或 FP16。通常包括图像输入的第一个卷积层重建图像的最后一个卷积层熵模型中输出概率分布的参数层。这些层数量不多但对视觉质量和码流大小影响巨大。固定它们的高位宽可以显著降低搜索难度。8.4 使用矩阵分解等技术降低 Hessian 评估成本如果逐参数计算 Hessian 仍然太贵可以用松耦合方式估计敏感度用少量样本计算每个通道的量化误差传播量再用该传播量代替 Hessian 值做排序。虽然不是严格的二阶信息但在工程上往往能逼近类似效果。Hessian-Aware 的价值在于提供判断依据而不在于每一步都必须使用最严格数学定义。8.5 安全与生产环境注意事项在生产环境做 PTQ 部署时务必遵循最小授权和灰度发布策略量化模型导出前在验证集上跑一遍完整指标记录 PSNR、MS-SSIM、bpp上线时先对小流量用户灰度对比全精度模型和量化模型的在线指标保留全精度模型备份方便快速回滚不要在生产环境中直接执行未经校验的量化脚本尤其是涉及模型文件覆盖、删除或替换的操作一定要先备份。9. 总结与下一步学习建议本文围绕 HAMP-LIC 的核心思路梳理了 Hessian-Aware Mixed-Precision Post-Training Quantization 在 Learned Image Compression 场景下的关键脉络从 LIC 模型部署痛点出发解释了量化为什么难做引入了 Hessian 作为敏感度度量并拆解了混合精度位宽分配的基本流程。通过伪代码示例和常见问题排查基本可以还原论文的方法主链路。如果接下来想深入研究可以从三个方向入手阅读 Hessian 近似估计的经典论文重点理解 Hutchinson 方法和 K-FAC 近似这是实现 Hessian-Aware 量化的基础工具阅读 LIC 模型的原始文献熟悉超先验模型、上下文模型和熵参数估计的实现细节否则很难针对不同模块做合理的位宽固定动手做一轮对比实验在同一 LIC 模型上分别用均匀 INT8、均匀 INT4、Hessian 引导混合精度做 PTQ对比重建图像和 bpp 的变化。这个实验体量不大但对理解“敏感度分析”的实际价值非常有帮助。最后补充一点量化方案没有“银弹”。HAMP-LIC 提供了一种理性、可量化的思路但真正落地时仍要结合硬件支持的位宽种类、算子库实现和业务对图像质量的要求做平衡。建议从本文的伪代码开始先跑通一个最小流程再逐步替换为更严格的 Hessian 估计方法和更精细的搜索策略。