跨多个数据集搜索基因以构建可变尺寸模型

# 通过跨多个数据集的 Learngene 搜索构建可变尺寸模型

来源：https://arxiv.org/html/2605.08209
###### 摘要

深度学习方法广泛应用于具有不同资源约束的各种场景中，从而产生了不同尺寸的模型，例如 Vision Transformer (ViT) 系列。部署这些模型通常需要从 scratch 进行预训练，然后进行微调，这一过程耗时且计算成本高昂。为了解决这些挑战，提出了 Learngene 范式。该范式从训练良好的大型模型（称为祖先模型，Ans-Net）中提取出具有强大学习能力的关键组件，称为 learngene，并利用这些 learngene 来初始化不同尺寸的模型，称为后代模型（Des-Nets）。之前的 learngene 提取方法仅关注单一数据集，这限制了提取的有效性，并导致下游任务中性能次优。为了克服这一局限性，我们提出了一种新的跨多个数据集提取 learngene 的方法，称为用于构建可变尺寸模型的跨多数据集 Learngene 搜索（Learngene Search Across Multiple Datasets for Building Variable-Sized Models, LSAMD）。LSAMD 通过引入数据集特定模块（dataset-specific blocks）和数据集适配器（Dataset Adapters, DADs）将 Ans-Net 扩展为可搜索的超级祖先模型（super Ans-Net）。数据集特定模块从输入数据集中学习，而 Ans-Net 中的原始模块（称为基础模块，base blocks）保持冻结以保留初始知识。DAD 根据输入数据集选择数据集特定模块与基础模块之间的传播路径。在训练过程中，LSAMD 从超级祖先模型中为每个数据集搜索最佳架构路径。然后，将这些路径中最常采用的基础模块提取为 learngene，用于初始化可变尺寸的 Des-Nets。在多个数据集上的广泛实验表明，LSAMD 不仅实现了与预训练-微调方法相当的性能，还显著降低了存储和训练成本。

## 引言

Vision Transformers (ViTs) (Dosovitskiy et al. 2020) 凭借高性能在计算机视觉领域取得了重大进展。这种能力使其能够部署在从资源受限的移动设备到资源丰富的计算中心的广泛场景中。资源设置的多样性导致了不同尺寸模型的产生。通常，这些模型会被随机初始化并从 scratch 开始训练以供部署。然而，随着部署场景数量的增加，这种方法相关的训练成本大幅增加。

为了解决这个问题，提出了预训练-微调（Pretrain-Finetune, PF）方法。该方法首先在一个大规模数据集（如 ImageNet-1K (Deng et al. 2009)）上训练源模型，然后在目标数据集上对该模型进行微调以导出目标模型。然而，这种方法需要将整个源模型转移到下游任务中，从而引入了一些挑战。对于每种资源约束，PF 都会训练特定尺寸的模型，并针对下游场景进行微调。随着资源约束数量的增加，PF 方法变得越来越耗时。此外，PF 需要存储不同尺寸的预训练模型以供微调，这导致了巨大的存储需求，并降低了对具有多样化资源约束的下游任务的灵活性。

如今，为了解决这些挑战，引入了一种名为 Learngene (Wang et al. 2022, 2023) 的新型学习范式。如图 1 (a) 所示，该范式主要包含两个步骤：(1) Learngene 首先从经过良好训练的大型模型（称为祖先模型，Ans-Net）中提取出一个虽小但关键的部分，称为 learngene；(2) 然后使用提取出的 learngene 来初始化不同尺寸的较小模型，称为后代模型（Des-Nets）。由于包含了来自 Ans-Net 的关键知识，learngene 能够初始化 Des-Nets 使其在下游数据集上表现出色。此外，learngene 紧凑的规模使其能够初始化大量多样的中型 Des-Nets，从而在适应具有不同资源约束的下游任务时表现出相当大的灵活性。由于 learngene 对于 Des-Nets 的发展至关重要，因此提取合适的 learngene 非常重要。为了实现这一目标，已有几项研究提出了从 Ans-Net 中提取 learngene 的方法。

