Initializing Models with Larger Ones

ICLR2024| Initializing Models with Larger Ones

• 论文链接：https://arxiv.org/pdf/2311.18823.pdf
• 源码链接：https://hub.yzuu.cf/OscarXZQ/weight-selection

简介

神经网络权重初始化对其优化至关重要。适当初始化有助于模型收敛，并防止梯度消失问题。两个著名的初始化技术，Xavier初始化和Kaiming初始化，在神经网络训练中发挥了重要作用。它们仍然是PyTorch等现代深度学习库中默认方法。
这些方法都是为从随机初始化开始训练神经网络而开发的。在当时，这是普遍的做法。然而形式已经发生了变化。由于社区的集体努力，各种预训练的模型现在很容易获得。这些模型在大型数据集上进行训练（如ImageNet-21K和LAION-5B），并经常由专家优化。因此，从这些预训练模型中进行微调通常被认为是当今的首选，而不是从头开始训练模型。
然而这些预训练大模型在其资源需求方面可能会令人望而却步，从而阻碍其广泛应用于资源受限的设置，例如在移动设备上。对于许多预训练的模型族，在某些情况下，即使是最小的模型实例也可以被认为是非常大的。例如，掩码自编码器和CLIP都提供了ViT Base。对于边缘设备上应用来说这已经太大了。最小的LLaMA模型甚至大100倍，具有7B参数。由于很少的小模型可用，开发者将不得不在目标数据集上从头训练满足它们需求。这种方法错过了利用大型预训练模型的机会，这些模型的知识是从对大型数据的广泛训练中学习的。
本文中通过引入一种权重初始化方法来解决这个问题，该方法使用大型预训练模型来训练小型模型。具体地，本文引入了权重选择，其中预训练的大模型中选择权重的子集来初始化较小的模型。这允许大模型学习的知识通过其权重转移到小模型。得益于现代神经网络的模块哈设计，权重选择只涉及三个简单步骤：层选择，组件映射和元素选择。此方法可以应用于与大模型相同的模型族中任何较小模型。使用权重选择来初始化一个小模型是简单的，与从头开始训练相比，不会增加额外的计算成本，它甚至对大模型训练也很有用。

模型选择

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

本文方法

现代神经网络架构通常遵循模块化方法：设计一个层并重复它构建模型。这一设计促进了可扩展性：模型可以通过增加嵌入维度或每个块中通道数量拓宽模型，并通过堆叠更多层加深模型。
层选择
从教师模型选择层是第一步。对于学生每一层，都会从教师那里选择一层作为初始化源。各向通性架构和分层架构的层选择程序略有不同。向通性架构是指每一层在整个模型中都表现出一致和统一的分层设计。ViT和MLP-mixer属于向通性架构。层次架构特征在于多尺度表示和嵌入维度的层次结构。分层架构通常具有不同规模和嵌入维度的阶段。例如像VGG这样经典卷积网络在增加通道的同时逐渐降低空间维度，捕捉多尺度特征。例如Swin Transformer和ConvNeXt的现代架构也使用层次架构。
对于向通性架构，从教师那里选择前N层，其中N表示学生层数，表示第一个N选择（First-N Selection）。当处理ConvNeXt等层次架构时，在每个单独阶段应用第一个N选择。另一种选择是均匀层选择，即教师中间隔均匀的分层。
组件映射
在第二步中，映射学生和教师之间的中间。从前面步骤中，获得了教师和学生的层映射。然后任务简化为用一个教师层初始化一个学生层。得益于现代神经网络设计采用的模块化方法，同一家族的模型中的层具有相同组件集，但宽度不同。匹配相应组件的过程因此是一个自然的一对一映射。
元素选择
建立组件映射之后，下一步是使用教师提供的较大组件初始化学生的组件。元素选择的默认方法是均匀选择，其中均匀间隔的元素是从教师的张量中选择。

元素选择算法

在这一部分中，制定了元素选择，并介绍了不同的选择标准。考虑来自学生的权重张量$W_s$，本文试图用教师权重行两$W_t$初始化。如果$W_r$具有形状$t_1,t_2,\dots ,t_n$，则与$W_t$具有相同部件类型的$W_s$也可以跨越n个维度。本文目标是选择$W_t$元素的子集初始化$W_s$。后文讨论一些元素选择可能的方法。
均匀选择
对于$W_t$每个维度i，从$t_i$中选择均匀间隔的$s_i$切片。例如要用形状为$4\times 6$的线性层初始化形状为$2\times 3$的线性层$W$，选择沿着第一位都的第一和第三切片，以及沿着第二维度的第一第三和第五切片。该算法从教师的权重张量$W_t$开始，并在$W_t$所有维度上迭代执行选择，以达到学生想要的形状。值得注意的是在包含分组组件的架构中，如ViT中多头注意力模块和ResNeXt中分组卷积，均匀选择吸收了所有组信息。例如，当应用ViTs时，均匀选择将从注意力块中所有头中选择参数，这将有利于从预训练ViTs中继承知识。
连续选择
对于$W_t$的每个维度i，从$t_i$中连续选择的$s_i$切片。与均匀选择相反，对于具有分组组件的架构，连续选择选择了一些完整的组，同时忽略对比。对于没有这样分组组件的架构，连续选择与均匀选择等价。
随机选择（有一致性）
对于所有权重张量，对于$W_t$每一个维度$i$，从$t_i$中选择相同的随机生成的$s_i$切片集合。根据经验实验，本文发现一致性（为所有权重矩阵选择相同索引）是权重选择达到最佳性能的关键。保持一致性的动机源于残差连接的存在–教师模型中添加的神经元应该具有他们操作保存在学生模型上。此外，由于只保留了一致的位置，因此在元素过程中保持了一致性可以保留完整的神经元。一致选择和连续一致本质上保持一致性，这都是具有一致性的随机选择的特殊情况。
随机选择（无一致性）
沿着$W_t$每个维度$i$，从$t_i$中随机选择$s_i$切片。与具有一致性的随机选择不同，该方法不需要为每个权重张量选择相同的索引。该设计是为了检验一致性的重要性。