[论文解读] LEMON: Lossless model expansion

Yite Wang, Jiahao Su+7 more|arXiv (Cornell University)|Oct 12, 2023

Advanced Neural Network ApplicationsComputer Science被引用 3

一句话总结

LEMON 是一种无损模型扩展框架，通过使用预训练的小型模型初始化更大的 Transformer 模型，实现在不损失性能的前提下加快训练速度。通过采用对称性破缺的权重初始化和优化的学习率调度策略，与从零开始训练相比，LEMON 在视觉 Transformer 和 BERT 上分别将计算成本降低了 56.7% 和 33.2%。

ABSTRACT

Scaling of deep neural networks, especially Transformers, is pivotal for their surging performance and has further led to the emergence of sophisticated reasoning capabilities in foundation models. Such scaling generally requires training large models from scratch with random initialization, failing to leverage the knowledge acquired by their smaller counterparts, which are already resource-intensive to obtain. To tackle this inefficiency, we present $ extbf{L}$ossl$ extbf{E}$ss $ extbf{MO}$del Expansio$ extbf{N}$ (LEMON), a recipe to initialize scaled models using the weights of their smaller but pre-trained counterparts. This is followed by model training with an optimized learning rate scheduler tailored explicitly for the scaled models, substantially reducing the training time compared to training from scratch. Notably, LEMON is versatile, ensuring compatibility with various network structures, including models like Vision Transformers and BERT. Our empirical results demonstrate that LEMON reduces computational costs by 56.7% for Vision Transformers and 33.2% for BERT when compared to training from scratch.

研究动机与目标

• 为解决从零开始训练大模型的低效问题，利用小型预训练模型的知识。
• 开发一种适用于多种 Transformer 架构（包括宽度和深度增量不可分的模型）的无损模型扩展方法。
• 在保持模型性能的同时，减少训练时间和计算成本。
• 优化扩展模型的训练方案，特别是通过量身定制的学习率调度策略。

提出的方法

• LEMON 采用对称性破缺的初始化策略，为扩展层中重复的神经元分配不相等的输出权重，以避免冗余并实现真正的容量提升。
• 引入平均扩展方法以处理宽度扩展过程中的 LayerNorm 层，确保即使在宽度增量不可分（例如 512 → 768）的情况下也保持兼容性。
• 通过适配架构差异，将扩展过程应用于多种 Transformer 变体，包括标准 Transformer 和 Pre-LN Transformer。
• 采用自定义的学习率调度器，在保持与从零开始训练时相同最大学习率的基础上，应用更快的衰减，从而提升扩展模型的收敛速度。
• 该框架支持将小型预训练模型的权重直接迁移至更大的目标模型，同时保持功能等价性。
• 实证验证表明，扩展后的模型在性能上与从零开始训练的模型相当，但训练时间显著缩短。

实验结果

研究问题

• RQ1 我们能否在不造成性能下降的前提下，将预训练的小型 Transformer 模型扩展为更大的模型，即使在宽度或深度增加不可分的情况下？
• RQ2 在模型扩展过程中，如何打破对称性，以确保更大模型的容量真正超越源模型？
• RQ3 针对扩展模型，最优的训练方案（尤其是学习率调度）是什么，以实现更快的收敛？
• RQ4 所提出的扩展方法是否可推广至多种 Transformer 架构，包括 ViT 和 BERT？
• RQ5 与从零开始训练相比，无损模型扩展在多大程度上降低了计算成本？

主要发现

• 与从零开始训练相比，LEMON 在视觉 Transformer 上将计算成本降低了 56.7%，在 BERT 上降低了 33.2%。
• 扩展后的 ViT（12,768）仅用 85 个周期即可达到与从零开始训练的模型相同的性能，训练时间减少了 28.3%。
• 该方法保持了与源模型完全的功能等价性，确保扩展后无性能下降。
• 优化后的学习率调度器（使用相同的最大学习率但更快衰减）显著提升了扩展模型的训练效率。
• 该方法兼容多种 Transformer 架构，包括 ViT、BERT 和 Pre-LN 变体，即使在宽度增量不可分的情况下也适用。
• 实证结果证实，扩展后的模型完整继承了源模型的所有知识且无损失，验证了扩展过程的无损性。