本文是一篇专注于减轻大型Transformer语言模型训练代价的工作。作者提出了一种基于低秩更新的ReLoRA方法。过去十年中深度学习发展阶段中的一个核心原则就是不断的“堆叠更多层（stack more layers），因此作者希望探索能否同样以堆叠的方式来提升低秩适应的训练效率，实验结果表明，ReLoRA在改进大型网络的训练方面更加有效。
在具体操作时，LoRA通常是加入新的可学习参数 和 来实现，这些参数可以在训练后合并回原始参数中。因此，尽管上述方程允许我们在训练阶段中实现总参数更新量高于任意单个矩阵的秩，但其仍然受到 的限制。因此作者想到通过不断叠加这一过程来突破限制达到更好的训练效果。这首先需要对LoRA过程进行重新启动，就可以在训练阶段不断合并每次得到的 和 来得到累加的权重更新量，计算公式如下：但是，想要对已经完成的LoRA过程重新启动并不容易，这需要对优化器进行精细的调整，如果调整不到位，会导致模型在重启后立即与之前的优化方向出现分歧。例如Adam优化器在更新时主要由先前步骤中所累积梯度的一阶矩和二阶矩引导。实际上，梯度矩平滑参数 和 通常非常高，因而在重新启动时的秩上界为 ，相应的梯度矩 和 都是满秩的，在合并参数后就会使用先前的旧梯度来优化 朝向与 相同的子空间方向。为了解决这个问题，作者提出了ReLoRA方法，ReLoRA在合并和重新启动期间可以对优化器进行部分重置，并在随后的预热中过程中将学习率设置为0。具体来说，作者提出了一种锯齿状学习率调度算法，如下图所示，在每次对ReLoRA参数进行重置时，都会将学习率设置为零，并执行快速（50-100 步）学习率预热使其回到与重置前相同的水平范围内。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在改进大型网络的训练方面更加有效，这与作者最开始研究探索一种改进大型网络训练方法的目标是一致的。此外，为了进一步判断ReLoRA是否能够通过迭代低秩更新来实现相比LoRA更高的秩更新训练，作者绘制了ReLoRA、LoRA和全秩训练的热启动权重与最终权重之间差异的奇异值谱。如下图所示，下图说明了LoRA和ReLoRA之间对于 、、 和 奇异值的显著差异，可以看到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.
