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},由输入空间 X ⊂ R c × h × w \mathcal X⊂\mathbb R^{c×h×w} X⊂Rc×h×w,标签空间 Y = { 1 , 2 , . . . , K } \mathcal Y = \{1, 2, ..., K\} Y={1,2,...,K}和 X × Y \mathcal X×\mathcal Y X×Y的联合分布 P X × Y P_{\mathcal X×\mathcal Y} PX×Y组成。给定一个在目标领域D上优化的教师模型 T ( x ; θ t ) T(x; θ_t) T(x;θt),KD通过直接最小化 original domain的经验风险( empirical risk),训练一个轻型学生模型来模仿教师的行为。
其中 l K L 和 l C E l_{KL}和l_{CE} lKL和lCE指的是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 ; θ g ) G(z; θ_g) G(z;θg)被训练来合成一个更有助于学生学习的分布 P G P_G PG。具体来说,MosaicKD是在 distributionally robust optimization分布稳健优化(DRO)框架的基础上发展起来的,该框架已被广泛用于解决领域转移问题[40, 16, 8]。给定一个预先定义的分布的距离度量 d ( ⋅ , ⋅ ) d(\cdot, \cdot) d(⋅,⋅),DRO框架的基本形式可以被形式化为以下内容。
在方程2中, l K L l_{KL} lKL表示学生学习的KL散度, d ( P G , P X ′ ) d(P_G, P_{X'}) d(PG,PX′)表示生成的样本和OOD数据之间的分布距离。超参数 ϵ \epsilon ϵ规定了以 P X ′ P_{X'} PX′为中心的球空间的半径。根据这个定义,DRO框架的目标是从搜索空间中寻找最坏情况分布,从而为搜索空间所覆盖的其他分布的经验风险建立一个上界。
理想情况下,如果原始训练数据的目标分布 P X P_X PX正好位于搜索空间中,其经验风险可以有效地被DRO框架所优化。然而,我们想说的是,这个前提在OOD设置中可能是有问题的,其中 X ≠ X ′ , Y ≠ Y ′ \mathcal X\ne \mathcal X',\mathcal Y \ne \mathcal Y' X=X′,Y=Y′。请注意,如果两个分布 P X 1 P_{X1} PX1和 P X 2 P_{X2} PX2在输入空间中在某些度量 d ( ⋅ ; ⋅ ) d(\cdot; \cdot) d(⋅;⋅)下是接近的,它们的标签空间 Y 1 \mathcal Y_1 Y1和 Y 2 \mathcal Y_2 Y2也应该是相似的[49]。
基于此,以OOD分布 P X ′ P_{X'} PX′为中心的小球空间内的分布很可能共享相同的标签空间,即 X ≈ X ′ , Y ≈ Y ′ \mathcal X≈\mathcal X',\mathcal Y≈\mathcal Y' X≈X′,Y≈Y′,这显然与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 ; θ g ) G(z; θ_g) G(z;θg)被训练为通过欺骗鉴别器网络 D ( x ; θ d ) D(x; θ_d) D(x;θd)来接近patch分布。请注意,我们的目标是合成完整的图像而不是patch的碎片,我们训练生成器G(x; θg)以产生完整分辨率的图像,并在生成的图像中制作patch(and craft patches on the fly)。让 C ( ⋅ ) C(\cdot) C(⋅)指的是裁剪操作,patch学习的目标可以表述如下。
其中 P X ′ P_{X'} PX′指的是OOD数据的分布, P z P_z Pz指的是隐变量z的先验分布。 C ( x ′ ) C(x') C(x′)和 C ( G ( z ) ) C(G(z)) 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 ; θ t ) = p ( y ∣ x , θ t ) T(x; θ_t) = p(y|x, θ_t) T(x;θt)=p(y∣x,θt),它对应于x属于第y类的信心。
为了对齐标签空间,一个天真的方法是最大化教师预测的置信度,即最小化熵项 H ( p ( y ∣ x , θ t ) ) H(p(y|x, θ_t)) H(p(y∣x,θt))。然而,先前的工作表明,这种简单的概率最大化可能只会导致一些 “垃圾样本”[14],没有太多有用的视觉信息用于学生训练。为了解决这个问题,我们提出了一个正则化的目标来对齐标签空间,其形式化为:
在式4中,第一项为式3中所述的discrimination loss,将局部模式正则化为可信。第二项是置信度最大化的熵损失,它对完整的图像进行处理,并将局部模式组合成所需的类别。同时利用式(3)对该目标进行优化,以保持局部patch的真实性。
DRO in MosaicKD. 如前文所述,公式3将局部模式正则化为可信,公式4将合成数据的标签空间与训练数据的标签空间进行对齐。他们共同构建了一个新的DRO框架搜索空间,如下所示:
其中 R ( G , D , T ) \mathcal R(G, D, T) R(G,D,T)是由方程3和4得出的正则化项。需要注意的是,优化生成式对抗网络如公式3相当于最小化两个patch分布的Jensen-Shannon散度,即 l J S D ( P X ′ p a t c h , P G p a t c h ) l_{JSD}(P^{patch}_{X'}, P^{patch}_G) lJSD(PX′patch,PGpatch),上述正则化可以写成:
正则化 R ( G , D , T ) \mathcal R(G, D, T) 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 ; θ g ) G(z; θ_g) G(z;θg)、判别器 D ( x ; θ d ) D(x; θ_d) D(x;θd)、固定的教师模型 T ( x ; θ t ) T(x; θ_t) T(x;θt)和学生 S ( x ; θ s ) S(x; θ_s) S(x;θs)在 L l o c a l \mathcal L_{local} Llocal、 L a l i g n \mathcal L_{align} Lalign和 L D R O \mathcal L_{DRO} LDRO的指导下被共同优化。
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 P_G PG为生成样本的分布。最小化Eq.(9)将强制类在合成过程中保持平衡。
9.1.2 Objectives of MosaicKD.
MosaicKD旨在解决一个分布式稳健优化(DRO)问题,具体如下
其中 R ( G , D , T ) ) ≤ ϵ \mathcal R(G, D, T))≤ \epsilon R(G,D,T))≤ϵ定义了搜索空间,即半径为满足R(G, D, T))=0的分布的球空间。中心分布的具体形式未知,但我们仍然可以训练一个生成器G来近似它。请注意,由于搜索空间的非微分条件,公式(5)是难以解决的。在拉格朗日对偶性的帮助下,我们可以将公式(5)的内部部分重新表达如下。
其中λ是拉格朗日乘数, λ ϵ λ\epsilon λϵ是一个常数项。如果 R ( G , D , T ) ) ≤ ϵ \mathcal R(G, D, T))≤ \epsilon R(G,D,T))≤ϵ,我们选择λ=0,即对R(G, D, T))不加限制,以获得最小成本。如果 R ( G , D , T ) ) > ϵ \mathcal R(G, D, T))> \epsilon 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中,搜索空间被 L l o c a l \mathcal L_{local} Llocal和 L a l i g n \mathcal L_{align} Lalign规范化,这使得生成的样本 locally authentic and globally legitimate。我们对上述正则化进行进一步研究,以显示它们对MosaicKD的意义。如图6所示,我们可视化了不同正则化下生成的样本。在图6(a)中,没有对生成器进行正则化处理,我们简单地将教师的信心最大化,这将导致一些劣质的样本[14]。在图6(b)中,判别器对完整的图像进行判断,在某种程度上,生成器将被OOD数据的类别语义所困,即合成一个像汽车的苹果或像马的枫叶。图(c)展示了MosaicKD的合成样本,它揭示了任务相关类的正确语义。