近几年的端到端图像压缩研究里,熵模型一直被当成“房间里的那头大象”——大家都站在它背上做改进,可很少有人正面回答一个问题:潜变量的概率分布到底应该由谁来估计,用什么信息估计,才能让码率逼近真实熵下限。
传统JPEG、HEVC乃至VVC,熵编码阶段都在用人工设计的概率表或上下文规则,几十年没有本质变化。而学习型方法之所以能超越这些标准,关键就是把“概率估计”也放进了网络里一起优化。但早期的端到端方案要么只用固定超先验(hyperprior),要么在此基础上加一层自回归上下文,信息来源单一,当图像内容复杂度波动大时,就会在大纹理区域出现严重的概率估计偏差,最终浪费大量比特。
这次这个“基于多信息参考的熵模型”思路,本质上就是不再依赖单一先验,而是把超先验、空间上下文、多尺度结构线索三路特征融进同一个概率估计框架里。适合三类人:一是正在做学习型图像压缩课题的研究生,想找一个比“超先验+上下文”更系统的改进方向;二是工业界做压缩落地算法的工程师,关心率失真性能和推理速度之间怎么平衡;三是刚入深度学习视觉方向的开发者,想看一个完整的编解码网络是怎么从头设计、训练并踩坑调通的。
下面按我自己复现和改造这套框架的经验,把设计思路、核心模块、训练细节和坑点逐层拆开讲。
1. 为什么非要在熵模型上动刀:传统压缩与学习型压缩的分水岭
1.1 从编解码流程看熵模型的位置
先理清最基本的链路。学习型图像压缩是一个典型的自编码器结构:
- 编码端 E: 原始图像 x 经过若干卷积下采样,得到潜变量 y;
- 量化 Q: y 被量化为 ŷ,这一步类似传统编码里的有损变换系数量化;
- 熵编码 CE: 把 ŷ 按照其概率分布 p(ŷ) 送入算术编码器,生成真正的比特流;
- 解码端 D: 从比特流中解码出 ŷ,再通过上采样网络重建出 x̂。
这里面的核心是:算术编码器本身并不产生压缩,它只负责按照给定概率把符号流“拧”成紧凑比特流。如果你给的概率分布和真实分布差距很大,就会有“冗余”,浪费比特。所以整个系统的压缩率上限,不取决于网络有多深,而取决于你对 ŷ 的分布估计得有多准。
传统方法这一步是靠人工设计的统计模型。JPEG 的霍夫曼表在统计上对自然图像做了一个平均拟合,HEVC 的 CABAC 引入了一系列上下文状态建模,但也是一种固定的规则。这些模型的问题在于:它们不知道图像里哪个区域是平滑的天空,哪个区域是布满文字的屏幕截图。而学习型方法把 p(ŷ) 的输出直接接到了网络末端,概率越准确,码率越低,重建质量越好,于是整个系统变成了一个可以端到端梯度优化的闭环。
1.2 从“超先验”到“上下文”的演进,卡在哪里
2016 年有人提出用自编码器直接做压缩时,用的还是一个固定的零均值高斯分布,效果自然一般。2018 年提出的超先验(hyperprior)思路是个标志性突破:把潜变量 y 再压一层得到 z,并假设 y 的均值、方差都由 z 来调制。这相当于给熵模型引入了一个“全局辅助参数”,在平滑区域方差小、纹理区域方差大,一下子把压缩效率拉高了一截。
但这还不够。超先验只能提供每个空间位置自身的统计特征,它看不见邻近位置之间的关系。于是后来有了使用掩码卷积(masked convolution)的自回归上下文模型:逐像素/逐块地利用已解码部分的信息来预测当前系数。这确实把压缩率又推进了一点,但代价非常大——自回归解码是严格串行的,解码速度慢得让人抓狂。
综合来看,现有方案里熵模型的信息路径往往是单一或双路的,它们没有真正利用“多信息参考”这样一个更完整的联合条件框架。具体来说,有三类信息是被低估的:第一,潜变量通道之间的相关性被忽略了,不同通道往往对应着不同频率或纹理敏感度,通道间的依赖对概率估计有很大帮助;第二,尺度上的结构参考(高层语义和低层细节)没有被显式建模;第三,解码端已知的重建图信息(条件生成式信息)在部分框架中被浪费了。
所以“多信息参考”的核心动机,不是简单堆模块,而是要把这些互补信号以一种可控的方式塞进熵模型里,让概率估计不再只“盯住局部”,也不再只“依赖全局”,而是同时看多个角度的证据。
2. 多信息参考熵模型的整体设计思路拆解
2.1 框架总览:一个“主干 + 三路参考”的结构
我复现时的整体结构可以拆成四个部分:
| 模块 | 输入 | 输出 | 在熵模型中的作用 |
|---|---|---|---|
| 主分析网络 E | 图像 x | 潜变量 y | 获得可量化编码的紧凑特征 |
| 超先验网络 H | y | 参数 z | 提供全局统计特征(均值/方差调制) |
| 上下文网络 C | 部分已知的 ŷ | 局部依赖特征 φ | 提供已解码区域的空间上下文 |
| 多尺度参考模块 P | y 的多层特征 | 多尺度结构特征 ψ | 提供不同粒度的纹理/语义先验 |
最后熵模型的实际输入是:
- 全局参考:由超先验 z 经过一个小网络产生均值 μ_g 和方差 σ_g;
- 局部参考:由上下文网络产生特征 φ,用于修正均值和方差;
- 多尺度参考:由特征金字塔结构产生 ψ,参与最终概率参数的综合。
这三路特征进入同一个“融合模块”——一般是一个轻量级注意力门控网络,用注意力权重把三路特征动态加权。融合后的参数再通过一个参数网络输出 ŷ 的高斯(或拉普拉斯)分布均值和方差,进而计算码率。
2.2 为什么“多路信息融合”能提升率失真性能:一个信息论的视角
整个过程可以用条件熵来理解。真实数据分布是 p(ŷ | x),但我们编码时能获得的参考条件只是一部分可计算信息。如果用单一超先验,我们只能近似 p(ŷ | z);如果加入上下文,是 p(ŷ | z, φ);再加入多尺度,是 p(ŷ | z, φ, ψ)。根据条件熵的链式法则:
H(ŷ | z, φ, ψ) ≤ H(ŷ | z, φ) ≤ H(ŷ | z)
也就是说,条件信息越多,理论上限的熵越低,也就意味着更低的比特率。当然,实际不会自动达到这个理论值,还得看每路信息怎么从数据中提取、怎么融合。但方向是对的:信息参考的维度决定了率失真性能的上限,而融合方式决定了实际能达到多接近上限。
这里还涉及一个关键问题:为什么不直接用一个大网络把所有条件信息拼起来?早期实验就是这么干的,但效果很差。原因在于不同信息源的量级、分布差异非常大:超先验 z 是低分辨率全局特征,上下文 φ 是逐位置局部特征,多尺度 ψ 又包含多个分辨率的特征。直接把通道拼接需要网络自己“学会”对齐这些特征的尺度差异,训练压力极大。所以现代方案更倾向于在通道维度上做轻量注意力融合,让网络先分别提取特征,再通过门控机制动态决定“以谁为主”。
2.3 与“超先验+上下文”基线架构的区别,不只是加法
很多人以为多信息参考就是“把之前的两路模型改成三路模型”,这个理解不准确。区别体现在两点:
第一是结合位置不同。超先验+上下文的经典做法是分别预测均值和方差,然后直接相乘相加,缺乏交互。多信息参考往往在特征层面就做了交叉注意力:多尺度特征会先修正超先验的特征,而上下文特则会反过来调制多尺度特征的空间权重,三者存在多次双向交互。
第二是“多尺度”本身的信息角色不同。以前的超先验本身也是多分辨率的,但它主要承载的是统计参数。而多尺度参考模块提取的是编码器中间层的结构信息——比如边缘、轮廓、重复纹理,这类信息对概率估计的价值不是统计方差能代表的。举个直观的例子:一张布满草地纹路的照片,如果只靠超先验的方差,草地部分的分布会被当成“高方差区”,码率自然高;但如果有尺度信息告诉网络“这里是周期性纹理”,就能用更紧凑的方式编码这种规律性,码率直接下降。
我自己的实验里,去掉多尺度参考后,同等码率点下 PSNR 会掉 0.1 到 0.2 dB,不算巨大,但在靠近曲线高码率端的区域,保留多尺度参考时的主观纹理细节明显更干净,因为重复纹理被估计得更准,伪影更少。
3. 核心模块逐层拆解与可复现的实操要点
3.1 主分析网络:下采样倍数、激活函数与通道数怎么定
主分析网络负责把图像压成潜变量。常见做法是下采样 4 次,每次 2 倍,总下采样因子 16。输入 256×256,潜变量是 16×16,然后在这个分辨率上做熵模型。
通道数方面,早期方案普遍是 128 或 192,后来为了让高码率点也能提供充足表征,主流会加到 192 甚至 256。我实际测试下来,通道数翻倍带来的率失真增益大约是 0.15 dB,但显存和计算量也几乎翻倍。所以工程项目建议定位在 192,研究原型可以尝试 256。
激活函数这里有个讲究。大部分视觉任务用 ReLU,但压缩网络特别钟爱 GDN(广义除法归一化)。原因在于:压缩网络的中间特征要承载大量空间相关性,ReLU 的零区间会导致梯度切断和特征坍缩,而 GDN 可以看作是一种可学习的、带局部竞争的归一化,能更好地保持特征的统计分布,让后续量化步骤更平滑。这个选择理论上和“独立性假设”有关:GDN 在特征层面最小化了通道间的统计相关,让熵模型的分布估计更容易收敛。
另外,主干网络里残差连接几乎必须用,原因很朴素:量化操作本身是不可微的,梯度要靠直通估计器(STE)回传,网络深度一旦超过 20 层,没有残差连接梯度就很难稳定传回编码器。我见过不少新手在骨干里堆纯卷积,结果训练两个 epoch 后码率 loss 崩掉,大多数情况就是这里埋的雷。
3.2 超先验旁路:均值方差预测的工程细节
超先验网络 H 对潜变量 y 再做一次下采样,得到 z。这个 z 有两个作用:一是作为 side information 进入熵模型,预测 ŷ 的均值和方差;二是 z 本身也要被量化、编码、传输,所以它也有自己的码率代价。
工程实现里要特别留意两点:
- 超先验的下采样倍数。一般取 4 倍下采样,这样 z 的空间尺寸非常小,它的码率占比通常控制在总码率的 5% 以下,避免 info 开销过大。
- 均值和方差输出头的初始化。方差头如果直接输出方差的原始值,训练初期很容易产生极小的方差导致熵值爆炸。标准做法是让网络输出 log(σ),并初始化成一个略为保守的常数(比如 log σ ≈ -1),等价于初始方差约 0.37,这个数值是无数实验试出来的稳定区间。
更实用的做法是把超先验解码出来的特征和上下文特征在通道维度上做一次 1×1 融合,再经过两个并行头分别输出均值和方差。不要直接接两个独立的全连接层,那样会丢失特征交互信息。
3.3 上下文建模:掩码卷积和自回归,速度瓶颈从这里来
上下文网络的典型实现是 masked convolution,也就是卷积核只覆盖当前像素之前的位置(比如左方和上方)。这样做是为了保证解码的时候,当前系数不会看到未来的信息,从而形成合法的自回归过程。
工程上常见的是 5×5 mask 卷积配合通道掩码。5×5 比 3×3 的接受域更大,解码品控更好,但代价是引入的自回归循环路径更复杂。我建议初版先用 3×3 的掩码卷积起步,等整个 pipeline 跑通了再升级到 5×5,不要一上来就在自回归结构上追求完美,因为它的 bug 调试成本是所有模块里最高的。
掩码卷积最隐蔽的坑在于:标准 PyTorch 的 Conv2d 没有内置 mask 支持,很多实现是在 weight 上乘一个 mask 矩阵。这个 mask 必须严格保证 “中心像素以及未来像素的权重被清零”,否则训练时网络直接“作弊”,用当前像素预测当前像素,看起来码率很低,实际上根本没法解码。排查这个问题的傻瓜办法是:随机初始化模型,直接跑一个前向,把潜变量 y 的全部值置成同一个常数,再观察输出的方差是否受扫描顺序影响。如果不受影响,大概率 mask 设计正确;如果第一条常数输入下不同位置输出方差仍有差异,那就要检查 mask 了。
自回归模型另一个问题是串行解码导致速度极慢。论文里通常只在 16×16 的低分辨率潜变量上做,逐位置解码 256 步。实际工程里,为了速度会退化为“块自回归”或“信息路径切分”,也就是一部分通道用上下文,另一部分通道不依赖上下文,并行度大幅提高。这在实现时属于压缩质量与解码速度的取舍,没有标准答案。
3.4 多尺度参考:如何在分支里保持结构信息不丢
多尺度参考最朴素的实现方法是从编码器主干的不同阶段引出特征金字塔:第 2 层(1/2 分辨率)、第 3 层(1/4 分辨率)、第 4 层(1/8 分辨率),分别通过 1×1 卷积对齐到固定通道数,再逐层上采样拼接。这段拼出来的特征和超先验特征做一次 cross-attention 或门控融合,形成最终的 ψ。
在训练中要注意梯度流向问题:多尺度参考的特征都是从主编码器引出来的,如果梯度直接回流到整个编码器主干,会造成主干和熵模型之间的耦合过强,训练极不稳定。一个有效的技巧是使用 stop-gradient 或 gradient scaling:编码主干部分的特征只在部分路径上回传梯度,熵模型侧的特征通过一个较浅的辅助网络单独更新。这样可以让两部分各司其职:主干专心提取图像重建特征,多尺度参考模块专心提取结构先验。这个技巧在很多内部实验里被证明是稳定训练的关键。
3.5 概率模型与率失真损失:不止高斯分布这一个选择
最终熵模型的输出是概率分布参数。最常见的是单高斯,即预测每个系数的 μ 和 σ,然后按高斯分布计算 p(ŷ)。这个方案自超先验工作以来最经典,但最近几年的改进包括:
- 非对称高斯/拉普拉斯:自然图像 DCT 系数和经验分布通常不是对称的,多一个形状参数能提升 1% 到 3% 的码率;
- 离散混合模型:把概率估计建模为多个离散分布的自适应加权(比如 3 个高斯混合),相当于给网络“多个可选的假设”,中码率点能获得更稳定的增益;
- 基于熵参数网络的均值中心化:先预测一个粗略均值,再用上下文信息做残差修正,这比直接从零预测 μ 要稳得多。
率失真损失的写法(以高斯为例):
L = Σ E[-log2 p(ỹ | 参考特征)] + Σ E[-log2 p(ẑ)] + λ · D(x, x̂)
其中:
- 第一项是潜变量的码率估计;
- 第二项是超先验的码率;
- 第三项是重建失真,通常是 MSE,也可以切换成 MS-SSIM;
- λ 是率失真权衡系数,λ 越大,码率越低但失真越高。
这个 λ 的选择直接影响训练出来的模型的码率点分布。实践中我一般会把 λ 在 0.001 到 0.05 之间按对数均匀采样,训练多个模型得到一个 rate-distortion 曲线。注意不要只训一个 λ,然后强行调整目标码率,那等于让网络在非目标区间内做外推,质量会明显下降。
3.6 算术编码器的对接:网络输出概率如何变成比特流
网络训练完只是第一步,部署时还需要把概率输出交给算术编码器。工业界通常用两种方式:一是整张潜变量图的概率参数一次性输出,然后逐符号调用一个 range coder;二是训练好的概率估计配合查表法对整数化的 ŷ 进行快速编码。
这里有三个工程要点:
- (\hat{y}) 是量化后的整数,但网络输出的 μ 和 σ 是浮点数。必须在编码端和解码端重复同样的量化方式,否则概率对不上,解码直接错。
- 算术编码器非常敏感,任何浮点计算不一致都可能导致整个比特流崩掉。解决办法是在最终部署时把概率参数缩放成整数查找表,并固定 softmax/exp 的实现方式。
- 为了能并行处理,许多现代实现会对多个潜变量行执行 batch 编码,但自回归部分的依赖关系会限制并行度,所以工程上常采用“分通道+分块”策略。
4. 训练流程与性能验证:怎样复现才能拿到好看的结果
4.1 数据准备与增广
压缩模型训练最好使用高分辨率图库,比如 ImageNet 的随机裁剪或专用自然图集。裁剪尺寸方面,我用过 256×256 和 384×384 两种。小 patch 训练快,但零散的裁剪会损失全局参考结构,多尺度模块在这种设置下不容易真正“学透”。
我的建议是:先用 256×256 跑通 baseline,确认所有模块能正常收敛,再升级到 384×384 做最终版本。如果直接上更大尺寸,遇到训练不稳定,都不好判断是代码问题还是数据问题。
增广方面,除了常用的水平翻转、随机裁剪,还有一个针对压缩任务的“剧毒”技巧:随机 JPEG 压缩重采样。把训练集以 50% 的概率先压一次再解出来,可以显著提升模型对不同压缩失真的鲁棒性。这个技巧是我实验里提升码率点稳定性的一个隐藏工作。
4.2 训练超参与收敛策略
训练参数基本可以照抄成熟设置:
- 优化器:Adam
- 学习率:初始 1e-4,每 30 个 epoch 乘以 0.1
- Batch size:16 左右,显存不够就降到 8,并用梯度累积补回来
- 训练总时长:单卡 80G 显存约 1 到 2 天可以训到 8 万步
我踩过最大的坑是梯度均衡问题。为了稳定,可以给损失的三项设置一个权重预热(warmup)。具体做法是:前 5000 步,只优化重建失真 D,不计算码率项,等主干网络迈出第一步后,再在 5000 步内把权重线性加到完整水平。为什么先不管码率?因为刚开始潜变量 y 的分布极不规律,如果熵模型一开始就强拟合,容易进入一个“坏局部最优”,后面很难逃出来。这个技巧能把曲线终点的码率点效果提升约 0.5 到 1 个百分点的 BD-Rate。
4.3 率失真曲线评测流程
评测阶段要准备一套测试集(常见的自然图片验证集),对每个 λ 训练出的模型做同一套流程:
- 对每张测试图做 padding,保证尺寸能被 16 整除;
- 编码端跑网络,提取 y, z, 概率参数;
- 调用算术编码器生成比特流,记录实际比特数;
- 解码端重建,计算 PSNR 和 MS-SSIM;
- 绘制 PSNR vs Bits-per-pixel (bpp) 的率失真曲线。
和传统编码器的对比,至少要拉上:BPG(HEVC 的静态图编码)、VTM/VVC 参考软件、以及超先验+上下文基线。表格化展示一个典型结果:
| 方法 | bpp=0.1 时 PSNR | bpp=0.2 时 PSNR | 解码速度 |
|---|---|---|---|
| BPG | 30.4 | 32.8 | 快 |
| 超先验(无上下文) | 29.8 | 32.1 | 中 |
| 超先验+上下文 | 30.5 | 33.3 | 慢 |
| 多信息参考熵模型 | 31.0 | 33.9 | 较慢 |
(具体数值会因数据集和训练配置浮动,这里展示相对趋势)
视化部分,我强烈建议除了率失真曲线,额外保存“熵图”:也就是每个潜变量位置实际比特数的热力图。平滑区域熵低、纹理区域熵高,这是符合预期的;如果你发现有些平滑区域反而熵极高,基本可以断定多尺度参考模块没有学到该学的结构先验。
5. 训练和部署中的常见问题与排查实录
5.1 量化阶段代码报错或码率不收敛
最常见的是训练阶段用“噪声代替量化”(给 y 加均匀噪声),推理阶段用真正取整量化,这中间如果接口不一致,训练和推理的效果会出现一条明显的鸿沟。排查方法是:训练到一半,直接把训练图同时过一遍训练模式和推理模式,对比两者的重建输出。如果噪声量化下重建保持平滑而整数量化后出现块效应,说明量化模拟不够顺滑。解法是采用带 STE 的取整近似或者在量化时混合使用更多噪声宽度。
5.2 自回归解码串行导致推理慢到不可接受
有些团队在多信息参考框架下加了自回归上下文,结果每张 256×256 的图解码要花 30 秒以上。解决方向有两个:一是把自回归改造为“通道级自回归”,即一次预测一群通道,而非一个系数;二是把多尺度参考里的大卷积核替换为深度可分离卷积,减少 80% 以上计算量。这属于工程侧调整,率失真损失通常能控制在 0.1 dB 内,但部署体验能提升一个量级。
其实我最终发布版本的推理速度,主要瓶颈并不是熵模型本身,而是上下文的串行部分。对于要求实时解码的产品,可以考虑完全去掉自回归分支,只保留超先验和多尺度参考两路,并行度会显著提升。和完整版相比,同等码率点 PSNR 损失约为 0.3 到 0.5 dB,但对很多实际应用来说,这个损失完全可以接受。
5.3 显存溢出:多基础模块堆叠时的显存爆炸
多信息参考框架最大的显存杀手有三个:超先验的第二个下采样分支、掩码卷积的因果 padding、多尺度特征金字塔的中间缓存。实际测试,240×240 的输入就可以把 24G 显存吃满。建议做法:
- 主干和熵模型共享部分特征,避免重复计算;
- 多尺度参考模块的回传使用 stop-gradient;
- 用 gradient checkpointing 把中间激活的显存消耗降下来。
5.4 融合模块的初始化到底怎么设
这是很多人忽视的细节点。通道注意力的门控初始值最好是“接近恒等映射”,也就是初始权重让融合后的特征约等于三路特征的平均值。如果初始随机化,融合模块会在大训练后期才学会合理加权,前期相当于在输出一堆放大比例的噪声特征,直接拖累收敛。我的经验是给融合模块最后加一个 tanh 或 sigmoid,并让 sigmoid 的初始 bias 设成合理数值,使门控比例初期保持在 0.5 附近。
5.5 训练集和测试集裁剪尺寸不匹配
如果你在 256×256 上训练,但测试时扔进去一张 1024×1024 的大图,率失真表现会下降。这是因为多尺度模块在训练时从未见过这么大的感受野。解决方案有两种:一是训练时就混合不同尺寸(multi-resolution training),让网络适应图像尺寸变化;二是测试时对图像做分块推理并保留重叠区,这种方法在工程上最稳妥。
6. 从这篇方法延伸出去:还能改什么、怎么落地
6.1 从图像到视频:多信息参考的自然迁移
这套框架天然可以迁移到视频压缩任务。视频压缩相当于把“图像内容结构”扩展成“时间+空间结构”,多信息参考从“多尺度空间特征”延伸到“时空多参考”——当前帧的潜变量不仅可以参考当前帧的超先验和上下文,还可以参考前一帧的重建特征和运动补偿特征。这个方向的增益会比单纯堆参数更明显,因为视频相邻帧之间的冗余远大于单帧内的空间冗余。
6.2 多光谱、深度图、医学影像:先验更稀缺的领域更受益
自然图像有很多通用先验,但多光谱遥感和医学影像数据量少、结构特殊。传统学习型压缩模型在这种数据上往往过拟合,因为它们假设自然图像的统计特征。多信息参考框架的价值在于:参考特征可以作为“领域知识注入”的入口。比如医学图像可以引入解剖结构先验特征,遥感图像可以引入地物分类特征,只要把多尺度参考模块替换成特定领域特征提取器,整个压缩模型就能在少量数据上获得明显增益。
6.3 部署到移动端的额外工程优化
这套方法原始版本的计算量不低,落地移动端需要两步优化:一是量化,把权重从 FP32 转到 INT8,通常带来的性能损失在 0.1 dB 以内;二是把熵模型的概率计算改成查表。注意算术编码器的整数范围选择一定要保证足够宽,否则概率值会截断导致错误。这些优化做完,模型体积能压缩 70%,点云和端侧实时编码基本可用。
写在最后的个人体会
我复现这类多信息参考框架时,最深的感受是:熵模型改进的核心不在“模型堆得多高”,而在“信息从哪里来,怎么让它们相互校准”。超先验给的是全局统计,上下文给的是局部依赖,多尺度给的是结构化语义,三路信息如果只是拼在一起,效果有限;只有在每一层都做了交互和门控,让网络自己决定当前系数应该信谁,才能真正把信息量压出来。
另外,这个方向目前远未到瓶颈期。比如自回归部分的串行问题和多尺度特征的冗余信息裁剪,都还有很大的研究空间。对刚上手的同学,我的建议是:先跑通一个简单的超先验 baseline,再一步步往上加参考分支;每加一步就记录率失真变化,便于定位改进来源。这种循序渐进的做法,比我一开始直接堆完整模型再痛苦排查要高效得多。希望这篇文章能帮你少踩几个我踩过的坑。