使用条件训练优化多项损失函数

文 / 研究员 Alexey Dosovitskiy，Google Research

在许多机器学习应用中，模型的性能无法归结于单个数字，而是取决于多个特性，有些特性甚至相互排斥。例如，一个训练得到的的图像压缩模型应该最大限度地降低压缩图像大小，同时尽可能提升图像质量。同时优化所有相关值往往不太可能，原因是有些值在根本上互相冲突（如上例中的图像质量和压缩比），或者仅仅是因为模型容量有限。因此在实际操作中，人们必须决定对于相关值应当如何取舍。

图像压缩中图像质量和文件大小之间的权衡：理想情况下，图像失真程度和文件大小都应达到最小，但这两者存在根本性冲突

训练模型时，如果必须在不同属性间取得平衡，常规方法是将属性相关的各项加权求和而成的损失函数最小化。例如，在图像压缩的示例中，损失函数包含两项，分别是图像重建质量和压缩比率。使用此损失函数训练出的模型会通过这些加权项的系数来决定重建的图像更加紧凑或质量更高。

如果需要在模型特性之间做出不同的取舍（例如图像质量与压缩比），标准做法是训练多个独立模型，并在每个模型的损失函数中使用不同的损失函数。这需要在训练和推理期间运行多个模型，效率十分低下。但是，所有这些独立模型都能解决一些相似的问题，这表明其可以共享某些信息。

在两篇同时收录在 2020 年国际学习表征大会 (ICLR 2020) 的论文中，我们提出了一种简单且广泛适用的方法，此方法可避免针对不同的损失权衡训练多个模型的低效性，改用单一模型解决所有问题。在《只需训练一次：深度网络的损失条件训练》(You Only Train Once: Loss-Conditional Training of Deep Networks) 一文中，我们给出了此方法的一个通用公式，并将其应用于多个任务，包括变分自编码器和图像压缩，而在《可调整的实时风格转化》(Adjustable Real-time Style Transfer) 中，我们深度探讨了风格转化方法的应用。

• 只需训练一次：深度网络的损失条件训练
  https://openreview.net/pdf?id=HyxY6JHKwr

• 可调整的实时风格转化
  https://openreview.net/pdf?id=HJe_Z04Yvr

损失条件训练

此方法背后的思路是训练一个涵盖损失项所有系数选项的单一模型，而不是为每组系数集训练一个模型。我们的实现方法是：

1. 按损失分布而非单个损失函数训练模型；

2. 根据损失项的系数向量调节模型输出。

这样一来，推理时的调节向量可能有所变化，我们因此能够遍历系数不同的损失函数所对应的模型空间。

下图说明了针对风格转化任务的这一训练流程。对于每个训练示例，首先对损失系数进行随机采样。然后，通过调节网络将其用于主网络，并计算损失。对整个系统进行联合端到端训练，即在训练模型参数的同时对损失函数进行随机采样。

以风格化为示例的方法概览：主风格化网络根据损失函数的随机采样系数进行调节，并按损失函数的分布进行训练，因此学习对整个损失函数系列进行建模

此方法在概念上简单易懂，因此仅需对现有代码库进行微小改动即可适用于诸多问题领域。在本文中，我们重点讨论图像压缩和风格转化这两类应用。

应用：可变比率图像压缩

作为此方法的第一个示例应用，我们将展示训练图像压缩的结果。压缩图像时，用户应该能够在图像质量和压缩比率之间按需取舍。您可利用经典图像压缩算法进行选择。然而，许多前沿的训练压缩方法需要为每个此类取舍训练一个独立模型，从而造成训练和推理时的计算量过大。我们的方法可针对此类问题提供一种简单方式，使用户无需针对不同损失优化一组模型，而是通过单一模型进行所有权衡，从而避免效率低下的情况。

• 图像压缩
  https://ai.googleblog.com/2016/09/image-compression-with-neural-networks.html

• 前沿的训练压缩方法
  https://openreview.net/forum?id=rkcQFMZRb

我们对 Balle 等人构建的训练图像压缩模型应用损失条件训练方法。这里的损失函数包含两项，即负责图像质量的重建项和负责压缩比率的压缩项。如下所示，我们的技术实现了训练单一模型即可覆盖一系列广泛的质量压缩权衡关系。

• Balle 等人
  https://openreview.net/forum?id=rkcQFMZRb

使用单一模型以不同的质量级别进行压缩，所有动画皆通过改变条件值，利用单一模型生成

应用：可调整的风格转化

我们演示的第二个应用是艺术风格转化，即将一张图像的内容根据另一张图像的风格来合成图像。最近的方法支持通过训练深度网络将图像实时风格化，并可变换多种风格。但是，对于每种指定的风格，这些方法均不支持用户对所合成输出的细节加以控制，例如，对图像风格化的程度，以及哪个风格特征需要进一步突出。如果风格化输入无法令用户满意，则必须使用不同的超参数训练多个模型，直到获得用户最喜爱的风格为止。

• 艺术风格转化
  https://openreview.net/pdf?id=HJe_Z04Yvr

• 最近的方法
  https://ai.googleblog.com/2016/10/supercharging-style-transfer.html

而借助我们提出的方法，训练单个模型即可覆盖一系列广泛的风格。此任务中，我们根据损失函数来调节模型，该损失函数具有对应五个损失项的系数，包括内容损失和针对风格化损失的其他四项。直观上说，内容损失决定风格化图像应该与原始内容的相似程度，而四个风格化损失则定义哪些风格特征将载入最终的风格化图像。我们在下方展示了所有这些系数发生变化时的单一模型输出：

可调整的风格转化：所有风格化皆通过改变调节值，从单一网络生成

很明显，模型会捕获每个风格中的多个变量，例如风格化程度、添加到图像的元素类型、其精确配置和位置等等。我们的网页以及交互演示中提供了更多示例。

• 网页
  https://sites.google.com/view/stochss/home

• 交互演示
  https://sites.google.com/view/stochss/adjusting-parameters

结论

我们提出了损失条件训练，这是一种简单通用的方法，只需训练单个深度网络即可完成之前需要大量独立训练网络才能解决的任务。虽然我们展示的是此方法在图像压缩和风格转化方面的应用，但其还可应用于许多领域。只要损失函数有可调整的系数，我们的方法就支持通过训练单一模型来对这些系数进行调整。

致谢

本文汇集了 Google Brain 团队多位研究员的工作成果：Mohammad Babaeizadeh、Johannes Balle、Josip Djolonga、Alexey Dosovitskiy 和 Golnaz Ghiasi。本篇博文的发布离不开他们所有人的重要贡献。来自 MS-COCO (http://cocodataset.org/#home)数据集和 unsplash.com 的图像用作演示目的。

— 推荐阅读 —