跳转至

Scalable Variational Bayesian Fine-Tuning of LLMs via Orthogonalized Low-Rank Adapters

会议: NeurIPS 2026
arXiv: 2604.03388
OpenReview: https://openreview.net/forum?id=q8fHjSdu4t
论文: https://arxiv.org/abs/2604.03388
领域: 模型压缩
关键词: 变分贝叶斯微调, 参数高效微调 (PEFT), 低秩自适应 (LoRA), 不确定性量化, 黎曼流形优化

一句话总结

针对传统贝叶斯低秩微调中 LoRA 秩崩溃损害距离感知几何特性以及推理端多次骨干前向蒙特卡洛采样开销巨大的缺陷,本文提出 PoLAR-VBLL 框架,通过 Stiefel 流形正交分解适配器与黎曼着陆场优化保持高稳定秩,并在最后一层执行闭式 Jensen 下界变分贝叶斯推断与可选 Laplace 精炼,在单次骨干前向开销下取得最优的泛化精度与不确定性标定。

研究背景与动机

在大语言模型(LLM)被加速部署于自动驾驶、医疗辅助诊断及高风险法律决策等安全关键场景的当下,模型自我评估预测置信度与不确定性量化(Uncertainty Quantification, UQ)的能力已成为保证系统可信度的基石。虽然海量预训练能够赋予基座模型初步的校准概率,但在下游特定领域通过参数高效微调(PEFT)适配有限样本后,模型极易出现严重的过度自信(Overconfidence)现象,不仅输出错误判断且赋予极高置信度,为严苛部署场景埋下了严重隐患。

为了恢复微调后 LLM 的不确定性感知能力,学术界引入了贝叶斯神经网络(BNN)体系,但既有路线均面临严峻的两难约束。以 Deep Ensemble 为代表的模型集成方案需要保存并运行多个微调模型副本,显存与计算开销过大;基于 Laplace 近似(如 Laplace-LoRA)的后验事后估计方法受限于单一模式周围的局部曲率,一旦确定性 MAP 训练收敛至不利于不确定性估计的平坦劣质区域便难以补救;而以 BLoB 为代表的变分贝叶斯 PEFT 方法虽然在训练期实现了均值与协方差的联合优化,但在推理阶段必须在全模型骨干网络(Backbone)中执行数十次完整的蒙特卡洛前向传播,推理延迟高出近一个数量级,根本无法适应低延迟在线服务。

更为深层的几何瓶颈在于,贝叶斯最后一层(BLL)等不确定性推理范式高度依赖特征提取器的「距离感知(Distance Awareness)」能力,即语义相异的输入必须在特征空间中保持充分的分离度。然而广泛采用的低秩自适应(LoRA)存在严重的有效自由度衰退——其权值增量在训练后期的稳定秩(Stable Rank)往往塌缩至理论下界 1.0,导致提取出的特征被剧烈压缩在极窄的各向异性子空间内,从根本上破坏了区分分布内与分布外(OOD)样本所依赖的几何间距。核心 idea:将施加 Stiefel 流形正交约束的极分解适配器(PoLAR)与轻量化变分贝叶斯最后一层(VBLL)深度耦合,利用高效黎曼着陆场替代高开销流形收缩算子,并推导闭式 Jensen-tightened ELBO 目标,实现单次骨干前向即可完成的鲁棒且距离感知的不确定性贝叶斯微调。

方法详解

整体框架

PoLAR-VBLL 的核心思想是将确定性特征抽取与随机参数不确定性建模解耦,同时利用结构化正交低秩分解保证特征空间的距离保持特性。给定预训练权重矩阵 \(\mathbf{W}_0\),微调权值更新建模为极分解形式 \(\mathbf{W} = \mathbf{W}_0 + \mathbf{U}\mathbf{\Lambda}\mathbf{V}^\top\),其中因子矩阵约束在 Stiefel 流形上以消除秩塌缩。输入样本经过 LLM 骨干网络抽取得到确定性特征表示 \(\boldsymbol{\phi}_{\mathbf{W}}(\mathbf{x}) \in \mathbb{R}^d\);模型的所有不确定性均被严格隔离并由随机分类层矩阵 \(\mathbf{\Theta} \in \mathbb{R}^{C \times d}\) 承担。在训练阶段,算法通过交替优化求解 PoLAR 权重的黎曼更新与 VBLL 变分参数更新;在预测阶段,骨干网络仅需单次前向,后续对极轻量的顶层线性层进行蒙特卡洛采样即可获取高标度预测分布。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入样本 x"] --> B["PoLAR 极分解正交适配器<br/>W = W0 + U Λ V^T"]
    B --> C["LLM 骨干网络单次前向<br/>提取高稳定秩特征 ϕ_W(x)"]
    C --> D["变分贝叶斯最后一层 (VBLL)<br/>分类权重 q(θ_c) ~ N(μ_c, S_c)"]
    D --> E["闭式 Jensen-tightened ELBO 目标<br/>联合约束对数似然下界与先验 KL 散度"]
    E --> F["交替优化求解<br/>着陆场更新 U, V 与变分梯度更新 μ, S"]
    F -->|训练收敛| G["可选事后 Laplace 协方差精炼<br/>Hessian 逆矩阵校准局部几何曲率"]
    G --> H["轻量化顶层蒙特卡洛预测<br/>单次骨干特征 + K 次线性分类采样"]

