自动多步蒸馏（AMD） | 突破大规模视觉模型压缩的效能界限！

AIGC 先锋科技

发布于 2024-07-20 10:19:56

560

发布于 2024-07-20 10:19:56

文章被收录于专栏：AIGC 先锋科技

基于 Transformer 的架构因其卓越的性能已成为各种视觉任务的默认标准模型。随着模型规模的不断扩大，模型蒸馏在各种实际应用中变得尤为重要，尤其是在计算资源受限的设备上。然而，当前流行的知识蒸馏方法在遇到教师模型和学生模型之间的容量差距较大时，例如10倍的压缩率，其有效性会大打折扣。在本文中，作者提出了一种名为自动多步蒸馏（AMD）的新方法，用于大规模视觉模型的压缩。特别是，作者的蒸馏过程分为多个步骤进行。最初，教师模型经过蒸馏形成一个中间的教师助手模型，随后该助手模型再进一步蒸馏到学生模型。作者引入了一种高效且有效的优化框架，来自动识别能够使学生模型性能最大化的最优教师助手。作者在多个图像分类数据集上进行了广泛的实验，包括CIFAR-10、CIFAR-100和ImageNet。研究一致发现，作者的方法优于几种已建立的 Baseline ，为未来大规模视觉模型上的知识蒸馏方法铺平了道路。

1 Introduction

视觉领域的基础模型（如 BiT ，ViT ，Swin，Florence ）因其在一众任务中的革命性表现而最近受到了广泛关注。特别是像ViT-Large（61.6 G FLOPS）和 Swin-Large （103.9 G FLOPS）这样的基于 Transformer 架构的模型，它们代表了一类通用的视觉基础模型，并取得了前所未有的成功，相比于卷积神经网络（CNNs）如ResNet18 （1.8 G FLOPS），它们的操作要复杂得多。然而，随着模型规模的不断扩大以提高性能，它们固有的高复杂度成为了在低功耗处理器和移动设备上部署的瓶颈，这些设备通常计算能力有限。

