跳转至

AdaPreLoRA: Adafactor-Preconditioned Low-Rank Adaptation

会议: NeurIPS 2026
arXiv: 2605.08734
OpenReview: https://openreview.net/forum?id=8o83y7bppR
论文: https://arxiv.org/abs/2605.08734
领域: 模型压缩
关键词: 低秩适应, Adafactor, 预条件优化器, 参数高效微调, 矩阵流形

一句话总结

针对 LoRA 生成映射雅可比矩阵秩亏导致权重空间预条件子无法直接拉回因子空间的根本阻碍,AdaPreLoRA 将现有优化器统一为带约束仿射逆解框架,采用低开销 Adafactor 对角 Kronecker 预条件子并在其赋范度量下求解因子不平衡最小化投影,以 \(\mathcal{O}((m+n)r)\) 显存闭式实现了全权重预条件更新的最优低秩逼近。

研究背景与动机

在大语言模型与扩散模型时代,参数高效微调(PEFT)已成为下游任务适配的标准范式。其中,低秩适应(LoRA)通过冻结预训练基座权重 \(\bm{W}_0\),将权重更新量重参数化为两个极低秩矩阵的乘积 \(\bm{W}=\bm{B}\bm{A}\)(\(\bm{B}\in\mathbb{R}^{m\times r}, \bm{A}\in\mathbb{R}^{r\times n}, r\ll\min(m,n)\)),成功将可训参数与优化器状态开销从 \(\mathcal{O}(mn)\) 骤降至 \(\mathcal{O}((m+n)r)\)。然而,直接在低秩因子空间 \([\bm{B}, \bm{A}]\) 中进行自适应优化(如原生 LoRA 的独立 AdamW 动量)存在根本性的结构缺陷:因为生成映射 \(\mathcal{G}: [\bm{B}, \bm{A}] \mapsto \bm{B}\bm{A}\) 存在固有规范冗余 \((\bm{B}\bm{C}, \bm{C}^{-1}\bm{A})\),其雅可比算子 \(J_{\mathcal{G}}\) 严重秩亏,其零空间 \(\ker(J_{\mathcal{G}})\) 维度高达 \(r^2\)。

这种秩亏性直接破坏了优化器预条件子的传递性。在自然梯度和预条件优化理论中,参数空间的 Fisher 信息矩阵或二阶动量应对应权重空间预条件子 \(\mathcal{F}_t\) 经由链式法则向因子空间的拉回算子 \(J_{\mathcal{G}}^* \mathcal{F}_t J_{\mathcal{G}}\)。然而,由于 \(J_{\mathcal{G}}\) 秩亏,该拉回算子必定奇异,不可直接求逆,使得权重空间的预条件更新方向无法唯一映射回因子更新。以往方法沿两条极端路线妥协:一类是廉价因子空间方法(如 Vanilla LoRA、LoRA+、Riemannian Preconditioned LoRA),它们完全退化为 \(\mathcal{F}_t=\bm{I}\) 或简单的块对角近似,彻底丢弃了权重空间的全局梯度二阶统计量;另一类是全权重投影或流形优化方法(如 LoRA-Pro AdamW、SOAP、RAdamW),前者在欧氏 Frobenius 空间求伪逆,与预条件矩阵内积度量失配,且显式维护 \(mn\) 维权重动量导致显存暴涨 \(2\times\),后者则依赖昂贵的收缩(Retraction)操作或流形 SVD 分解。在 \(\mathcal{O}((m+n)r)\) 显存预算内结合二阶梯度统计量与封闭解析解始终是一个空白。

本文的切入点在于发现:虽然拉回算子 \(J_{\mathcal{G}}^* \mathcal{F}_t J_{\mathcal{G}}\) 奇异,但正规方程系统 \(J_{\mathcal{G}}^* \mathcal{F}_t J_{\mathcal{G}} [\bm{\Delta}_{\bm{B}}, \bm{\Delta}_{\bm{A}}] = J_{\mathcal{G}}^*(\bm{G}_t)\) 永远是代数相容的,其解集构成一个非空的 \(r^2\) 维仿射子空间。核心 idea:选用极轻量的 Adafactor 行列行向量 Kronecker 分解 \(\mathcal{H}_t = \bm{L}_t^{1/2} \otimes \bm{R}_t^{1/2}\) 作为权重预条件子,并在该度量诱导的加权范数下引入因子不平衡最小化准则,推导出仅需 \(\mathcal{O}((m+n)r)\) 显存与 \(\mathcal{O}((m+n)r^2 + r^3)\) 计算开销的严格正交投影闭式更新解。

方法详解

