Mosaicking to Distill Knowledge Distillation from Out-of-Domain Data

Mosaicking to Distill: Knowledge Distillation from Out-of-Domain Data

在本文中，我们试图解决一项雄心勃勃的任务，即域外知识蒸馏（OOD-KD），它允许我们只使用可以以非常低的成本轻易获得的OOD数据来进行KD。诚然，由于不可知的领域差距，OOD-KD本质上是一项极具挑战性的任务。为此，我们介绍了一种方便而又令人惊讶的有效方法，被称为MosaicKD。

MosaicKD背后的关键在于，来自不同领域的样本有共同的局部模式（local patterns），即使它们的全局语义可能有很大的不同；这些共同的局部模式反过来可以被重新组合，以接近领域内的数据，并进一步缓解域差异。在Mosaic-KD中，这是通过一个四人的min-max游戏来实现的，在这个游戏中，在一个预先训练好的老师的指导下，一个生成器、一个鉴别器、一个学生网络以对抗的方式被共同训练。我们在各种基准的分类和语义分割任务中验证了MosaicKD，并证明它在OOD数据上产生的结果远远优于最先进的同行。Our code is available at https://github.com/zju-vipa/MosaicKD.

1 Introduction

我们提出了一种新assembling-by-dismantling approach，称为MosaicKD，它使我们能够利用OOD数据进行KD。我们的动机源于这样一个事实：即使来自不同领域的数据表现出不同的全局分布，它们的局部分布，如图像中的patch，可能会彼此相似。这一观察进一步激发了我们利用OOD和目标域数据所共享的局部模式来解决OOD-KD中的域转移问题。因此，MosaicKD的核心思想是合成域内数据，其中的局部模式模仿来自真实世界的OOD数据，而由局部分布组成的全局分布则有望骗过预先训练的教师模型。如图1所示，共享的局部模式被从OOD数据中提取出来，并重新组合成域内数据。

图1：自然图像有共同的局部模式。在MosaicKD中，这些局部模式首先从OOD数据中分解出来，然后集合起来合成域内数据，使OOD-KD成为可行的。

具体来说，在MosaicKD中，我们将OOD-KD问题设定为一个新的四人min-max游戏，涉及一个生成器、一个判别器、一个学生和一个老师，其中前三个需要学习，而最后一个是预先训练的，因此是固定的。生成器，就像之前的GANs一样，将随机噪声向量作为输入，在其他三个参与者的监督下，学习合成具有locally-authentic and globally-legitimate分布的合成域内样本。

另一方面，判别器学习区分从真实世界的OOD数据和合成样本中提取的局部patch。整个合成图像被反馈给预先训练的教师和将要训练的学生，在此基础上，教师为数据合成提供类别知识，学生模仿教师的行为以进行KD。四个参与者以对抗的方式相互协作加强，共同完成学生的训练。

简而言之，我们的贡献是对OOD-KD的第一次专门尝试，这是一个非常实用但在很大程度上被忽视的问题，通过一个新颖的方案，将领域内的数据镶嵌在一起实现。通过一个四人的min-max的游戏产生的合成样本，享有现实的局部结构和合理的全局语义，为从预先训练的教师那里提炼出可靠的知识奠定了基础。我们对各种基准的分类和语义分割任务进行了实验，并证明MosaicKD产生了真正令人鼓舞的结果，比其最先进的竞争对手在OOD数据上的结果要好得多。

Domain adaptation and generalization. 大多数学习算法都强烈依赖于这样一个前提，即用于训练的源数据和用于测试的目标数据是独立和相同分布的[50]，忽略了现实世界应用中经常遇到的OOD问题。在文献中，OOD问题通常通过domain generalization（DG）或adaptation（DA）来解决[3, 64]。Adaptation是一种流行的调整源域和目标域的技术[42, 11, 48]，它通常要求目标域在训练期间是可以访问的。近年来，领域适应已被扩展到开放集设置，其中训练和测试数据的标签空间是不同的[37]。相比之下，领域泛化与领域适应相似，但不需要目标领域的信息[3]。域泛化只对源域进行一次模型训练，并直接将模型应用于目标域[28, 12, 51, 16]。尽管DA和DG在监督学习中取得了成功，但在知识蒸馏的背景下，OOD问题仍然没有得到充分研究。