关键设计

1. 正交极分解适配器与流形几何保持:解决传统 LoRA 秩崩溃导致的特征空间坍缩

针对标准 LoRA 中矩阵分解 \(\Delta \mathbf{W} = \mathbf{B}\mathbf{A}\) 出现的稳定秩塌缩至 1.0 的缺陷,方法引入极分解低秩自适应(PoLAR)重参数化方案: $$ \Delta\mathbf{W} = \mathbf{U}\mathbf{\Lambda}\mathbf{V}^\top $$ 其中输出投影基 \(\mathbf{U} \in \text{St}(m,r)\) 与输入投影基 \(\mathbf{V} \in \text{St}(n,r)\) 均严格约束在 Stiefel 流形 \(\text{St}(m,r) := \{\mathbf{M} \in \mathbb{R}^{m \times r} \mid \mathbf{M}^\top\mathbf{M} = \mathbf{I}_r\}\) 上,即列向量两两正交且模长归一;核心转换矩阵 \(\mathbf{\Lambda} \in \mathbb{R}^{r \times r}\) 保持非约束自由状态以保障连续优化表达力。该参数化在代数上强制将低秩增量分散在各个正交子空间基底上,阻断了低秩矩阵优化中各主成分朝单一主导奇异方向合流的现象。论文理论与实测证明,PoLAR 能将适配增量的稳定秩从 LoRA 的约 1.53 提升至 2.86,确保了各向同性特征扩散与几何距离敏感度,为贝叶斯判决面准确识别分布偏移奠定了特征基础。

2. 黎曼着陆场(Landing Field)无收缩优化:摆脱高开销 SVD/QR 的可扩展流形训练

直接在 Stiefel 流形上进行黎曼优化需要使用收缩算子(Retraction)将切空间切向量投影回流形本身,然而标准收缩算子高度依赖奇异值分解(SVD)或 QR 分解,在 GPU 上处理高维大模型时具有极高的显存搬运与延迟开销。PoLAR-VBLL 引入了非可行流形轨迹演化方案,利用着陆场(Landing Field)代替硬性投影更新。以矩阵 \(\mathbf{U}\) 的更新为例,其更新方向被构造为反对称流形切向量梯度与非可行惩罚项的线性和: $$ \boldsymbol{\Gamma}(\mathbf{U}) = \boldsymbol{\psi}(\mathbf{U})\mathbf{U} + \lambda \nabla N(\mathbf{U}) $$ 其中 \(\boldsymbol{\psi}(\mathbf{U}) = \text{Skew}(\nabla_{\mathbf{U}}\mathcal{L} \cdot \mathbf{U}^\top) = \frac{1}{2}(\nabla_{\mathbf{U}}\mathcal{L} \cdot \mathbf{U}^\top - \mathbf{U}\nabla_{\mathbf{U}}\mathcal{L}^\top)\) 代表广义黎曼切空间梯度的反对称操作,而 \(N(\mathbf{U}) = \|\mathbf{U}^\top\mathbf{U} - \mathbf{I}_r\|_F^2\) 为正交违规偏离损失,对应的惩罚场梯度为 \(\nabla N(\mathbf{U}) = 4\mathbf{U}(\mathbf{U}^\top\mathbf{U} - \mathbf{I}_r)\)。超参数 \(\lambda > 0\) 约束违规扩张,使得迭代轨迹无需任何 SVD 即可在渐近收敛时精确落回 Stiefel 正交流形,在 GPU 并行加速中带来了 3× 至 18× 的反向传播耗时削减。

3. 闭式 Jensen 紧致化 ELBO 目标:实现免蒙特卡洛采样的快速变分推断