整体框架

AdaPreLoRA 将 LoRA 参数更新过程建模为在权重空间预条件方向下的加权最小二乘切空间正交投影。整个流程由三个核心阶段组成:首先在反向传播中采集因子梯度并维护 \(\mathcal{O}(m+n)\) 维度的 Adafactor 行列二阶统计向量;接着在加权内积空间中求解关于因子更新的相容仿射方程族;最后利用二阶加权平衡准则消去零空间冗余,闭式导出低秩因子的最终更新步。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["反向传播因子梯度<br/>G_B = G_t A^T, G_A = B^T G_t"] --> B["Adafactor 行列二阶矩统计<br/>维护 l_t 与 r_t 向量 (O(m+n))"]
    B --> C["加权切空间投影求解<br/>极小化 H_t 加权残差系统"]
    C --> D["H_t 加权因子平衡选择<br/>消除 r^2 维零空间规范冗余"]
    D --> E["闭式解析因子更新输出<br/>ΔB_t 与 ΔA_t 落地更新"]

关键设计

1. Adafactor 行列二阶矩统计:极低显存的权重空间预条件子

为了在不破坏 LoRA 显存优势的前提下引入全局梯度二阶曲率,方法避免了维护 \(\mathcal{O}(mn)\) 维度的逐元素二阶矩阵,转而引入 Adafactor 秩-1 对角 Kronecker 分解算子 \(\mathcal{H}_t = \bm{L}_t^{1/2} \otimes \bm{R}_t^{1/2}\)。通过指数移动平均(EMA)追踪梯度外积矩阵的行和与列和向量 \(\bm{l}_t \in \mathbb{R}^m\) 和 \(\bm{r}_t \in \mathbb{R}^n\):

\[\bm{l}_t = \beta_1 \bm{l}_{t-1} + (1-\beta_1) \sum_{j=1}^n (\bm{G}_t \odot \bm{G}_t)_{i,j}, \quad \bm{r}_t = \beta_2 \bm{r}_{t-1} + (1-\beta_2) \sum_{i=1}^m (\bm{G}_t \odot \bm{G}_t)_{i,j}\]

归一化后构造对角算子 \(\bm{L}_t = \mathrm{diag}(\bm{l}_t / \sqrt{\|\bm{l}_t\|_1})\) 与 \(\bm{R}_t = \mathrm{diag}(\bm{r}_t / \sqrt{\|\bm{r}_t\|_1})\)。该算子作用于任意矩阵 \(\bm{Y} \in \mathbb{R}^{m\times n}\) 定义为 \(\mathcal{H}_t \bm{Y} = \bm{L}_t^{1/2} \bm{Y} \bm{R}_t^{1/2}\),诱导了加权内积 \(\langle \bm{Y}, \bm{Z} \rangle_{\mathcal{H}_t} = \langle \bm{L}_t^{1/2} \bm{Y} \bm{R}_t^{1/2}, \bm{Z} \rangle\)。这一设计将预条件状态的显存复杂度严格锁死在 \(\mathcal{O}(m+n)\),完美匹配 LoRA 的轻量特性。

2. 加权切空间投影求解:绕过拉回奇异性的相容最小二乘系统

由于雅可比算子 \(J_{\mathcal{G}}\) 的秩为 \((m+n)r - r^2\),直接对算子 \(J_{\mathcal{G}}^* \mathcal{H}_t J_{\mathcal{G}}\) 求逆必定崩溃。本文将其转化为等价的加权最小二乘投影问题,目标是寻找在 \(\mathcal{H}_t\) 范数下距离预条件权重梯度 \(\widetilde{\bm{G}}_t = \mathcal{H}_t^{-1}\bm{G}_t = \bm{L}_t^{-1/2} \bm{G}_t \bm{R}_t^{-1/2}\) 最近的低秩切空间切向量:

\[\min_{\bm{\Delta}_{\bm{B}_t}, \bm{\Delta}_{\bm{A}_t}} \big\| J_{\mathcal{G}}[\bm{\Delta}_{\bm{B}_t}, \bm{\Delta}_{\bm{A}_t}] - \mathcal{H}_t^{-1} \bm{G}_t \big\|_{\mathcal{H}_t}^2\]

