平衡的艺术: 多目标优化如何提升学习型图像压缩性能

在数字媒体的世界里，我们总是面临着一场拉锯战。一方面，我们想要清晰逼真的高质量图像 (低失真) 。另一方面，我们希望文件足够小，能够即时流式传输、存储和分享 (低码率) 。这种权衡是图像压缩的核心。

JPEG 或 HEVC 等传统编解码器通过手工调优的工程设计来解决这个问题。但最近, 学习型图像压缩 (Learned Image Compression, LIC)——即利用深度神经网络来压缩图像——已经开始超越这些传统方法。LIC 模型从数据中学习如何最好地表示图像。

然而，训练这些模型十分棘手。你必须教网络同时最小化文件大小和最大化图像质量。通常，这两个目标在训练过程中会相互冲突，导致产生一个并不是最优的“平衡”模型。

在本文中，我们将深入探讨一篇最新的研究论文，它提出了一种更智能的网络训练方法。通过将压缩视为一个多目标优化 (Multi-Objective Optimization, MOO) 问题，研究人员开发了一种动态平衡“码率”和“失真”目标的方法，在不改变模型架构本身的情况下，显著提升了性能。

问题所在: 当目标发生冲突时

要理解解决方案，我们需要先了解学习型图像压缩模型通常是如何训练的。

标准 R-D 损失

在标准设置中，训练 LIC 模型的目的是最小化一个结合了两方面因素的损失函数:

  1. 码率 (\(\mathcal{L}_R\)): 存储压缩代码所需的比特数。
  2. 失真 (\(\mathcal{L}_D\)): 原始图像与重建图像之间的差异 (通常通过均方误差测量) 。

它们通过一个权衡参数 \(\lambda\) (lambda) 结合在一起。损失函数如下所示:

Standard Rate-Distortion Loss Equation

在这里，\(\lambda\) 就像一个音量旋钮。调大它，模型就会专注于质量 (高比特率) 。调小它，模型就会专注于压缩 (低比特率) 。

不平衡

理想情况下，随着模型的训练，码率和失真都会同时改善。实际上，梯度——告诉神经网络如何更新其权重的信号——对于码率和失真来说可能有巨大的差异。

有时失真的梯度很大，而码率的梯度很小 (或者反之亦然) 。或者它们可能指向相互冲突的方向。当我们简单地将它们相加时，较强的信号会淹没较弱的信号。

这导致了不平衡的优化 。 模型可能会花费几十个 epoch 来修复微小的失真错误，而忽略了节省比特的机会，反之亦然。结果就是模型满足于一个“足够好”的解，而不是最佳解。

我们可以在下图中清楚地看到这一现象:

Comparison of loss trends and improvement speeds

在 图 1 (a) 中，请看代表“标准 (Standard)”方法的虚线。注意失真损失是如何下降的，但 bpp (每像素比特数) 损失几乎没有变化——甚至略有上升。目标并没有协同移动。

现在看看实线 (即提出的“平衡 (Balanced)”方法) 。两条曲线都平滑且一致地下降。通过平衡优化，模型学会了同时提高质量和减小文件大小。

解决方案: 多目标优化 (MOO)

研究人员建议放弃简单的损失求和。相反，他们将率失真优化构建为一个多目标优化问题。

MOO General Formulation

目标是找到一个能同时改善两个目标的参数更新。如果我们将码率和失真视为两个独立的任务，我们希望找到一个梯度更新方向，尽可能平等地让两个任务受益。

最大化改进速度

我们如何定义“平等”？我们看改进速度 。 如果我们在某个方向 \(d_t\) 上更新模型参数 \(\theta\)，损失相对于其当前值下降了多少？

任务 \(i\) (\(i\) 为码率或失真) 的相对改进速度定义为:

Improvement speed equation

为了确保平衡学习，我们要最大化最小改进速度。把这想象成一支船队；船队的速度取决于最慢的那艘船。我们要确保“最慢”的目标 (即落后的那个) 得到最多的关注。

这引出了一个“极小极大 (minimax)”优化问题。我们希望找到一个方向 \(d_t\)，它能最大化码率和失真之间最坏情况下的改进:

Saddle point optimization problem

别被数学公式吓倒。这个方程只是在说: “找到更新方向 \(d_t\)，以保证无论哪个目标当前最困难，都能得到最好的改进。”

对偶问题: 重新加权梯度

在高维神经网络 (拥有数百万参数) 中直接求解更新方向 \(d_t\) 在计算上是非常昂贵的。然而，研究人员利用了一个叫做拉格朗日对偶性的数学技巧。

他们不是直接寻找方向向量，而是求解应用于梯度的权重 (\(w_R\) 和 \(w_D\)) 。

事实证明，寻找最佳方向 \(d_t\) 在数学上等同于寻找最小化组合梯度幅度的权重 \(w\)。

Minimizing the weighted gradient norm

这里，\(J_t\) 代表码率和失真损失的梯度。我们想要找到权重 \(w_t\) (其中 \(w_R + w_D = 1\)) ，使得组合梯度向量最小化。这个特定的加权梯度指向满足帕累托平稳性条件的方向——即有效地同时改善两个目标。

一旦找到这些最佳权重，我们就通过加权和来计算最终的更新方向 \(d_t\):

Balanced gradient calculation

这里有一个实际问题: 损失值 \(\mathcal{L}\) 可能会变得非常小，导致数值不稳定。为了解决这个问题，研究人员使用了一个归一化常数 \(c_t\):

