Scalable Bayesian Low-Rank Adaptation of Large Language Models via Stochastic Variational Subspace Inference
通过随机变分子空间推理实现大型语言模型的可扩展贝叶斯低秩自适应
https://www.arxiv.org/pdf/2506.21408
![]()
![]()
摘要
尽管大语言模型(LLMs)被广泛使用,但它们常会生成错误信息(即“幻觉”),且置信度校准性能较差。这使得对这些模型进行不确定性量化变得至关重要,尤其是在自动驾驶、医疗等高风险领域。先前的研究通过在微调模型的低秩适应(LoRA)参数上进行贝叶斯深度学习推理,使该问题更具可行性。尽管这些方法有效,但由于相比标准LoRA还需引入更多额外参数,因此在扩展到更大的LLM时面临困难。在本研究中,我们提出了 基于随机变分子空间推理的可扩展贝叶斯低秩适应方法 (ScalaBL)。我们在一个r维子空间中进行贝叶斯推理(r为LoRA的秩)。通过将LoRA参数重新用作投影矩阵,我们能够将该子空间中的样本映射到LLM的完整权重空间中。这使得我们可以使用随机变分推理来学习本方法的所有参数。尽管子空间维度较低,我们的方法仍能与当前最先进的方法达到相当的性能,且仅需约1000个额外参数。此外,该方法使我们能够扩展到迄今为止最大的贝叶斯LLM,其基础模型参数数量是此前工作的四倍。
1 引言
大语言模型(LLMs)的使用已遍及众多领域,包括医疗健康 [Clusmann 等, 2023]、科学发现 Zhang 等 [2024]、信息物理系统 [Cobb 等, 2023]、代码生成 [Jiang 等, 2024] 以及日常通用场景 [Anil 等, 2023]。因此,确保这些模型的可靠性与可信度比以往任何时候都更加重要。然而,众所周知,LLMs 会以“幻觉”形式输出错误信息 [Huang 等, 2024],且其置信度往往校准不佳 [Zhu 等, 2023, Spiess 等, 2024]。
为解决这些问题,一个研究方向聚焦于量化 LLM 输出的不确定性。针对这一任务,已有多种事后(post-hoc)方法被提出,例如口头化置信度 [Tian 等, 2023, Xiong 等, 2023]、量化词元级别不确定性 [Kuhn 等, 2023, Farquhar 等, 2024],或使用共形预测(conformal prediction)[Kaur 等, 2024]。
相比之下,贝叶斯深度学习(BDL)为深度模型的不确定性量化提供了一种原理性更强的方法。在这类方法中,通过直接对模型权重上的分布进行推理来实现不确定性量化 [Gal 和 Ghahramani, 2016, Blundell 等, 2015, Lakshminarayanan 等, 2017]。在此框架下,我们通过贝叶斯法则,对参数后验分布(记为 P(W∣D))进行积分边缘化,来估计模型对测试样本 x的预测不确定性,即 P(y∣x,D),其形式如下:
![]()
其中,D是训练(或微调)数据集,W是模型参数。然而,在将这类方法扩展到大语言模型(LLMs)时,由于 W的维度极高,对这一不可行积分进行良好近似变得愈发困难。因此,近期的研究考虑在一些流行的参数高效微调(PEFT)方法所学习的较小参数子集上进行贝叶斯推理 [Fu 等, 2023]。
在 Hu 等人 [2022] 提出的广泛使用的低秩适应(LoRA)技术中,仅有少量参数被更新,相比更新全部参数可节省大量资源,同时仍能保留基础模型的大部分性能。方便的是,这些参数的低维度也使其非常适用于贝叶斯深度学习(BDL)方法。然而,Yang 等人 [2024a] 和 Wang 等人 [2024] 指出,直接在 LoRA 参数上应用诸如深度集成(Deep Ensembles)[Lakshminarayanan 等, 2017] 或蒙特卡洛 Dropout [Gal 和 Ghahramani, 2016] 等 BDL 方法,相较于最大似然估计(MLE)或最大后验估计(MAP)等简单的微调方法,在不确定性量化指标上的提升十分有限。
该领域首个成功尝试来自 Yang 等人 [2024a],他们在 MAP 微调后对参数后验分布进行拉普拉斯近似。而 Wang 等人 [2024] 提出的当前最优方法则采用随机变分推理,他们称之为“通过反向传播实现贝叶斯 LoRA”(BLoB)。尽管该方法性能优于此前所有方法,但其所需参数数量比标准 LoRA 多出约 40%,这在高风险、资源受限的部署场景中可能成为严重的内存瓶颈,因为在这些场景中,计算贝叶斯模型平均本身就已经接近可用内存的极限 [Vadera 等, 2022]。
在本研究中,我们提出了基于随机变分子空间推理的可扩展贝叶斯低秩适应方法(ScalaBL)。如图 1 所示,我们在完整权重空间 W的一个更小的子空间内进行贝叶斯推理,该子空间的维度等于 LoRA 的秩 r。我们展示了如何将 LoRA 参数 A和 B重新用作投影矩阵,将来自低维子空间的样本映射到完整的权重空间 W中。随后,我们使用随机变分推理来学习本方法的所有参数。
![]()
我们方法的一个主要优势是:对于每个LoRA层,仅需学习 2r 个额外的变分参数,而BLoB方法则需要 rd 个(其中 d是大语言模型的嵌入维度)。例如,在对一个拥有70亿参数、嵌入维度 d=3584的大模型进行微调时,若LoRA秩 r=8,BLoB需要数百万个额外参数,而ScalaBL仅需约1000个。此外,只要秩 r保持不变,无论基础LLM的嵌入维度如何,我们每层所需的额外参数数量都保持恒定。因此,我们能够将该方法扩展到一个拥有320亿基础参数、d=5120的模型,相比之下,Yang等人 [2024a] 和 Wang等人 [2024] 的先前工作仅适用于70亿参数的模型。
通过大量实验,我们表明ScalaBL在分布内和分布外的多种常识推理基准任务上,其性能与当前最先进的基线方法相当甚至更优。我们的主要贡献总结如下:
我们提出了ScalaBL,一种在低维子空间内进行随机变分推理的贝叶斯LoRA方法;
与先前工作相比,ScalaBL具有显著的参数效率,额外参数减少约2000倍;
ScalaBL在不确定性量化指标上的性能达到或优于当前最优方法,同时使用更少的参数;
我们的工作首次将贝叶斯LoRA方法扩展到拥有320亿基础参数的预训练模型,远超此前70亿参数的规模。
本文结构如下:第2节介绍与本工作相关的主要前期研究;第3.1节阐述如何构建参数高效的子空间;第3.2节介绍如何在该子空间中使用随机变分推理训练概率模型;第4节展示广泛的实验结果;第5和第6节讨论局限性并总结全文。更多细节和实验结果见附录。我们的代码已开源:github.com/SRI-CSL/BayesAdapt。
2 相关工作
本节介绍我们方法所基于的前期研究。
2.1 低秩适应(Low-Rank Adaptation)
![]()
![]()
2.2 拉普拉斯LoRA
Yang 等人 [2024a] 的工作是将不确定性量化技术应用于 LoRA 层的首个研究实例,其方法是对低秩参数应用拉普拉斯近似。他们将微调后得到的最大后验(MAP)估计视为一个多维高斯分布的均值,并通过逆Hessian矩阵来推导其协方差。然而,即使将拉普拉斯近似限制在 LoRA 参数上,计算Hessian矩阵在实际中仍不可行。因此,Yang 等人 [2024a] 采用克罗内克分解(Kronecker factorization)[Ritter 等, 2018, Daxberger 等, 2021] 对Hessian矩阵施加结构化假设。这些克罗内克因子仍然占用大量内存,因此 Yang 等人 [2024a] 不得不进一步通过迭代截断奇异值分解的方法进行近似。该拉普拉斯近似是在 LoRA 参数微调完成后进行的,属于事后(post-hoc)处理。另一个局限是,在测试时,他们需要反向传播整个模型以构建近似的协方差矩阵,这限制了该方法在资源受限环境下的可扩展性和实际应用。
2.3 BLoB
目前该领域的最先进方法是通过反向传播实现贝叶斯低秩适应(Bayesian Low-Rank Adaptation by Backpropagation, BLoB)。BLoB摒弃了拉普拉斯LoRA的两阶段方法,转而对LoRA参数 A直接进行随机变分推理。更具体地说,他们采用了Blundell等人 [2015] 提出的“贝叶斯反向传播”(Bayes by Backprop)方法。即将 A重新视为一个低秩高斯分布的均值,记为 Aμ,并学习一组方差参数 Aσ。通过重参数化技巧(reparameterization trick)[Kingma 和 Welling, 2013],他们从该低秩分布中采样,并将样本投影到完整的权重空间:
![]()
其中,εt∼N(0,1)。Wang 等人 [2024] 通过实验表明,他们的方法在性能上优于拉普拉斯LoRA。然而,拉普拉斯近似的一个显著优势是无需学习任何额外参数。而由于引入了方差参数 Aσ,BLoB 所需学习的参数数量比基础的LoRA微调过程多了约40%。即使对于较小的70亿参数模型,这也意味着要额外增加数百万个参数。
2.4 贝叶斯子空间推理
![]()
![]()
正如 Izmailov 等人 [2020] 所强调的,该模型并不是原始参数后验分布的一种重新参数化,因为向完整参数空间的投影是不可逆的。然而,这种方法的优势在于,能够在子空间中执行贝叶斯推理,从而使得许多原本难以实现的常见贝叶斯推理技术变得可行。据我们所知,此类子空间推理技术此前从未被应用于大语言模型(LLMs)。
3 方法
在本节中,我们详细介绍所提出的方法。我们首先说明如何构建完整权重空间 W中的一个 r维子空间,然后讨论如何在该子空间中使用随机变分推理来训练一个概率模型。
3.1 子空间构建
![]()
![]()
![]()
![]()
我们注意到,BU是两个线性参数矩阵的乘积,其表示能力与单独使用 B相同。此外,由于 V的维度与 A相同,我们直接将 V重命名为 A。由此得到如下子空间:
![]()
其中 f是一个投影函数,为记法方便而定义。直观上,我们将 LoRA 参数重新用作投影矩阵,用于构建一个位于 A和 B之间的 r维子空间。
3.2 变分子空间推理
接下来,我们在该子空间中构建一个概率模型,其数据似然函数为:
![]()
![]()
![]()
4 实验
在本节中,我们对 ScalaBL、若干标准基线方法以及当前最先进的方法进行了实验比较。
4.1 数据集
我们遵循 Yang 等人 [2024a] 和 Wang 等人 [2024] 的实验协议,使用一组常识推理数据集(见表1)对我们的方法进行微调和评估。这些数据集以多项选择题的形式呈现。给定一个包含问题的输入提示,我们获取大语言模型(LLM)在下一个词元上的 softmax 分布,然后提取每个可能答案(例如 A、B、C、D)对应的 logits,并重新归一化。通过这种方式,我们将这些常识推理任务转化为分类任务,从而便于计算标准的不确定性指标。具体而言,我们报告分类准确率(ACC)、期望校准误差(ECE)[Guo 等, 2017] 以及正确类别的负对数似然(NLL)。有关这些指标的更多细节见附录第7.1节。
![]()
4.2 基线方法
我们与一系列标准基线方法进行了比较。首先,我们考虑标准的LoRA训练流程,包括不带权重衰减(标记为 MLE)和带权重衰减正则化(标记为 MAP)两种情况。接着,我们与两种经典的贝叶斯深度学习基线进行比较:深度集成(Deep Ensembles)[Lakshminarayanan 等, 2017] 和蒙特卡洛Dropout(Monte Carlo Dropout)[Gal 和 Ghahramani, 2016]。最后,我们还与两种最新的最先进方法进行了对比:Yang 等人 [2024a] 提出的拉普拉斯近似方法,以及 BLoB [Wang 等, 2024]。
4.3 实现细节
我们基于 Wang 等人 [2024] 开发的 bayesian-peft 库构建我们的方法。该库提供了标准基线方法以及 BLoB 的实现。对于拉普拉斯近似方法,我们使用了 Yang 等人 [2024a] 提供的官方代码。与 Yang 等人 [2024a] 和 Wang 等人 [2024] 不同,我们使用更新的 Qwen2.5 模型系列 [Yang 等, 2024b],而非此前工作中使用的较早的 Llama-2-7b 模型 [Touvron 等, 2023]。为了便于比较,使用 Llama-2-7b 的结果见附录第8.5节。
我们遵循 Yang 等人 [2024a] 和 Wang 等人 [2024] 的做法,将 LoRA 应用于每个自注意力层的查询(query)和值(value)参数,以及大语言模型的 softmax 输出头,LoRA 秩设为 r=8。我们采用与 BLoB 相同的训练流程和超参数设置。所有方法均使用 AdamW 优化器训练 5000 步。对于 70亿参数模型,训练批大小为 4;对于 320亿参数模型,批大小为 2。与 Wang 等人 [2024] 不同,我们对冻结的模型参数使用 16 位精度训练,而非 8 位量化,而可学习参数仍保持 32 位精度。所有实验均在单块 80GB 的 NVIDIA A100 GPU 上完成。
对于 ScalaBL,我们使用与 BLoB 相同的 KL 散度加权调度策略,其最大值设为 β=0.1。我们未采用 BLoB 中使用的 Flipout 技术 [Wen 等, 2018],因为我们发现其对性能无明显影响。这使得我们的方法实现比 BLoB 更加简洁。与 BLoB 一样,我们使用标准的 N(0,Ir)作为先验 P(s)。我们通过对一个随机初始化的矩阵进行奇异值分解(SVD)来初始化 sμ和 A,由于 LoRA 矩阵的低秩特性,该操作非常快速。类似地,sσ的方差参数也像在 BLoB 中一样,初始化为较小的均匀随机值。我们对 sσ使用对数参数化,以确保方差始终为正。基于 sμ类似于 A的奇异值的直觉,我们也对其使用对数参数化以保证其正值。
对于变分推理方法(BLoB 和 ScalaBL),我们在评估时使用 N=10个后验权重样本,这是 Wang 等人 [2024] 发现性能最佳的设置。该超参数的影响在附录第8.1节中进一步探讨。同样地,MC-Dropout 基线也进行 10 次前向传播。对于深度集成(Deep Ensembles),我们使用集成大小为 3。
4.4 分布内结果
在表2中,我们展示了使用 Qwen2.5-7B 大语言模型在标准分布内设置下的测试集结果。我们首先注意到,简单的最大似然估计(MLE)微调方法在所有数据集上均实现了较高的准确率,但通常会出现过拟合,表现为校准效果较差(ECE值高)。MAP 方法等价于在 MLE 基础上加入 10−2的权重衰减正则项,仅略微改善了最终的校准性能。
![]()
当转向蒙特卡洛 Dropout 和深度集成时,ECE 和 NLL 有轻微改善,其中深度集成在标准基线方法中表现最佳,但代价是显著更高的资源消耗。
验证了 Yang 等人 [2024a] 和 Wang 等人 [2024] 的结论,我们发现最近的最先进方法在 ECE 和 NLL 指标上明显优于基线方法,且分类准确率几乎没有下降。此外,我们观察到 ScalaBL 的性能始终与 BLoB 相当,甚至在 Winogrande-Medium 数据集上以 ECE 指标达到了当前最优性能。
不出所料,BLoB 在所有方法中通常表现最好,且经常以微弱优势优于 ScalaBL。然而,由于其参数数量更多,BLoB 的表示能力严格强于 ScalaBL 或拉普拉斯方法。与 MLE 相比,BLoB 需要额外约 1.4 倍的参数,而 ScalaBL 仅需约 1.0001 倍。对于本实验所用的大模型和秩设置,BLoB 比 MLE 多出约 160 万个参数,而 ScalaBL 仅增加 912 个参数。在此背景下,ScalaBL 以远低于 BLoB 的成本实现了极具竞争力的性能。例如,在 ARC-Challenge 数据集上,BLoB 的 ECE 性能比 ScalaBL 提高约 1.3 倍,准确率相近,但 BLoB 所需的额外参数数量是 ScalaBL 的 1792 倍。
4.5 分布外结果
接下来,我们进行一项分布外实验:模型在 OpenBookQA(OBQA)数据集上进行训练,该数据集包含小学水平的多项选择科学问题。首先,我们在 ARC 数据集上评估该微调后的模型,ARC 同样为小学水平的多选题,代表较小的分布偏移。接着,我们进一步研究更大的分布偏移,通过在更具挑战性的 MMLU-Chemistry 和 MMLU-Physics 数据集上进行评估,这两个数据集分别包含大学本科水平的化学和物理多选题。所有方法在该实验中的结果如表3所示。
![]()
我们再次观察到,最新的最先进方法在不确定性量化方面优于标准基线方法,同时保持了相当的准确率。我们注意到,所有方法在面对较大的分布偏移时,校准性能均有所下降。此外,我们特别指出拉普拉斯方法在 MMLU 数据集上的准确率表现较差。我们提出的方法表现出强劲性能:在多个数据集上,ScalaBL 在 ECE 指标上优于 BLoB 和拉普拉斯方法。在小和大的分布偏移下,ScalaBL 在所有指标上均能实现与 BLoB 相当的性能。
4.6 扩展到更大模型
Yang 等人 [2024a] 和 Wang 等人 [2024] 的先前工作存在一个局限性:他们仅使用了相对较小的、仅有70亿参数的大语言模型。这使得其结论是否适用于当前广泛使用的更大规模模型尚不明确 [Anil 等, 2023]。为此,我们将我们的方法扩展到迄今为止最大的贝叶斯大语言模型——Qwen2.5-32B,其基础参数数量是此前工作的四倍。我们进行了与之前相同的分布内实验,测试集结果如表4所示。需要说明的是,我们未报告拉普拉斯基线的结果,因为其事后(post-hoc)计算过程即使在使用8位参数且测试时批大小为1的情况下,仍超出了我们80GB A100 GPU的内存容量,这凸显了该方法在可扩展性方面的严重不足。
![]()
与之前在小模型上的结果不同,当使用更大的基础模型时,标准基线方法的表现显著提升。我们发现,即使是 MLE 或 MAP 这样简单的方法,所得到的模型校准效果也远优于其在小模型上的对应版本。这一现象在先前的研究中已有观察 [Xiong 等, 2023, Spiess 等, 2024]。此外,我们注意到,深度集成(Deep Ensembles)在所有三个指标上通常表现最佳,但其代价是显著更高的资源消耗。
我们观察到,所提出的 ScalaBL 方法在与各基线方法(包括 BLoB)的对比中仍表现出强劲的竞争力。在 ECE 和 NLL 指标上,它通常排名第二,分类准确率与 BLoB 相当。当使用更大的基础模型时,我们方法的效率和可扩展性优势更加突出。从 Qwen2.5-7B 升级到 Qwen2.5-32B,模型的嵌入维度从 3584 增加到 5120,并额外增加了 12 层。由于 BLoB 的方差参数数量随嵌入维度增长,其相比 MLE 额外增加了约 520 万个参数。相比之下,ScalaBL 的额外参数数量仅与秩 r相关,而 r在此大模型上并未改变,因此 ScalaBL 仅需额外增加 2064 个参数。事实上,在此 LLM 和秩的设置下,BLoB 所需的额外参数数量是 ScalaBL 的 2560 倍。因此,我们认为 ScalaBL 是目前唯一有望扩展到当前前沿大模型(基础参数已超万亿 [Anil 等, 2023])的方法。
5 局限性
尽管 ScalaBL 具有参数效率优势,但对投影后的权重样本进行贝叶斯模型平均的计算在运行时间上与 BLoB 相同。本工作与 Yang 等人 [2024a] 和 Wang 等人 [2024] 共同的一个局限是:我们仅在多项选择分类数据集上进行了评估。这凸显了未来研究中对开放式生成任务进行不确定性量化的迫切需求。
6 结论
在本研究中,我们提出了基于随机变分子空间推理的可扩展贝叶斯低秩适应方法(ScalaBL)。我们在一个 r维子空间上进行贝叶斯推理,并将 LoRA 的 A和 B参数重新用作投影矩阵,将该低维子空间中的样本映射到大语言模型(LLM)的完整权重空间 W中。我们展示了如何通过随机变分推理学习该方法的所有参数。由于子空间维度极低,我们的方法在参数效率上远超先前工作,同时在多种常识推理基准任务上仍能实现与当前最先进方法相当的性能。正因如此,我们的工作首次将贝叶斯 LoRA 方法成功扩展到 320 亿参数的模型,且所需的额外参数比此前工作少几个数量级。
![]()
![]()
![]()
8.2 子空间的选择
在本节中,我们考虑了方法中所使用子空间的不同选择。在表7中,我们展示了使用公式(8)定义的SVD子空间的实验结果。此外,我们还包含了一项实验结果:在微调过程中冻结 A矩阵,这类似于 Izmailov 等人 [2020] 提出的随机子空间方法。
我们首先注意到,SVD子空间与ScalaBL所用子空间之间的性能差异可以忽略不计。这并不令人意外,正如论文正文所述,额外引入的参数矩阵 U并不会改变模型的表示能力。使用随机子空间的一个有趣优势是,它进一步减少了需要学习的参数数量。我们发现,在某些数据集上,随机子空间的性能与参数更多的子空间相当。然而,在某些数据集上(例如 Winogrande-Medium),使用随机子空间会导致分类准确率显著下降。
![]()
![]()
![]()
![]()
![]()
![]()
原文链接:https://www.arxiv.org/pdf/2506.21408
特别声明:以上内容(如有图片或视频亦包括在内)为自媒体平台“网易号”用户上传并发布,本平台仅提供信息存储服务。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.