3 Out-of-Domain Knowledge Distillation

在不失一般性的情况下，我们在图像分类任务的背景下研究OOD问题。数据集的基础领域被定义为一个三triplet D = { X , Y , P X × Y } \mathcal D =\{\mathcal X ,\mathcal Y, P_{\mathcal X×\mathcal Y} \} D={X,Y,PX×Y​}，由输入空间$\mathcal{X}\subset {\mathbb{R}}^{c\times h\times w}$，标签空间 Y = { 1 , 2 , . . . , K } \mathcal Y = \{1, 2, ..., K\} Y={1,2,...,K}和$\mathcal{X}\times \mathcal{Y}$的联合分布$P_{\mathcal{X}\times \mathcal{Y}}$组成。给定一个在目标领域D上优化的教师模型$T(x;{\theta }_t)$，KD通过直接最小化 original domain的经验风险( empirical risk)，训练一个轻型学生模型来模仿教师的行为。

其中$l_{KL}和l_{CE}$指的是KL散度和交叉熵损失。然而，当原始训练域D不可用，而使用另一个域 D ′ = { X ′ , Y ′ , P X × Y } D'=\{\mathcal X',\mathcal Y', P_{\mathcal X×\mathcal Y} \} D′={X′,Y′,PX×Y​}的一些替代数据进行训练时，如果域的差距很大，方程1可能会出现问题。在这项工作中，我们专注于知识提炼中的out-of-domain问题，描述如下。

在OOD- kd中，由于OOD数据与原始训练数据的领域分歧，一些重要的模式可能会缺失，教师对这些模式的相关知识也可能没有恰当地传递给学生。为了解决OOD问题，我们提出了一种由域外样本到域内样本的l assembling-by-dismantling方法，有效地缓解了转移集与不可用训练集之间的域内差距，使KD在域外数据上也能应用。

4 Proposed Method

在没有原始训练数据X的情况下，由于数据域的分歧，直接在OOD集X’上最小化风险会有问题。在这项工作中，我们为OOD-KD引入了一种生成方法，被称为MosaicKD，其中一个生成器$G(z;{\theta }_g)$被训练来合成一个更有助于学生学习的分布$P_G$。具体来说，MosaicKD是在 distributionally robust optimization分布稳健优化（DRO）框架的基础上发展起来的，该框架已被广泛用于解决领域转移问题[40, 16, 8]。给定一个预先定义的分布的距离度量$d(\cdot ,\cdot )$，DRO框架的基本形式可以被形式化为以下内容。

在方程2中，$l_{KL}$表示学生学习的KL散度，$d(P_G,P_{X^{\prime }})$表示生成的样本和OOD数据之间的分布距离。超参数$ϵ$规定了以$P_{X^{\prime }}$为中心的球空间的半径。根据这个定义，DRO框架的目标是从搜索空间中寻找最坏情况分布，从而为搜索空间所覆盖的其他分布的经验风险建立一个上界。

理想情况下，如果原始训练数据的目标分布$P_X$正好位于搜索空间中，其经验风险可以有效地被DRO框架所优化。然而，我们想说的是，这个前提在OOD设置中可能是有问题的，其中$\mathcal{X}\ne {\mathcal{X}}^{\prime }，\mathcal{Y}\ne {\mathcal{Y}}^{\prime }$。请注意，如果两个分布$P_{X1}$和$P_{X2}$在输入空间中在某些度量$d(\cdot ;\cdot )$下是接近的，它们的标签空间${\mathcal{Y}}_1$和${\mathcal{Y}}_2$也应该是相似的[49]。

