跳转至

RoSA: Rotational Sparse Adaptation for Memory-Efficient Fine-Tuning

会议: NeurIPS 2026
arXiv: 2610.06243
OpenReview: https://openreview.net/forum?id=QZlXiCc3kK
代码: https://github.com/xxx
领域: 其他
关键词: 大模型微调, 显存优化, 旋转块训练, 稀疏优化, 梯度截断

一句话总结

RoSA 提出一种新型显存高效微调范式,通过固定浅层前缀、常驻输出锚点并让单一可训练块在深层之间随训练周期旋转移动,结合块内稀疏优化实现反向传播深度截断与优化器状态开销的大幅缩减。

研究背景与动机

随着大语言模型规模的持续扩展,下游任务全参数微调面临的主要硬件瓶颈往往并非静态模型权重本身,而是优化器状态与中间激活值所带来的海量显存开销。以 Llama-2-7B 等经典模型为例,模型权重在推理阶段仅需十数吉字节显存,但若采用 AdamW 进行全参数微调,优化器一阶与二阶动量状态必须为所有可训练参数常驻分配显存,导致峰值显存飙升至近 80 GiB,使得普通商用或单卡 40 GB 显存的算力环境根本无法承担全参数稠密适配。

传统的参数高效微调(PEFT)方法如 LoRA、DoRA 与各类 Adapter 主要通过冻结绝大多数主干参数、仅在低秩子空间或瓶颈层内引入微量额外参数来削减显存。然而,这种策略本质上是将模型更新容量死死束缚在一个固定的可训练参数子空间内,缺乏全网络不同层级特征协同演化的灵活性。另一方面,诸如 BAdam 的块坐标优化虽然分块更新,但其遍历全深度且各块均为稠密更新;LISA 随机重要性采样则无法界定固定的冻结前缀,因而无法利用 Transformer 的深度序列特性缩减反向传播路径。

现有方案的核心矛盾在于:如何在维持严格受限显存预算的同时,打破固定子空间的表征束缚,让更大容量的网络参数获得更新机会。本文发现模型在微调期间的梯度幅值沿深度分布呈现显著非均匀性——浅层表示极度稳定且梯度较弱,而中深层承担了绝大多数适配信号。核心 idea:显式冻结浅层前缀并保持输出锚点常驻激活,让单一可训练层块在各训练轮次间沿中深层轮转迭代,同时在激活块内引入稀疏优化,从层级与参数级双重压缩瞬时优化器状态并截断反向传播路径。

方法详解

整体框架

RoSA 针对 Transformer 架构的序列深度特性,将模型参数解耦为三部分:恒定冻结的前缀层集合 \(\mathcal{F}\)、输出端始终常驻激活的锚点集合 \(\mathcal{A}\),以及在训练过程中动态轮换的可训练层集合 \(\mathcal{R}\)。在任意训练步中,仅有一个包含连续层的小规模层块 \(G_k\) 处于激活状态,与输出锚点 \(\mathcal{A}\) 共同参与参数更新与梯度反向传播,其余所有层均被严格冻结。

由于浅层前缀 \(\mathcal{F}\) 在整个微调生命周期中从未发生改变,计算图在反向传播时无需穿透该前缀,反向计算一旦到达当前激活块的起始层 \(l_{\text{start}}\) 即可提前截断,显著消除了早期的梯度与激活存储,并使得越深层的激活块迭代速度越快。同时,RoSA 在激活块内部无缝集成稀疏优化器(如 NanoAdam),仅对块内绝对值最小的参数子集追踪动量状态,从而实现层级旋转与参数级稀疏的双重正交显存压缩。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Tokens"] --> B["恒定冻结前缀<br/>Embeddings + Layers 0–7"]
    B --> C["可训练层轮转集合<br/>Layers 8–30 (划分为 K 个连续分组)"]
    subgraph S1 ["当前激活周期 e (单块激活)"]
        direction TB
        D["当前轮转激活块 G_k<br/>仅维持激活块参数梯度"] --> E["块内稀疏优化<br/>NanoAdam 仅追踪 bottom-ρ 动量"]
    end
    C --> S1
    S1 --> F["常驻输出锚点<br/>Layer 31 + LM Head"]
    F --> G["任务预测输出与损失计算"]
    G -.->|"反向传播截断<br/>至激活块起始层 l_start 终止"| S1

关键设计

1. 深度感知层级旋转与锚点常驻:解耦瞬时显存开销与累积更新容量

