论文笔记：Collaborative Distillation for Ultra-Resolution Universal Style Transfer

摘要

（1）通用的风格迁移方法通常利用对大量图像集合上进行预训练的Deep Convolutional Neural Network（CNN）模型（例如VGG-19）的丰富表示形式。尽管有效果，但由于内存有限，其大尺寸模型无法处理 ultra-resolution图像，因此其应用受到了严重限制。

（2）作者提出了一种新的知识蒸馏方法（称为协作蒸馏），用于基于编码器-解码器的神经网络风格转移，以减少卷积滤波器。主要思想是通过发现（编码器-解码器对构成了排他性的协作关系，这被视为样式迁移模型的一种新型知识）。

（3）此外，为了克服在应用协作蒸馏时出现的特征尺寸不匹配的问题，引入了线性嵌入损失，以驱动学生网络学习教师特征的线性嵌入。 大量的实验表明，即使模型尺寸减小了15.5倍，我们的方法在应用于不同的通用样式转换方法（WCT和Ada​​IN）时仍然有效。

（4）大量的实验表明新方法适用于不同的通用风格迁移方法，而且在具有压缩模块的WCT下，首次在12GB的GPU上实现了超过4000万像素的 ultra-resolution的通用风格迁移。

图1

1. Introduction

通用神经样式转换（neural style transfer, NST）致力于将内容图像与来自任何参考图像的新样式进行合成。 这通常需要具有相当大能力的模型来提取有效的表示形式，以捕获任意样式的统计数据。

• 基于神经网络的最新通用样式迁移方法一致地表明，使用像VGG19 这样的预训练深层神经网络提取的表示，可以同时实现视觉上令人愉悦的转移结果以及对任意样式图像的泛化能力。

• 但是由于硬件上的显存有限，VGG-19的大模型尺寸极大地限制了输入图像的分辨率。到目前为止，当前的通用样式迁移方法仅在具有12GB内存的单个GPU上的结果大约1百万像素（例如1024×1024像素）。

近年来在模型压缩领域已经看到了快速的发展，其目的是在不损失大量性能的情况下减少大型CNN模型的参数。尽管取得了进展，但大多数模型压缩方法仅着眼于high-level任务，例如分类和检测。对于low-level任务的压缩模型仍然较少探索。

• 知识蒸馏（Knowledge distillation, KD）是一种有前途的模型压缩方法，它可以将大型网络（称为教师网络teacher）的知识转移到小型网络（称为学生网络student），在这种情况下，knowledge可以被softened（可以反映固有的类相似性结构（称为dark knowledge）或样本关系（可以反映不同样本之间的相似性结构）。这些knowledge可作为one-hot标签的补充信息，因此可以提高学生网络student的表现。但是，这些额外信息主要取决于标签，因此几乎不适用于low-level任务。

• 在低级视觉任务（例如，神经风格迁移 neural style network）中，什么是dark knowledge仍然是一个悬而未决的问题。 同时，基于编码器-解码器的模型被广泛用于神经样式转换中，其中通常通过编码器的知识来训练解码器。值得注意的是，它们共同在风格化过程中构建了排他性的协作关系，如图2所示：
  

图2

• 由于解码器$D$被训练为只能与编码器$E$一起使用，因此如果另一个编码器$E^{\prime }$也可以与D一起使用，则意味着$E^{\prime }$可以在功能上扮演E的角色。基于此思想，我们提出了一个新的knowledge来蒸馏神经样式迁移(neural style trransfer NSP)中的深层模型：编码器和解码器之间的协作关系。
• 给定一个冗余的大型编码器（例如VGG-19），我们提出了两步压缩方案：
  首先，为编码器训练协作网络，即我们语境中的解码器;
  其次，将大型编码器替换为小型编码器，然后在固定协作器的情况下训练小型编码器。
  由于小型编码器通常具有较少的通道，因此其输出特征的维度比大型编码器小。 因此，小型网络无法直接与协作者合作。 为了解决这个问题，作者限制student学习teacher输出的线性嵌入，以便在将student的输出馈入协作者之前，可以通过student输出的简单线性组合来重建teacher的输出。
• 我们在提出的方法中不限制特定的协作形式。而且它可以应用于两种不同的最新通用的样式迁移方案：WCT （其中协作（collaboration）是图像重建）和AdaIN （其中协作（collaboration）是样式转移）。

2. Related Work

2.1 Style transfer