基于此，以OOD分布$P_{X^{\prime }}$为中心的小球空间内的分布很可能共享相同的标签空间，即$\mathcal{X}\approx {\mathcal{X}}^{\prime }，\mathcal{Y}\approx {\mathcal{Y}}^{\prime }$，这显然与OOD的设置相冲突。为此，原始训练数据的目标域可能不被搜索空间所覆盖，不能被DRO框架所约束。对这个问题的补救措施是使用一个足够大的半径。不幸的是，这只会导致难以解决的搜索空间，充斥着无意义的分布。

4.1 Mosaicking to Distill

如上所述，建立在OOD数据上的搜索空间不足以建立一个可靠的优化上界。为了解决这个问题，MosaicKD引入了一种新的方式来构建基于局部patch的搜索空间。我们的动机源于这样一个事实：自然图像的模式通常是分层组织的，其中高层模式是由低层模式组合而成的。尽管原始训练数据X和OOD数据X’的领域是不同的，但它们的局部模式仍然可能相互相似。例如，"毛皮 "的图案可以由不同领域的不同动物物种共享。需要注意的是，每个图像都是由局部斑块组装而成的，我们提出了一个通过assembling-by-dismantling的策略来重新组织共享的局部patch并合成域内数据用于训练。

Patch Learning. 实现MosaicKD的第一步是从OOD数据X’中提取局部模式，并估计生成patch分布。给定OOD数据集 X ′ = { x 1 ′ , x 2 ′ , . . . , x N ′ ; x i ′ ∈ R H × W × 3 } X'=\{x'_1, x'_2, ..., x'_N; x'_i∈\mathbb R^{H×W×3}\} X′={x1′​,x2′​,...,xN′​;xi′​∈RH×W×3}，我们通过L×L裁剪获得patch，产生patch数据集 C = { c 1 , c 2 , . . . , c M ; c i ∈ R L × L × 3 } C=\{c_1, c_2, ..., c_M; c_i∈\mathbb R^{L×L×3}\} C={c1​,c2​,...,cM​;ci​∈RL×L×3}。

patch大小L是MosaicKD的一个重要超参数。例如，如果L=W=H，每个patch将覆盖一个完整的图像，其中包含原始图像的所有高级特征。当我们把patch的大小减少到L=1时，那么每个patch只包含低层次的颜色信息。显然，小尺寸的L比大尺寸的L能导致更普遍的模式，这些模式更有可能被不同领域所共享。此外，增加patch大小将引入更多的结构信息，使patch的分布更接近于完整图像的分布。

在这项工作中，我们将局部patch学习建模为一个生成问题，其中一个生成器$G(z;{\theta }_g)$被训练为通过欺骗鉴别器网络$D(x;{\theta }_d)$来接近patch分布。请注意，我们的目标是合成完整的图像而不是patch的碎片，我们训练生成器G(x; θg)以产生完整分辨率的图像，并在生成的图像中制作patch(and craft patches on the fly)。让$C(\cdot )$指的是裁剪操作，patch学习的目标可以表述如下。

其中$P_{X^{\prime }}$指的是OOD数据的分布，$P_z$指的是隐变量z的先验分布。$C(x^{\prime })$和$C(G(z))$指的是来自OOD数据和生成数据的裁剪patch。公式(3)与vanilla GANs[13]中的目标的主要区别在于patch级的判别，MosaicKD只将局部模式正则化，让全局结构不受限制。如上所述，全局模式可以由局部模式组装而成，MosaicKD通过标签空间对齐，将这些assembled组装起来，合成域内数据。

Label Space Aligning. 由于公式（3）中没有引入patch间的限制(As no inter-patch restriction is introduced)，生成器可能只生成全局语义无意义的图像，尽管其局部模式是可信的。在这一步，我们转向教师模型，以获得更多的域内数据综合信息。在KD中，教师模型是在原始训练数据X上训练的，其输出是条件概率$T(x;{\theta }_t)=p(y|x,{\theta }_t)$，它对应于x属于第y类的信心。

为了对齐标签空间，一个天真的方法是最大化教师预测的置信度，即最小化熵项$H(p(y|x,{\theta }_t))$。然而，先前的工作表明，这种简单的概率最大化可能只会导致一些 “垃圾样本”[14]，没有太多有用的视觉信息用于学生训练。为了解决这个问题，我们提出了一个正则化的目标来对齐标签空间，其形式化为：

