[论文解读] Collaborative Distillation for Ultra-Resolution Universal Style Transfer

Huan Wang, Yijun Li+3 more|arXiv (Cornell University)|Mar 18, 2020

Generative Adversarial Networks and Image Synthesis参考文献 63被引用 7

一句话总结

本文提出协同蒸馏（Collaborative Distillation），一种新颖的知识蒸馏方法，通过利用编码器-解码器对之间的独特协作关系，压缩通用风格迁移中的VGG-19编码器。通过训练学生编码器学习教师特征的线性嵌入，该方法可在12GB GPU上实现超分辨率风格迁移（最高达4000万像素以上），在参数量减少15.5倍的同时保持视觉质量，达到当前最优性能，适用于多种风格迁移框架。

ABSTRACT

Universal style transfer methods typically leverage rich representations from deep Convolutional Neural Network (CNN) models (e.g., VGG-19) pre-trained on large collections of images. Despite the effectiveness, its application is heavily constrained by the large model size to handle ultra-resolution images given limited memory. In this work, we present a new knowledge distillation method (named Collaborative Distillation) for encoder-decoder based neural style transfer to reduce the convolutional filters. The main idea is underpinned by a finding that the encoder-decoder pairs construct an exclusive collaborative relationship, which is regarded as a new kind of knowledge for style transfer models. Moreover, to overcome the feature size mismatch when applying collaborative distillation, a linear embedding loss is introduced to drive the student network to learn a linear embedding of the teacher's features. Extensive experiments show the effectiveness of our method when applied to different universal style transfer approaches (WCT and AdaIN), even if the model size is reduced by 15.5 times. Especially, on WCT with the compressed models, we achieve ultra-resolution (over 40 megapixels) universal style transfer on a 12GB GPU for the first time. Further experiments on optimization-based stylization scheme show the generality of our algorithm on different stylization paradigms. Our code and trained models are available at https://github.com/mingsun-tse/collaborative-distillation.

研究动机与目标

• 解决因VGG-19模型过大和内存消耗过高导致的超分辨率神经风格迁移限制。
• 克服压缩编码器-解码器风格迁移模型时的特征维度不匹配挑战。
• 识别并利用编码器与解码器之间的独特协作关系，作为知识蒸馏的新形式。
• 在消费级硬件（如12GB GPU）上实现高效、高分辨率的风格迁移，同时不损失视觉质量。
• 在不同风格迁移范式（包括前馈与基于优化的方法）中验证该方法的泛化能力。

提出的方法

• 提出两阶段压缩方案：首先为大型编码器训练一个协作者（解码器），然后在保持协作者固定的情况下，用更小的编码器替换大型编码器。
• 引入线性嵌入损失，将学生编码器的低维输出与教师编码器的高维输出对齐，解决特征尺寸不匹配问题。
• 将编码器与解码器之间的独特协作关系视为一种新型知识形式，区别于基于标签或激活的蒸馏。
• 将该方法应用于两种最先进的通用风格迁移框架：WCT（基于图像重建）和AdaIN（基于风格迁移）。
• 在训练过程中结合感知损失、内容损失与协同蒸馏损失，以保持风格迁移质量。
• 将压缩后的编码器应用于基于优化的风格迁移（Gatys et al.），以验证其在不同风格迁移范式中的泛化能力。

实验结果

研究问题

• RQ1 神经风格迁移中编码器与解码器之间的独特协作关系能否作为知识蒸馏的新知识形式加以利用？
• RQ2 在压缩编码器的同时，如何有效解决其与预训练解码器之间的兼容性问题？
• RQ3 通过协同蒸馏实现的模型压缩在超高分辨率风格迁移（如4000万像素以上）中，能在多大程度上保持视觉质量？
• RQ4 所提出的方法是否能在不同风格迁移框架中实现泛化，包括前馈与基于优化的方法？
• RQ5 压缩后的模型是否在视觉质量与分辨率可扩展性方面优于现有压缩方法（如滤波器剪枝）？

主要发现

• 所提出的协同蒸馏方法首次实现在单张12GB GPU上对超分辨率神经风格迁移（最高达10240×4096像素）进行处理，耗时约31秒完成风格迁移。
• 在模型参数量减少15.5倍的情况下，压缩后的模型在多个数据集和风格迁移方法中均实现了与原始VGG-19相当的视觉效果。
• 用户研究表明，压缩模型在偏好度上优于基于滤波器剪枝的模型（如FP），且与原始模型的偏好得分相近。
• 该方法在基于优化的风格迁移（Gatys et al.）中也表现出良好泛化能力，压缩后的编码器生成的结果几乎与原始模型无法区分。
• 对解码器进行蒸馏会降低视觉质量，表明由于解码器受到像素损失与感知损失的强监督，应保持其完整性。
• 实验表明，标准知识蒸馏方法（如软标签蒸馏）在风格迁移中效果不佳，因为分类准确率的微小提升无法转化为感知质量的改善。