神经样式迁移（neural style transfer），它利用了来自预训练的VGG19模型的深层特征，并通过匹配生成图像和给定样式图像之间的二阶统计量来实现样式化。但是，所有这些基于神经网络的方法的一个共同局限性在于，在内存有限的情况下，它们无法处理 ultra-resolution内容和样式图像。
一些方法通过学习用于特定样式示例或分类的小型前馈网络，获得高分辨率的风格化结果（最多10兆像素，例如3000×3000像素），但它们并不能概括为其他unseen的样式。 相反，作者的目标是仅使用一个模型就可以实现通用样式的 ultra-resolution图像样式转换。

2.2 Model compression

模型压缩和加速最近也引起了很多关注，其目的是在不显着影响性能的情况下获得更小更快的模型。现有方法大致可分为五类，即低阶分解，修剪，量化，知识蒸馏和紧凑型架构重新设计或搜索。但是，这些方法主要在high-level视觉任务中进行探索，通常是分类和检测。 很少有方法关注诸如样式迁移之类的low-level视觉任务，其中许多方法也受到CNN巨大模型尺寸的限制。

与仅需要维护全局语义信息以保持精度的用于high-level视觉的CNN压缩不同，用于low-level视觉的模型压缩的额外挑战可能是如何维护局部结构，例如局部样式转移中的纹理和颜色多样性。

• 作者开发了一种深度监督的知识蒸馏（deeply-supervised knowledge distillation）方法，以从预训练的冗余VGG-19中学习更小的模型。 压缩后的模型带来了超过15倍的参数和减少的计算量。 更重要的是，模型尺寸的减小使得可以在 ultra-resolution图像上进行通用样式转换。

据我们所知，只有最近的一项工作使用GAN来学习ultra-resolution图像上的非配对样式迁移网络。 但是，他们通过处理图像子样本然后将它们合并回整个图像来实现此目的。相比之下，我们的方法从根本上降低了模型复杂度，可以直接处理整个图像 。

3. Proposed Method

3.1 Collaborative distillation

风格不可知的风格化方法通常采用编码器-解码器方案，以学习用于风格渲染的深层表示，然后将其转换回风格化图像。 由于样式信息未在模型中直接编码，因此编码器部分必须足够表达，用俩提取通用样式的信息表示。

考虑到其庞大的容量和层次结构，现有方法通常选择VGG-19作为编码器。至于解码器，取决于不同的风格化方案，它可以与编码器具有不同的实体协作关系。

这里讨论了两种SOTA迁移方法，即WCT和AdaIN。
（1）对于WCT：样式化过程是使用样式特征的二阶统计量将白化和着色变换应用（ Whitening and Coloring Transform）于内容特征。 然后，解码器将变换后的内容特征转换为图像。 因此，解码器训练不直接涉及样式化。 WCT中的合作关系本质上是图像重建。
（2）对于AdaIN：与WCT不同，其解码器训练直接涉及样式化。 将两个图像（内容和样式）馈入编码器，然后在特征空间中，内容特征被样式特征的统计信息（均值和方差）所渲染进行呈现。最后，解码器将渲染的内容特征反转回风格化图像。 就内容（或样式）的距离而言，样式化的图像应该与内容（或样式）接近。 因此，AdaIN的协作关系是样式转移。

尽管以上两种方案在范式上有所不同，但它们都是基于编码器-解码器的，并且解码器是通过编码器的知识来训练的。这意味着，在训练解码器期间，编码器的知识会泄漏到解码器中。

假定并凭经验确认，解码器D只能与匹配的编码器E一起工作。 对于另一个编码器$E^{\prime }$，即使它具有与E相同的体系结构，D和$E^{\prime }$也不能一起工作（见图2）。这种排他性表明解码器具有一些特定于其编码器的固有信息。 如果我们也找到使网络$E^{\prime }$与D兼容的方法，则意味着$E^{\prime }$可以在功能上替代原始编码器E。

如果$E^{\prime }$的模型远小于E，则可以达到模型压缩的目标。 基于此思想，作者提出了一种新的，专门针对（neural style transfer， NST）的蒸馏方法，称为协作蒸馏，包括两个步骤：

• 第一步，我们为大型编码器训练一个协作网络。如图3所示：

图3

• 图3显示了：协作蒸馏框架（最佳观看颜色）。(a)和(b)描绘了两种用于通用的neural style transfer的编码器/解码器协作关系：分别用于WCT的图像重建和用于AdaIN 的样式迁移。 蓝色箭头表示训练协作者网络（即解码器）时的前进路径。 当训练小型编码器（“ SEncoder”）在功能上替换原始编码器（“ Encoder”）时，绿色箭头显示前进路径。©显示了提出的线性嵌入方案，以解决特征尺寸不匹配的问题，加强对小型编码器SEncoder中间层的监督。