理论证明该系统的最优解存在且代数相容,等价于使得复合更新量 \(\bm{\Delta}_{\bm{B}_t} \bm{A}_t + \bm{B}_t \bm{\Delta}_{\bm{A}_t}\) 精确等于 \(\widetilde{\bm{G}}_t\) 在正切流形子空间 \(\mathbb{T}_t = \mathrm{range}(J_{\mathcal{G}})\) 上的 \(\mathcal{H}_t\) 正交投影 \(\widetilde{\mathcal{P}}_{\mathbb{T}_t}(\widetilde{\bm{G}}_t)\)。所有可解的因子更新向量构成一个参数化解族,包含任意自由矩阵 \(\bm{X}_t \in \mathbb{R}^{r\times r}\) 引起的零空间平移 \((\bm{B}_t \bm{X}_t, -\bm{X}_t \bm{A}_t)\)。

3. \(\mathcal{H}_t\) 加权因子平衡选择:消去规范冗余的解析闭式更新

解族中不同的 \(\bm{X}_t\) 虽诱导相同的权重空间步长,却决定了因子空间的不同演化轨迹。为防止某一因子的范数发生爆炸或崩溃,本文在 \(\mathcal{H}_t\) 赋范度量下引入因子更新不平衡极小化准则:

\[\min_{\bm{X}_t \in \mathbb{R}^{r\times r}} \frac{1}{2} \big\| \bm{\Delta}_{\bm{B}_t} \bm{A}_t - \bm{B}_t \bm{\Delta}_{\bm{A}_t} \big\|_{\mathcal{H}_t}^2\]

对目标函数求导并令其为零,可以解析求出最优平衡偏置矩阵 \(\bm{X}_t^{\mathrm{opt}} = -\frac{1}{2}(\bm{B}_t^\top \bm{L}_t^{1/2} \bm{B}_t)^{-1} \bm{B}_t^\top \bm{G}_t \bm{A}_t^\top (\bm{A}_t \bm{R}_t^{1/2} \bm{A}_t^\top)^{-1}\)。将其回代解族后,惊喜地化简为仅依赖于因子梯度 \(\bm{G}_{\bm{B}_t} = \bm{G}_t \bm{A}_t^\top\) 与 \(\bm{G}_{\bm{A}_t} = \bm{B}_t^\top \bm{G}_t\) 的封闭解析式:

\[\bm{\Delta}_{\bm{B}_t}^{\mathrm{opt}} = \big(\bm{I} - \tfrac{1}{2} \widetilde{\bm{P}}_{B_t}\big) \bm{L}_t^{-1/2} \bm{G}_{\bm{B}_t} \big(\bm{A}_t \bm{R}_t^{1/2} \bm{A}_t^\top\big)^{-1}\]
\[\bm{\Delta}_{\bm{A}_t}^{\mathrm{opt}} = \big(\bm{B}_t^\top \bm{L}_t^{1/2} \bm{B}_t\big)^{-1} \bm{G}_{\bm{A}_t} \bm{R}_t^{-1/2} \big(\bm{I} - \tfrac{1}{2} \widetilde{\bm{Q}}_{A_t}\big)\]

其中 \(\widetilde{\bm{P}}_{B_t} = \bm{B}_t (\bm{B}_t^\top \bm{L}_t^{1/2} \bm{B}_t)^{-1} \bm{B}_t^\top \bm{L}_t^{1/2}\),\(\widetilde{\bm{Q}}_{A_t} = \bm{R}_t^{1/2} \bm{A}_t^\top (\bm{A}_t \bm{R}_t^{1/2} \bm{A}_t^\top)^{-1} \bm{A}_t\) 分别为加权投影矩阵。公式中精准出现的系数 \(\frac{1}{2}\) 正是二阶平衡准则的数学印记。该闭式求解仅涉及 \(r\times r\) 小矩阵求逆,单步附加计算量仅为 \(\mathcal{O}((m+n)r^2 + r^3)\)。

损失函数 / 训练策略

在实际训练中,该更新规则支持与一阶动量和解耦权重衰减(AdamW 形式)结合。定义一阶偏置校正动量项后,因子更新形式为:

\[\bm{B}_{t+1} = (1 - \lambda \eta_t)\bm{B}_t - \eta_t \frac{\sqrt{1-\beta_1^t}}{1-\beta_3^t} \bm{\Delta}_{\bm{B}_t}^{\mathrm{opt}}, \quad \bm{A}_{t+1} = (1 - \lambda \eta_t)\bm{A}_t - \eta_t \frac{\sqrt{1-\beta_1^t}}{1-\beta_3^t} \bm{\Delta}_{\bm{A}_t}^{\mathrm{opt}}\]

