Pixel Distillation: Cost-flexible Distillation across Image Sizes and Heterogeneous Networks【小白读论文】

摘要：
以往的知识蒸馏（KD）方法主要集中在压缩网络架构上，但由于传输带宽和成像设备等成本与图像大小有关，因此在部署上不够彻底。 因此，我们提出像素蒸馏，将知识蒸馏扩展到输入级别，同时打破架构限制。
我们首先提出了一种输入空间表示蒸馏（ISRD）机制，将空间知识从大图像转移到学生的输入模块，这可以促进 CNN 和 ViT 之间的稳定知识转移。
进一步建立了教师-助理-学生（TAS）框架，将像素蒸馏分解为模型压缩阶段和输入压缩阶段，显着降低了像素蒸馏的整体复杂度和中间知识蒸馏的难度。
最后，我们通过 TAS 的对齐特征保存（AFP）策略将像素蒸馏应用于目标检测，该策略通过操纵助手的特征和锚点来对齐每个阶段检测器的输出尺寸。 图像分类和目标检测的综合实验证明了我们方法的有效性。

本文的贡献有四个方面：
• 我们提出了一种称为像素蒸馏的新蒸馏方案，它为小输入尺寸的边缘设备建立灵活的 KD 方案提供了早期尝试。像素蒸馏将知识蒸馏的思想推广到输入级别，其中大输入用于重型教师模型，小尺寸用于轻量级学生模型。
• 我们提出了一种输入空间表示蒸馏机制，可适应不同小尺寸的输入图像，并可应用于常见的网络架构，如CNN 和ViT。
• 我们提出了一种教师-助理-学生蒸馏框架，该框架降低了学生在像素蒸馏中的学习难度，并在学生的输入大小减小时启用特征蒸馏。
• 我们将像素蒸馏应用于目标检测，为辅助网络制定对齐的特征保留策略，以应对由于图像分辨率不同而导致输出尺寸不一致的挑战。
• 图像分类和目标检测的大量实验证明了所提出的蒸馏方案的有效性和效率。

方法：
3.1 图像分类中 KD
传统的知识蒸馏方法通过学习教师模型的信息来训练学生模型。 根据获取监督信息的方式，我们将图像分类中的先前工作分为两类：基于预测的方法和基于特征的方法。 基于预测的蒸馏方法通过使用老师预测的类别结果来训练学生。 一种基本方法是将教师模型的预测logits视为学生的软目标。 损失为：

$p_t=softmax(\frac{x_t}{T})$ 是教师预测的类别分数，y 是真实值，$p_s=softmax(\frac{x_s}{T})$，$X_t$ 和 $X_s$ 分别是教师和学生模型的预测类别分数，T 是温度参数，α 是平衡超参数，分类损失 $L_{cls}$ 和 KL 散度损失 $L_{kl}$。
基于特征的蒸馏方法使用教师模型的中间表示来指导学生模型。这些方法的学习过程可以表示为：

其中 $Fs=F_{s,1},F_{s,2},...,F_{s,M}$ 表示学生的中间特征，而 $Ft=F_{t,1},F_{t,2},...,F_{t,M}$ 表示教师的特征。 变量 M 表示网络中的块数。 B是所选特征的集合，对于不同的方法它有所不同。 gt(·) 和 gs(·) 是从中间特征中提取信息的函数。 δ(·)是距离度量函数。 β是一个超参数，用于平衡分类损失和特征蒸馏损失。
 
3.2 Baseline
在本文中，我们的目标是在教师模型的帮助下训练学生模型，其中网络架构和输入大小都不同。 教师模型采用大图像并使用重型网络，而学生模型采用小图像作为输入并使用轻量级网络。像素蒸馏方法应遵循两个标准：首先，它应适用于教师和学生模型的各种架构，包括不同的 CNN 和 ViT。 其次，该方法应该适应学生不同大小输入量。
提出vanilla PD 作为baseline，如图 3(a) 所示，所提出的基线由两个蒸馏过程组成：
1）遵循基于预测的蒸馏方法，我们使用教师的 logits 作为对学生的监督。基于预测的蒸馏方法不受网络架构和输入的影响。
2）提出了输入空间表示蒸馏（ISRD）机制，让学生的输入模块从教师的大量输入中学习有价值的空间知识。
vanilla PD 的损失定义为：

其中 $L_{pkd}$ 是基于预测的蒸馏损失，$L_{isrd}$ 是 ISRD 的损失。 γ是平衡$L_{pkd}$和$L_{isrd}$的超参数。

图3

 
ISRD：
考虑到像素蒸馏中的小输入会导致学生浅层特征的信息不足，因此在 CNN 或 ViT 的输入卷积之后使用所提出的 ISRD。 如图3a所示，ISRD是一个自动编码器，以学生的输入特征作为输入并输出大图像。 ISRD的编码器将CNN和ViT的空间信息变换为相同的形式，ISRD的解码器通过变换后的特征来预测大图像。 在本文中，我们计算伪大图像 $I_{hr}^{\prime }$ 和真实大图像 $I_{hr}$ 之间的 $L1$ 损失作为 ISRD 模块的损失。假设学生的输入是宽度为W像素、高度为H像素的图像，教师的输入是宽度为KW像素、高度为KH像素的图像，则定义ISRD的损失为：

