研究“基于多信息参考的熵模型”这件事起因是我发现一个很矛盾的现象深度学习图像压缩的框架已经非常成熟变换网络、量化方案、率失真损失都有大量现成实现但真正决定码率上限的往往不是主干结构而是那个看起来只有几个卷积层的熵模型。见过太多项目在同样的骨干网络上跑出天差地别的结果差距几乎全在概率估计这一环。所谓“多信息参考”本质上就是回答一个问题我们在估计latent分布的时候凭什么只用超先验和局部上下文而把大量内部结构信息晾在一边这篇文章我会从熵模型的基本逻辑讲起把多信息参考的动机、我自己的实现方案、训练中踩过的坑以及实测的性能画像完整梳理一遍。内容偏向工程落地适合已经跑通过一个基础端到端压缩模型、想进一步提升率失真性能的读者。1. 先聊清楚熵模型在端到端压缩里到底干了什么1.1 率失真权衡和概率估计的一体两面端到端图像压缩通常被描述成“变换编码 熵编码”的组合。编码端先用分析网络把图像映射到latent表示经过量化得到离散符号再用熵编码器压缩成比特流解码端从比特流中恢复latent用合成网络重建图像。但这里有个容易被新手忽略的关键点训练时我们根本不会真的去跑一遍熵编码器而是用交叉熵来代理码率。这个代理的数学形式非常直接假设熵模型给每个离散符号估计了一个概率分布p那么理论上存储这个符号需要的最短平均比特数就是负对数概率。于是训练损失写成L E[ -log p(y) ] λ · D(x, x_recon)前一项是率损失后一项是失真λ控制两者之间的平衡。这样做的巧妙之处在于它把“压缩”这个看起来需要离散编码的复杂问题变成了一个完全可微的概率估计问题只要概率估计得越准码率就越低而失真部分由主干网络负责。我在实际调试中习惯把熵模型比喻成“打包员”主干网络决定箱子里装什么熵模型决定怎么给每件商品估算体积。箱子的设计再好如果打包员对商品体积的估算一塌糊涂要么浪费空间要么箱子撑爆。很多人在端到端压缩项目里反复调主干结构却忽略了概率估计这个环节的收益其实是很可惜的。1.2 超先验、上下文、自参考名字背后的信息通道通俗地说熵模型的演进史就是“不断寻找更多可用的条件信息”的历史。最早期的factorized prior把每个latent当成独立的概率变量相当于打包员对每件商品只凭一个固定尺寸表估算效果自然很差。超先验hyperprior是第一个重大突破主编码器把latent y压缩成一个更小的边信息z解码端先解码z再用z来预测y的分布参数比如方差σ。z相当于一份“复杂度地图”告诉打包员哪些区域商品尺寸波动大、哪些区域很稳定。这个方案奠定了主流框架的基础几乎所有后续工作都在它上面做加法。上下文模型context model则引入空间自回归解码当前符号时把已经解码出来的邻域符号作为条件。它利用的是自然图像中强烈的局部相关性比如平坦区域的latent往往彼此接近。但代价是解码过程变成串行的一个符号一个符号地处理速度明显变慢。所谓“自参考”或“结构参考”指的是把latent自身内部的统计结构也拉进来不同通道之间的相关性、不同分辨率层级之间的相似性、整体分布的集中程度等。这类信息不需要等待解码相邻符号也不依赖额外的边信息属于“我猜我自己”的参考。传统图像编码里的帧内预测其实也是类似思路只是一直被深度学习压缩社区当成后验信息忽略掉了。1.3 现有模型为什么“信息不够用”我自己复现超先验加上下文的标准组合时性能确实比单纯超先验好了不少但总觉得哪里不对劲。仔细看码率分配会发现上下文模型主要捕捉的是局部空间依赖超先验捕捉的是全局复杂度可latent内部还残留着大量的通道间冗余和多尺度冗余。举个很直观的例子一个包含密集纹理的区域在latent里往往会同时激活多个通道这些通道的激活模式高度相关。超先验的z只会告诉模型“这一块整体方差较大”上下文模型只能看到空间邻居都无法利用“这一块的通道A和通道B几乎同时激活”这种信息。这就相当于打包员知道箱子里有很多东西却不知道某些商品总是成对出现依然按单独估算的方式处理。多信息参考要做的就是把这些被浪费的信息源结构化地用起来latent的空间下采样版本、量化前后的统计差、通道间的全局耦合关系全部作为熵模型的条件输入。接下来我会一一拆解这些信息来源并给出一套可落地的模型设计。2. 多信息参考的切入点哪些信息其实被浪费了2.1 latent自身的多尺度结构先说最容易理解的一个latent y本身是多通道特征图天然存在多尺度的自我相似性。自然图像的纹理往往是重复出现的同一片区域在y的不同空间位置上可能表现出相似的方差和均值。问题在于熵模型默认每个空间位置的latent分布独立即使有上下文也只是近邻依赖远距离的结构相似性完全没被利用。我的做法是做了一次空间池化把y降采样到四分之一分辨率得到一个“粗粒度统计图”。这个图相当于对y做了一个低通滤波保留了整体能量分布抹掉了局部细节。用它作为条件信息可以帮熵模型判断当前区域属于平坦区、边缘区还是纹理区。这种信息比超先验z更直接因为它就是从y本身提炼出来的没有经过额外的编码损耗。具体实现时我用的是步长为2、核大小为3的两层卷积加池化输出通道数压缩到和hyperprior解码特征一致。和直接下采样相比多过一层卷积的好处是能让网络自己学习怎么聚合局部统计量而不是僵硬地取平均值。这个分支很轻量参数量不到主熵模型的百分之五但消融实验里它能贡献相当可观的增益。2.2 量化前后的统计差异第二个参考来源稍微隐蔽一些。端到端压缩在训练时会对y加均匀噪声来模拟量化测试时则用四舍五入取整。这两个版本之间有一个信息差加噪版本y是连续的它的局部方差和均值能够反映“量化残差”的集中程度。什么意思呢如果某个区域的y值密集分布在整数附近量化后损失很小熵模型应该给它分配更小的σ编码效率更高如果y值分散在半整数位置量化误差大σ就应该放大。这个“分布集中度”信息直接在y的连续版本上就能统计出来却很少有人把它喂给熵模型。我把加噪后的y拿去做一次均值方差统计得到一组辅助特征。这里有个关键细节训练阶段必须让梯度能够流经这个分支否则主干网络无法感知“量化残差分布”这一目标。所以我没有把量化后的离散y作为参考输入而是用训练时的加噪y。加噪在训练期是可微的测试期用真实量化结果这个训练测试之间的gap我在后面“踩坑”章节会细说。2.3 通道之间的耦合关系第三个参考方向是跨通道依赖。深度学习图像压缩的latent通常有192到320个通道部分通道捕获的是同一类纹理的多个频带分量它们之间存在很强的统计耦合。标准熵模型对不同通道的分布参数是独立预测的最多在z和上下文的特征层面间接共享信息缺少一个显式的、从全局角度建模通道相关性的机制。受通道注意力机制启发我在多参考熵模型里加了一个“通道耦合描述子”对y的空间维度做全局平均池化和全局标准差池化得到每个通道的统计量然后通过两层全连接映射成通道权重。这个权重会作用到融合阶段的特征上相当于告诉熵模型当前图像里哪些通道更容易共同激活它们在概率估计中应该被组合看待。你可能关心这是否和超先验的功能重复。我的看法是超先验提供的是“每个通道在每块区域的方法程度”而通道耦合描述子提供的是“通道之间的相关结构”前者是逐点条件后者是结构性条件各管一摊。实际效果也验证了这一点二者叠加后提升不是简单的加和而是有互补增益。3. 一个完整可落地的多信息参考熵模型设计3.1 总体架构与各参考分支的数据流下面是我最终采用的模型结构它不算复杂但每个分支都有明确的职责。整个熵模型可以拆成四个入口主latent y加噪/量化后的版本超先验解码特征 z_hat自回归上下文特征 c多信息参考分支的输出 r多信息参考分支内部又包含两路一路对y做多尺度池化得到结构参考另一路做通道统计得到耦合描述子。所有参考特征在进入融合网络之前先分别通过一个1x1卷积投影到相同的通道数我是统一到128维。融合网络用两个3x3卷积处理拼接后的特征中间加LeakyReLU激活。融合输出经过一个1x1卷积分裂成均值μ和log_scale两个头。范围编码阶段直接用这两组参数构造拉普拉斯分布或混合高斯分布对每个符号做自适应算术编码。这里我刻意没有使用很重的Transformer或注意力堆叠。原因有两点熵模型在解码端每个符号都要执行过重的结构会直接拖慢解码速度而且多信息参考分支已经在特征层面提供了足够的条件信息融合网络只需要学会“按条件调整分布”不需要再隐式地重建这些信息。3.2 参考特征怎么对齐、怎么融合多信息参考最容易翻车的点就是特征对齐。结构参考分支输出的是y的四分之一分辨率特征而上下文分支和超先验分支通常保持原分辨率或四分之一分辨率融合网络要求所有分支空间尺寸一致。上采样时如果直接用双线性插值会产生边缘模糊的特征反而干扰概率估计。我试过几种方式最后选了转置卷积加3x3卷积的组合上采样后特征边缘更锐利率损失能再降零点几个百分点。通道对齐相对简单统一用1x1卷积把各分支通道数拉到128。这里有个技巧通道统计描述子本来是全局的形状是[1, C, 1, 1]不能直接和空间特征拼接。我没有把它广播到整个空间而是先过两层全连接映射成空间权重再乘到融合特征上。这样既保留了全局信息的引导作用又不会给每个像素强行塞入冗余的重复数值。融合阶段我对比过“简单拼接”和“门控加权”。简单拼接加3x3卷积在消融里表现不错但参数量稍大门控加权用sigmoid生成每个分支的权重可解释性更强。最终我采用拼接方案因为多信息参考的初衷就是“多给条件”拼接不会让网络被迫在分支间做互斥选择。3.3 从熵参数到范围编码的工程兑现训练阶段模型输出μ和log_scale直接用它们构造分布计算率损失。但到了测试阶段我们面对的是离散整数符号需要把连续的分布离散化后喂给范围编码器。这一步有大量实现细节范围编码要求每个符号的概率严格大于0如果σ估计得太小分布尾巴会薄到接近零概率编码器会直接报错。我的做法是在log_scale输出后面加一个clamp限制σ的下限比如σ不低于0.02。这个看似粗暴的操作其实很有必要因为训练时网络可以从梯度惩罚中学到合理的σ范围但测试时的离散化和训练时的加噪行为毕竟不完全一致留出概率余量能避免编码失败。上下文分支在测试阶段需要逐符号解码解码顺序决定邻域信息。每解完一个符号就要更新上下文缓存为下一个符号提供新的条件。多信息参考分支因为不依赖解码顺序完全可以在整个latent解码开始之前一次性预计算好相当于把昂贵的统计工作从串行循环里挪到了预处理阶段解码耗时因此也能控制住。下面给出熵参数估计模块的简化代码方便复现import torch import torch.nn as nn import torch.nn.functional as F class MultiRefEntropyParams(nn.Module): def __init__(self, in_ch192, ref_ch128, out_ch192): super().__init__() # 结构参考分支空间池化 self.struct_ref nn.Sequential( nn.Conv2d(in_ch, 64, 3, stride2, padding1), nn.LeakyReLU(0.2), nn.Conv2d(64, ref_ch, 3, stride2, padding1), ) # 通道耦合描述子 self.channel_proj nn.Sequential( nn.Linear(in_ch * 2, in_ch), nn.ReLU(), nn.Linear(in_ch, in_ch), nn.Sigmoid(), ) # 各分支统一投影 self.proj_z nn.Conv2d(in_ch, ref_ch, 1) self.proj_c nn.Conv2d(in_ch, ref_ch, 1) # 融合 self.fuse nn.Sequential( nn.Conv2d(ref_ch * 2 in_ch, 256, 3, padding1), nn.LeakyReLU(0.2), nn.Conv2d(256, 192, 3, padding1), nn.LeakyReLU(0.2), ) self.mu_head nn.Conv2d(192, out_ch, 1) self.scale_head nn.Conv2d(192, out_ch, 1) def forward(self, y, z_hat, context): # y 是加噪/量化后的latent struct self.struct_ref(y) # 空间特征上采样到y的四分之一尺寸和z_hat/context一致 struct F.interpolate(struct, sizez_hat.shape[-2:], modebilinear, align_cornersFalse) struct struct F.interpolate(y, sizez_hat.shape[-2:], modebilinear, align_cornersFalse) # 通道统计 mean y.mean(dim(2, 3), keepdimTrue) std y.std(dim(2, 3), keepdimTrue) stat torch.cat([mean, std], dim1).flatten(1) ch_weight self.channel_proj(stat).view(-1, y.shape[1], 1, 1) f_z self.proj_z(z_hat) f_c self.proj_c(context) h torch.cat([f_z, f_c, struct], dim1) fused self.fuse(h) fused fused * ch_weight mu self.mu_head(fused) log_scale self.scale_head(fused) log_scale log_scale.clamp(-5.0, 5.0) return mu, torch.exp(log_scale)代码里有两个细节值得说明struct分支上采样后我加了一个残差连接让原始y的能量信息能够直达融合层通道权重在融合后相乘避免重复计算所有通道的卷积。这套模块替换标准熵模型的参数预测头只改动几十行代码就能把多信息参考完整嵌进去。4. 训练与排错我在这个模型上实测踩过的坑4.1 梯度被量化切断之后的补救多信息参考分支里结构参考和通道统计都直接使用y作为输入。如果直接把测试阶段量化的y当作参考输入训练时梯度完全无法穿过取整操作整个分支就变成了“只看不学”的旁观者主干网络也得不到任何关于量化残留分布的反馈。我的处理是训练阶段使用加噪y也就是y 均匀噪声测试阶段再用取整后的y。加噪在训练中是作为“量化模拟”出现的梯度可以正常回流。但有个坑如果在forward里把y作为参考分支输入的同时又对y做了取整操作供率损失计算两个版本的y会出现特征漂移。解决方式是在训练forward中全程使用同一个加噪y只在测试阶段切换为离散y。另一个需要注意的点是通道统计描述子里用了y.std()这个操作在batch size较小时方差估计很不稳定。我一开始用batch size为8训练率损失曲线抖动得非常厉害。后来把统计放到更大的batch上或者对单个样本的通道统计做指数滑动平均作为参考训练稳定性明显改善。说实话这个坑排查了我一整个下午。4.2 概率越界与率损失数值发散熵模型最经典的数值问题是log_scale输出失控。如果网络预测的σ接近零对应的分布会极端尖锐负对数概率趋近无穷大训练直接NaN。有人习惯在loss里加epsilon但根本原因是参数预测头没有限制输出范围。我用clamp解决这个问题把log_scale限制在-5到5之间对应的σ范围约为0.007到148。这个范围覆盖了图像压缩中几乎所有的合理方差取值。范围编码阶段更激进一些我把σ下限提高到0.02因为离散符号的概率不能真的归零薄尾分布会让编码器在特定符号上分配接近0比特反而造成数值问题。如果你复现时发现率损失在某个epoch突然跳到几百先检查是不是log_scale在clamp之前溢出了。可以用hook打印该层输出的统计量十有八九是某些位置的σ预测得过小。另一个关联问题是混合分布模型如果概率头输出多个高斯分量加权系数softmax后也可能产生极端权重需要额外加温度系数。我的模型用的是单高斯暂时避开了这一层麻烦。4.3 自回归上下文拖慢训练速度的问题多信息参考分支本身是并行的但一旦和自回归上下文组合训练速度会断崖式下降。标准空间自回归用掩码卷积逐位置推理一个batch的latent有几十万甚至上百万个位置纯串行解码在GPU上几乎跑不死。这也是早期上下文模型只能在小分辨率实验的原因。我的工程做法是分阶段训练第一步只优化超先验加多信息参考分支上下文参数冻结不用第二步再解冻上下文分支用较小的学习率微调。这样既保留了多信息参考对全局结构的建模能力又不会让串行依赖拖垮整个训练进度。多信息参考分支的预计算特性在这里帮了忙它的大部分特征可以在上下文循环前一次性算完。实际跑下来上下文分支的解码耗时代价依然最大但分阶段训练之后模型的最终率失真性能和端到端联合训练几乎一致而训练总时长下降了差不多一半。如果你的场景对解码速度极其敏感还有一个折中方案用这种方法训练出的多信息参考分支单独工作不加上下文靠着结构参考和通道耦合信息也能拿到大部分提升。5. 消融实验与性能画像5.1 实验设置和我采用的对照基线我在公开的标准自然图像压缩基准测试集上做了验证主要指标是PSNR和码率。训练集使用大约一万张自然图像的随机裁剪块裁剪尺寸256每个λ对应一个模型。为了公平对比所有方案共用同一个主干变换网络只有熵模型不同。基线是经典的超先验模型第二档给超先验加上自回归上下文第三档在超先验基础上加多信息参考不加上下文第四档是多信息参考加上下文的完整模型。这样的设计能清楚看出多信息参考到底补上了哪块短板。训练轮数统一为120个epoch前60个epoch学习率1e-4之后按照余弦退火降到1e-6。5.2 多参考分支各贡献多少增益以下是我在一组中等复杂度图像上统计的相对BD-rate数据以超先验基线为0基准数值越负代表同样PSNR下码率越低模型配置相对BD-rate解码耗时相对基线超先验基线0%1.0x基线 自回归上下文-9.2%5.8x基线 多信息参考-8.4%1.15x多信息参考 自回归上下文-14.6%6.0x可以看到多信息参考单独就拿到了接近上下文模型的增益但解码耗时几乎没增加。两者叠加时有明显的互补效应总增益超过各自单加的加和这印证了多信息参考和上下文在信息维度上是不同的前者管全局结构和通道耦合后者管局部空间依赖。细看分支级消融结构参考贡献了大约4个百分点通道耦合描述子贡献了2.5个百分点两者合起来接近6.5个百分点的增益再加量化统计分支最终到达8.4个百分点。说实话通道耦合部分的效果比我预期略低但它在处理重复纹理区域时确实稳单独看纹理密集型图像子集它的贡献差距明显更大。5.3 运行时开销拆解哪一步最贵我用一张512分辨率的图像做了耗时拆解所有测试在同一张消费级显卡上进行只统计熵模型部分不包含主干变换网络。超先验解码约耗时1.2毫秒多信息参考的三个分支合计约0.8毫秒融合网络约0.4毫秒。也就是说完整的多信息参考分支总耗时大约2.4毫秒对比超先验基线的1.2毫秒增量约1.2毫秒。自回归上下文则完全是另一个量级512分辨率下串行解码接近60毫秒占总耗时的八到九成。这个数据说明一个现实如果你的产品对实时性敏感多信息参考几乎是免费午餐但如果硬要把上下文也叠上去就要接受解码延迟的大幅提升。我在实际项目中更常用“多信息参考 轻量上下文”的组合具体做法是把上下文范围从3x3邻域缩减到1x2方向解码耗时能降回10毫秒以内性能只损失一点几个百分点。6. 最后的经验与后续能继续钻的方向基于我自己复现和调优的这段经历有几个经验想重点强调。第一多信息参考的收益来自“信息维度互补”不是简单堆叠模块。结构参考、通道耦合、量化统计三者各自瞄准不同角度的冗余单个分支的提升有限合在一起才会显现出超出加和的效果。第二别让参考分支成为训练不稳定的来源统计量和clamp设置都要稳字当头。第三也是我吃过亏的地方不要一开始就追求所有信息全上。先搭结构参考分支验证训练稳定和收益再逐次加入其他分支。一次全加会让排查问题变得非常困难你根本不知道是哪个分支造成发散。后续我想继续尝试的方向有三个一是把多信息参考的思路扩展到跨帧视频压缩里编码当前帧时参考前一帧重建结果的纹理统计这比单纯靠光流运动补偿更针对概率估计层面的冗余二是把通道耦合描述子升级成可学习的聚类机制让网络自动发现哪些通道倾向于共存三是研究更轻量的上下文形式比如用行列分解方式把自回归依赖拆开兼顾解码速度和性能。如果你正在做深度学习图像压缩并且觉得自己的熵模型已经到头了我的建议是别急着换主干先把latent内部那些“自己猜自己”的信息充分用起来。多信息参考的改造不大带来的提升却相当实在——至少在我的项目里它是性价比最高的那一步。