在保持有限计算资源的同时利用大型模型优越性能的一种传统方法是知识蒸馏。这个过程将这些大型教师模型学到的知识转移到紧凑且可部署的学生模型中。尽管它们在各种应用中都很有效，但最近的研究[揭示了当教师和学生模型之间存在巨大的容量差距时，传统蒸馏方法会遭受严重的性能退化。因此，提出了多步蒸馏方法，首先将教师模型蒸馏到一个中间的教师助手模型。然后，教师助手作为替代教师更好地向学生传递知识。尽管多步蒸馏通常能提高学生的性能，但作者的初步研究（图1）表明教师助手的选择对学生性能有显著影响。然而，由于缺乏关于教师助手规模/大小的选择指导，通常需要尝试性训练，即遍历所有可能的规模，以寻找最优的教师助手，这导致了高昂的计算成本。

针对这一需求，作者提出了一种用于大型视觉模型压缩的自动多步蒸馏（AMD）方法。AMD通过包括三个阶段的级联策略进行操作。

结构剪枝：采用一种刻度和剪枝算法，生成一系列跨不同规模的教师助手架构；

联合优化：开发了一个联合优化框架，在单次运行中高效地获得所有规模上的最佳教师助手；

最优选择：根据负性能-规模导数（NPSD）指标（见SS3.2）从所有候选人中选择最优的教师助手。这个指标旨在衡量候选人的性能-规模最优性。

AMD是一种高效且强大的大型视觉模型知识蒸馏方法。实验证明，在基于 Transformer 架构的 ViT 和 Swin 上，AMD在top-1准确度上超过了同期单步和多步知识蒸馏方法，即在CIFAR-10 上**1.79-15.91%，CIFAR-100 [43]上1.81-24.25%，ImageNet 上2.77-4.03%**。更令人印象深刻的是，与其它多步蒸馏 Baseline 相比，它的训练速度要快得多（例如，快10倍）。关于AMD在学生模型可扩展性、单一教师助手效率、训练计划、采样率和损失组件方面的全面消融研究进一步证明了其鲁棒性和有效性。作者还将在补充材料中按照常见做法将AMD扩展到基于传统CNN的架构，表现出与现有技术水平相当的性能。总之，考虑到视觉模型压缩大型差距方面的研究有限，所呈现的结果尤为重要。作者相信这项工作有助于对相关领域的根本理解做出贡献。

2 Related Work

借鉴自然语言处理（NLP）领域transformer的显著进展，研究行人随后将transformer架构应用于各种视觉任务中，包括图像分类，图像分割，目标检测等。与卷积神经网络（CNNs）[29, 36]相比，transformers因其卓越的性能和可扩展性而在视觉相关学科中占据主导地位。然而，transformers的显著复杂性开销也阻碍了其在资源有限的现实场景中的部署。例如，著名的基于transformer的架构ViT-Base[16]的计算开销是ResNet101 的2.23倍（即17.6G FLOPS对比7.9G FLOPS）。尽管人们渴望利用大规模视觉模型的卓越能力，但现有研究主要集中在传统卷积神经网络的 knowledge distillation 上，忽视了基于transformer的架构中规模缩减的日益增长的需求。鉴于这一观点，作者提出了AMD，一种知识提炼方法，以减少大规模视觉模型的开销，同时保持令人满意的性能。

模型剪枝 模型剪枝可以大致分为无结构剪枝和结构剪枝。这两种方法的基本区别在于，结构剪枝通过物理移除成组参数来修改神经网络架构，而无结构剪枝在不改变网络固有结构的情况下零化选定权重[18]。尽管无结构剪枝能享受更细粒度的剪枝，但它只适用于专用设备。另一方面，结构剪枝不需要专用加速器或软件框架就能实现内存利用和计算开销的减少[11, 84]。因此，它更适合更广泛的实际应用。在基于transformer的架构中，一系列工作已经提出以确定所需规模的设计结构（例如，动态搜索[32]，层丢弃[17]和剪枝[54]）。在作者的工作中，作者采用结构剪枝来推导候选人的结构，因为它有利于提炼。

知识蒸馏在视觉领域模型压缩方面，知识蒸馏[31]已经引起了广泛关注。其基本概念是将较大规模模型（即教师模型）的知识压缩到计算效率更高的较小模型（即学生模型）中。近年来已经进行了大量研究，可以广泛分为两类：单步蒸馏和多步蒸馏。单步方法首先受到[31]设计的启发，并将教师的预测作为“软”标签，直接模仿教师模型的最终预测[21]。除了基于对数的方法，各种研究还应用了特征模仿方法，直接或间接地利用空间知识。还有基于关系的研究进一步探索不同层或数据样本之间的关系，使用来自FSP矩阵[90]的知识、实例关系、相似性矩阵、互信息流等。另一方面，当研究者注意到大尺度网络与小型学生模型之间的显著模型容量差异会降低知识转移时，多步方法显著出现。然而，确定学生模型的最佳教师助手呈现了一个微妙的挑战。如[71]的方法通过迭代指导多个教师助手到每一个较小教师助手，并引入“随机丢弃”以提高教师或教师助手的效率；[37]在学习过程中构建了一个逐渐模仿的序列。尽管这些方法减轻了在蒸馏过程中的性能退化，但由于与教师助手相关的计算开销增加，特别是随着教师模型容量的增加，它们并没有得到广泛接受。在目标检测中的学习容量由[6]涵盖，提供了一条通过多教师渐进式蒸馏的可行策略。然而，他们应对这一挑战的方法似乎还处于初级阶段（见SS3.3），忽略了通过共享参数的模型剪枝进一步缩小容量差距的潜力。此外，当前方法很少考虑大规模视觉模型，这些模型相对于当代卷积架构往往需要更高的压缩比。最近的研究（例如，CSKD [101]，DeiT [75]，CviT [66]，DearKD [10]）取得了值得称赞的性能。

然而，他们的重点是通过知识蒸馏来提高大规模模型的效能，忽略了在计算受限的设备上进行模型压缩和部署的必要性。相比之下，作者的方法注重计算效率和高压缩比，同时保持了教师模型的高性能。

3 Methodology

在本节中，作者首先正式定义了多步蒸馏问题及其表示法。随后，作者引入了负性能-规模导数（NPSD）指标，该指标用于评估中间教师-助教模型的最优性。作者的初步研究证明了NPSD作为评估教师-助教性能的指标的有效性。接着，作者提出了所提出的自动多步知识蒸馏方法。

Problem Formulation

在多步骤蒸馏范式内（如图2所示），主要目标是确定一个最优的教学助手，记作，这将促进从完全规模的教师模型（其中）到学生模型的知识蒸馏。这里和分别表示模型的性能和规模。这个过程旨在通过利用适当的教学助手（即，）来提高学生模型的性能。在实践中，教师助手的规模和性能受到两个边界的限制：对于可扩展性，S_{s}<s_{ta}<s_{t}p_{s}<p_{ta}<p_{t}；对于性能，p_{s}<p_{ta}<p_{t}

Negative Performance-Scale Derivative

在本工作中，作者提出为多步蒸馏识别一个合适的教师助手测量方法，考虑到性能与规模之间的权衡。同时优化这些因素的一个适当的目标函数并非易事。理想情况下，作者希望找到一个性能最大、规模最小的教师助手。然而，这种理想情况往往并不实际。一方面，当教师助手的规模接近学生模型的规模时，其性能往往会下降，这使得将全部知识转移到学生模型变得具有挑战性。另一方面，规模更大的教师助手与学生模型之间仍可能存在显著的规模差距，导致在蒸馏过程中性能大幅下降。因此，识别一个在性能与规模之间达到最优平衡的教师助手至关重要。

为了应对这一挑战，作者引入了负性能-规模导数（NPSD）作为优化目标。NPSD评估了教师助手在匹配学生规模方面的好处与因性能下降而产生的成本之间的平衡。较高的NPSD表明更有利的好处-成本比。作者的初步研究证实，NPSD与所有任务和 Backbone 网络的下游学生性能呈正相关（参见图1）。基于这一见解，作者的目标是在可管理的时序复杂度内最大化NPSD。

NPSD的定义。形式上，NPSD表示性能对规模的负推导，定义为：

其中和分别代表教师和教师助手。和代表两个模型的性能，而和表示它们对应的模型规模。需要注意的是，在给定教师模型的情况下，和在此背景下是常数。直观地说，性能高且规模小的教师助手会导致高的NPSD值。尽管NPSD计算没有明确涉及学生，但教师助手的性能（即）受到学生模型的影响，因为这是通过与学生的联合优化实现的。换句话说，不同的学生规模将为相同规模的教师助手产生不同的最优解。因此，教师助手的NPSD值在理论上受到学生的限制。作者进一步使用 -NPSD度量的变体进行实验，即，明确考虑学生模型（见表6）。很明显，与方程1中的原始NPSD相比，-NPSD并没有带来更好的性能，但具有额外的复杂性（即超参数）以及额外的蒸馏步骤以获得学生的性能。

初步结果。为了评估NPSD作为最优性测量的有效性，作者在CIFAR-10 和 CIFAR-100 上使用各种教师模型进行了多步蒸馏实验。这些教师模型经过蒸馏，通过不同规模的{10%、20%，...，80%，90%}的教师助手转变为10%的学生。每个教师助手模型的架构是通过结构剪枝从教师模型中剪枝掉重要性较低的参数得到的。

此外，作者观察到NPSD度量与学生模型的性能之间存在正相关关系。当以最大的NPSD值进行蒸馏时，学生达到了最高性能，这表明了NPSD度量的有效性。SS4.2中进一步的结果加强了作者的观察。

Automatic Multi-step Distillation

上述目标因此转化为一个优化问题，即寻找具有最大NPSD值的最佳教师助手。然而，确定最优的教师助手存在几个关键挑战。首先，考虑到连续的搜索空间，潜在的教师助手数量是无限的，即使对于固定规模的模型，其架构也可能表现出显著的变体。其次，在为每个教师助手进行多步蒸馏时，寻找所有规模上性能最高的教师助手的计算成本变得极其高昂。最后但同样重要的是，人们需要在所有候选人中选择最佳的教师助手。

为了应对这些挑战，作者将问题分解为三个阶段。在结构剪枝阶段，通过网格化和剪枝技术生成一系列不同规模的教师助手候选人。然后引入联合优化框架，以在单个优化中获取所有规模上的最高性能教师助手。前两个阶段共同将教师助手的无限搜索分解为一个有限集合，其中所有规模的教师助手都在单个试验中进行优化。最后，在最优选择阶段，选择具有最高NPSD的最佳教师助手。

结构剪枝。作者采用网格化和剪枝来构建不同规模的教师助手。具体来说，首先引入网格化以允许有限数量的候选人，作者将连续的搜索空间转换为离散空间。理想情况下，为了连续找到“完美”的教师助手，人们需要在每个可能规模上生成候选人。然而，在连续空间中搜索所有可能的候选人是既不可能（即无限细分）又效率低下（即共同优化大量教师助手）的。因此，作者将规模均匀地分为部分，其中每个候选者之间的间隔为，在教师和学生之间产生个候选人。接下来，为了获得不同规模上每个候选人的架构，可以应用各种剪枝技术。在这项工作中，作者采用了结构剪枝方法[54, 84]，因为它在知识蒸馏[86]中的已知优势。它根据在学习可 Mask 变量引入的教师推理阶段获得的重要性分数，逐渐剪除最不重要的网络参数。由此产生的教师助手候选者表示为 = 1到。更多细节在附录材料中提供。

联合优化。寻找最佳教师助手的直观方法涉及最大化每个候选个体的性能，然后使用NPSD来衡量每个候选者的最优性。然而，在评估所有候选者的性能时，由于候选者众多，计算开销仍然很大。为了进一步减轻计算复杂性，作者探索了在不同尺度上教师助手候选者之间的参数共享，利用从结构剪枝得到的候选者所固有的增量属性。本质上，增量属性指出，对于两个候选者和，如果 S_{i}<s_{j}\mathcal{m}_{i}="" 的参数是="" \mathcal{m}_{j}="" 的一个子集。这确保了可以通过不断剪枝较不重要的参数，将较大的候选者转换为较小的候选者，从而在候选者之间实现参数共享。<="" p="">

通过参数共享，待优化的总参数数量等于最大候选者的数量。随后，作者设计了一个联合优化框架，该框架可以一次性同时优化从教师到所有候选者的蒸馏过程。这样，作者有效地将所有候选者的内存占用减少到与一个特定候选者相当的水平。通过参数共享优化，计算成本也显著降低。更重要的是，作者提出了一种自然解决[6]中提出的问题的方法，研究行人发现，减少学生模型与教师模型之间的容量差距可以促进更有效的知识转移。与现有方法不同，现有方法需要从零开始初始化并教导一个独立的学生模型，无论是在同构蒸馏（例如，具有不同层数的相似架构）还是异构蒸馏（例如，不同的网络设计）中，作者的方法是正交的，并引入了联合优化，使得学生模型可以直接从同一单一架构设计中的教师对应模型继承参数。这种方法建立了一种有效的机制来弥合学生模型与教师模型之间的容量差距，从而实现最优的蒸馏性能。总之，AMD的总体目标包括交叉熵损失、基于逻辑的损失和基于特征的损失：

= 是交叉熵损失，= 是软化解码器输出之间的Kullback-Leibler散度损失[77, 41]，而是最后一层隐藏状态之间的均方误差。分别代表标签、输出逻辑值、温度值和最后一层隐藏状态。

最优选择. 然后从所有候选者中选择具有最高NPSD值的最佳教师助手。随后，根据方程式2中的相同训练目标，在选定的教师助手和学生之间进行额外的蒸馏。应该认识到，在NLP领域内，有一系列研究[47, 84]通过结合超参数调整的训练曲线的实证分析来描述训练目标。相比之下，作者的方法关注教师和教师助手之间的关系，促进在多种数据集和模型上的鲁棒性探索，从而对知识蒸馏动态有了更细致的理解。

4 Experiment

以下是第4节实验部分的开始部分。

Experimental Setup

数据集。 评估分别在CIFAR-10 [43]，CIFAR-100 [43] 和 ImageNet [68] 数据集上进行。CIFAR-10 包含 K（K/K 用于训练/测试）的个类别的彩色图像。CIFAR-100 数据集包含个类别，每个类别有张训练图像和张测试图像。ImageNet 包含 M 张图像用于训练和 K 张图像用于验证的 1K 个类别。请注意，在所有基于 Transformer 的架构实验中，图像分辨率为 224 224。在补充材料中，当考虑基于 CNN 的架构时，作者改为在 CIFAR-10 和 CIFAR-100 上应用 32 32 的分辨率以保持通用性。

网络与 Baseline 。 在本文中，作者主要关注大规模视觉模型，并使用不同的网络进行实验，即Vision Transformer (ViT) [16] 和 Swin Transformer (Swin) [51]，具有不同的模型大小（ViT-Tiny，ViT-Small，ViT-Base，Swin-Base）。当前的视觉方法主要围绕基于CNN的蒸馏（见SS2），将它们特别设计的方法转移到Transformer架构中存在挑战。在大规模视觉模型背景下，作者的研究复现并比较了提出的高级模仿蒸馏（AMD）与几种广泛采用的知识蒸馏方法。这些包括四个单步 Baseline ：KD [31]，CRD [73]，DKD [100] 和 ADKD [84]，以及两个多步 Baseline ：TAKD [55] 和 DGKD [71]。对于这两种多步方法，作者使用了与AMD相同的教师助手集合。在DGKD的情况下，作者用两个中间教师助手进行了实验。为了完整性，作者在补充材料中进一步设计了基于CNN架构的广泛研究。

实现细节。 AMD的实现设置遵循了常见的实践[71, 100, 73]，包括预处理、优化、训练计划和学习率等。作者使用标准化的预处理[16]，包括归一化和重缩放，_不_应用其他预处理技术。作者将初始学习率设置为0.003，并使用AdamW优化器[52]，通过余弦调度进行3个 Warm up 周期的安排。训练周期设置为160，批量大小为128，与当前研究[71, 55]一致。作者将比例采样率设置为9，这平衡了性能和计算成本（见表4），在结构剪枝阶段后得到候选比例集。对于公式中的超参数和。

在第二节中，作者遵循常见做法[84]，分别调整从和从。在SS4.3中彻底研究和后，作者为所有实验设置选择=0.2，=100。作者使用默认的温度=1。对于本研究中比较的其他单步方法，作者遵循[31, 100]，并应用它们的默认训练周期以进行公平比较（即，在CIFAR-10和CIFAR-100上训练240个周期，在ImageNet上训练100/120个周期）。作者利用公开可访问的代码库，仅对教师和学生网络进行修改，融入基于 Transformer 的架构。在训练过程中，作者注意到具有阶梯学习率计划的随机梯度下降（SGD）常常导致 Transformer 的准确度较低。因此，作者采用AdamW优化器以确保稳定训练。

可复现性。 AMD是用Pytorch[63]实现的。实验在NVIDIA A100-40GB GPU上进行。

Main Results

在CIFAR-10和CIFAR-100上的结果。表1报告了在CIFAR-10和CIFAR-100数据集上五次运行的最高1准确度得分。从这些结果中可以得出几个关键观察。首先，AMD能够大幅超越当前的知识蒸馏方法，在CIFAR-10和CIFAR-100数据集上都取得了显著的性能差距。例如，在CIFAR-10上，作者的模型在ViT-Base知识蒸馏上的最高1准确度提高了1.34-15.89%；其次，鉴于基于对数概率的方法（如KD、DKD和TAKD）在多种规模和数据集上展现出的次优性能，可以合理推理这些方法可能并不代表基于 Transformer 架构的知识蒸馏的理想解决方案。

相比之下，通过利用特征空间监督，CRD、ADKD和作者提出的AMD取得了令人称赞的性能。这一趋势表明，在知识蒸馏的背景下，进行特征模仿研究是必要的，特别是在 Transformer 架构的深度和宽度扩展时，为了保持蒸馏性能的完整性和有效性；

第三，与多步方法相比，作者的AMD实现了更快的训练速度，表明结构剪枝和联合优化的有效性；

最后，作者的方法在各种基于 Transformer 的架构设计中一致优于其他知识蒸馏方法，为未来的模型部署展示了一条有希望的途径。

在ImageNet上的结果。为了进一步证明作者方法的有效性和泛化能力，作者遵循常见做法[48, 55, 71]，并将AMD扩展到ImageNet[68]，结果如表2所示。由于作者在表1的CIFAR-10和CIFAR-100数据集中已经表明，基于对数概率的方法在 Transformer 架构上的性能不佳，因此在ViT和Swin上的这次实验中，作者仅与特征模仿方法CRD进行了比较。如所见，作者的方法在ViT-Base上比CRD取得了超过**2.61%**的最高1准确度提升。

Diagnostic Experiments

作者设计了几项消融研究来评估作者提出的AMD的鲁棒性和有效性。大多数实验是在CIFAR-100上训练的ViT上进行的。为了研究单个教师辅助的充分性，作者还在CIFAR-10上训练的ViT上报告了额外结果以保证完整性。[方法]需要3.68G FLOPs，在CIFAR-100上达到77.94% [4] 的top-1准确率。作者有意不在本文中将学生模型扩展到更大规模，因为本文的主要目标是保持紧凑的学生模型。因为学生模型在规模上的扩展会导致推理成本显著增加。这种增强不仅对部署灵活性提出了重大挑战，还加剧了对计算资源的必要需求。可以看出，AMD在各种模型规模上均取得了持续强劲的性能。

单个教师助手的充足性。 为了验证使用单个教师助手的效率，作者进一步在方法中设置了关于多个教师助手的消融研究。在图3(a)中，作者首先展示了没有教师助手会导致性能明显下降（即79.77%），这与表1中的结果一致。作者还发现，当有多个教师助手时，性能提升是微小的（即80.19% 80.24%），而GPU小时数却显著增加（即当有两步和三步教师助手时分别增加3倍和4倍）。在ViT-Base CIFAR-10上进行的训练得到的结果与观察一致（见图3(b)）。因此，在AMD中作者选择了一个教师助手，既保证了训练效率又满足了性能充足性。关于ViT和Swin的额外结果将在补充材料中提供。

训练时间表的好处。正式地，作者遵循常见做法（参见SS4.1），并将160个周期设置为标准训练时间表。然而，作者注意到AMD在早期周期内快速收敛。例如，在图4中，仅用40个周期，作者的AMD与其他流行方法（见表1）相比，展现出具有竞争力的性能（即CIFAR-100上的70.82%）。请注意，尽管KD [31]在不同周期间表现出较小的性能差距（即7.05%），但它很难在基于 Transformer 的架构设置下达到令人满意的结果。另一方面，即使周期数减少，AMD的表现也始终良好。这强调了作者的方法作为在有限训练时间表下训练的有效策略的潜力。

候选采样率的影响。 作者进一步研究通过改变采样候选数量来变化候选采样率的影响。的值越高，表明采样率的精细度越高。这种增加的精细度与训练时间的延长直接相关。表4报告了GPU小时数及相应的学生表现。作者设置以在性能和计算开销之间达到满意的折中。提高采样率不可避免地导致训练时间更长，这带来了性能的轻微提升。例如，当时，作者观察到可以获得的性能提升，但GPU小时数增加了。作者认为这对于训练计划来说是不高效的。

不同超参数的影响。在方程2中，作者引入了两个超参数和来平衡交叉熵损失、基于逻辑的损失和基于特征模仿的损失。作者在图5中展示了不同和值对性能的影响。结果显示对这些超参数的波动具有很大的鲁棒性，标准差低至0.42。在实验中，不同的数据集上观察到了一致的规律：较大的与对教师助手逻辑响应的影响减小相关，这反过来又导致了次优的性能结果。此外，作者注意到随着值的增加，监督效果逐渐增强。与之前在NLP领域的研究建议根据不同任务采用不同的值[84]不同，在作者的实验中，在几个数据集上相对稳定。总的来说，对于最佳性能，作者得到了超参数组合和，这些值用于所有实验。