为了避免在训练期间对全模型进行耗时的前向重参数化采样,模型假设分类顶层权重 \(\mathbf{\Theta} = [\boldsymbol{\theta}_1, \ldots, \boldsymbol{\theta}_C]^\top\) 满足各类别独立的高斯先验 \(p(\boldsymbol{\theta}_c) = \mathcal{N}(\mathbf{0}, \sigma_\theta^2 \mathbf{I}_d)\),并采用变分后验 \(q(\mathbf{\Theta}) = \prod_{c=1}^C \mathcal{N}(\boldsymbol{\theta}_c; \boldsymbol{\mu}_c, \mathbf{S}_c)\) 进行逼近。针对 softmax 中难以直接解析积分的对数配分函数 \(-\log \sum_{c=1}^C \exp(z_{n,c})\),利用 Jensen 不等式构建可导下界: $$ \mathbb{E}{q(\mathbf{\Theta})}\left[-\log\sum)\right] $$ 利用高斯变量矩母函数展开,可得单类别的解析指数期望: $$ \mathbb{E}[\exp(z_{n,c})] = \exp\left(\boldsymbol{\mu}}^C \exp(z_{n,c})\right] \ge -\log \sum_{c=1}^C \mathbb{E}\left[\exp(z_{n,cc^\top \boldsymbol{\phi}}}(\mathbf{xn) + \frac{1}{2}\boldsymbol{\phi}}}(\mathbf{x}_n)^\top \mathbf{Sc \boldsymbol{\phi}_n)\right) $$ 配合先验与变分高斯分布之间的解析 KL 散度,将原本需要多次随机模拟的证据下界转化为严格的闭式数值可导损失函数,消除了训练期的随机采样方差,同时将全周期训练耗时控制在与确定性微调相近的量级。}}(\mathbf{x

4. 变分初始化驱动的事后 Laplace 精炼:全局趋势捕获与局部精确几何曲率的互补

通过变分训练求解出的均值 \(\hat{\boldsymbol{\mu}}_c\) 虽能精准定位高概率全局后验模态,但在一阶梯度优化下得到的对角或低秩方差 \(\hat{\mathbf{S}}_c\) 易受优化轨迹平滑影响,无法充分反映局部鞍点与最优极大值处的实际极小曲率。因此,PoLAR-VBLL 在模型收敛后设计了可选的事后 Laplace 精炼机制:以变分训练求得的高质量均值 \(\hat{\boldsymbol{\mu}}_c\) 为中心基底,计算对数后验关于分类权重的经验 Fisher 信息(精确 Hessian 矩阵): $$ \mathbf{H}c = -\nabla}_c}^2 \left( \log p(\mathcal{D} \mid \boldsymbol{\thetac, \hat{\mathbf{\Psi}}}}) + \log p(\boldsymbol{\thetac) \right) \Big\rvert $$ 通过求逆 }_c = \hat{\boldsymbol{\mu}}_c\(\mathbf{\Sigma}_c = \mathbf{H}_c^{-1}\) 重构精炼后的后验协方差矩阵 \(\tilde{q}(\boldsymbol{\theta}_c \mid \mathcal{D}) = \mathcal{N}(\boldsymbol{\theta}_c; \hat{\boldsymbol{\mu}}_c, \mathbf{\Sigma}_c)\)。与传统在确定性微调劣质解周围直接套用 Laplace 近似不同,VBLL 提供的优质初始化使 Laplace 近似无需遍历大范围搜索即可稳健捕获清晰的局部几何曲率,进一步强化了尾部分布与高频漂移数据的标定表现。

损失函数 / 训练策略

训练阶段最大化全样本解析 ELBO 目标函数 \(\mathcal{L}^{\text{ELBO}}(\mathbf{\Psi}, \mathbf{W}; \mathcal{D})\): $$ \begin{aligned} \mathcal{L}^{\text{ELBO}} = \sum_{n=1}^N \Bigg[ &\sum_{c=1}^C y_{n,c} \boldsymbol{\mu}c^\top \boldsymbol{\phi}}}(\mathbf{x}_n) - \text{LSE}_c \left( \boldsymbol{\muc^\top \boldsymbol{\phi}}}(\mathbf{xn) + \frac{1}{2}\boldsymbol{\phi}}}(\mathbf{xn)^\top \mathbf{S}c \boldsymbol{\phi}}}(\mathbf{xn) \right) \Bigg] \ &- \sum{2\sigma}^C \left[ \frac{1\theta^2}\left(\text{tr}(\mathbf{S}_c) + \boldsymbol{\mu}_c^\top\boldsymbol{\mu}_c\right) - \frac{1}{2}\log|\mathbf{S}_c| \right] + \text{const} \end{aligned} $$ 训练采用交替迭代更新:PoLAR 矩阵通过着陆场梯度下降迭代,学习率 \(\eta_{\text{polar}} = 10^{-4}\),约束违规惩罚系数 \(\lambda = 1.0\);变分均值与协方差矩阵 \(\boldsymbol{\mu}_c, \mathbf{S}_c\) 使用 AdamW 优化器进行参数更新,并将正则化先验权重设置为 \(\lambda_{\text{reg}} = 1/N\)。预训练语言模型线性头权重直接作为变分均值的初始值 \(\boldsymbol{\mu}_{c,0} = \mathbf{w}_{\text{pretrain},c}\),协方差初始赋值为各向同性阵 \(\mathbf{S}_{c,0} = \sigma_\theta^2 \mathbf{I}_d\)(默认设置 \(\sigma_\theta^2 = 1.0\))。