在式4中，第一项为式3中所述的discrimination loss，将局部模式正则化为可信。第二项是置信度最大化的熵损失，它对完整的图像进行处理，并将局部模式组合成所需的类别。同时利用式(3)对该目标进行优化，以保持局部patch的真实性。

DRO in MosaicKD. 如前文所述，公式3将局部模式正则化为可信，公式4将合成数据的标签空间与训练数据的标签空间进行对齐。他们共同构建了一个新的DRO框架搜索空间，如下所示:

其中$\mathcal{R}(G,D,T)$是由方程3和4得出的正则化项。需要注意的是，优化生成式对抗网络如公式3相当于最小化两个patch分布的Jensen-Shannon散度，即$l_{JSD}(P_{X^{\prime }}^{patch},P_G^{patch})$，上述正则化可以写成：

正则化$\mathcal{R}(G,D,T)$迫使生成器利用OOD数据的局部模式进行数据合成，这导致了一个定义在所有可能的patch组合方案上的特殊搜索空间。 与传统的DRO不同，MosaicKD使用小半径进行稳健优化，目标域可以被搜索空间覆盖。我们放宽了方程5的正则化，得到了一个可优化的DRO训练目标，形式化为：

4.2 Optimization

Patch Discriminator. 为了提高训练效率，方程3中的判别器可以作为一个Patch GAN[19]来实现，它具有精心设计的感受野和补丁重叠(patch overlap)。具体来说，我们将几个卷积层堆叠起来，建立一个全卷积网络，其输出是一个 score map，而不是一个单一的真假标量。每个score unit接受一个L×L的局部斑patch进行判别。我们在 score map上应用一个额外的步长为s的降采样，以控制patch之间的重叠。一个大的步长s会导致更多的独立斑patch，从而有效地减少OOD图像中的结构限制。

Full Algorithm. MosaicKD的完整算法总结在Alg. 1，其中一个生成器 $G(z;{\theta }_g)$、判别器$D(x;{\theta }_d)$、固定的教师模型$T(x;{\theta }_t)$和学生$S(x;{\theta }_s)$在${\mathcal{L}}_{local}$、${\mathcal{L}}_{align}$和${\mathcal{L}}_{DRO}$的指导下被共同优化。

5.2 Results of Knowledge Distillation

CIFAR-100. 表1报告了对CIFAR-100数据集的知识提炼结果。这里我们使用CIFAR-10、ImageNet、Places365和SVHN作为OOD数据来评估MosaicKD的OOD设置。我们将提出的MosaicKD与各种基线进行比较，包括无数据KD方法（DAFL[7]、ZSKT[33]、DeepInv. [61], DFQ[8]）和从最先进的KD方法（BKD[18], Balanced[35], FitNet[41], RKD[38], CRD[47]和SSKD[54]）天真地适应的OOD-KD方法。

如表1所示，尽管分布不匹配，但传统KD方法仍然可以从OOD数据中学习到一些有用但不完全的知识(即比随机猜测的性能明显优越)，这说明OOD数据和训练数据之间存在共享模式。此外，还进行了一些探索，以评估类别平衡和代表性转移对OOD-KD的重要性。

首先，我们根据教师的预测，通过重新抽样稀缺类别来平衡OOD数据。但是结果表明，平衡OOD数据并不能帮助学生学习正确的班级信息，因为OOD数据中的大部分样本只是错误分类的离群值。在OOD设置的上下文中，balance操作可能会导致对离群值的过拟合，从而进一步降低学生在测试集中的表现。

如前所述，OOD数据和原始数据可能共享一些局部模式，这些模式可以通过网络的浅层来提取。我们应用四种表示方法，即FitNet、RKD、CRD和SSKD来研究它们在OOD-KD中的作用。与注重instance relation的RKD相比，我们发现response-based methods，如Fitnet，可以在OOD环境中转移更多的有用信息，学生直接模仿教师的中间输出的教师。一般来说，转移低层次的特征有时对OOD-KD有帮助。然而，请注意，CRD是在倒数第二层提取的高级表征上工作的，转移这些知识可能不适合OOD-KD，因为高级特征可能与目标任务不相关。

