PYRA：超轻量级ViT适应&推理高效微调模块

论文：https://arxiv.org/abs/2403.09192

最近，Transformer的规模迅速增长，这向基础模型的下游任务适应中引入了相当大的训练开销和推理效率方面的挑战。现有的工作，即参数高效微调（Parameter-Efficient Fine-Tuning，PEFT）和模型压缩，分别研究了这两部分挑战。然而，PEFT不能保证原始骨干网络的推理效率，特别是对于大规模模型。模型压缩需要显著的训练成本进行结构搜索和重新训练。因此，它们的简单组合不能保证以最小的成本同时实现训练效率和推理效率。在本文中，我们提出了一种新颖的并行生成重激活（Parallel Yielding Re-Activation，PYRA）方法，以应对训练-推理高效任务适应的挑战。PYRA首先利用并行生成的自适应权重来全面感知下游任务中的数据分布。之后，PYRA应用一种重激活策略进行token调制，以合并tokens，从而校准token特征。广泛的实验表明，PYRA在低压缩率和高压缩率下均优于所有竞争方法，证明了其在大模型微调当中的训练效率和推理效率方面的有效性和优越性。

简介

视觉Transformer（ViT）模型[1]在各个计算机视觉领域产生了深远的影响，例如图像分类、对象检测、图像分割等。近年来，视觉Transformers的规模已发展至十亿参数级别[2]。因此，将如此大规模的模型适应到下游任务中，呈现出越来越复杂的挑战，特别是在实际部署场景中。当将大规模Transformer应用于下游应用时，研究者们广泛认为有两个关键问题阻碍了这一点：（1）在下游任务上进行微调时的训练开销，以及（2）模型部署后的推理效率。

具体来说，首先，传统的微调方法（即全参数微调，Full Fine-Tuning）需要调整模型的所有参数，鉴于基础模型的广泛规模，这导致了GPU资源和训练时间的巨大消耗。近些年，针对第一个问题，研究人员已深入研究参数高效微调（Parameter-Efficient Fine-Tuning，PEFT）[3]算法，这些算法通常冻结预训练模型并且只调整额外的小参数，导致训练时间和存储开销的大幅度减少。第二个问题涉及到推理效率，要求部署的模型能够迅速处理输入数据。模型的计算复杂性显著影响到达满意的推理吞吐量。代表性的解决方案包括模型压缩方法，包括模型剪枝、知识蒸馏、模型量化等。

在现有工作当中，这两个问题被分别进行了研究：（1）针对第一个问题，现有的PEFT方法要么识别骨干中可调参数的子集以进行微调（BitFit [4]），要么在微调过程中向冻结的骨干引入可学习参数（LoRA [5], Consolidator [6]）。虽然有效地减少了训练成本，但这些方法中的大多数不可避免地增加了计算复杂性，导致推理效率低下。（2）针对第二个问题，模型压缩方法经常需要大量的计算资源来识别修剪的最优结构。修剪后，使用大量数据的全面重新训练过程对于防止性能显著下降至关重要。因此，模型压缩方法在训练效率方面通常是低效的。

这些观察自然引出了一个问题：我们能否同时实现下游任务的训练效率和推理效率？ 我们将这一挑战称为训练-推理高效任务适应（Training-Inference Efficient Task Adaptation）。探索这个问题可以使我们方便地以最小的成本部署先进的大规模基础模型到现实世界的下游应用中，这对于基础模型的广泛实现是极为必要的。一个直接的解决方案是结合PEFT和模型压缩。然而，对于高效的任务适应，一个完整的重新训练（Re-Training）阶段是不可承受的。因此，简单地结合PEFT和模型压缩容易遭受显著性能下降。例如，我们可以将常用且高效的LoRA [5]，与ToMe [7]，一个针对视觉transformers的无参数模型压缩技术结合起来（实际上，ToMe+LoRA是一个简洁而极为强大的解决方案，这一点在实验部分有所说明）。如图2所示，在较低的压缩率（大约1.7）下，两个Transformer骨干（ViT-L/16和ViT-B/16）的性能相较于直接在做PEFT，表现出轻微下降（<1%），这表明尽管方案需要进一步的性能改进，但ToMe+LoRA组合在较低压缩率范围内可以作为一个基础解决方案。在高压缩率（>3.0）下，ToMe+LoRA的性能迅速下降，其性能甚至不如直接对相应吞吐量的小规模Transformer进行直接PEFT。我们将这种现象称为逆压缩。这两种现象表明，直接结合现有工作无法有效应对新挑战。

本文提出了一种针对视觉Transformer进行训练-推理高效任务适应设计的并行生成重激活方法（Parallel Yielding Re-Activation，PYRA）。PYRA采用token合并范式以提高推理效率。为了进行有效的任务适应，我们从token特征和通道特征两方便分别进行特征调制，并通过轻量级的的调制权重生成器并行生成待合并tokens的调制权重。这些并行生成的权重可以全面感知下游任务中的数据分布。然后，通过重激活，将它们应用于特征调制，从而实现自适应的token调制。得益于这样一种token调制策略，PYRA能够以低计算复杂度自适应地校准下游任务的特征分布学习，最终实现有效的训练和高效的推理。

实验证明，在低倍压缩率下，PYRA的性能和压缩前的模型相差无几，在高倍压缩率下，PYRA可以达到和对应小模型PEFT相当的性能，彻底消除逆压缩现象。在实际应用中，下游任务所支持的模型大小很可能并不能找到对应大小的模型架构进行微调，因此用更大的模型做高倍压缩到对应FLOPs具有很大的应用价值。消除逆压缩的现象证明PYRA的训练-推理高效路径切实可行。

我们在不同大小模型、不同预训练范式的模型、不同模型架构上的实验证明，PYRA是一种广泛有效的方法。

方法

Preliminaries

ViT模型。本文主要关注ViT模型 [1]的训练-推理高效任务适应。一个ViT模型包含个相同的Encoder，每个Encoder由一个多头自注意力（MHSA）模块和一个前馈网络（FFN）组成。形式上，输入图像被重塑并线性投影为个维的tokens 。为简化，我们省略了分类token（[CLS]）和蒸馏token。对于Encoder，我们将输入表示为，输出表示为。对于MHSA，输入tokens首先被三个FC层处理以生成,和矩阵，输出通过计算后被另一个FC层投影。对于FFN，tokens被两个FC层投影。我们的方法主要关注在输入tokens 送入MHSA模块之前。

LoRA。LoRA [5] 是一种广泛使用的PEFT方法。ViT包含大量的密集参数矩阵。当适应特定任务时，对这些矩阵的更新位于小的子空间中，可以用低秩分解来建模。LoRA在微调过程中只训练分解后的矩阵。具体来说，对于密集矩阵和输入，更新后的的修改前向传播为：

其中且。在推理过程中，可以与合并，不产生额外的计算开销。

Token合并。Token合并 [7-9] 是一种针对ViT的无参数压缩技术。这类方法与Transformer结构正交，且能够灵活地改变压缩率。具体来说，第个ViT块的输入tokens 在MHSA之前被随机分成两组：