实验关键数据

主实验

在 LLaMA-3.1-8B 骨干网络上对 6 项常识推理分布内(ID)数据集(Winogrande-Small/Medium、ARC-Challenge/Easy、OpenBookQA、BoolQ)以及 4 项分布外(OOD)迁移任务(以 OBQA 为微调源域,迁移至 MMLU Chemistry、Physics 等复杂大跨度偏移任务)进行全方位对比测评(微调 5000 步,Rank \(r=8\))。

评测维度 任务数据集 指标 MLE 基线 Laplace-LoRA BLoB (SOTA) PoLAR-VBLL (本文) 增益 / 优势
ID 分布内 WG-S ACC (↑) / ECE (↓) 77.92% / 21.11% 77.92% / 16.41% 72.36% / 5.78% 77.26% / 3.89% 准确率反超 BLoB 4.9%,ECE 最佳
ID 分布内 ARC-C ACC (↑) / ECE (↓) 81.05% / 17.95% 81.05% / 9.72% 79.42% / 7.34% 81.79% / 4.92% 准确率最高,校准误差降低 2.42%
ID 分布内 ARC-E ACC (↑) / ECE (↓) 90.66% / 8.95% 90.66% / 4.51% 90.16% / 5.66% 91.38% / 3.71% 准确率与校准误差双指标最优
ID 分布内 OBQA ACC (↑) / ECE (↓) 88.30% / 8.33% 88.30% / 7.40% 87.53% / 5.30% 88.43% / 2.34% 准确率达到顶峰,ECE 显著压缩
OOD 小漂移 OBQA→ARC-C ACC (↑) / ECE (↓) 79.33% / 13.89% 79.33% / 5.85% 79.10% / 4.87% 81.11% / 4.55% 跨域泛化精度提升 1.78%
OOD 大漂移 OBQA→Chem ACC (↑) / ECE (↓) 48.00% / 28.73% 48.00% / 11.69% 45.67% / 13.23% 49.30% / 10.30% 大偏移下校准表现远超竞品
OOD 大漂移 OBQA→Phy ACC (↑) / ECE (↓) 43.33% / 37.02% 43.33% / 13.02% 45.67% / 11.31% 48.91% / 11.12% 准确率较 SOTA 提升 3.24%

消融实验

基于 LLaMA-2-7B 在 Winogrande-Small 数据集上对各技术组件进行剥离评估,验证 PoLAR 正交性、VBLL 变分机制以及事后 Laplace 的分立贡献。

实验配置 适配器形式 贝叶斯机制 ACC (%) ↑ ECE (%) ↓ NLL ↓ 说明
确定性基线 PoLAR 无 (MLE) 70.33 29.83 3.17 确定性微调存在严重概率过度自信
全局事后近似 PoLAR 全层 Laplace (PoLAR-LA) 70.33 12.16 0.69 事后逼近受限,校准精度有限
顶层事后近似 PoLAR 最后一层 Laplace (LA-LL) 70.33 14.63 0.71 劣质模式下仅修顶层导致校准劣化
全层采样变分 PoLAR 核心阵变分 (PoLAR-BLoB) 70.39 12.06 0.73 显存与采样开销大且精度偏低
变分基础版 PoLAR VBLL (无 Laplace 精炼) 71.62 8.26 0.66 变分训练本身是校准的核心驱动力
完整框架 (Full) PoLAR VBLL + 事后 Laplace 精炼 71.62 7.31 0.60 全局模式引导 + 局部曲率精细刻画