在这项工作中，我们将OOD-KD问题作为一个生成性问题来处理，而不是直接使用OOD数据进行训练。所提出的方法利用OOD数据的局部模式进行数据合成，其中一些任务相关的模式将从共享的局部patch中 "assembled "出来。结果表明，这些重新组合的数据可以有效地将知识从教师传递给学生。在表1中，我们还将我们的方法扩展到不同类型的OOD集。我们发现，MosaicKD的性能与OOD数据和原始数据之间的领域分歧程度有关。例如，ImageNet是一个物体识别数据集，而Places365是一个场景分类数据集。结果显示，对于目标数据CIFAR-100，与Places365相比，MosaicKD可以在ImageNet上取得更好的性能。

5.3 Quantitative Analysis

Data balance and FID. 图3提供了OOD数据和生成样本的一些统计信息，包括教师预测的类别平衡和每类的FID分数。类别是根据它们的百分比来排列的。请注意，原始的CIFAR-10数据集只包含10个类别，与CIFAR-100的100个类别相比，这是非常有限的。如图3（a）所示，CIFAR-10中缺少一些CIFAR-100的类别。此外，OOD数据和原始训练数据之间较大的FID也表明，即使一些样本被老师归类到某些类别，它们仍然可能属于异常值。相比之下，我们的方法成功地平衡了不同的CIFAR-100类别，缓解了领域差距（较低的类别FID），特别是对于不平衡的类别。

表4:不同patch大小的学生测试准确度(%)。同时报告了OOD数据与原始数据之间的Patch FID评分。结果表明，我们的方法需要更小的patch大小来处理严重的域差异。

The influence of patch size. patch大小在我们的方法中起着至关重要的作用，它决定了数据合成的灵活性。最佳patch大小实际上取决于OOD数据和域内数据之间的差异。对于域内差异较大的OOD数据，由于局部相似度有限，通常需要小的patch。如表4所示，我们用不同的patch尺寸评估了我们的方法，并报告了学生模型的测试精度以及patch的FID。根据表4，我们发现对于像CIFAR-10和Places365这样的OOD数据，大的patch尺寸（如18）可以用于学生学习。然而，对于SVHN数据集来说，较小的patch尺寸会更合适，因为SVHN与CIFAR-100有严重的分歧。

Figure 4: 有和没有补丁学习的合成数据的可视化。没有补丁学习的GANs会被OOD数据所困，无法为不同的类别提供正确的语义（蓝色强调）。在我们的方法中，语义可以正确地与目标领域对齐。

Relation to data-free knowledge distillation.

无数据算法和MosaicKD的共同点在于，它们都是通过数据合成来解决KD问题。然而，无数据KD利用一些简单的先验条件，如类别置信度[7]和高斯假设[61]，忽略了自然图像中的结构细节。相比之下，MosaicKD以一种通过assembling-by-dismantling的方式实现数据合成，其中自然模式可以被利用来提高合成质量。不同方法的训练曲线可以在图5中找到。

9 Appendix

在这份文件中，我们提供了由于篇幅限制而无法放入主稿的细节和补充材料。具体来说，我们在第9.1节提供了MosaicKD的优化细节，在第9.2节提供了实验设置，在第9.3节提供了更多的实验结果。

9.1 Optimization Details

9.1.1 Alleviating Mode Collapse.

在这项工作中，我们部署了一个生成器来合成知识蒸馏的transfer set。然而，众所周知，GANs会受到模式崩溃的影响，不能产生多样化的模式。为此，我们利用OOD数据和合成数据来训练我们的学生模型，这样生成器就不需要为KD合成所有样本。此外，我们还部署了一个额外的平衡损失，以缓解训练期间的模式崩溃，定义为：

式中p(y|x， θt)为softmax后的概率预测，$P_G$为生成样本的分布。最小化Eq.(9)将强制类在合成过程中保持平衡。