具体而言，Vanilla Learngene (Van-LG) (Wang et al. 2022) 基于 Ans-Net 在源自单一数据集（如 Cifar100 (Krizhevsky 2009) 或 ImageNet100 (Deng et al. 2009)）的各种任务训练期间的梯度信息来提取 learngene。此外，Auto-Learngene (AL) (Wang et al. 2023) 训练一个元网络来计算 Ans-Net 每一层输出与 Cifar100 (Krizhevsky 2009) 或 ImageNet100 (Deng et al. 2009) 上辅助模型之间的相似度，并选择 Ans-Net 中最相似的层作为 learngene。PEG (Wang et al. 2024) 采用概率混合方法，从 Ans-Net 中采样多头自注意力和前馈层子集作为 learngene，利用 ImageNet-1K (Deng et al. 2009)。LearngenePool (Shi et al. 2024)、TLEG (Xia et al. 2024a)、SWS (Xia et al. 2024c) 和 LeTs (Xia et al. 2024b) 使用单一数据集 ImageNet-1K (Deng et al. 2009) 将 Ans-Net 蒸馏为一个或多个较小的辅助模型以提取 learngene。然而，这些研究存在以下局限性：它们仅从单一数据集中提取 learngene。一个数据集通常遵循特定的数据分布，这限制了从 Ans-Net 中定义具有通用知识的关键部分，并导致由 learngene 初始化的 Des-Nets 性能降低。

受此问题启发，我们提出了一种新方法，称为跨多数据集 Learngene 搜索（Learngene Search Across Multiple Datasets, LSAMD），用于从 Ans-Net 在多个数据集上提取 learngene。所提出的 LSAMD 涉及以下步骤。如图 1 (b) (i) 所示，我们首先将 Ans-Net 的架构扩展为超级祖先模型（super Ans-Net）。在超级祖先模型的每一层中，它额外包含一个数据集特定模块和一个数据集适配器（DAD）。数据集特定模块使用与 Ans-Net 原始模块（称为基础模块）相同的参数进行初始化。在训练超级祖先模型时，DAD 决定基础模块和数据集特定模块之间的前向传播路径。然后，如图 1 (b) (ii) 所示，我们在多个数据集上训练超级祖先模型。在此过程中，数据集特定模块被更新，而基础模块保持固定，这使得每一层既可以学习数据集特定的知识，也可以利用 Ans-Net 现有的能力。训练后，DAD 从超级祖先模型中为每个数据集搜索特定的传播路径。然后，我们对所有搜索路径中基础模块的使用频率进行逐层统计分析。基于此，Ans-Net 中的每个基础模块被归类为两种情况之一：重度使用或轻度使用。重度使用的基础模块被识别为跨数据集共享的关键模块，并提取为 learngene，如图 1 (b) (iii) 所示。最后，我们利用提取的 learngene 初始化可变尺寸的模型，用于具有各种资源约束的下游任务。

我们的主要贡献如下：(1) 我们将 learngene 提取从单一数据集扩展到多个数据集。据我们所知，这是 Learngene 范式中第一个考虑多数据集场景的工作。(2) 我们引入了一种称为 LSAMD 的方法，用于从扩展后的 Ans-Net 中跨多个数据集搜索 learngene。(3) 广泛的实验表明，LSAMD 可以灵活地构建可变尺寸的模型，在使用更少存储的情况下实现与预训练-微调方法相当的性能，并降低下游任务的训练成本。

## 相关工作

### 混合专家（Mixture of Experts, MoE）

混合专家（MoE）技术最初由 Jacobs 等人 (Jacobs et al. 1991) 提出，通过组合多个子模型并执行条件计算来提高效率。最近在自然语言处理 (Shazeer et al. 2017; Chen et al. 2024)、视觉 (Riquelme et al. 2021; Wu et al. 2022; Liu et al. 2024) 和多模态 (Lin et al. 2024; Akbari et al. 2023) 应用中的研究采用了稀疏 MoE 技术，以降低计算成本并训练大规模模型。与传统专家模型不同，现代 MoE 方法如记忆混合专家（Memory Mixture of Experts, MMoE）(Ye and Xu 2023) 在多任务训练期间动态组装特定于任务的特征，显著增强了多任务预测能力。Switch Transformer (Fedus, Zoph, and Shazeer 2022) 利用了分布式训练机制和专家共享技术。在本文中，我们从 MoE 方法中获得灵感，提出了一种新的 learngene 提取方法，以决定数据集相关模块与基础模块之间的传播路径。

