想象一下,你正在试图教导一位庞大而聪慧的学生(一个大语言模型)像人类一样写作。为此,你需要向他们展示数十亿页文本。问题在于?这位学生的“大脑”(模型)如此巨大,以至于需要超级计算机来容纳所有信息,且整个过程耗时数月,并耗费巨额电费。
这篇论文介绍了CR-Net,一种训练这些巨型模型的新方法,它就像给这位学生提供了一套智能捷径,同时不会让他们忘记任何重要内容。
以下是其工作原理,通过日常类比进行分解:
1. 问题:“沉重的背包”
目前,训练这些模型就像要求一名学生背着一个装满他们读过的每一本书的背包,再加上一个记录他们每一个想法的笔记本。
- 问题所在:背包太重了(内存占用过大)。学生花费大量时间仅仅是在背负重量,导致他们无法以应有的速度学习。
- 旧方案:以前的方法试图通过扔掉书籍(减少参数)或压缩书籍来减轻背包重量。但往往这会让学生的智力下降(性能变差),或者压缩/解压的过程如此缓慢,以至于并未节省任何时间。
2. 发现:“邻居效应”
研究人员发现了这些模型思考方式中一个有趣的现象。他们发现,模型中某一层产生的“想法”(激活值)与紧邻的前一层的“想法”非常相似。
- 类比:想象一场接力赛。第二跑道的选手不需要从头开始;他们只需要知道第一棒选手的位置与他们需要到达的位置之间的微小差异。
- 洞察:模型不需要从头计算整个新想法,而只需计算当前想法与前一想法之间的微小差异。关键在于,研究人员发现这些“差异”非常简单且易于描述(在数学上,它们是“低秩”的)。
3. 解决方案:CR-Net(“智能接力”)
CR-Net 改变了模型的架构以利用这一洞察。
- 工作原理:与其让每一层都从头构建一堵全新的、沉重的砖墙,CR-Net 会说:“直接沿用下一层的墙,只需在上面添加一层薄薄的、轻量级的纸,以做出必要的修改。”
- 结果:模型使用更少的材料(参数)就能建成同样的墙。它将“沉重”的、全强度的砖块保留给最底层(以确保基础稳固),而对其他所有部分则使用这些轻量级的“纸张”。
4. 内存技巧:“重做按钮”
即使背包变轻了,模型仍然需要记住它的步骤以便从错误中学习(在训练的“反向”传播阶段)。通常,这需要存储海量数据。
- 创新点:论文引入了一种特殊策略,模型不再存储所有内容。相反,它只存储几个关键检查点。如果它需要回忆某个未保存的步骤,它会利用上述“薄纸”逻辑快速重新计算该特定步骤。
- 类比:与其为了日后回忆而写下长篇故事的每一个字,你只需写下章节标题和每章的第一句话。如果你忘记了中间的某个细节,只需快速重读相关段落。由于这些“纸张”如此简单,重读它们既极快又极便宜。
5. 结果:更快、更便宜、更智能
该论文在从小型(6000 万参数)到大型(70 亿参数)的模型上测试了这种方法。
- 性能:CR-Net 的学习效果与笨重的全尺寸模型一样好,有时甚至更好。
- 效率:它使用的内存显著减少(使其能在更少或更小的计算机上运行),并且所需的计算能力更少。
- 速度:由于需要移动的数据更少,训练速度更快。
总结:
CR-Net 就像这样教导一位巨型学生:“不要一遍又一遍地背诵整本百科全书。只需记住你读过的最后一页,并只写下你今天学到的新单词。”这使得学习过程对计算机而言更快、更便宜、更省力,同时不会损失学生的任何智力。
技术摘要:CR-Net——利用跨层低秩结构扩展参数高效训练
1. 问题陈述
随着大语言模型(LLM)的参数量从百万级扩展至十亿级,其预训练面临的计算和内存需求呈指数级增长。尽管低秩架构已成为降低参数复杂度和内存占用的主要解决方案,但现有方法存在三个关键缺陷:
- 性能次优:低秩参数化(如 LoRA)往往会损害模型容量,因为 Transformer 权重通常表现出接近满秩的特性,这对于知识表示至关重要。
- 计算瓶颈:依赖低秩梯度的方法(如 GaLore)由于涉及奇异值分解(SVD)或随机投影等操作,引入了显著的计算开销,降低了训练吞吐量。
- 激活内存节省有限:虽然现有方法减少了参数、梯度和优化器状态的内存占用,但它们往往未能解决激活存储(反向传播所需的中间变量)带来的巨大内存负担,其大小可达模型参数规模的 1 倍至 4 倍。
2. 方法论
核心洞察:跨层低秩残差
作者提出了一个新颖的观察:LLM 中相邻层激活值之间的差异具有强烈的低秩结构。在 LLaMA-3 和 GPT-2 等模型上的实证评估表明,利用前一层的激活值(Yl−1)加上一个低秩差异项来近似当前层的激活值(Yl),其相对重构误差显著低于仅对 Yl 进行直接低秩近似。
CR-Net 架构
基于这一洞察,作者引入了CR-Net(跨层低秩残差网络),这是一个具有双路径架构的参数高效框架:
- 第一层:保留满秩参数(W1)以保留高秩信息并确保稳定的初始化。
- 后续层(l≥2):用跨层残差结构替换满秩权重矩阵。激活值的计算方式为:
YlP=βlPYl−1P+XlPAlPBlP
其中:
- Yl−1P 是前一层的激活值。
- AlP 和 BlP 是可学习的低秩矩阵(r≪min(hin,hout))。
- βlP 是一个可学习的缩放因子,用于动态平衡历史信号(高秩)与低秩残差的贡献。这使得模型能够在浅层表达层和深层低秩转换之间进行插值,而不会坍缩到低维子空间。
激活高效的重计算策略
为了解决激活内存开销,作者开发了一种专为 CR-Net 架构定制的重新计算策略。
- 挑战:将标准梯度检查点(GCP)应用于 CR-Net 会产生 O(L2) 的开销,因为由于跨层依赖关系,重建某层的激活值需要所有前驱层的完整前向传播。
- 解决方案:该策略存储一部分线性激活值(检查点)和低秩输出(XlAl)。对于未检查点的层,通过跨层残差连接的逆运算来恢复激活值:
Yl−1P=sign(βl+1P)(∣βl+1P∣+ϵ)1(YlP−XlPAlPBlP)
这种方法避免了对前驱层进行完整的前向传播,显著降低了重计算成本,同时保持了内存效率。
3. 主要贡献
- 新颖的基础原理:发现了层间激活残差具有低秩特性,这与 prior works 所利用的梯度或参数的低秩特性不同。
- 参数高效框架(CR-Net):一个利用跨层低秩激活差异的预训练框架。它通过利用满秩第一层和可学习残差缩放,避免了标准基于 LoRA 的方法中重复低秩近似所带来的信息损失。
- 激活高效的重计算:一种定制策略,消除了反向传播过程中存储大多数层激活值的需求,在不产生朴素检查点应用所带来的 O(L2) 开销的情况下,解决了激活内存瓶颈。
- 实证验证:大规模预训练实验表明,CR-Net 在验证性能上优于最先进的低秩框架,同时需要更少的计算资源和内存。
4. 实验结果
实验在参数量从 60M 到 13B 的 LLaMA-2 模型上进行,使用了 C4-en 数据集。
- 性能与参数:CR-Net 在参数数量对齐的情况下,始终优于参数高效基线(LoRA、ReLoRA、SLTrain、CoLA、LORO)。例如,在 10 亿参数模型上,CR-Net 实现了比满秩训练更好的验证困惑度,同时将参数复杂度降低了56.5%,每步计算量降低了63.2%。
- 性能与内存:在内存开销对齐的情况下,CR-Net 超越了优化器高效的方法(GaLore、RSO、Apollo),特别是对于大于 10 亿参数的模型。
- 吞吐量:CR-Net 展现出优越的训练和推理吞吐量。在 LLaMA-2 10 亿参数预训练中,即使考虑数据并行通信开销,其吞吐量也比最先进的方法提高了6% 以上。
- 扩展至 13B:在 130 亿参数模型上,CR-Net 实现了超过 50% 的参数减少,而验证性能相比满秩训练仅下降了2%。
- 消融研究:
- 秩选择:中间层使用较高的秩、侧边层使用较低的秩可获得最佳性能。
- 可学习缩放:与固定缩放因子相比,使用可学习的 βlP 显著提高了数值稳定性和收敛性。
- 第一层:保持满秩第一层至关重要;将其替换为低秩参数会导致困惑度下降 3.5%。
5. 意义与主张
本文将 CR-Net 定位为一个原则性框架,而非技术的临时组合。其意义在于:
- 弥合差距:它成功解耦了参数效率与模型性能之间的权衡,以低秩参数预算实现了类似满秩的能力。
- 系统效率:通过解决常被忽视的激活内存瓶颈,它使得在现有硬件约束下训练更大模型成为可能,而无需承担传统梯度检查点的计算惩罚。
- 稳定性:该设计确保了与满秩方法相似的稳定训练动态,克服了低秩预训练通常伴随的不稳定性。
作者得出结论,CR-Net 为高效的 LLM 预训练提供了一条可扩展的路径,在保持或提升模型能力的同时,降低了内存消耗和计算成本。确定的未来工作方向包括集成混合精度训练以进一步缓解内存开销,并将该框架推广到替代的注意力架构。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。