9.1.2 Objectives of MosaicKD.

MosaicKD旨在解决一个分布式稳健优化（DRO）问题，具体如下

其中$\mathcal{R}(G,D,T))\le ϵ$定义了搜索空间，即半径为满足R(G, D, T))=0的分布的球空间。中心分布的具体形式未知，但我们仍然可以训练一个生成器G来近似它。请注意，由于搜索空间的非微分条件，公式（5）是难以解决的。在拉格朗日对偶性的帮助下，我们可以将公式（5）的内部部分重新表达如下。

其中λ是拉格朗日乘数，$\lambda ϵ$是一个常数项。如果$\mathcal{R}(G,D,T))\le ϵ$，我们选择λ=0，即对R(G, D, T))不加限制，以获得最小成本。如果$\mathcal{R}(G,D,T))>ϵ$，则应采用大的λ作为惩罚。根据公式(11)的推导，我们得到了难以解决的公式(5)的一个宽松版本。(5)的放松版本，表示如下:

9.2 Experimental Settings

Datasets. 所提出的方法在两个主流视觉任务上进行了评估，即图像分类和语义分割，在四个用于教师训练的标记数据集和四个用于学生学习的OOD数据上，如表6所总结。请注意，CIFAR-100、ImageNet和Places365可能包含域内类别。如算法9.2所述，我们通过选择预测置信度低的样本，从完整的ImageNet和Places365数据集中精心制作OOD子集。这些OOD子集可以被看作是CIFAR-100的域外数据。此外，我们将OOD数据的大小调整为与域内数据相同的分辨率，例如，CIFAR-100为32×32，细粒度数据集为64×64，而NYUv2为128×128。

Network Training. 在这项工作中，所有的教师模型都是使用表6中列出的域内数据集进行训练，并使用交叉熵损失。我们使用SGD优化器，{lr = 0.1, weight_decay =1e - 4, momentum = 0.9}，用余弦退火调度器训练每个模型200 epoch。在知识蒸馏中，学生模型是用无标签的数据集训练，只利用了教师的 soft targets。我们使用与教师训练相同的训练协议，并在测试集上报告最佳的学生准确性。我们使用Adam进行优化，用超参数{lr = 1e - 3, β1 = 0.5, β2 = 0.999}作为生成器和判别器。

Generator and Discriminator 表7和表8说明了CIFAR-100数据集的GAN架构。对于CUB-200（64×64）和NYU（128×128），我们增加了更多的卷积层和上采样层以生成高分辨率的图像。

9.3 More Experimental Results

9.3.1 Patch Overlap

考虑到固定的patch大小，patch之间的重叠在patch学习中起着重要作用。重叠是由patch判别器中的区间采样控制的。请注意，判别器产生一个预测图来预测原始图像上的每个小区域，这意味着预测应该共享较少的信息。我们在最后的判别中加入一个prediction stride，以控制patch重叠。表9显示了在不同patch重叠的情况下获得的学生准确率，其中较大的stride对应于较小的重叠量。结果显示，增加stride对学生的准确性没有好处。

请注意，我们使用patch GAN架构进行patch学习，它包含了discriminator内部的stride操作。这些stride操作已经为patch 学习提供了适当的重叠。此外，stride越大，训练样本越少，这可能对GAN训练有害。

9.3.2 DRO Regularization

在MosaicKD中，搜索空间被${\mathcal{L}}_{local}$和${\mathcal{L}}_{align}$规范化，这使得生成的样本 locally authentic and globally legitimate。我们对上述正则化进行进一步研究，以显示它们对MosaicKD的意义。如图6所示，我们可视化了不同正则化下生成的样本。在图6(a)中，没有对生成器进行正则化处理，我们简单地将教师的信心最大化，这将导致一些劣质的样本[14]。在图6(b)中，判别器对完整的图像进行判断，在某种程度上，生成器将被OOD数据的类别语义所困，即合成一个像汽车的苹果或像马的枫叶。图（c）展示了MosaicKD的合成样本，它揭示了任务相关类的正确语义。