通过阶段性权重共享探索 Learngene 以初始化可变尺寸模型

在实践中，我们通常需要构建可变尺寸的模型，以适应不同应用场景下的多样资源约束，其中权重初始化是训练前的重要步骤。最近提出的 Learngene 框架首先从训练良好的大模型中学习一个被称为 learngene 的紧凑部分，随后将 learngene 扩展以初始化可变尺寸的模型。在本文中，我们首先分析了指导训练良好的 learngene 层扩展的重要性，这启发了一种简单但高效的 Learngene 方法的设计，称为 SWS (阶段性权重共享)，其中 learngene 层及其学习过程都为初始化不同尺度的模型提供了关键的知识和指导。具体而言，为了学习 learngene 层，我们构建了一个包含多个阶段的辅助模型，其中每个阶段的层权重是共享的，随后我们通过蒸馏对其进行训练。接着，我们在相应的阶段扩展这些包含阶段信息的 learngene 层，以初始化可变深度的模型。在 ImageNet-1K 上的大量实验表明，与许多从头开始训练的模型相比，SWS 取得了持续更好的性能，同时将总训练成本降低了约 6.6 倍。在某些情况下，SWS 仅在 1 个 epoch 的微调后就表现更好。在初始化适应不同资源约束的可变尺寸模型时，与预训练加微调方法相比，SWS 取得了更好的结果，同时将用于初始化这些模型的存储参数减少了约 20 倍，预训练成本减少了约 10 倍。