在实现中,全权重梯度 \(\bm{G}_t\) 可通过反向传播 Hook 获取或直接采用切空间代理 \(\bm{G}_{\bm{B}_t}\bm{A}_t + \bm{B}_t\bm{G}_{\bm{A}_t}\) 计算,动量缓冲直接挂载于低秩因子梯度上以避免开辟全权重显存缓存。当矩阵 \(\bm{B}_t^\top \bm{L}_t^{1/2} \bm{B}_t\) 接近奇异时,添加 \(\epsilon \bm{I}\)(\(\epsilon = 10^{-6}\))进行阻尼正则化。

实验关键数据

主实验

论文在两类主流 7B 级大语言模型(Mistral-7B 和 Qwen2-7B,秩 \(r=8\))上系统评估了下游 NLU(GLUE 中的 RTE, CoLA, MRPC)、多步推理(ARC)及数学推理(GSM8K)表现。所有基线均进行了独立的学习率网格搜索。

模型 基准任务 Vanilla AdamW Scaled AdamW LoRA-Pro AdamW SOAP AdaPreLoRA (本文)
Mistral-7B GLUE-RTE (%) 89.4 89.1 88.8 88.9 89.5
Mistral-7B GLUE-CoLA (%) 69.4 71.5 68.5 68.3 71.4
Mistral-7B GLUE-MRPC (%) 89.5 89.7 84.8 86.3 90.0
Qwen2-7B GLUE-RTE (%) 90.6 90.6 88.1 86.6 91.0
Qwen2-7B GLUE-MRPC (%) 90.0 87.0 89.2 89.2 90.4
Qwen2-7B ARC Reasoning (%) 84.3 85.3 80.9 80.9 85.6
Qwen2-7B GSM8K Math (%) 75.1 74.2 75.7 73.2 76.4

在 Mistral-7B 任务上,对计算效率和显存开销进行了实测统计(单步用时与峰值显存):

优化器方案 理论显存复杂度 实测峰值显存 (GB) 单步用时 (s/step) 显存相对开销
Vanilla AdamW \(\mathcal{O}((m+n)r)\) 25.8 0.24 \(1.00\times\) (基准)
Scaled AdamW \(\mathcal{O}((m+n)r)\) 26.0 0.36 \(1.01\times\)
SOAP \(\mathcal{O}(m^2+n^2)\) 26.4 0.65 \(1.02\times\)
LoRA-Pro AdamW \(\mathcal{O}(mn)\) 50.4 1.35 \(1.95\times\)
AdaPreLoRA SGD (本文) \(\mathcal{O}((m+n)r)\) 21.5 0.46 \(0.83\times\)
AdaPreLoRA AdamW (本文) \(\mathcal{O}((m+n)r)\) 26.0 0.99 \(1.01\times\)

消融实验与多任务泛化

论文在 GPT-2(E2E 生成任务,\(r=4\))以及基于 Stable Diffusion 的 Mix-of-Show 个性化文生图任务上探究了优化器消融与跨领域泛化性。

在 GPT-2 (small, \(r=4\)) E2E 自然语言生成任务上的优化器对比:

优化器算法族 具体方法 BLEU ↑ NIST ↑ METEOR ↑ ROUGE-L ↑ CIDEr ↑
SGD 族 (无一阶动量) 原生 SGD 54.8 4.56 34.0 63.3 1.29
SGD 族 Scaled GD 68.5 8.72 45.5 69.4 2.40
SGD 族 LoRA-Pro SGD 68.4 8.72 45.5 69.6 2.43
SGD 族 AdaPreLoRA SGD (本文) 69.5 8.77 46.5 71.5 2.50
AdamW 族 原生 AdamW 69.1 8.75 46.0 70.5 2.47
AdamW 族 Scaled AdamW 69.5 8.80 46.2 70.9 2.48
AdamW 族 LoRA-Pro AdamW 69.2 8.73 45.9 70.8 2.47
AdamW 族 AdaPreLoRA AdamW (本文) 70.0 8.84 46.3 71.3 2.50

在文生图 Mix-of-Show 个性化微调任务上的表现(LoRA Scaling = 1.0):

优化器方法 CLIP 文本对齐度 ↑ FID 图像分布距离 ↓ 说明
原生 SGD 31.40 40.95 无二阶预条件与平衡
Scaled GD 30.60 29.62 块对角雅可比近似
LoRA-Pro SGD 30.48 29.19 欧氏伪逆投影
AdaPreLoRA SGD (本文) 31.58 28.18 加权切空间投影
原生 AdamW 30.68 27.80 因子独立坐标缩放
Scaled AdamW 24.51 34.18 在扩散任务上出现退化
LoRA-Pro AdamW 30.60 28.18 需维护全权重动量
AdaPreLoRA AdamW (本文) 30.73 27.13 最低 FID 与最优生成质量

