论文:
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。对于Encoderl, 我们将输入表示为
,输出表示为
。对于MHSA, 输入 tokens首先被三个FC层处理以生成
,
和
矩阵, 输出通过
计算后被另一个FC层投影。对于FFN, tokens被两个FC层投影。我们的方法主要关注在输入 tokens
送入MHSA模块之前。
LoRA。
LoRA [5] 是一种广泛使用的PEFT方法。ViT包含大量的密集参数矩阵。当适应特定任务时,对这些矩阵的更新位于小的子空间中, 可以用低秩分解来建模。LoRA在微调过程中只训练分解后的矩阵。具体来说, 对于密集矩阵
和输入
, 更新后的
的修改前向传播为:
其中
且
。在推理过程中,
可以与
合并, 不产生额外的计算开销。
Token合并。
Token合并 [7-9] 是一种针对ViT的无参数压缩技术。这类方法与Transformer结构正交,且能够灵活地改变压缩率。具体来说, 第
个ViT块的输入 tokens
在MHSA之前被随机分成两组:
然后, 通过余弦相似度, 为每个
token匹配最相似的
token。之后, 从
中, 选择与
中 tokens有最相似连接的
个tokens形成token对
, 其中
, 且
。形式上可以表示为:
当前工作 [7-9] 通常通过平均池化合并tokens。
我们采用上述技术作为我们的基线方法,在其中我们在微调LoRA进行任务适应时附加了token合并。选择这些方法是因为它们的优势与训练-推理高效任务适应高度兼容。首先,token合并是无参数和无需训练的,并且不改变模型结构,这保留了PEFT的存储高效优势。此外,我们选择LoRA因其受欢迎度、简单性以及合并性(在推理期间不引入额外的FLOPs)。实验结果表明ToMe+LoRA是一个强大的基线。
PYRA: Parallel Yielding Re-Activation
传统的微调方法通过大量微调参数, 指导模型动态地与下游任务的目标数据分布对齐。然而, 在训练-推理高效任务适应的背景下, 只能微调一部分参数, 这在准确捕捉数据分布的细微差别方面提出了重大挑战。虽然通过token合并减少模型复杂性在提高推理效率上卓有成效, 但它引入了在ViT的逐层处理过程中的信息丢失风险。由于在高效任务适应场景中高效微调限制了数据分布的理解, 这种损失难以纠正。因此, 直接结合token合并和PEFT算法以实现训练-推理高效任务适应可能不会产生最优结果。我们提出PYRA, 以自适应调制token特征, 在token合并过程中增强对数据分布的感知。具体来说, 在PYRA内部, 首先通过每个ViT块中的一对轻量级可学习向量以并行方式生成自适应合并的权重。然后, 这些生成的权重通过重激活应用于待合并的tokens以进行调制。PYRA使得以低计算复杂度自适应校准学习到的特征分布成为可能。
(1) 并行生成的自适应调制参数
我们旨在优化每对选定token对的合并过程。受到特征调制[9-10]的启发, 我们在合并前调制token特征。形式上, 对于具有
对待合并的
维 tokens的 Encoder
, 我们将
和
tokens 分组为 token矩阵
和
, 其中
。我们学习一个调制矩阵
, 以自适应地在每个通道的粒度上调制tokens。我们强调, 直接学习
是冗余的, 并且不能自适应地满足不同图像和token对的条件。因此, 我们进一步将目标具体化为分别学习特征通道的解耦权重
和特征 tokens的解耦权重
, 其中
。我们以并行生成的方式生成
和
。具体来说, 对于层
中要合并的
个tokens, 我们在 Encoder块内部创建两个可学习向量作为调制权重生成器:
和
。为了保证
和
从token对中的两个tokens提取特征信息, 我们首先计算token信息矩阵:
我们通过利用LayerNorm
操作来归一化
tokens的分布, 以在训练
和
时实现更平滑的梯度。有了token信息矩阵, 我们之后并行生成自适应权重
和
:
(2)重激活Token调制
通过矩阵乘法简单生成
和
仍然面临一些可能的问题。首先, 没有措施确保
和
保持在正常范围内。其次, 将权重解耦到特征tokens和特征通道导致调制权重矩阵
具有低秩性, 这显示出有限的表达能力, 因此可能无法在复杂的数据分布中最优地调制tokens。为了应对这些问题, 我们采用重激活策略进行token调制。具体来说, 我们首先将
广播到
并在其上进行
激活
, 然后调制
以获得中间调制结果: