用于可扩展贝叶斯推断的压缩活跃子空间

# 面向可扩展贝叶斯推断的压缩主动子空间
来源：https://arxiv.org/html/2609.19539
###### 摘要

主动子空间方法通过识别并沿着对模型输出影响最大的参数方向进行推断，为量化高维模型中的预测不确定性提供了一种框架。然而，构建主动子空间需要存储大量全维度模型梯度，这随着模型规模的增长变得不可行。我们通过提出**压缩主动子空间（CAS）** 来解决这一限制，这是一种可扩展的方法，首先使用结构化等距嵌入将模型参数映射到压缩空间，然后在该约化参数化中构建主动子空间。我们的方法显著减少了构建主动子空间所需的内存，并使得对标准主动子空间方法难以处理的大型模型进行贝叶斯推断成为可能。我们在规模递增的神经网络上演示了CAS的可扩展性，同时保持了预测性能和稳健的不确定性估计。

###### 索引术语：

主动子空间、贝叶斯深度学习、随机投影、不确定性量化（UQ）

††地址：⋆计算与数据科学部门，布鲁克海文国家实验室，美国纽约州尤普顿
†电气与计算机工程系，德克萨斯A&M大学，美国德克萨斯州大学城
‡数学与计算机科学部，阿贡国家实验室，美国伊利诺伊州莱蒙特

## 1 引言

深度学习模型的贝叶斯推断为量化预测不确定性提供了一个有原则的框架。然而，其计算成本随着参数空间维度的增加而迅速增长。主动子空间通过识别由对模型输出影响最大的参数方向所张成的低维线性子空间，提供了一种高效应对这一挑战的方法\[1 (https://arxiv.org/html/2609.19539#bib.bib1)\]。最近的研究表明，限制在此类主动子空间上的贝叶斯推断可以提供与全网络贝叶斯推断相当的预测不确定性估计，同时保持计算效率\[1 (https://arxiv.org/html/2609.19539#bib.bib1)\]。利用主动子空间已被证明能够实现对相对较大的深度学习模型的不确定性量化（UQ）\[2 (https://arxiv.org/html/2609.19539#bib.bib2)\]和不确定性引导的微调\[3 (https://arxiv.org/html/2609.19539#bib.bib3)\]。

主动子空间方法的一个关键步骤是识别模型梯度协方差矩阵 \(\mathcal{C}\) 的主导特征向量：
\[
\mathcal{C}=\mathbb{E}\left[\nabla_{\boldsymbol{\theta}}f_{\boldsymbol{\theta}}(\boldsymbol{x})\nabla_{\boldsymbol{\theta}}f_{\boldsymbol{\theta}}(\boldsymbol{x})^{T}\right].\tag{1}
\]
其中，期望是在模型参数 \(\boldsymbol{\theta}\) 和输入 \(\boldsymbol{x}\) 上进行的。遵循 \[1 (https://arxiv.org/html/2609.19539#bib.bib1)\]，我们考虑围绕预训练模型参数 \(\boldsymbol{\theta}_{0}\) 的参数扰动 \(\boldsymbol{\theta}\sim\mathcal{N}\left(\boldsymbol{\theta}_{0},\sigma^{2}_{0}\mathbf{I}\right)\)，以及从训练数据分布中采样的 \(\boldsymbol{x}\sim p\)。\(\mathcal{C}\) 中与最大特征值对应的特征向量标识了模型输出对这些扰动最敏感的参数方向。

在实践中，\(\mathcal{C}\) 的主导特征向量可以通过对采样的模型梯度进行截断奇异值分解（SVD）来近似。令 \(\theta_{1},\ldots,\theta_{M}\) 表示来自 \(\mathcal{N}\left(\boldsymbol{\theta}_{0},\sigma^{2}_{0}\mathbf{I}\right)\) 的独立同分布样本，\(\boldsymbol{x}_{1},\ldots,\boldsymbol{x}_{M}\) 表示来自训练数据分布的样本。我们可以构建梯度矩阵：
\[
\mathcal{G}=[\boldsymbol{g}_{1},\ldots,\boldsymbol{g}_{M}]\in\mathbb{R}^{N\times M},\quad\boldsymbol{g}_{i}=\nabla_{\boldsymbol{\theta}}f_{\boldsymbol{\theta}_{i}}(\boldsymbol{x}_{i}),\tag{2}
\]
其中 \(N\) 是模型参数数量。\(\mathcal{G}\) 的前 \(K\) 个左奇异向量提供了 \(\mathcal{C}\) 的主导特征向量的近似。虽然这种方法在神经网络上已显示出有希望的结果 \[1 (https://arxiv.org/html/2609.19539#bib.bib1),2 (https://arxiv.org/html/2609.19539#bib.bib2),3 (https://arxiv.org/html/2609.19539#bib.bib3)\]，但其应用于大型模型受到存储 \(\mathcal{G}\) 所需内存的限制。特别是，每个额外的模型梯度样本都需要存储另一个 \(N\) 维向量，导致内存消耗为 \(O(NM)\)。

我们通过首先在构建主动子空间之前压缩模型参数空间来解决这一限制。我们提出**压缩主动子空间（CAS）**，它使用一种高效的等距嵌入引入模型的固定低维参数化。然后在该压缩参数空间内进行主动子空间构建，显著减少了存储采样模型梯度所需的内存。生成的主动子空间可以映射回原始模型参数空间，以使用 \[1 (https://arxiv.org/html/2609.19539#bib.bib1)\] 的贝叶斯推断框架量化预测不确定性。

我们在这项工作中的主要贡献如下：
- 我们提出了压缩主动子空间，用于在高维参数空间模型中进行可扩展的主动子空间构建。CAS使用结构化等距嵌入，利用压缩模型梯度构建主动子空间。
- 我们在规模递增的神经网络上演示了CAS的可扩展性，并使用基于主动子空间的贝叶斯推断评估其预测不确定性估计。

**相关工作**。在低维子空间上进行贝叶斯推断提供了一种高效的替代方案，用以替代在整个神经网络参数空间上进行推断。先前的方法使用SGD轨迹或神经网络损失景观中的学习区域构建低维参数子空间 \[4 (https://arxiv.org/html/2609.19539#bib.bib4),5 (https://arxiv.org/html/2609.19539#bib.bib5),6 (https://arxiv.org/html/2609.19539#bib.bib6)\]。最近，子空间推断与参数高效表示相结合，使得在大型语言模型中进行贝叶斯推断成为可能 \[7 (https://arxiv.org/html/2609.19539#bib.bib7)\]。

主动子空间起源于计算机模型文献，其中梯度信息用于识别输入空间中捕获模型输出最大变异性的方向 \[8 (https://arxiv.org/html/2609.19539#bib.bib8),9 (https://arxiv.org/html/2609.19539#bib.bib9)\]。最近的研究将此方法扩展到神经网络参数空间，用于深度学习和生成模型中的不确定性量化 \[1 (https://arxiv.org/html/2609.19539#bib.bib1),2 (https://arxiv.org/html/2609.19539#bib.bib2),3 (https://arxiv.org/html/2609.19539#bib.bib3)\]。其他方法通过稀疏子网络选择提高了贝叶斯神经网络的可扩展性 \[10 (https://arxiv.org/html/2609.19539#bib.bib10),11 (https://arxiv.org/html/2609.19539#bib.bib11),12 (https://arxiv.org/html/2609.19539#bib.bib12)\]。

随机投影已用于数值线性代数和机器学习中的高效计算。在随机线性代数中，它们构成了低秩近似、随机SVD和预处理方法的基础 \[13 (https://arxiv.org/html/2609.19539#bib.bib13),14 (https://arxiv.org/html/2609.19539#bib.bib14),15 (https://arxiv.org/html/2609.19539#bib.bib15)\]。随机压缩也已用于分布式机器学习中，以减少与高维模型梯度相关的通信和存储成本 \[16 (https://arxiv.org/html/2609.19539#bib.bib16),17 (https://arxiv.org/html/2609.19539#bib.bib17),18 (https://arxiv.org/html/2609.19539#bib.bib18),19 (https://arxiv.org/html/2609.19539#bib.bib19)\]。

低秩参数化广泛用于大型神经网络的参数高效微调。低秩适应（LoRA）\[20 (https://arxiv.org/html/2609.19539#bib.bib20)\] 将模型更新限制在低秩参数化中，而最近的方法通过结构化投影和参数共享进一步减少了可训练参数 \[21 (https://arxiv.org/html/2609.19539#bib.bib21)\]。在LoRA参数空间中也研究了贝叶斯推断和不确定性量化 \[22 (https://arxiv.org/html/2609.19539#bib.bib22),23 (https://arxiv.org/html/2609.19539#bib.bib23),24 (https://arxiv.org/html/2609.19539#bib.bib24)\]。CAS中使用的投影与Uni-LoRA \[21 (https://arxiv.org/html/2609.19539#bib.bib21)\] 中提出的等距构造以及CountSketch矩阵 \[25 (https://arxiv.org/html/2609.19539#bib.bib25),26 (https://arxiv.org/html/2609.19539#bib.bib26)\] 密切相关。与Uni-LoRA构造相比，我们在非零条目中包含了随机符号，并使用所得的嵌入来支持*压缩主动子空间*构建。

## 2 预备知识

**贝叶斯子空间推断**。令 \(\mathcal{D}=\{(\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n}\) 表示训练数据，\(\boldsymbol{\theta}_{0}\in\mathbb{R}^{N}\) 表示预训练的神经网络模型参数。对完整模型进行贝叶斯推断需要推断后验分布 \(p(\boldsymbol{\theta}|\mathcal{D})\)，这随着 \(N\) 的增长而变得昂贵。子空间推断则将参数变化限制在围绕 \(\boldsymbol{\theta}_{0}\) 的 \(K\) 维线性子空间内 \[1 (https://arxiv.org/html/2609.19539#bib.bib1),4 (https://arxiv.org/html/2609.19539#bib.bib4)\]：
\[
\boldsymbol{\theta}=\boldsymbol{\theta}_{0}+\boldsymbol{V}\boldsymbol{z},\qquad\boldsymbol{z}\in\mathbb{R}^{K},\tag{3}
\]
其中 \(\boldsymbol{V}\in\mathbb{R}^{N\times K}\) 包含子空间的基向量。然后，对约化参数 \(\boldsymbol{z}\) 而不是完整模型参数 \(\boldsymbol{\theta}\) 进行贝叶斯推断。可以使用高斯先验 \(p(\boldsymbol{z})=\mathcal{N}(\boldsymbol{\mu}_{\boldsymbol{z}},\tilde{\sigma}^{2}\mathbf{I}_{K})\)，并通过使用来自 \(p(\boldsymbol{z}|\mathcal{D})\) 的样本进行贝叶斯模型平均来获得后验预测分布。对 \(\boldsymbol{z}\) 的后验推断可以使用标准马尔可夫链蒙特卡洛（MCMC）采样或变分推断方法，如 \[1 (https://arxiv.org/html/2609.19539#bib.bib1)\] 中所述。

**算法1 标准主动子空间构建**
1: **输入** 模型 \(f\)，权重 \(\boldsymbol{\theta}_{0}\)，#梯度样本 \(M\)，AS维度 \(K\)，扰动方差 \(\sigma_{0}^{2}\)
2: **for** \(m=1,2,\ldots,M\) **do**
3: 采样 \(\boldsymbol{x}_{m}\sim p\)
4: 采样 \(\boldsymbol{\theta}_{m}\sim\mathcal{N}\left(\boldsymbol{\theta}_{0},\sigma_{0}^{2}\mathbf{I}\right)\)
5: 计算梯度 \(\boldsymbol{g}_{m}\leftarrow\nabla_{\boldsymbol{\theta}}f_{\boldsymbol{\theta}_{m}}(\boldsymbol{x}_{m})\)
6: **end for**
7: 令 \(\mathcal{G}=[\boldsymbol{g}_{1},\ldots,\boldsymbol{g}_{M}]\in\mathbb{R}^{N\times M}\)
8: 计算SVD \(\mathcal{G}=\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{W}^{T}\)
9: **return** 前 \(K\) 个左奇异向量 \(\boldsymbol{U}_{K}=[\boldsymbol{u}_{1},\ldots,\boldsymbol{u}_{K}]\)

**主动子空间推断**。对于主动子空间推断，(3) 式中的 \(\boldsymbol{V}\) 的列是 (2) 式中采样梯度矩阵 \(\mathcal{G}\) 的前 \(K\) 个左奇异向量。这些方向对应于 (1) 式中梯度协方差矩阵 \(\mathcal{C}\) 的主导特征向量。算法1 (https://arxiv.org/html/2609.19539#alg1) 总结了标准的主动子空间构建。CAS通过首先将模型参数映射到压缩参数空间，然后再计算主动方向来修改此构建。

## 3 压缩主动子空间

在压缩主动子空间中，我们引入一个固定的等距嵌入 \(\boldsymbol{P}\in\mathbb{R}^{N\times R}\)，其中 \(R\ll N\)。约化参数 \(\boldsymbol{\phi}\in\mathbb{R}^{R}\) 通过以下方式映射到完整模型参数空间：
\[
\boldsymbol{\theta}=\boldsymbol{\theta}_{0}+\boldsymbol{P}\boldsymbol{\phi}.
\]
我们使用下面描述的结构化构造随机采样 \(\boldsymbol{P}\)，并在整个主动子空间构建过程中保持其固定。然后我们在约化模型 \(f_{\boldsymbol{\theta}_{0}+\boldsymbol{P}\boldsymbol{\phi}}(\boldsymbol{x})\) 上进行主动子空间构建。这将每个采样模型梯度的维度从 \(N\) 降低到 \(R\)。

我们对 \(\boldsymbol{P}\) 的选择基于计算效率和存储成本。一个通用的 \(N\times R\) 矩阵需要 \(O(NR)\) 存储，并且对于高维参数空间模型会引入另一个大矩阵。我们转而使用归一化的类CountSketch构造 \[25 (https://arxiv.org/html/2609.19539#bib.bib25),26 (https://arxiv.org/html/2609.19539#bib.bib26)\]。每个模型参数被独立且均匀地分配到 \(R\) 个类别之一。令 \(M_{i}\in\{1,\ldots,R\}\) 表示分配给参数 \(i\) 的类别，\(n_{j}\) 表示分配到类别 \(j\) 的参数数量。我们还为每个参数采样独立的随机符号 \(s_{i}\in\{-1,+1\}\)。\(\boldsymbol{P}\) 的条目则为：
\[
P_{i,j}=\mathbf{1}_{j=M_{i}}\frac{s_{i}}{\sqrt{n_{j}}}.\tag{4}
\]
由于不同列具有不相交的支持集且每列已归一化，所得的嵌入满足 \(\boldsymbol{P}^{T}\boldsymbol{P}=\mathbf{I}_{R}\)。因此，\(\boldsymbol{P}\) 在嵌入完整参数空间时保留了约化参数空间的欧几里得几何结构。

(4) 式中的构造与Uni-LoRA \[21 (https://arxiv.org/html/2609.19539#bib.bib21)\] 中使用的等距投影密切相关。主要区别在于应用于每个非零条目的随机符号 \(s_{i}\)，类似于CountSketch矩阵。由于 \(\boldsymbol{P}\) 的每一行只包含一个非零条目，该嵌入可以使用 \(N\) 个类别分配和随机符号来表示，而无需存储一个稠密的 \(N\times R\) 矩阵。同样，使用 \(\boldsymbol{P}\) 或 \(\boldsymbol{P}^{T}\) 进行乘法实现的压缩和解压缩也可以在 \(O(N)\) 操作内完成。

对于CAS，我们采样约化扰动为 \(\boldsymbol{\phi}_{m}\sim\mathcal{N}\left(\boldsymbol{0},\sigma_{0}^{2}\mathbf{I}_{R}\right)\)，并在 \(\boldsymbol{\theta}_{m}=\boldsymbol{\theta}_{0}+\boldsymbol{P}\boldsymbol{\phi}_{m}\) 处评估模型。相应的梯度为：
\[
\boldsymbol{q}_{m}=\nabla_{\boldsymbol{\phi}}f_{\boldsymbol{\theta}_{0}+\boldsymbol{P}\boldsymbol{\phi}_{m}}(\boldsymbol{x}_{m})=\boldsymbol{P}^{T}\boldsymbol{g}_{m}\in\mathbb{R}^{R}.
\]