ISRD 编码器由学生输入卷积层和 GSFP 操作组成。 如图 3b 所示，GSFP 操作是无参数操作，学生输入卷积层是编码器唯一可学习的参数。 CNN和ViT都使用卷积层将输入图像映射到特征空间，但它们的输入特征的形式有很大不同，因此我们需要将CNN和ViT的特征变换为相同的形式，以实现广义蒸馏。
ISRD 解码器，如图3c所示，给定广义输入空间特征Fg，我们需要将其体积扩展为与大输入$I_{hr}$相同。使用1×1卷积层来实现通道扩展。 然后使用 pixel shuffle 将扩展特征转换为3通道的特征图。
 
vanilla PD的总结：为图像分类任务的像素蒸馏方案提出了一个基线，称为 vanilla PD，具体来说，假设像素蒸馏中的小输入图像会导致学生浅层特征的空间信息不足，提出了一种输入空间表示蒸馏（ISRD）机制，从大输入中提取知识，以帮助学生的输入模块获得更丰富的表示。如图3b所示，考虑到我们的学生网络不仅包括CNN，还包括ViT，我们设计了一个通用空间特征预处理（GSFP）模块作为编码器，将CNN和ViT的输入空间特征转换为相同的形式。此外，如图3c所示，ISRD的解码器可以将任意体积的编码特征转换为伪大图像，这使得ISRD机制能够在学生的任何输入大小下使用。通过将 ISRD 与之前的预测蒸馏方法相结合，我们可以获得一种简单的单阶段训练基线像素蒸馏方法，即 vanilla PD。

 
3.3 教师-助理-学生框架
在知识蒸馏方案的经典师生框架中，教师和学生具有相同的输入大小和不同的网络架构。 然而，在像素蒸馏方案中，教师和学生具有不同的输入大小和网络架构，这使得学生更难成功模仿教师。
为了降低学生像素蒸馏的学习难度，本文在经典师生框架中引入辅助网络，将像素蒸馏过程解耦为模型压缩阶段和输入压缩阶段。 如图5所示，助理网络保持与教师模型相同的大输入大小，并具有与学生相同的轻量级网络架构。 在模型压缩阶段，助理模型被视为学生模型，从具有复杂架构的教师模型接收知识。 在输入压缩阶段，助理模型被视为教师模型，将空间知识从大量输入转移到学生模型中。 整个框架称为教师-助理-学生（TAS）框架。
所提出的TAS框架将为像素蒸馏带来三重好处：1）TAS是一个两阶段学习框架，每个阶段的学习难度都小于单阶段师生框架，这将为像素蒸馏带来更高的性能增益。 2）TAS框架的模型压缩阶段可以使用现成的知识蒸馏方法。 3）在输入压缩阶段，由于助理模型和学生模型具有相同的网络架构，因此它们的特征形式和数量相同，仅通道（CNN）或补丁（ViT）的数量不同，这使得更容易使用特征蒸馏策略

在本文中，我们设计了一种简单的特征蒸馏策略，它将上采样操作应用于assistant模型中的特征。对于来自 CNN 的特征图，我们直接在空间维度上对特征图进行上采样。 对于 ViT 的注意力图，我们首先将空间维度扩展为二维，然后应用上采样操作，以便学生在输入压缩阶段的特征蒸馏。TAS 中输入压缩的总体损失为：

 
3.4 目标检测中的像素蒸馏
在本文中，我们还将像素蒸馏扩展到目标检测领域。输入分辨率显着影响目标检测器的有效性，降低输入分辨率不仅会影响模型性能，还会改变输出特征，例如锚框的数量，从而使蒸馏过程中从教师检测头传输知识的过程变得复杂。
用于保存的对齐特征：如前所述，降低输入分辨率会影响输出特征，例如锚框的数量，如果我们想保存来自教师预测头的知识，这使得在一个阶段中执行像素蒸馏变得困难。 为了解决这个问题，我们在本研究中完善了两阶段的教师-助理-学生框架。 设计了一个对齐特征保存（AFP）策略来对齐每个阶段检测器的输出尺寸，由于在推理过程中未使用assistant网络，因此可以对其架构进行调整以抵消分辨率降低对检测器的影响。 具体来说，在第一个输入压缩阶段，教师和助理都使用相同的骨干网络。 助理的特征被上采样以匹配教师特征的空间分辨率，我们将其称为“助理↑”。 这种方法允许辅助检测器利用与教师相同的anchoring schema，从而促进在预测级别保存知识。 该对象检测输入压缩阶段的相关损失函数详细如下：

其中 Lpkd 和 Lfkd 分别表示预测蒸馏和特征蒸馏的损失。
在第二阶段的模型压缩阶段，为教师和学生提供相同的输入，这允许删除辅助检测器内的特征上采样操作。因此，助理现在采用与学生相同的anchoring schema。 这种一致性使得采用标准知识蒸馏技术成为可能。 综上所述，通过有效地操纵辅助网络的特征和锚点配置，我们的像素蒸馏方法被扩展到目标检测任务。 此外，基于预测和基于特征的蒸馏技术都可以在两个阶段中的每个阶段中使用。 值得注意的是，输入压缩阶段改变了辅助网络的结构并带来了额外的计算成本。 因此，教师-助理-学生（TAS）框架中两个阶段的顺序对于对象检测任务是不可互换的。 输入压缩阶段必须先于模型压缩阶段。