### Learngene

Learngene 涉及两个关键阶段：从 Ans-Net 中浓缩 learngene，使用这些 learngene 构建 Des-Net，并在最少的步骤内对下游数据集上的 Des-Net 进行微调。Vanilla Learngene (Van-LG) (Wang et al. 2022) 根据 Ans-Net 在多个任务上训练期间的梯度信息，将高层识别为 learngene。这些 learngene 与可变数量的随机初始化低层结合，以创建不同尺寸的 Des-Nets。Auto-Learngene (Wang et al. 2023) 训练伪 Des-Nets，并采用元网络自动选择与伪 Des-Nets 最相似的 Ans-Net 层作为 learngene。Learngene Pool (Shi et al. 2024) 将 Ans-Net 蒸馏为多个较小的模型，使用它们的层作为 learngene 实例来构建一个池，从中构建 Des-Nets。TLEG (Xia et al. 2024a) 将 Ans-Net 蒸馏为线性扩展以创建不同深度的 Des-Nets 的层。PEG (Wang et al. 2024) 使用概率方法从 Ans-Net 中采样多头自注意力和前馈层作为 learngene，并通过非线性映射扩展它们以形成 Des-Nets。SWS (Xia et al. 2024c) 将 Ans-Net 蒸馏为具有多个阶段、共享层权重的单一模型，其中每个阶段的层作为 learngene 指导 Des-Net 的扩展。然而，这些方法从单一数据集中提取 learngene，限制了其有效性，并最终影响扩展后的 Des-Nets 的性能。

## 方法论

为了克服这些局限性，我们引入了一种称为 LSAMD 的新方法，详见图 2。在本节中，我们通过关注两个方面来概述 LSAMD：使用多个数据集从 Ans-Net 中搜索 learngene，以及基于 learngene 构建可变尺寸的 Des-Nets。

### 跨多个数据集的 Learngene 搜索

**向 Ans-Net 添加数据集特定模块。** 我们首先采用 DeiT-Base (Touvron et al. 2020) 作为 Ans-Net，因为其较大的参数数量使其能够从预训练数据集中获取广泛的知识。随后，我们使用具有不同数据分布的 M 个数据集 $T_1, T_2, \dots, T_M$ 来提取 learngene。为了使 Ans-Net 在学习数据集特定知识的同时保留其原有能力，我们在 Ans-Net 的每一层中引入了一个额外的组件，称为数据集特定模块（$h_i^D$），它与称为基础模块（$h_i^B$）的原始模块并行操作，如图 2 (a) 所示。由于与训练良好的基础模块相比，随机初始化的数据集特定模块学习能力有限，这种设置可能会在训练期间对数据集特定模块不利。为了解决这个问题，我们在 Ans-Net 的每一层中使数据集特定模块和基础模块共享相同的架构和初始参数。在训练期间，对于每个数据集，我们更新数据集特定模块的参数，同时保持基础模块冻结。这种方法确保基础模块不受数据集特定知识的影响。此外，为了防止两个模块之间的干扰，我们实施了一种机制，根据输入数据每层仅选择一个模块进行传播。

**向 Ans-Net 添加数据集适配器。** 为了在这两个模块之间进行选择，受 Mod-Squad 方法 (Chen et al. 2022) 的启发，我们引入了一个类似路由器的模块，称为数据集适配器（DAD）。DAD 通过在数据集特定模块和基础模块之间进行选择来确定传播路径，如图 2 (a) 所示。对于第 $i^{th}$ 层，给定输入数据 $x_t \in \mathbb{R}^{N \times D}$