引言

如果你曾经尝试在本地机器上微调大型语言模型 (LLM) ,你很可能遇到过令人畏惧的“CUDA Out of Memory” (显存不足) 错误。像 LLaMA-3 这样的现代模型虽然能力惊人,但也极其庞大。即使出现了像低秩自适应 (LoRA) 这样的参数高效微调 (PEFT) 方法,其显存需求通常仍超过标准消费级硬件 (如 NVIDIA RTX 3090 或 4090) 的可用容量。

目前的黄金标准 LoRA 通过冻结模型权重并训练小型低秩适配器矩阵来工作。它显著减少了需要保存的参数数量。然而,参数数量只是战斗的一半。训练过程中真正的显存杀手通常是反向传播所需的中间状态 (激活值) 的存储。

今天,我们将深入探讨一篇名为 “From Weight-Based to State-Based Fine-Tuning” (从基于权重到基于状态的微调) 的新研究论文,它挑战了 LoRA 的基础观点。研究人员提出了一种从调整权重转向控制状态的转变。他们介绍了一种名为 并行控制 (Parallel Control) 的方法,实现了一项了不起的成就: 它允许在单张 24GB 消费级 GPU 上微调 LLaMA-3-8B 模型,且无需对模型进行量化,同时还能减少计算时间并保持性能。

现状: 基于权重的微调

要理解这篇论文的创新之处,我们需要先看看目前的微调方式。主导范式是 基于权重的微调 (Weight-Based Fine-Tuning, Weight-FT)

LoRA 如何工作