常规微调在每一轮优化中必须为所有可调层分配显存,导致瞬时占用极大。RoSA 将待微调深度层集 \(\mathcal{R}\) 划分为 \(K\) 个互不重叠的连续子块 \(G_1, G_2, \dots, G_K\)(每个子块大小为 \(b\) 层)。在第 \(e\) 个周期内,激活层组为 \(\mathcal{S}_e = \mathcal{A} \cup G_k\),其中 \(k = 1 + ((e - 1) \bmod K)\)。当一个子块离开激活状态时,其一阶与二阶动量状态被立即从显存中清除释放;当下一次周期重新访问该块时,执行冷启动重新初始化动量状态。实验表明这种冷启动不会引发训练损失震荡,反而促使模型平滑收敛。通过将输出层与预测头作为常驻锚点 \(\mathcal{A}\),模型保留了对最终任务目标的连续映射能力,使得瞬时显存仅对应 \(b+1\) 层,而在全周期内累积更新参数量可扩展至全模型的 18.8% 以上。

2. 前缀冻结与反向传播截断:利用序列结构降低计算与激活负担

基于全参数微调时浅层梯度接近平坦的经验观察,RoSA 永久冻结前 \(L_{\text{prefix}}\) 层(在 Llama-2-7B 上为嵌入层与第 0–7 层)。由于前缀参数恒定不变,无需为其分配参数梯度与优化器状态。更关键的是,设当前激活块的起始层为 \(l_{\text{start}}\),则根据链式法则,损失函数对低于 \(l_{\text{start}}\) 的层参数不存在反向传播需求:

\[\nabla_{\theta_{\mathcal{F}}} \mathcal{L} = 0 \quad (\text{无需向浅层传递激活梯度})\]

因此,反向传播图在 \(l_{\text{start}}\) 处直接终止。随着激活块向更深层旋转移动,反向传播经过的层数越来越少,深层激活块的单步迭代耗时可缩短高达 27%。

3. 块内正交稀疏优化:双重因子压榨优化器状态显存

单纯减少激活层数仍可能在模型维度较大时承受可观的优化器内存。RoSA 在激活块内引入参数级稀疏更新机制(默认采用 NanoAdam)。在当前激活参数集合 \(\theta_{\mathcal{S}_e}\) 中,依据权重绝对值动态维护一个比例为 \(\rho\) 的稀疏掩码 \(M_t = \text{Select}_\rho(\theta_{\mathcal{S}_e})\),仅对绝对值最小的底端参数追踪一阶与二阶动量并执行自适应更新,其余参数不维护任何动量状态。掩码每隔 \(\Delta t\) 步周期性刷新,确保参数具备进出优化子集的动态自由度。若全模型参数量为 \(\Psi\)、单位参数动量状态字节数为 \(K_{\text{opt}}\),则 RoSA 优化器显存理论复杂度缩减为:

\[\text{Memory}_{\text{opt}} = \left(\frac{l}{L}\right) \cdot \Psi \cdot \rho \cdot K_{\text{opt}}\]

其中 \(l/L\) 为激活层比例,\(\rho\) 为稀疏密度。这一设计将优化器状态显存压低至极致,配合层旋转实现了正交乘积级的显存缩减。

损失函数 / 训练策略

RoSA 沿用下游微调任务的标准交叉熵因果语言建模损失 \(\mathcal{L}(f(x;\theta), y)\)。在优化策略方面,激活块内部参数采用 NanoAdam 优化器更新,设置初始学习率 \(1 \times 10^{-4}\),动量超参 \(\beta_1 = 0.9, \beta_2 = 0.95\),稀疏密度 \(\rho = 0.2\)。掩码更新间隔 \(\Delta t = 600\) 步。训练调度采用循环双周期(2 cycles,共 6 epochs),使得每个旋转块被完整访问两次。当轮转切换发生时,旧激活块的一阶与二阶矩被丢弃,新激活块分配初始显存;模型权重在各个块之间平滑传递且无梯度检查点依赖。

实验关键数据

主实验

在 Llama-2-7B 模型及 8 个常识推理基准(BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC-Easy, ARC-Challenge, OpenBookQA)上的评测结果表明,RoSA 成功将峰值显存压制在 40 GB 硬件上限以内,同时性能极其逼近全参数稠密微调。

方法 轮数 (Epochs) 峰值显存 (GiB) 是否适配 40GB 平均准确率 (%) 训练耗时
FFT (AdamW) 3 78.5 否 79.99 4.37h
Full NanoAdam (\(\rho=0.1\)) – 46.3 否 81.71 –
LoRA 3 – 是 77.61 –
DoRA 3 – 是 79.70 –
RoSA (1 cycle) 3 34.9 是 79.10 3.5h
RoSA (2 cycles) 6 34.9 是 79.61 6.5h

消融实验

在 Llama-2-7B 上的静态块对比消融(保持锚点常驻,NanoAdam \(\rho=0.2\),训练 3 epochs)验证了层级动态轮转相比固定局部层微调的不可替代性;同时多模型泛化实验证明了该方案在 Qwen 与 Mistral 架构上的通用性。