关键发现

  • 显存与效果的解耦:全权重预条件方法(如 LoRA-Pro AdamW)在 7B 模型上引入了近 \(2\times\) 的显存开销(从 25.8GB 飙升至 50.4GB),但在复杂推理任务(如 ARC 80.9% vs AdaPreLoRA 85.6%)上表现反而滞后。这表明不当的欧氏度量投影无法充分发挥权重曲率优势,而 AdaPreLoRA 在完全维持原生 LoRA 显存级(26.0GB)的前提下全面胜出。
  • SGD 场景下优势最为显著:在 GPT-2 的无动量 SGD 组中,原生 LoRA 表现崩塌(BLEU 54.8),而 AdaPreLoRA 凭借内在的二阶正交投影直接冲至 69.5,甚至超越了原生 AdamW(69.1),证明了算法所导出的加权几何正交步具有极其强悍的优化导向力。
  • 跨模态稳健性:在扩散模型微调中,部分依赖经验性缩放的方案(如 Scaled AdamW)遭遇了优化退化(FID 激增至 34.18),而 AdaPreLoRA 保持了严格的数学闭式推导,在多尺度下均取得最低 FID。

亮点与洞察

  • 严格自洽的几何视角:指出了 LoRA 优化器领域的统一本质是相容线性系统的广义逆选择,指出 LoRA-Pro 等前序工作在加权度量与欧氏度量间的混淆,构建了第一个兼顾全局梯度统计量与因子空间投影的优化器。
  • 极具美感的闭式推导:将 \(r^2\) 维零空间规范冗余转化为最小化因子不平衡问题,推导出带有特征性 \(\frac{1}{2}\) 投影系数的闭式更新,无需任何数值迭代循环。
  • 真正的实用价值:无需修改模型架构,单步耗时适度(0.99s vs 0.36s),显存开销与标准 AdamW 完全一致,具备极高的大模型微调即插即用落地价值。

局限与展望

  • 反向梯度的计算开销:严格计算 Adafactor 统计量需要权重梯度的行和与列和,虽然可用切空间代理简化,但单步时间开销仍约为原生 AdamW 的 4 倍(0.99s vs 0.24s),后续可通过周期性异步更新预条件统计量来进一步平摊耗时。
  • 极端低秩假设:闭式解假设低秩因子 \(\bm{B}_t\) 列满秩、\(\bm{A}_t\) 行满秩,在微调初始化或奇异值崩塌时可能出现数值病态,需要依赖 \(\epsilon \bm{I}\) 正则阻尼。
  • 拓展应用场景:作者指出未来应进一步向 MoE 适配器(多 Expert 独立因子追踪)、量化 LoRA(QLoRA 反量化梯度重构)以及基于 DiT 的扩散 Transformer 模型拓展。

相关工作与启发

  • vs Vanilla LoRA [14]: 原生 LoRA 在因子空间直接运行坐标级 AdamW,完全忽视了生成映射 \(J_{\mathcal{G}}\) 的非线性与奇异性;AdaPreLoRA 在权重空间曲率度量下投影,赋予了因子更新真实的全局曲率感知。
  • vs LoRA+ [11] & Riemannian Preconditioned LoRA [37]: 后两者仅针对拉回算子做简单的标量非对称缩放或静态块对角近似,本质上取 \(\mathcal{F}_t = \bm{I}\) 并丢弃二阶梯度统计;AdaPreLoRA 融合了真实的动态 Adafactor 梯度外积统计量。
  • vs LoRA-Pro [33]: LoRA-Pro 在全权重空间计算一阶/二阶矩并求解欧氏 Frobenius 伪逆,导致显存翻倍且投影度量失配;AdaPreLoRA 在加权内积空间求解闭式投影,显存与计算开销压至低秩水平。
  • vs 黎曼优化类方法 (Riemannian Muon [5] / RAdamW [4]): 黎曼流形方法必须在正切空间更新后执行复杂的流形收缩(Retraction)或高开销正交分解;AdaPreLoRA 直接在因子坐标系中解析闭式求解,规避了繁琐的几何投影算子。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从代数相容解空间与加权度量投影的角度统一 LoRA 优化器,推导优美且理论自洽。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖从 124M/355M GPT-2、7B Mistral/Qwen2 到 Stable Diffusion 文生图,指标完备且显存开销对比详实。
  • 写作质量: ⭐⭐⭐⭐⭐ 符号系统严谨,附录数学证明清晰完备,设计动机与痛点推演层层递进。
  • 价值: ⭐⭐⭐⭐⭐ 为大模型的高性能、低显存参数高效微调提供了兼具理论保证与实用价值的全新基石优化器。