【论文笔记】Initializing Models with Larger Ones

Abstract

介绍权重选择，一种通过从预训练模型的较大模型中选择权重子集来初始化较小模型的方法。这使得知识从预训练的权重转移到更小的模型。
它还可以与知识蒸馏一起使用。

权重选择提供了一种在资源受限的环境中利用预训练模型力量的新方法，希望能够成为大模型时代训练小模型的有用工具。

OscarXZQ/weight-selection (github.com)

1 Introduction

神经网络权重的初始化对于其优化至关重要，正确的初始化有助于模型收敛并防止梯度消失等问题。
本文引入一种权重初始化的方法解决该问题，使用大型预训练模型来训练小型模型。

引入了权重选择(weight selection)，选择预训练大型模型的权重子集来初始化较小的模型。
使用权重选择来初始化小模型非常简单，与从头开始训练相比，不会增加额外的计算成本。

3 Weight Selection

给定预训练模型，目标是为同一模型系列中较小尺寸的模型获得有效的权重初始化。借用知识蒸馏中的术语，将预训练模型称为教师，将我们旨在初始化的模型称为学生。

3.1 Approach

模型可以通过增加嵌入维度或每个块中的通道数量来加宽，并通过堆叠更多层来加深。
它还使我们能够按照三个步骤执行权重选择：层选择、组件映射和元素选择。

图2：权重选择(Weight Selection)，为了初始化预训练模型的较小变体，从预训练模型的相应组件中统一选择参数。

3.1.1 Layer Selection

第一步是从教师模型中选层。对于学生的每一层，教师模型中被选中的层都是用于初始化的数据来源。对于各个同性架构和分层架构，层选择的过程略有不同。

各向同性架构是指神经网络的每一层在整个模型中都表现出一致且统一的分层设计。

分层结构的特点是多尺度表示和嵌入维度的层次结构。分层结构通常具有不同规模和嵌入维度的阶段，如VGG这样的经典卷积网络逐渐减少空间维度，同时增加通道维度，捕获多尺度特征。

对于各向同性架构，我们选择教师的前$N$层，$N$为学生模型的层数，记作前$N$层选择(first-$N$ selection)。
处理分层结构时，在每一个独立的阶段单独应用前$N$层选择。

另一种方法是统一层选择(uniform layer selection)，即选择教师中均匀间隔的层。

3.1.2 Component mapping

第二步，在学生和教师之间映射组件。上一步获得了从教师到学生的层的映射，将任务简化为用一个教师层初始化一个学生层。
得益于现代神经网络设计采用的模块化方法，同一系列模型中的层具有一组相同的组件，只是宽度不同。
因此，匹配相应组件的过程是自然的一对一映射。

组件(component)：在人工智能领域中，神经网络的component通常指的是网络的组成部分，也可以理解为神经网络的各个层（layers）或模块（modules）。神经网络由许多不同类型的层组成，每一层都有自己特定的功能和参数。这些层可以分为输入层、隐藏层和输出层，其中隐藏层可以进一步细分为卷积层、池化层、全连接层等。

3.1.3 Element selection

建立组件映射后，下一步是用教师的较大对应组件来初始化学生的较小组件，使用uniform selection，从教师tensor中均匀选择间隔的元素，具体方法在下一部分介绍。

3.2 Methods for element selection

本节制定元素选择并介绍不同的选择标准。

考虑一个学生的权重tensor $W_s$，想找到用于初始化的教师权重tensor $W_t$。如果$W_t$具有shape：$t_1,t_2,\cdots ,t_n$，那么具有相同组件类型的$W_s$，在选择上也要跨过这$n$个维度。
目标是选择$W_t$的子集来初始化$W_s$。

下面讨论几种可能的元素选择方法，并在4.3节中对其表现进行比较。只要保持了一致性（如具有一致性的随机选择段落中所述），权重选择就可以达到类似性能的水平。推荐实践中使用uniform selection作为权重选择的默认值。

一致性(consistancy)：通常指的是在模型训练、推断或应用过程中，确保系统行为和输出的稳定性和一致性。

1. 数据一致性：确保数据在不同阶段（采集、预处理、训练、测试）的一致性，相同的数据在不同情况下应有相同的表现和输出。
2. 模型一致性：确保模型在不同环境或输入下的行为一致，在不同硬件、软件或数据集上训练的模型应该产生相似的结果。
3. 结果一致性：确保模型的输出在不同运行中保持一致，给定相同的输入或条件，模型应该产生相同的输出或行为。

3.2.1 Uniform selection (default)

对于$W_t$的第$i$维，从$t_i$中均匀间隔地选择$s_i$。
比如，想使用$4\times 6$的线性层$W_t$来初始化$2\times 3$的线性层$W_s$，沿着第一个维度选择$1,3$数据（均匀从$1,2,3,4$中选择2个），沿着第二个维度选择$1,3,5$切片（均匀从$1,2,3,4,5,6$中选择3个）。

3.2.2 Consecutive selection

对于$W_t$的第$i$维，从$t_i$中选出连续的$s_i$条数据。
对于具有分组组件的架构，连续选择会选择一些整个组，同时忽略对比。
对于没有这种分组组件的架构，连续选择相当于均匀选择。

3.2.3 Random selection (with consistency)

对于所有权重tensor，对于$W_t$的第$i$维，选择$s_i$大小的$t_i$的随机子集。
通过验证实验得到，一致性（为所有权重矩阵选择相同的指数）是权重选择达到最佳性能的关键。
保持一致性的动机源自于残余连接的存在——添加到教师模型中的神经元应该将其操作保留在学生当中。

此外，保持一致性可以在元素选择过程中保留完整的神经元，因为只选择了一致的位置。
均匀选择和连续选择本质上保持了一致性，它们都是具有一致性的随机选择的特殊实例。

3.2.4 Random selection (without consistency)

对于所有权重tensor，对于$W_t$的第$i$维，选择$s_i$大小的$t_i$的随机子集。

该方法不需要为每一个权重tensor选择相同的索引。设计这个方法用来测试保持一致性的重要程度。

4 Experiments

表3：与经典初始化方法的比较，使用Cifar-100数据集测试多种初始化方法的准确度，具有一致性的权重选择方法优于经典的初始化方法。

5 Analysis

5.1 Reduction in training time

weight selection可以显著减少训练时间。直接测量了使用不同数量的epochs，使用weight selection训练ViT-T省去的时间，将结果展示在图4(a)。有weight selection后，只需要1/3的epoch数量即可打到传统初始化方法的表现。

图4：更快地训练。和随机初始化比较，在Cifar-100上，ViT-T只需要1/3的epoch就可以达到其相同的表现。和在ImageNet-1K上预训练+微调对比，weight selection可以在预训练的第60个epoch达到其表现，省区6.12倍的时间。

5.2 Comparison with transfer learning

通过实验来找到预训练所需的训练预算，以匹配权重选择的准确性。在本实验中，我们在ImageNet-1K上以不同的epoch数训练ViT-T，然后在Cifar-100上微调300个 epoch。如图 4b 所示，在 ImageNet-1K 上进行 60 轮预训练才能在Cifar-100上达到相同的性能。在此设置下，与通过预训练达到相同性能相比，权重选择速度提高了6.12倍，且无需访问用于预训练的数据集。