HAMP-LIC:基于Hessian的混合精度训练后量化,破解图像压缩模型部署难题
HAMP-LIC 这类工作,本质上是在解决一个工程上非常现实的问题:学习型图像压缩模型在实验室里性能很好,但真要部署到手机、摄像头、服务器上时,模型体积和算力开销往往让人下不了手。直接量化又会掉精度,重新训练成本又高。HAMP-LIC 的核心思路是把 Hessian 信息引入训练后量化,通过计算每一层对量化误差的敏感度,给不同模块分配不同位宽,从而在不重训的情况下尽量保住率失真性能。如果你正在做神经图像压缩的落地、模型压缩,或者只是想知道“训练后量化遇到精度崩塌时该怎么分析”,这篇文章值得花几分钟看完。下面我按论文思路结合常见工程实践,把这个方案的原理、流程、可复现性和坑点拆开讲。
1. 先看清楚 HAMP-LIC 解决的是学习型图像压缩落地的哪个卡点
1.1 学习型图像压缩从学术 Demo 到产品部署差在哪
学习型图像压缩(Learned Image Compression)在公开数据集上的率失真性能已经做得相当好,结构上通常包含编码器、解码器、超先验网络、上下文模型和熵模型,典型如超先验架构、自回归上下文模型、可加性增强等变体。这类网络在 GPU 上推理时效果不错,但放到实际产品里就暴露几个问题:
第一是模型参数量不小。一个完整的 LIC 模型往往有几十到上百兆参数,对端侧存储和内存压力很大。第二是计算量不低。编码器和解码器内部有大量卷积、Masked Convolution、注意力模块,某些环节还有循环或逐块处理,CPU 上跑得很慢。第三是部署设备对权重精度有硬限制。很多 NPU、DSP、FPGA 平台对单层权重位宽有限制,或者低比特推理收益远高于高比特。
这时量化就变成一个绕不开的工程步骤。量化并不是新方法,但在 LIC 里做量化,比在分类网络上做量化更容易出现性能波动。原因是图像压缩的输出是一整张重建图像,任何一层量化误差都可能被后续网络放大,最后表现为视觉伪影、块效应、色偏或 RD 曲线明显变差。
1.2 为什么要选训练后量化,而不是量化感知训练
量化感知训练(QAT)能在量化后把精度拉回来,但要重新训练模型,对图像压缩网络来说尤其昂贵。LIC 模型的训练成本本身就比分类模型高,而且涉及码率估计、熵编码和失真项的联合优化,重训一遍需要大量算力和调参时间。对于很多团队来说,更希望拿到一个已经训好的预训练模型后,只做轻量的后端处理就完成部署。
Post-Training Quantization(训练后量化)的优势正好在这里:不需要重训,只需要少量校准数据,算一下每一层的量化参数,再做推理验证。缺点是量化误差难控制。HAMP-LIC 的核心方向,就是把训练后量化从“一刀切统一位宽”拉到一个更精细的层面:先分析每一层对量化误差的敏感度,再按敏感度分配不同的位宽。这个思路本身并不难理解,难点在于如何准确评估敏感度。论文选择使用 Hessian 信息来做这件事,所以叫 Hessian-Aware Mixed-Precision Post-Training Quantization。
注意:如果是新项目从零开始训练,QAT 仍然值得考虑;但如果你需要快速部署已有模型,PTQ 路线是更务实的选择。HAMP-LIC 属于后者。
2. 直接量化容易崩,核心原因是不同层对误差的敏感度差异极大
2.1 图像压缩网络里的数值敏感性来自哪里
在分类网络里,量化的主要影响是最终分类精度,网络结构对误差有一定抑制能力。但在图像压缩网络里,问题更复杂。
第一层量化误差会直接影响编码器提取的隐表示,也就是送入熵模型和上下文模型的中间特征。中间特征的分布直接影响码率估计。如果量化改变了特征分布,码率估计就会偏离实际,导致熵编码的压缩效率下降。更麻烦的是,解码器端会把这些被量化的特征重新映射成图像像素,如果误差在空间维度上聚集,就会出现明显的视觉瑕疵。
不同模块对量化的容忍度完全不一样。有些卷积层权重非常集中,小位宽也能保持输出分布;有些层的权重范围大、奇异值分散,量化后直接就破坏了输出分布;某些 Masked Convolution 和解码器尾部层,对重建图像的敏感度极高,一旦量化就容易造成色块或纹路异常。统一位宽方案为了保住最敏感的那一层,只能把全局位宽设得偏高,但这会让模型体积和计算量都降不下来。
2.2 Hessian 信息为什么能成为敏感度判断的依据
量化一个层之后,损失变化可以用泰勒展开来近似。假设某个层的权重是 w,量化扰动是 Δw,那么损失变化大致可以写成:
ΔL ≈ gᵀ Δw + 1/2 Δwᵀ H Δw
在训练完成的收敛点附近,梯度 g 通常很小,所以第一项可以忽略,主要影响来自第二项,也就是 Hession 矩阵 H 决定的二次项。H 越大,在同样大小的 Δw 下,损失变化就越剧烈,说明这层越“敏感”。
Hessian 矩阵是一个二阶导数矩阵,维度等于参数数量。实际计算完整 Hessian 几乎不可能,图像压缩网络哪怕一层也经常有几十万参数,直接用全矩阵会直接内存爆炸。所以实际做法是用 Hessian 的近似形式,比如对角 Hessian、Kronecker-Factored 近似、Hutchinson 无偏迹估计或使用少量校准样本的损失对每个权重求二阶导数。HAMP-LIC 这一类工作通常会在近似精度和计算成本之间取舍,使用的就是这种基于二阶信息的敏感度评估。
理解了这一点,你对论文的核心贡献会有一个很清楚的把握:不是发明了量化本身,而是把 Hessian 信息引入到 LIC 的混合精度位宽分配里,让每一层都能得到“相对合理”的位宽,而不是靠人工经验去试。
2.3 常见误区:敏感度不等于层大小,也不等于输出分布方差
有人会以为参数越多或特征分布越分散的层越敏感,实际不是这样。敏感度必须结合任务损失来看。比如某个层参数很多,但输出被后续的归一化或截断操作限制住了范围,量化它对整体率失真损失可能并不大;反之,某个小卷积层卡在关键路径上,权重幅度很小,量化噪声却会直接扩散到解码器多个阶段,反而非常敏感。
这也是为什么直接看权重直方图来决定位宽是不够的。Hessian 信息更接近“当前任务在当前模型状态下,这个参数变化对最终损失的影响程度”。实际测试时,你会发现有些层用 2 bit 或 4 bit 也能稳住,有些层必须给 8 bit 甚至保留浮点。
3. 混合精度位宽分配:把总比特预算花在最值得保护的层上
3.1 统一位宽为什么不足够好
如果所有层都用同一个位宽,比如 8 bit,那敏感层和非敏感层没有区别,模型压缩比例上不去。如果全部压到 4 bit,非敏感层也许没问题,但敏感层的误差会被放大,最终重建图像和码率都会恶化。
混合精度的目标是:在总体比特预算不变的前提下,尽量把低位宽分配给不敏感层,把高位宽留给敏感层。这样既控制了模型体积,又保住了率失真性能。可以把这理解成资源分配问题:总预算是一个定值,每一层需要“花费”多少比特,取决于它对率失真损失的影响程度。
3.2 HAMP-LIC 的位宽分配可以抽象成什么形式
从方法学上看,HAMP-LIC 的位宽分配大致可以做这样的抽象:
- 假设模型有 N 个层或 N 个量化单元。
- 每个单元可以选择若干候选位宽,比如 2、4、6、8 bit。
- 每个候选位宽对应一个模型大小或计算量开销。
- 每个单元在位宽 b 下的敏感度,用 Hessian 信息计算出的损失变化量来衡量。
- 最终需要在总比特数不超过预算的条件下,最小化整体量化损失。
这个组合优化问题如果穷举,选择空间会非常大,因为 N 可能几十上百,候选位宽每多一个,组合数就会爆炸。所以常见做法是用贪心策略、动态规划或者一些启发式搜索来逼近最优解,而不是逐层全搜索。
需要注意,HAMP-LIC 的位宽分配对象不是只有主编码解码网络,还应该包括超先验网络、熵参数网络、上下文特征提取等模块。如果你复现时只量化主干卷积,忽略超先验,效果往往不理想,因为超先验网络输出的是一组边信息,它的失真会被传递到主网络的码率估计和重建过程里。
3.3 衡量位宽分配好不好,看什么指标
衡量方案好坏的常用工具有三个:
- 模型体积或平均位宽,直接看压缩后的模型大小。
- RD 曲线,也叫率失真曲线,就是码率与重建质量的关系曲线。量化后的曲线越贴近原始全精度曲线,说明量化损失越小。
- BD-Rate,用来表示在相同质量下码率增加多少百分比,负值说明码率还降低了,正值说明码率升高。量化后 BD-Rate 增加越小越好。
我在实际对比中一般会同时保留全精度模型、统一 8 bit 模型和混合精度模型三组结果,先看统一 8 bit 能不能跑通,再看混合精度是否比 8 bit 进一步压缩而不显著增加 BD-Rate。这样比单独看一组数据更容易判断方案价值。
4. 论文方案的流程拆解:从预训练模型到量化模型大致怎么走
4.1 整体流程大致分为四段
虽然原始论文的具体实现细节没有放在输入材料里,但按照这类 Hessian-Aware 混合精度 PTQ 方法的通用做法,流程可以拆成四段,便于理解和复现:
- 准备预训练模型和校准数据。
- 计算每一层的敏感度。这里会用到 Hessian 信息的近似值。
- 根据敏感度和比特预算,做位宽分配。
- 对量化后的模型做校准、对齐或短时微调,然后用测试集评估。
这个过程不需要对编码器、解码器做大规模重训,所以整体成本比 QAT 低很多。但“训练后量化不重训”不代表什么都不做,量化缩放因子、截断范围、通道对齐这些细节仍然直接影响最终效果。
4.2 校准数据的选择:校准集太小,Hessian 信息就不稳定
Hessian 敏感度是用校准数据算出来的,校准集的分布和数量直接影响位宽分配结果。实操中我一般会用训练集或验证集中随机抽取的一定数量图像,而不是直接用测试集。原因很简单:避免把测试集统计信息泄漏到量化决策里。如果你用测试集去算敏感度,再在测试集上汇报指标,结果会偏乐观。
校准图数量需要根据网络复杂程度调整。常见的范围是几十到几百张图。太少会把噪声当成真信号,分配结果不稳定;太多计算成本高了,但敏感度估计会稳一些。如果你做实验时发现两次抽样算出来的位宽分配差异很大,说明校准集数量不够,或者图像内容多样性不足,应该加量或换更分散的样本。
输入端的大小也会影响 Hessian 计算。LIC 模型通常需要把图切成块或者缩放到固定分辨率,用不同分辨率跑出来的敏感度不太一样。建议在校准时先固定一个和测试集接近的输入分辨率,或者使用多样性更强的批处理,避免训练和测试条件不一致导致位宽分配失真。
4.3 Hessian 计算会遇到哪些工程障碍
计算 Hessian 时最常见的三个工程问题是内存、时间和数值稳定性。
内存方面,Hessian 向量积是实际可行的一种方式,不需要把完整矩阵存下来,只需要在反向传播过程中计算 Hv。如果你的层参数是几百万维,完整 Hessian 矩阵是不可行的。使用 Hessian 对角近似或块对角近似也可以大幅降低内存开销,但会损失一些相关性信息。HAMP-LIC 这类工作通常会在近似方式上做实验,我建议你在复现时先看官方代码用的是什么近似,不要盲目追求精确 Hessian。
时间方面,每层都算一遍 Hessian 相关统计,在超大网络上会非常耗时。更稳妥的做法是先跑一个小子集,比如 8 到 16 张图,确认耗时和内存符合预期,再逐步加数据量。
数值稳定性方面,Hessian 的特征值可能跨度很大,计算过程中容易出现 NaN 或异常值。如果发现某层的 Hessian 统计量异常大或异常小,不要急着把它作为最终依据,先检查该层的梯度计算、激活值范围和输入图像是否有异常像素值。
4.4 位宽搜索的约束条件和工程落点
位宽搜索时,通常要设置一个总比特预算,比如目标是原模型大小的 1/4 或 1/8。注意这个预算不只是数学上成立就行,还要看部署平台对位宽的支持情况。很多 NPU 只支持 4、8、16 bit 或对应的混合配置,不支持 6 bit;这种情况下你把 6 bit 分配给敏感层,实际部署时只能向上取整到 8 bit,等于预算白算。
部署层面还要考虑每层算子是否真正按你分配的位宽执行。有些框架对卷积、全连接、矩阵乘的量化支持程度不同,某些层权重是 4 bit,但激活必须 8 bit,最终整体收益会打折。所以论文里看着漂亮的压缩比例,落地时还要结合平台算子库再看一次。
5. 如果要在自己的压缩模型上复现或验证这套思路,可以这样安排
5.1 环境准备:并不要把要求想得太高
这类实验通常用 PyTorch 就能跑。建议准备一张 12 GB 以上的 GPU,因为 LIC 模型本身推理时需要把整张图像经过编码器和解码器,显存需求比普通分类网络高。校准集中如果包含较大分辨率图像,要先用裁剪或缩放控制输入尺寸。
依赖方面主要就是 PyTorch 和基本的图像处理库。原始论文是否有配套代码,要以你查到的版本为准;如果暂时没有,你可以先在自己熟悉的 LIC 仓库和训练后量化工具上跑通一个简化版本。关键是先把“能不能正确完成一次量化、解码并测出 BD-Rate”这件事跑通,再考虑完整复现论文实验。
5.2 最小验证流程建议
我一般会把流程拆成下面几步,每步确认无误后再进下一步:
- 加载预训练 LIC 模型,跑若干张验证图,得到全精度 RD 点。
- 直接做统一 8 bit 量化,重跑验证图,对比重建图像和 RD 点。
- 实现或接入一个 Hessian 敏感度估计模块,抽样校准图算出每层敏感度。
- 在总预算限制下生成混合精度位宽配置。
- 按照位宽配置逐层量化,再跑验证图和指标。
- 对比全精度、统一 8 bit、混合精度三组结果。
从最短路径看,第一步和第二步通常一两天内能完成。第三步到第五步是实现论文核心的时间大头,会涉及不少代码细节。建议先把第六步的测试脚本写好,后面每次改位宽配置都能快速看到效果,而不是反复手工导出图像再算指标。
5.3 实验表格怎么设计才有效
对比实验建议至少包含这几列:
| 模型配置 | 平均位宽 | 模型大小 | 码率 | PSNR | MS-SSIM | 备注 |
|---|---|---|---|---|---|---|
| 全精度 | 32 bit | 基准 | 基准 | 基准 | 基准 | 无 |
| 统一 8 bit | 8 bit | 明显缩小 | 略增或持平 | 略降 | 略降 | 作为稳定基线 |
| 统一 4 bit | 4 bit | 显著缩小 | 可能大增 | 明显下降 | 明显下降 | 观察最坏情况 |
| HAMP-LIC 混合精度 | 按配置 | 接近统一 4 bit | 略增 | 尽量接近全精度 | 尽量接近全精度 | 核心对比 |
不需要把所有码率点的图像都贴出来,但至少选一两张有代表性的重建图,方便看视觉差异。指标只看 PSNR 容易忽略码率变化,所以 BD-Rate 更需要优先计算。如果你身边工具没有现成 BD-Rate 脚本,可以用编码面积比较方法,也可以手动实现常用版本。
5.4 怎么判断结果算好还是算坏
判断标准不是“有没有掉点”,而是“掉多少点、换回多少压缩收益”。理想情况下,混合精度模型和全精度模型之间 BD-Rate 差距在可接受范围内,比如几个百分点以内,同时模型大小显著小于统一 8 bit。如果混合精度几乎没有比统一 8 bit 小多少,或者精度变化很明显,就需要回头检查敏感度估计和位宽分配过程是否合理。
如果量化后 RD 曲线出现拐点异常,比如低码率点变好、高码率点变差,这种非单调变化通常说明某个敏感模块在高码率分支下被过度量化了,需要针对该码率分支单独调整位宽。
6. 复现中的高频问题排查:按这个顺序查更快
6.1 量化后 BD-Rate 上升严重,先别调位宽搜索,先查这两点
第一是校准集和验证集是否存在分布偏移。我用过很多次,最经常出现的问题是校准图分辨率偏小,而验证图分辨率偏大,导致敏感度估计偏低;模型在视觉质量要求更高的区域提前被压坏。第二是激活层是否也做了量化。有些工具默认只量化权重,激活留在浮点;还有的混合精度方案只约束了权重位宽,但激活位宽仍由框架统一决定,这样也会导致实际效果和论文不一致。
检查方法很简单:把位宽配置输出逐层打出来,对照模型结构图看哪些层被分到了低比特。如果发现超先验或熵模型相关层大面积是低比特,而主干解码层大多是高比特,也不一定错;但如果解码器尾部低比特层过多,重建图像出现大块异常,那就要优先降低这些层的压缩程度。
6.2 Hessian 计算时 OOM 或速度过慢
先减少输入批次大小,再把用于估计的采样数量降低。不要一开始上 64 张全分辨率图,先用 4 到 8 张缩略图把代码跑通。如果是反传过程中梯度保存过多导致 OOM,可以在计算每层 Hessian 相关量时用 torch.no_grad 或按层分离,逐步释放中间变量。还可以用更轻量的 Hessian 近似,例如基于梯度平方的 Fisher 近似,虽然和严格 Hessian 有差异,但有时候已经足以给出稳定的敏感度排序。
6.3 混合精度位宽搜索的组合空间太大
如果候选位宽集合是 {2,4,6,8},层数 50,穷举完全不现实。常见做法是先做统一 8 bit 量化,再用贪心策略从 8 bit 逐层降到 4 bit,每次选择掉点最少的那一层继续降。这样每一轮只需要评估 N 种候选,一共最多跑 N 步,可执行性高很多。贪心不一定最优,但对工程任务通常已经够用。如果你希望更快,可以先用更粗的敏感度指标预筛,比如按 Hessian 迹求和排序,再对候选集合做精确评估。
6.4 量化后模型部署到实际平台,速度提升不明显
先确认实际推理框架有没有按你给的位宽执行。很多框架需要在转换阶段显式指定层位宽,或者只在特定后端算子下才支持混合精度加速。如果你的硬件不支持某位宽,框架会静默把该层还原成 FLOAT16 或 INT8,最终模型体积可能小了,但推理速度没有变化。另一个要确认的是内存带宽瓶颈,某些层参数量不大但数据传输频繁,降低位宽也只能带来少量收益,这时候要结合层耗时分析看瓶颈在哪。
提示:如果你的目标是端侧部署,最好在一开始就把目标框架的算子列表拿过来,看它支持哪些位宽和哪些算子,然后反向约束你的位宽候选集。这样论文里算出来的总比特预算,才真正能落到产品里。
7. 几点实践建议,以及这个方向后续还能怎么做
7.1 先把单模型量化流程调稳定,再扩大验证范围
不要一上来就在多个数据集、多个码率点上全面测试。先挑一个中低码率模型,跑通统一的量化管线,确认指标能复现、日志完整、位宽分配可解释,然后再推广到更多码率点和更多模型架构。否则一旦出现量化后效果波动,你很难判断是敏感度计算的问题、位宽分配的问题,还是测试脚本的问题。
7.2 训练后量化不是终点,配合短时微调效果更好
虽然叫 Post-Training Quantization,但实际工程里可以在量化之后加一个非常短时的微调,比如冻结大部分层,只更新量化缩放因子或者部分解码器层。这种做法的训练成本远低于完整 QAT,但能把重建质量拉回不少。HAMP-LIC 作为纯 PTQ 方案,优点是快速,缺点是性能上限受限于量化本身;如果你想追求更高精度,可以在它给出的位宽分配基础上,再做浅层微调,这样相当于把两种方法的优势组合起来。
7.3 对后续工作可以关注什么
从技术方向看,Hessian-Aware 思路还能延伸到激活量化、通道粒度量化、在线自适应量化等角度。图像压缩网络的特征图分布随时间、内容变化较大,静态 PTQ 很难在所有内容上都做到最优,所以加入少量在线校准指令或者元学习位宽分配,可能会成为下一步探索方向。
对工程团队来说,更值得关注的不是某个论文指标的绝对值,而是它的方法能不能落到你们自己的模型和算子库上。如果你们平台只支持统一的 8 bit,那混合精度再强也无法落地;如果平台支持多种位宽,那 HAMP-LIC 这类方案就非常值得引入,它能帮你把总比特预算花到刀刃上。
从我个人经验看,这类训练后量化方法最核心的收益不是“省了多少比特”这一个数字,而是它提供了一套用数据和损失函数说话的分析机制。有了 Hessian 敏感度之后,你不再需要靠猜来决定哪层保护、哪层压缩,而是能给出可以被讨论和验证的排序。这个思维方式的转变,比某次实验多压缩 10% 模型大小更有长期价值。
如果你正在复现 HAMP-LIC 或者类似方案,我建议先跑通最小流程,确认敏感度和位宽分配能稳定复现,再去看最终 RD 指标。这个过程中遇到问题不要急着换网络结构或重训模型,优先检查校准数据、量化范围、位宽候选集和部署后端的一致性。踩过几个坑之后你会发现,真正限制量化的往往不是算法本身,而是那些藏在实验设置里的细节。
