突破限制：ReLoRA技术让大型神经网络训练效率飞跃

LLM.

已于 2024-05-05 13:37:49 修改

阅读量1.2k

点赞数 14

文章标签：神经网络人工智能深度学习 python tensorflow 机器学习

于 2024-03-20 10:43:56 首次发布

本文链接：https://blog.csdn.net/2301_81888214/article/details/136869378

版权

概述

ReLoRA，即“High-Rank Training Through Low-Rank Updates”，该篇文章基于LoRA，提出了一种通过低秩更新训练高秩网络的方法，从而能够降低训练模型所需要的算力。自2011年以来，随着深度学习技术的飞速发展，模型规模不断扩大，从百万参数到现在的数百亿参数。然而，这种规模的增长也导致计算成本不断升高，使得训练过程变得日益复杂。

ReLoRA提出了一种新的训练技术，旨在通过低秩更新来高效地训练高秩网络，不仅能够节省资源，还能使得模型性能和常规训练相同。

核心理念

ReLoRA的核心在于利用矩阵的秩性质，通过一系列低秩的更新来聚合成一个高秩的网络。这种方法的灵感来源于LoRA（Low-Rank Adaptation of Large Language Models），但ReLoRA进一步优化了训练过程，使其更加高效。

方法细节

ReLoRA的训练过程包括以下几个关键步骤：

全秩预热：在训练开始阶段，使用全秩的方式对网络进行热启动，以确保模型有一个良好的起点。（其实就是进行一定程度的预训练。）
低秩更新：在预热之后，ReLoRA通过应用低秩更新来逐步调整模型参数，这些更新是通过特定的矩阵分解技术实现的。
参数合并与重置：在训练过程中，定期将低秩参数合并回主网络，并重置优化器状态，这一步骤有助于模型探索更广阔的参数空间。
学习率调整：ReLoRA采用了一种特殊的学习率调度策略，即锯齿形余弦调度，以适应模型在不同阶段的训练需求。

1. 全秩预热（Full-Rank Warm Start）

在训练的初始阶段，ReLoRA首先使用全秩的方式对神经网络进行预热————网络的所有参数都以常规的方式进行初始化和训练。简而言之就是对于整个网络进行预训练：即所有的参数都是可训练的，并且使用标准的优化器（如Adam）进行更新。

预热周期：预热所需要的epoch数目并不是固定的，而是一个超参，实际的预热周期可能会根据模型的表现和收敛速度进行调整。
预热后的状态：在全秩预热阶段结束后，模型应该已经学习到了一些有效的表示，这为后续的低秩更新奠定了基础。在这个阶段，模型可能还没有达到最优性能，但已经足够成熟。
过渡到低秩更新：一旦完成了预热步骤，模型就会过渡到ReLoRA的低秩更新阶段。在这个阶段，模型的线性层会被转换为ReLoRA的形式，即通过低秩矩阵乘积来更新权重。

2. 低秩更新（Low-Rank Updates）

通过LoRA（Low-Rank Adaptation）技术，我们可以知道，一个模型中的权重更新值ΔW\Delta WΔW，通常是一个巨大的低秩矩阵。而这就意味着：权重更新ΔW\Delta WΔW可以被分解为两个较小秩的矩阵WAW_AWA和WBW_BWB的乘积，即ΔW=ABT,A∈Rn×d,B∈Nd×m,d≪n\Delta W=A B^T, A \in \mathbb{R}^{n \times d}, B \in \mathbb{N}^{d \times m}, d \ll nΔW=ABT,A∈Rn×d,B∈Nd×m,d≪n， ddd 就是 ΔW\Delta WΔW 这个参数矩阵的秩（Rank，lora_dim）。由于下游细分任务的域非常小，所以ddd可以取得很小，很多时候我们可以取1。因此我们可以将原本的线性层转化为参数量非常小的矩阵乘积。

为大家举个直观的例子，我们假设原来的ΔW\Delta WΔW是1001024的参数矩阵，那么参数量为102400，LoRA模型将ΔW\Delta WΔW矩阵拆成了两个矩阵相乘，如果我们设置Rank=8，那么就是1008的B矩阵与8*1024的A矩阵做矩阵乘法，参数量为800+8192=8992，整体参数量下降了约11.39倍。

在低秩更新阶段，WAW_AWA和WBW_BWB是新增的可训练参数，它们在训练后可以合并回原始权重矩阵中。

3. 参数合并与重置（Parameter Merging and Reinitialization）

在训练过程中，ReLoRA会定期将低秩更新参数WAW_AWA和WBW_BWB合并回网络的主参数中。这一步骤通过简单的矩阵加法来完成，即Wi=(Wi+WAiWBi)W_i = (W_i + W_A^iW_Bi)Wi=(Wi+WAiWBi)，其中WiW_iWi是原始权重，WAiW_A^{iWAi和WBiW_B}iWBi是低秩更新参数。合并后，WAW_AWA和WBW_BWB会被重新初始化，WAW_AWA使用Kaiming初始化，而WBW_BWB则初始化为零。这一过程有助于模型在保持之前学习到的知识的同时，继续探索新的参数空间。