配置 / 对比项 峰值显存 (GiB) 平均准确率 (%) 说明
RoSA 1 cycle (旋转 8–30 层) 34.9 79.10 完整轮转机制
静态微调 Layers 8–15 26.2 78.32 仅微调最具活力的中间块
静态微调 Layers 16–23 25.9 30.07 仅微调中深层严重欠拟合
静态微调 Layers 24–30 25.5 33.68 仅微调靠近输出深层失效
静态微调 Layers 8–23 (翻倍层数) 31.5 78.07 静态增大层数仍不及动态旋转
LISA (8 active layers, 基线) 30.71 70.05 随机层采样性能大幅落后 RoSA 8.45%

关键发现

  • 显存缩减归因分解:在 Llama-2-7B 峰值显存 34.9 GiB 中,静态模型权重占 12.55 GiB,激活块参数梯度占 3.64 GiB,激活层及显存缓存占约 16.2 GiB,而优化器状态仅占 2.52 GiB。相较全模型 NanoAdam (46.3 GiB),层轮转机制独立砍掉了 12.4 GiB 显存开销。
  • 轮转超越静态容量:即便静态微调包含 16 层的中层区域(Layers 8–23),其最终精度 (78.07%) 依然显著低于每次仅激活 8 层但遍历深度的单周期 RoSA (79.10%),证实模型各层特征协同在时域上的动态重排优于局部静态容量。
  • 极佳的正交可组合性:RoSA 可作为即插即用的层调度调度器与 PEFT 复合。在 Qwen-2.5-3B 上,RoSA-DoRA 相比原始 DoRA 将显存由 16.42 GB 降至 12.87 GB(后序轮转更降至 9.44 GB),精度保持在 81.79%(vs. 81.58%);在 GSM8K 上,RoSA-LoRA 以 14.6 GB 显存取得 35.25% 准确率,大幅压制 BAdam (24.48%, 16.7 GB)。

亮点与洞察

  • 软硬件深度协同的截断反向传播:RoSA 并不只是在数学上将参数分块,而是敏锐地利用了因果 Transformer 沿前向深度的拓扑依赖,通过固定前缀将反向计算图直接截断在激活块前沿,带来高达 27% 的后序阶段步速提升。
  • 瞬时开销与累积容量的精巧置换:传统微调将硬件显存死死绑定在瞬时全网络状态上,而 RoSA 揭示出优化器状态可以“在不同时间服务于不同深度”,使累积更新参数量达 12.6 亿(18.8% 模型规模),远超传统 LoRA 的表征空间。
  • 动态冷启动的训练平滑性:在层块切换并丢弃动量后,新激活层重设冷启动未引发任何训练发散或损失阶跃,各周期 transition 处损失单调稳定下降,证明了大模型局部残差更新具备优良的层间表征兼容性。

局限与展望

  • 前缀激活未实现持久化缓存:当前开源实现仍需进行全深度的完整前向传播;由于未将恒定前缀层的中间激活落盘或缓存,导致前向计算依然存在冗余,未来需依赖高效的 INT8/FP8 激活压缩与跨设备流水线优化。
  • 深度超参需要启发式设计:前缀层数、轮转块跨度及稀疏掩码刷新频率目前依赖经验规则与前期梯度探针扫描,在更大规模多模态或超长上下文模型中尚缺乏自动化最优分配理论。
  • 训练总时间与硬件时限的折中:虽然单步计算与显存开销压低至商用 GPU 可承受范围,但为了达到接近全参数微调的表现,通常需要 2 个完整轮转周期,使得总壁钟训练时间相较于拥有海量显存直接全速微调的场景略有增加。

相关工作与启发

  • vs LoRA / DoRA: 传统 PEFT 方案在全训练周期内将可训练参数冻结在固定的低秩适配层内,无法更改原始权重全貌;RoSA 改变的是“何时更新哪一层”,二者正交且可直接叠加(如 RoSA-DoRA),进一步压缩 PEFT 的显存占用。
  • vs BAdam: BAdam 是在全网络深度内进行稠密块坐标下降且各块均完整稠密微调,缺乏对前缀的恒久冻结与输出常驻锚点;RoSA 引入深度感知的截断反向传播与块内稀疏优化,在显存更低的同时显著提升了下游推理任务准确率(GSM8K 上 35.25% vs 24.48%)。
  • vs LISA: LISA 依赖随机重要性采样动态挑选部分层微调,由于激活层随机分布,反向传播无法稳定截断在前缀浅层;RoSA 维持严格的连续块结构,反向路径长度确定可控,同等显存下准确率高出 LISA 8.45 个百分点。

评分

  • 新颖性: ⭐⭐⭐⭐ [打破固定子空间微调定势,将深度感知层轮转与参数级稀疏优化结合]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 7B/3B 多模型、常识推理与数学生成多任务,包含详尽消融与显存成分拆解]
  • 写作质量: ⭐⭐⭐⭐⭐ [动机叙述逻辑链条紧凑,软硬件协同分析透彻且图表论证充分]
  • 价值: ⭐⭐⭐⭐⭐ [为消费级与受限 40GB 显存设备进行大模型近全参数质量微调提供了极具实用价值的范式]