在标准的神经网络层中,我们有一个预训练的权重矩阵 \(W_0\)。当我们进行微调时,我们要找到一个更新量 \(\Delta W\),使新权重变为 \(W' = W_0 + \Delta W\)。

在全量微调中,\(\Delta W\) 的大小与 \(W_0\) 相同,这是巨大的。LoRA 提出我们可以使用两个小得多的低秩矩阵 \(A\) 和 \(B\) 来近似这个更新。

LoRA 权重更新方程。

在这里,\(W_0\) 是冻结的。只有 \(A\) 和 \(B\) 被训练。由于秩 \(r\) 很小,可训练参数的数量大幅下降。

隐藏的瓶颈

虽然 LoRA 减少了参数,但它并没有完全解决显存问题。为什么?因为神经网络是很深的。为了计算网络较前层中 \(A\) 和 \(B\) 的梯度,系统必须在正向传播期间存储每一层的输入和输出 激活值 (状态) 。

即使你冻结了一个巨大的前馈网络 (FFN) 块,你仍然需要存储它的输入和输出,以便计算连接到它的 LoRA 适配器的梯度。随着序列长度或批次大小 (batch size) 的增加,这些激活值会消耗数 GB 的 GPU 显存——通常远超权重本身。

范式转变: 从权重到状态

研究人员认为,仅将微调视为“权重修改”限制了我们优化效率的能力。相反,他们建议通过 控制理论 (Control Theory) 的视角来看待神经网络。

作为动力系统的神经网络

不要将神经网络视为一堆矩阵,而应将其视为数据随时间 (层) 流动的动力系统。

  • 系统: 预训练模型。
  • 状态 (\(x_t\)): 第 \(t\) 层的特征 (激活值) 。
  • 控制: 我们为引导模型完成下游任务而做出的调整。

在传统的 Weight-FT 中,我们试图通过重建引擎 (改变 \(W\)) 来引导系统。更新看起来像这样:

带有非线性激活的基于权重的更新。

在这里,修改量 \(\Delta W_t\) 被困在非线性激活函数 \(f_t\) 内部。这使得动力学变得复杂且“非仿射”。

基于状态的微调 (State-FT)

如果我们不改变引擎,而是在侧面加一个推进器会怎样?这就是 基于状态的微调 的核心思想。我们不再修改函数内部的参数 \(W\),而是直接向状态流中注入控制信号。

在数学上,这看起来像一个仿射控制系统:

仿射控制系统方程。

在这个框架中,我们将预训练层视为一个固定函数 \(f_t\)。然后我们添加一个并行控制项 \(G(t)K(t)x(t)\)。这将“微调”与“预训练动力学”解耦了。

研究人员将其推广为一个基于图的框架。如果神经网络是一个状态图,我们可以通过添加基于其前序状态 \(x_u\) 的扰动来修改任何状态 \(x_v\) (节点) :

通用状态更新方程。

这里,\(g_v^u\) 是 控制函数 。 它接收输入状态 \(x_u\) 和一个可学习的控制矩阵 \(M\),计算修正量,并将其添加到原始冻结块的输出中。

解决方案: 并行控制

理论框架引出了一种名为 并行控制 (Parallel Control) 的实用架构。目标是通过跳过昂贵的中间状态存储来最大化显存节省。

设计控制块

现代 Transformer (如 LLaMA 或 ViT) 建立在 残差块 (Residual Blocks) 之上。一个块通常包含一个注意力机制或前馈网络 (FFN) ,后面跟着一个残差连接 (将输入加到输出上) 。

作者建议将 整个残差块 (例如,包括 LayerNorm、扩展层、激活函数和投影层的整个 MLP 部分) 视为单个“黑盒”单元。

让我们看看这与 LoRA 有何不同:

LoRA、Control 和 DoubleControl 架构的比较。

  • 图 (a) LoRA: 适配器被注入到块内部的线性层中 (例如 Q, K, V) 。这意味着块的内部结构在梯度计算图中仍然是活跃的。
  • 图 (b) Control: 在整个 FFN 块的平行位置添加了一条新的“绿色”路径。原始 FFN 块被视为一个固定函数。
  • 图 (c) Double Control: 这一概念被扩展到同时覆盖注意力块和 FFN 块。

并行控制的数学原理

对于特定块 (比如 FFN 块) ,输入是 \(x_u\)。原始块执行复杂的变换 \(f_v^u(x_u)\)。并行控制添加了一个修正项 \(g_v^u\)。

状态的更新规则变为:

并行控制更新方程。

控制函数 \(g_v^u\) 可以很简单。为了保持低参数量,作者使用了类似于 LoRA 的低秩瓶颈结构:

控制函数定义。

这看起来像 LoRA,但 位置 不同。它是与整个块并行,而不是在某个线性层内部。

为什么这能节省大量显存

这是论文最关键的部分。为什么将适配器从“层内部”移到“与块并行”能节省显存?

在标准 LoRA (应用于 MLP 块) 中,你是在微调块内的权重。为了计算这些权重的梯度,你必须存储 MLP 第一个扩展层之后的激活值。在 LLM 中,这个扩展层将隐藏维度投影到 \(4\times\) 大小 (例如,从 4096 到 16384) 。为批次中的每个 token 存储这个巨大的张量是非常昂贵的。

在并行控制中:

  1. 原始 MLP 块是冻结的。
  2. 我们不需要更新其内部的任何参数。
  3. 因此, 我们不需要存储其内部的中间状态来进行反向传播。
  4. 我们可以以“推理模式” (不跟踪梯度) 执行原始块,并立即丢弃其内部数据。
  5. 我们只需要存储输入 \(x_u\) 和低秩控制路径的小型中间状态。

这种差异的影响如下图所示:

显存消耗分析: LoRA vs. 并行控制。

看右边的柱状图。

  • LoRA (Top): 橙色条 (States) 非常大。这代表了必须保持存活的内部激活值 (\([x_t^1]_L\))。
  • Control (Bottom): 显存占用极小。巨大的内部状态消失了,取而代之的是控制路径微不足道的显存成本。

理论验证

一个自然的担忧是,这种“旁路”是否具有与修改内部权重相同的能力。作者提供了理论支持,以确保我们没有失去表达能力。

表达能力

对于深度线性网络,作者证明了只要总秩保持一致,并行控制在数学上与 LoRA 在表达能力上是等价的。

秩不等式定理。

这个不等式意味着存在一个控制路径中的权重矩阵,可以匹配内部低秩适应的效果。

更好地处理奇异点

有趣的是,State-FT 框架在某些非线性场景中实际上可能 更好。如果预训练模型有“死”神经元或奇异点,即梯度 \(\nabla f(x_t)\) 为零,标准的 Weight-FT (依赖于乘以该梯度) 就会卡住。

带有控制的状态更新。

因为并行控制添加了一个解耦于冻结块内部非线性的加法项 \(x(t)u(t)\),它可以绕过这些死区并继续有效地引导状态。

实验结果

理论听起来很扎实,但在实践中效果如何?研究人员在视觉 Transformer (ViT)、RoBERTa 和 LLaMA 模型上测试了该方法。

1. ViT 测试

使用视觉 Transformer 在 CIFAR-100 上,他们对比了 LoRA 和并行控制。

ViT 性能表。

结果非常明显:

  • 显存: 从 18 GB 降至 12 GB。
  • 时间: 训练时间从 4小时42分 降至 3小时24分。
  • 准确率: 于 LoRA。

2. GLUE 基准 (RoBERTa)

在自然语言理解任务上,这种趋势得以延续。

GLUE 基准测试结果。

Control 方法在几乎所有任务上都优于 LoRA 和 DoRA (LoRA 的一种近期改进版) 。值得注意的是“Avg”得分: Control 为 86.14 , 而 LoRA 为 85.62

3. “圣杯”: 在消费级硬件上运行 LLaMA

这是会让学生和业余爱好者兴奋的结果。研究人员采用了 LLaMA-2-7B 和 LLaMA-3-8B,并尝试在单张 NVIDIA RTX 3090 (24GB VRAM) 上训练它们。

通常,以 16 位精度加载一个 8B 模型大约占用 16GB。LoRA 的梯度和优化器状态很容易吃掉剩下的 8GB,导致 OOM 错误,除非你使用 4-bit 或 8-bit 量化 (这可能会降低性能) 。

使用 Double Control (对 Attention 和 MLP 块都应用控制) ,他们实现了以下成果:

在 Nvidia-3090 上训练 7B/8B 模型。

  • 显存使用: LLaMA-2-7B 约为 20.6 GB,LLaMA-3-8B 约为 22.1 GB。
  • 可行性: 它可以舒适地放入 24GB 的显卡中,且 无需量化
  • 性能: 在常识推理任务上的平均准确率保持竞争力 (例如,LLaMA-3-8B 上 DoubleControl 的平均分为 84.7) 。

这意味着微调最先进模型的显存门槛已显着降低。你不再需要 A100 或激进的量化来有效地微调这些模型。

结论

基于权重基于状态 的微调转变,代表了我们对如何适应大型神经网络的根本性反思。通过接受预训练模型作为固定的动力系统并对其状态应用 并行控制 , 我们可以将预训练模型的显存成本与微调过程解耦。

关键要点:

  1. 状态优于权重: 控制数据流 (状态) 比修改引擎 (权重) 更简洁且显存效率更高。
  2. 并行控制: 为残差块添加旁路允许我们在训练期间丢弃巨大的内部激活状态。
  3. 普及化: 这种方法使得在消费级硬件 (RTX 3090) 上全精度微调 7B 和 8B 参数模型成为可能,这是以前不进行量化很难实现的壮举。

这项研究为更易于访问的 AI 研究打开了大门。它表明,高效训练的未来可能不仅仅在于更少的参数,还在于对计算图本身的更智能控制。对于计算预算有限的学生和研究人员来说,“并行控制”是一项非常值得实施的技术。