简而言之，相对于LoRA，ReLoRA其实就是在训练过程中，周期性的采用LoRA的方式进行训练，将模型中的线性层分解为两个低秩矩阵，而后直接更新低秩矩阵，从而减少需要计算的参数量。而后周期性的合并低秩矩阵到线性层中，实现权重的彻底更新。

但是，我们在训练的时候，是需要记录优化器状态和调整学习率的，那么当我们重新初始化低秩矩阵后，我们该如何处理已计算的优化器状态和自动计算的学习率呢？

4. 优化器状态重置（Optimizer State Reset）

很多优化器，如Adam，其在训练时会保留关于模型梯度的累计信息。在合并参数后，为了确保模型不会受到之前训练步骤中累积的梯度信息的影响，ReLoRA 执行了优化器状态的部分重置。

优化器（如 Adam）维护了两组矩阵来跟踪梯度的一阶矩（mean）和二阶矩（unbiased variance）。在 ReLoRA 中，这两组矩阵通常被称为 M（一阶矩）和 V（二阶矩）。为了重置优化器，ReLoRA 执行了以下操作：

**基于阈值的剪枝：基于M 和 V 矩阵中元素的绝对值大小，保留较大值的梯度信息，丢弃较小的梯度信息，从而将 M 和 V 矩阵中的大部分元素设置为零。**这个过程可以通过以下伪代码表示：

scss
复制代码
FOR each matrix M and V in the optimizer state
    PRUNE(M, threshold) // 保留幅度大于阈值的元素，其余设置为零
    PRUNE(V, threshold)

由于 M 和 V 矩阵的大部分元素被剪枝，优化器的状态实际上是被部分重置了。这样做有助于模型在新的参数空间中探索，而不是被旧的梯度信息所束缚。

5. 学习率调整（Learning Rate Scheduling）

在优化器状态重置之后，ReLoRA 会调整学习率，以帮助模型在新的参数配置下稳定训练。

ReLoRA采用了一种锯齿形余弦调度（jagged cosine scheduler）来调整学习率。这种调度策略在每次优化器重置后将学习率设为零，然后通过一个预热阶段（例如，50-100步）逐渐恢复到之前的学习率。而后通过余弦退火方法计算新的学习率。

如下图所示：

每一次将学习率设置为0之后，会快速把学习率拉回设置之前的值，而后再改用余弦退火方法。

余弦退火是一种常见的学习率调度方法，它模拟了余弦函数的周期性变化。在这种方法中，学习率从一个较高的初始值开始，随着训练的进行逐渐减小。余弦退火的基本公式如下：

ηt=12η0(1+cos⁡(T2πt))\eta_t = \frac{1}{2} \eta_0 (1 + \cos(\frac{T}{2} \pi t))ηt=21η0(1+cos(2Tπt))

其中，ηt\eta_tηt 是在时间步 ttt 的学习率，η0\eta_0η0 是初始学习率，TTT 是总的训练步数。

6. 局部低秩训练（Locally Low-Rank Training）

ReLoRA的实验分析表明，尽管预训练的神经网络在长期轨迹上表现出高秩更新，但对于足够小的轨迹，训练可以通过低秩更新有效近似。这意味着网络训练在局部上是低秩的，这一观察结果直接激励了ReLoRA的设计。

实验结果

实验设置

实验中，ReLoRA 被应用于不同规模的 Transformer 语言模型，包括 60M、130M、250M、350M 和 1.3B 参数的模型。所有实验都在没有数据重复的情况下进行单次训练（single epoch），并且在计算最优的数据量上进行训练，这是根据 Scaling Laws 估计的。

训练性能

实验结果显示，ReLoRA 在所有测试的模型规模上都显著优于 LoRA 训练。具体来说，ReLoRA 在 250M 模型上的性能甚至接近了全秩训练的性能。这表明 ReLoRA 能够有效地通过低秩更新来近似高秩训练的效果。

内存和速度提升

ReLoRA 在内存和训练速度方面都取得了显著的提升。例如，在 1.3B 参数的模型上，ReLoRA 节省了高达 5.5Gb 的 GPU RAM，并且根据模型大小和硬件设置，训练速度提高了 9-40%。这说明 ReLoRA 是一种高效的训练方法，可以在保持性能的同时减少资源消耗。

高秩训练的证据

通过对学习到的权重更新的奇异值谱进行分析，实验结果支持了 ReLoRA 能够执行高秩更新的观点。ReLoRA 的奇异值谱与全秩训练的更为相似，而与仅使用 LoRA 的有显著差异。这表明 ReLoRA 通过多次低秩更新，能够有效地实现高秩训练的效果。

结论

综合实验结果，ReLoRA 证明了其在训练大型神经网络时的有效性。它不仅能够节省大量的内存资源，还能在不同的硬件配置上提高训练速度。更重要的是，ReLoRA 能够在不牺牲模型性能的前提下，实现与全秩训练相似的结果。这些发现表明 ReLoRA 是一种有前景的高效训练方法，特别适合于资源受限或需要快速迭代的场景。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以点击下方链接免费领取或者V扫描下方二维码免费领取🆓

在这里插入图片描述