层回收和微调效率

Yusuf Sarıgöz

2022年8月23日

Layer Recycling and Fine-tuning Efficiency

Allen AI 最近发表的一篇论文引起了自然语言处理界的关注，他们通过在训练和推理阶段缓存某个中间层的输出来实现约83%的速度提升，而模型性能损失可以忽略不计。这项技术与Quaterion 中的缓存机制非常相似，但后者适用于任何数据模态，而前者仅专注于语言模型，尽管他们从实验中提出了一些重要见解。在这篇文章中，我将分享我们的发现以及这些见解，希望为社区提供关于层回收的更广阔视角。

层回收的工作原理

层回收的主要思想是通过避免同一数据对象重复通过冻结层来加速训练（和推理）。相反，可以将对象仅通过这些层一次，缓存输出，并在未来的 epoch 中将它们用作未冻结层的输入。

在论文中，他们通常缓存50%的层，例如在一个12块编码器中缓存第6个多头自注意力块的输出。然而，他们发现这并不适用于所有任务。例如，对于问答任务，回收50%的层会导致性能出现更显著的下降，他们选择将这一比例降至25%，因此他们建议根据具体任务确定缓存级别。他们还指出，对于较大的模型和低端机器，缓存能提供更显著的速度提升。

在层回收中，缓存命中是针对完全相同的对象。这在文本数据中很容易实现，因为文本很容易哈希，但当你想将这项技术推广到不同类型的数据时，你可能需要更高级的技巧来生成缓存键。例如，对 PyTorch 张量进行哈希可能不会像你预期的那样工作。Quaterion 带有一个智能键提取器，可以应用于任何数据类型，但也允许通过传递一个可调用对象作为参数进行自定义。得益于这种灵活性，我们能够在不同的设置下进行各种实验，我相信这些发现对您未来的项目会有所帮助。

实验

我们进行了不同的实验来测试性能，包括

在相似汽车搜索示例中回收不同数量的层。
在相似汽车搜索的训练和微调数据集中使用不同数量的样本。
在问答示例中回收不同数量的层。

使用 Quaterion 轻松实现层回收

在 Quaterion 中缓存层的最简单方法是组合一个带有冻结的可训练模型 (TrainableModel) 和未冻结的编码器 (Encoder) 的编码器头部 (EncoderHead)。因此，我们在示例中修改了 TrainableModel 如下

class Model(TrainableModel):
    # ...


    def configure_encoders(self) -> Union[Encoder, Dict[str, Encoder]]:
        pre_trained_encoder = torchvision.models.resnet34(pretrained=True)
        self.avgpool = copy.deepcopy(pre_trained_encoder.avgpool)
        self.finetuned_block = copy.deepcopy(pre_trained_encoder.layer4)
        modules = []

        for name, child in pre_trained_encoder.named_children():
            modules.append(child)
            if name == "layer3":
                break

        pre_trained_encoder = nn.Sequential(*modules)
        
        return CarsEncoder(pre_trained_encoder)

    def configure_head(self, input_embedding_size) -> EncoderHead:
        return SequentialHead(self.finetuned_block,
        self.avgpool,
        nn.Flatten(),
        SkipConnectionHead(512, dropout=0.3, skip_dropout=0.2),
        output_size=512)


    # ...

这个技巧让我们可以在微调过程中将基础模型中的一个或多个层作为 EncoderHead 的一部分进行训练，同时仍然受益于缓存为冻结的 Encoder 提供的速度提升。

实验 1：回收的层数百分比

论文指出，与完全微调相比，回收50%的层几乎不会导致性能损失。在这种设置下，我们比较了四种方法的性能

冻结整个基础模型，仅训练 EncoderHead。
将四个残差块中的一个移到 EncoderHead，并与头部层一起训练，同时冻结其余部分（75%层回收）。
将四个残差块中的两个移到 EncoderHead，同时冻结其余部分（50%层回收）。
将整个基础模型与 EncoderHead 一起训练。

注意：在这些实验中，我们使用 ResNet34 而非 ResNet152 作为预训练模型，以便在完全训练中使用合理的批处理大小。使用 ResNet34 的基准分数是 0.106。

模型	RRP
完全训练	0.32
50% 回收	0.31
75% 回收	0.28
仅头部	0.22
基准线	0.11

如表中所示，50% 层回收的性能与完全训练非常接近。此外，在性能仅小幅下降的情况下，50% 层回收仍然可以带来可观的速度提升。虽然 75% 层回收优于仅训练 EncoderHead，但与 50% 层回收和完全训练相比，其性能下降很快。

实验 2：可用数据量

在第二个实验设置中，我们比较了不同数据集大小下的微调策略性能。我们随机抽取了训练集的 50%，同时仍然在整个验证集上评估模型。

模型	RRP
完全训练	0.27
50% 回收	0.26
75% 回收	0.25
仅头部	0.21
基准线	0.11

这个实验表明，可用数据集越小，我们在完全训练、50% 和 75% 层回收中观察到的性能下降越明显。另一方面，与其它方法相比，仅训练 EncoderHead 的性能下降幅度非常小。当我们进一步减少数据集大小时，完全训练在某个点变得不可训练，而我们仍然可以通过仅训练 EncoderHead 来超越基准线。

实验 3：问答任务中的层回收

我们还希望在不同的领域测试层回收，因为论文最重要的结论之一是层回收的性能取决于任务。为此，我们使用基于相似性学习的问答教程中的代码设置了一个实验。

模型	RP@1	RRK
完全训练	0.76	0.65
50% 回收	0.75	0.63
75% 回收	0.69	0.59
仅头部	0.67	0.58
基准线	0.64	0.55

在此任务中，与完全训练相比，50% 层回收仍能取得不错的性能，且性能下降很小。然而，性能下降程度小于相似汽车搜索示例。这可能归因于预训练模型质量、数据集大小和任务定义等多种因素，这也可能是一个更详尽和全面的研究项目的主题。另一个观察结果是，75% 层回收的性能更接近于仅训练 EncoderHead，而不是 50% 层回收。

结论

我们设置了几个实验来测试不同约束条件下的层回收，并确认层回收在不同任务和领域中表现各异。最重要的观察之一是，与完全训练相比，层回收的性能下降程度是亚线性的，也就是说，我们损失的性能百分比小于我们回收的百分比。此外，仅训练 EncoderHead 对小数据集大小更具鲁棒性。甚至存在一个临界大小，低于该大小，完全训练根本无法进行。性能差异问题表明，关于层回收仍有进一步研究的空间，幸运的是 Quaterion 足够灵活，可以快速进行此类实验。我们将继续报告我们在微调效率方面的发现。

有趣的事实：本文的预览图片由 Dall.e 使用以下提示创建：“逼真的机器人使用调音叉调整钢琴。” 点击这里查看全尺寸图片！

学习