Normalization constant

这确保了更新步长不会爆炸或消失。

两种平衡策略

论文介绍了两种解决这个加权权问题的不同方法。一种是适用于从头开始训练的近似方法，另一种是非常适合微调的精确解析解。

方案 1: 梯度下降 (由粗到细)

当从头开始训练模型时，损失函数的景观变化很快。我们不一定需要每一步都有完美的权重；我们只需要大致朝正确的方向移动。

方案 1 将权重 \(w_t\) 本身视为可学习的参数。在每一步中，它对权重执行快速梯度下降更新，以最小化我们上面定义的目标。

权重是迭代更新的:

Weight update rule

为了确保权重保持为正且和为 1 (有效的概率分布) ，研究人员通过 Softmax 函数映射它们:

Softmax mapping

他们还在更新规则中添加了一个衰减项，这有助于稳定训练，防止权重因单批数据而过度震荡:

Update with decay

最适合: 从头开始训练新的 LIC 模型。它的计算成本比方案 2 低，并且随着模型的学习逐步适应。

方案 2: 二次规划 (解析法)

有时，我们需要精度。如果我们正在微调一个已经收敛的模型，我们需要精确的最佳权重来挖掘最后一点性能。

我们前面讨论的最小化问题实际上是一个具有特定约束 (权重和必须为 1) 的二次规划 (QP) 问题。

Quadratic Programming Formulation

因为 Hessian 矩阵 \(Q\) (关联码率和失真的梯度) 是正定的，所以这个问题有一个封闭形式的解析解。我们不需要迭代；我们可以直接使用矩阵代数计算答案。

使用拉格朗日乘数法，论文推导出了最佳权重的精确公式:

Analytical solution for weights

最后，为了确保数值稳定性和非负性，他们使用 Softmax 将此解投影到单纯形上:

Softmax projection for QP solution

最适合: 微调现有模型。它需要计算和求逆矩阵 (\(Q^{-1}\))，这增加了计算成本，但它为每次更新提供了数学上的最佳平衡。

实验结果

那么，平衡梯度真的有效吗？研究人员在标准数据集 (Kodak, Tecnick, CLIC2022) 上使用流行的架构 (如 “M&S Hyperprior”、“ELIC” 和 “TCM-S”) 测试了他们的方法。

R-D 性能

结果显示出明显的改善。仅仅通过改变模型的优化方式 (而不改变模型架构本身) ，他们就实现了更好的压缩性能。

R-D Curves Comparison

在 图 2 中，我们看到了率失真曲线。横轴是比特率 (bpp，越低文件越小) ，纵轴是 PSNR (质量，越高越好) 。

  • 实线 (方案 1 和 2) 始终高于 虚线 (标准训练) 。
  • 这意味着在相同的文件大小下，平衡方法提供了更高的质量。或者，在相同的质量下，它产生的文件更小。

下表使用 BD-Rate 量化了这些收益，这是一个衡量同等质量下平均码率节省的指标。负数表示文件大小更小。

BD-Rate Table

提出的方法实现了大约 2% 到 3% 的 BD-Rate 降低 。 在图像压缩领域，算法们为了零点几个百分点的提升而争斗，纯粹靠优化策略获得 2-3% 的提升是一个巨大的胜利。

计算成本

天下没有免费的午餐。计算最佳权重需要在训练期间进行额外的计算。

Complexity Table

  • 方案 1 增加了约 20% 的训练时间。
  • 方案 2 增加了约 50% (由于矩阵运算) 。

然而，必须注意的是, 这个成本仅适用于训练阶段。 推理时间 (模型训练后压缩/解压图像的时间) 保持完全相同 。 最终模型的架构是一样的；只是权重更好。

消融研究: 细节为何重要

研究人员进行了消融研究，以证明他们具体的设计选择是必要的。

Ablation Studies

  • 图表 (a) - 归一化: 橙色和红色线显示了没有归一化常数 \(c_t\) 时的性能。它们的表现明显更差。归一化梯度对于收敛至关重要。
  • 图表 (b) - 权重衰减: 在方案 1 中，使用衰减项 (\(\gamma\)) 有助于平滑更新。蓝色方块 (\(\gamma=0.001\)) 与其他值相比产生了最佳性能。
  • 图表 (c) - 从头训练 vs. 微调: 这证实了预期的用例。方案 1 (蓝色) 非常适合从头开始训练。方案 2 (橙色) 也不错，但较慢。有趣的是，使用方案 1 进行微调 (绿色) 不如方案 2 有效。

结论

论文《Balanced Rate-Distortion Optimization in Learned Image Compression》强调了 AI 训练中的一个基本问题: 仅仅因为你将两个损失函数相加，并不意味着它们得到了平等的优化。

通过将图像压缩重新构建为多目标优化问题，作者提供了一个强大的框架来处理码率和失真这两个相互竞争的目标。

  • 方案 1 提供了一种实用的、迭代的方法来从头开始训练平衡模型。
  • 方案 2 提供了一个数学上精确的工具，用于将模型微调到其理论极限。

对于深度学习领域的学生和研究人员来说，这项工作提醒我们: 有时最大的收益并非来自新的网络架构，而是来自简单地帮助网络更有效地学习。随着 LIC 的不断成熟，像这样的优化技术可能会成为标准做法，以缩小与——并最终超越——传统编解码器的差距。