本文是一篇专注于减轻大型Transformer语言模型训练代价的工作。作者提出了一种基于低秩更新的ReLoRA方法。过去十年中深度学习发展阶段中的一个核心原则就是不断的“堆叠更多层（stack more layers），因此作者希望探索能否同样以堆叠的方式来提升低秩适应的训练效率，实验结果表明，ReLoRA在改进大型网络的训练方面更加有效。
论文链接：
https://arxiv.org/abs/2307.05695
代码仓库：
https://github.com/guitaricet/peft_pretraining
一段时间以来，大模型（LLMs）社区的研究人员开始关注于如何降低训练、微调和推理LLMs所需要的庞大算力，这对于继续推动LLMs在更多的垂直领域中发展和落地具有非常重要的意义。目前这一方向也有很多先驱工作，例如从模型结构上创新的RWKV，直接替换计算量较大的Transformer架构，改用基于RNN范式的新架构。还有一些方法从模型微调阶段入手，例如在原有LLMs中加入参数量较小的Adapter模块来进行微调。还有微软提出的低秩自适应（Low-Rank Adaptation，LoRA）方法，LoRA假设模型在任务适配过程中对模型权重的更新量可以使用低秩矩阵进行估计，因而可以用来间接优化新加入的轻量级适应模块，同时保持原有的预训练权重不变。目前LoRA已经成为大模型工程师必备的一项微调技能，但本文作者仍然不满足于目前LoRA所能达到的微调效果，并进一步提出了一种可叠加的低秩微调方法，称为ReLoRA。
本文来自马萨诸塞大学洛厄尔分校的研究团队，作者团队将ReLoRA应用在具有高达350M参数的Transformer上时，展现出了与常规神经网络训练相当的性能。此外，本文作者还观察到ReLoRA的微调效率会随着模型参数规模的增加而不断提高，这使得其未来有可能成为训练超大规模（通常超过1B参数）LLMs的新型手段。
一、引言
虽然目前学术界和工业界都在不断推出自家的各种基座模型，但不可否认的是，完全预训练一个具有初等推理能力的LLMs仍然需要非常庞大的算力，例如大家熟知的LLaMA-6B模型[1]就需要数百个GPU才能完成训练，这种规模的算力已经让绝大多数学术研究小组望而却步了。在这种背景下，参数高效微调（PEFT）已经成为了一个非常具有前景的LLMs研究方向。具体来说，PEFT方法可以在消费级GPU（例如RTX 3090或4090）上对十亿级语言或扩散模型进行微调。因此本文重点关注PEFT中的低秩训练技术，尤其是LoRA方法。作者思考到，过去十年中深度学习发展阶段中的一个核心原则就是不断的“堆叠更多层（stack more layers）”，例如ResNet的提出可以使我们将卷积神经网络的深度提升到100层以上，并且也获得了非常好的效果。因此本文探索能否同样以堆叠的方式来提升低秩适应的训练效率呢？
本文提出了一种基于低秩更新的ReLoRA方法，来训练和微调高秩网络，其性能优于具有相同可训练参数数量的网络，甚至能够达到与训练100M+规模的完整网络类似的性能，对比效果如上图所示。具体来说，ReLoRA方法包含（1）初始化全秩训练、（2）LoRA 训练、（3）参数重新启动、（4）锯齿状学习率调度（jagged learning rate schedule）和（5）优化器参数部分重置。作者选择目前非常火热的自回归语言模型进行实验，并且保证每个实验所使用的GPU计算时间不超过8天。
二、本文方法
ReLoRA通过序列叠加的方式仅训练一小组参数就可以实现与全秩训练相当的性能，并且遵循LoRA方法的基础原则，即保持原始网络的冻结权重并添加新的可训练参数。乍一看，这种方式可能显得计算效率低下，但我们需要清楚的是，这种方法可以通过减小梯度和优化器状态的大小，来显著提高显存效率。例如Adam优化器状态消耗的显存通常是模型权重占用的两倍。通过大幅减少可训练参数的数量，ReLoRA可以在相同的显存条件下使用更大的batchsize大小，从而最大限度地提高硬件效率，ReLoRA的整体操作细节如下图所示。
三、实验效果
为了清晰的评估ReLoRA方法的性能，作者将其应用在各种规模大小（60M、130M、250M 和 350M）的Transformer模型上，并且都在C4数据集上进行训练和测试。为了展现ReLoRA方法的普适性，作者重点考察NLP领域的基础语言建模任务。模型架构和训练超参数设置基本与LLaMA模型保持一致。与LLaMA不同的是，作者在实验中将原始的注意力机制（使用float32进行 softmax计算）替换为了Flash注意力[2]，并且使用bfloat16精度进行计算，这样操作可以将训练吞吐量提高50-100%，且没有任何训练稳定性问题。此外，使用ReLoRA方法训练的模型参数规模相比LLaMA要小得多，最大的模型参数才仅有350M，使用8个RTX4090上训练了一天时间就可以完成。
下图展示了本文方法与其他方法的性能对比效果，可以看到ReLoRA显着优于低秩LoRA方法，证明了我们提出的修改的有效性。此外，ReLoRA还实现了与满秩训练（Full training）相当的性能，并且我们可以观察到，随着网络规模的增加，性能差距逐渐缩小。有趣的是，ReLoRA 唯一无法超过的基线模型是仅具有60M参数的最小模型。这一观察结果表明，ReLoRA在改进大型网络的训练方面更加有效，这与作者最开始研究探索一种改进大型网络训练方法的目标是一致的。
四、总结
本文是一篇专注于减轻大型Transformer语言模型训练代价的工作，作者选取了一条非常具有前景的方向，即低秩训练技术，并且从最朴素的低秩矩阵分解 (LoRA) 方法出发，利用多个叠加的低秩更新矩阵来训练高秩网络，为了实现这一点，作者精心设计了包含参数重新启动、锯齿状学习率调度算法和优化器参数重置等一系列操作，这些操作共同提高了ReLoRA算法的训练效率，在某些情况下甚至能够达到与全秩训练相当的性能，尤其实在超大规模的Transformer网络中。作者通过大量的实验证明了ReLoRA的算法可行性和操作有效性，不知ReLoRA是否也会成为大模型工程师一项必备的算法技能呢？
参考
[1] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
[2] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Re. Flashattention: Fast and memory-efficient exact attention with IO-awareness. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
作者：seven_
快抢沙发！