关键发现

  • VBLL 是不确定性校准的核心驱动引擎:对比 PoLAR-VBLL (w/o LA) 与 PoLAR-LA-LL,在二者均仅针对最后一层进行概率建模的前提下,纯变分驱动的 VBLL 将 ECE 从 14.63% 大幅压低至 8.26%,表明动态变分优化能够主动引导模型收敛至利于不确定性刻画的高质量后验模式,而事后 Laplace 仅起到边际增益(降至 7.31%)的修饰作用。
  • 破解精度与不确定性的权衡悖论:传统 BLoB 类方法存在明显的性能妥协(随着推理采样次数从 0 增加至 10,校准指标改善但准确率出现 1-2% 的不可逆下降);而 PoLAR-VBLL 依托正交高保真特征表示,在提升校准能力(ECE 降至 2-4%)的同时,所有任务的绝对判决准确率反超确定性 MLE 与 BLoB。
  • 推理延迟与显存吞吐量形成代际差距:在 ARC-E 推理基准测试中,BLoB 与 TFB 由于需要通过 LLM 骨干网执行数十次蒙特卡洛前向传播,耗时高达 80s 至 90s;而 PoLAR-VBLL 仅需单次骨干前向,测试总耗时缩减至 12s(提速近 7×),同时显存消耗仅 18.4GB,较全网 Laplace(40.7GB)节省超过 54% 显存。

亮点与洞察

  • 正交分解破解低秩几何退化:巧妙指出标准 LoRA 诱导的有效秩塌缩(稳定秩低至 1.0)破坏了特征空间的距离敏感度,首次将具有保角与正交约束的 PoLAR 引入贝叶斯微调,建立了特征空间距离保持度与贝叶斯 OOD 探测能力之间的直接因果链路。
  • 黎曼优化与变分训练的高效无缝融合:采用着陆场方法回避了高维 Stiefel 流形更新中的密集 SVD/QR 分解,同时推导了基于 Jensen 紧致不等式的解析 ELBO 闭式解,使得完整的贝叶斯微调能够以几乎等同于常规确定性微调的计算速度平滑收敛。
  • 可复用的「先验模式搜索 + 事后曲率对准」范式:将变分训练视作优质模式发现器、将 Laplace 近似视作局部高阶几何显微镜的分步设计,可无缝推广至视觉 Transformer、多模态判决头等各类需要低开销可靠置信度的深度推理体系。

局限与展望

  • 分类任务与生成任务的场景适用边界:当前基于 BLL 的变分推断主要聚焦于判别式离散选择与多分类常识推理,尚未直接覆盖自回归序列级自由文本生成过程中的逐 token 贝叶斯解码与动态认知不确定性建模。
  • 各类别独立后验假设:变分后验假设各类别权重相互解耦独立,忽略了多类复杂标签体系下类别之间的交叉先验关联,未来可探索 Kronecker 分解结构以进一步捕获类别间协方差。
  • 大秩扩展下的显存权衡:尽管着陆场优化极大地降低了正交计算复杂度,但当适配秩 \(r\) 显著增大时,核心矩阵 \(\mathbf{\Lambda}\) 与分类协方差 \(\mathbf{S}_c\) 的存储开销仍呈二次方增长,需结合低秩协方差分解进一步压缩。

相关工作与启发

  • vs BLoB (Wang et al., 2024):BLoB 在全网 LoRA 参数上施加变分推断,推理时需对整个 LLM 骨干网络执行多次蒙特卡洛采样导致延迟暴增,且受制于 LoRA 内部秩崩溃;本文通过 PoLAR 恢复特征多样性,并将不确定性集中于顶层,实现单次骨干前向的高速低延迟推理。
  • vs Laplace-LoRA (Yang et al., 2024):Laplace-LoRA 依赖确定性训练终点,易陷入不利于不确定性度量的次优解模式;本文由 VBLL 主动引导参数收敛至高品质变分后验模式,仅将 Laplace 作为局部曲率的精密微调工具。
  • vs C-LoRA (Rahmati et al., 2025) & TFB (Shi et al., 2024):相较于额外训练上下文 MLP 或对固定权重进行启发式事后网格搜索,本文给出了严谨的解析变分下界优化方案,具备更完备的贝叶斯理论支撑与更稳定的跨域迁移鲁棒性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [首次将极分解正交低秩适配器与变分贝叶斯最后一层结合,并引入无收缩黎曼着陆场优化]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖两代 LLaMA 骨干模型、10 项分布内/外评测基准,并完备给出了稳定秩、Jensen 紧致性与显存时延消融]
  • 写作质量: ⭐⭐⭐⭐⭐ [理论推导严谨自洽,从几何距离感知到变分解析解的逻辑链条紧凑清晰]
  • 价值: ⭐⭐⭐⭐⭐ [切实解决了大模型贝叶斯微调长期面临的「高计算延迟」与「过度自信」兼顾难题,极具落地部署实用性]