对于WCT，训练解码器以使其特征尽可能地忠实于输入图像（即图像重建），同时采用像素重建损失和感知损失（ the pixel reconstruction loss and the perceptual loss）。

其中$k\in 1,2,3,4,5$表示VGG-19的第k个阶段; $F(i)$表示ReLU_i_1层的特征图；${\lambda }_p$是权衡感知损失和图像重建损失的权重；$I_o$和$I_r$分别表示原始图像和重构图像。

Perceptual Loss 讲解

对于AdaIN ，解码器直接参与样式迁移。 因此，其解码器的损失函数由内容损失和样式损失组成：

其中$G$是描述style的Gram矩阵，${\lambda }_s$是权衡style loss和content loss的权重； 下标“ st”，“ c”和“ s”分别表示样式化图像，内容图像和样式图像。

有了合作者后，

• 算法的第二步是用小的编码器$E^{\prime }$替换原始编码器E。 为简单起见，在本文中，我们将$E^{\prime }$与E的体系结构相同，但每层中的过滤器较少。 我们期望小型编码器$E^{\prime }$在功能上可以与原始编码器E等效，如图3（a）和（b）所示。 与第一步类似，根据特定的协作任务，表示为$L_{collab}$的the collaborator network loss可以采取不同的形式。

本文的任务，对于WCT，$L_{collab}=L_r$；对于AdaIN，$L_{collab}=L_{st}$。

3.2 Linear embedding

在提出的协作蒸馏方法中，小型编码器与解码器网络连接。 在他们的接口处出现了特征尺寸不匹配的问题。具体来说，如果原始编码器输出尺寸为$C\times H\times W$的特征，则解码器的输入也被假定为尺寸为$C\times H\times W$，但是，由于小型编码器的滤波器数量较少，它将输出尺寸为$C^{\prime }\times H\times W(C^{\prime }<C)$的特征，解码器不能接受这样的特征。
为了解决这个问题，我们首先看一下channel numbers如何在样式化过程中发挥作用。图像的样式由从VGG-19提取的深层特征的gram矩阵描述。

其中，$F$是从VGG-19的某些卷积层提取的$C\times HW$大小的深层特征；$G$表示大小为$C\times C$的Gram矩阵；$T$代表矩阵转置。由于我们旨在压缩这些特征（即它们被认为是多余的），因此可以将其表示为$F$是较低维某些特征基向量的线性组合。

其中$Q$是大小为$C\times C^{\prime }$的变换矩阵，$F^{\prime }$是大小为$C^{\prime }\times HW$的特征基础矩阵，可以将其视为原始深层特征$F$的线性嵌入。这样，新特征$F^{\prime }$的Gram矩阵$G^{\prime }=F^{\prime }\cdot F^{\prime T}$具有与原始Gram矩阵$G$相同数量的特征值。换句话说，如果采用$F^{\prime }$代替原始冗余的$F$，同样可以保持style description power。在我们的方法中，F是原始编码器的输出，$F^{\prime }$是小型编码器的输出。 通过没有非线性激活功能的全连接层学习变换矩阵$Q$，以实现线性假设。 因此，线性嵌入损失可以表示为：

• 以上提出的解决方案不限于最终输出层。 对于小型编码器SEncoder的中间层，它也可以找到一个应用。具体而言，我们将线性嵌入应用于原始编码器和小型编码器之间的其他四个中间层(ReLU_k_1，$k\in 1,2,3,4$)，如图3（c）所示。

为此，我们有两个动机。

• 首先，在提出的方法中，当使用SGD算法训练小型编码器时，其唯一的梯度源将是解码器，该解码器通过全连接层Q。但是，Q没有太多参数，因此它将实际上形成了信息瓶颈，减慢了Student的学习速度。 而将这些分支插入网络的中间层后，它们将为Student注入更多的梯度，从而促进其学习，尤其是对于那些倾向于梯度消失的深层网络。
• 其次，在神经样式传递中，图像的样式通常由许多中间层的特征来描述。

因此，有必要对这些层增加更多的监督，以确保它们不会在随后的样式转移中失去太多的样式描述能力。
因此，作者提出的算法中训练小型编码器的总损失可总结为：

其中$\beta$是权衡两个损失项的权重因子。