跳转至

Flow-Transformed Implicit Processes for Function-Space Variational Inference

会议: NeurIPS2026
arXiv: 2606.01954
代码: https://github.com/Ludvins/ImplicitProcessZoo
领域: 优化/理论
关键词: 隐式随机过程、函数空间变分推断、正规化流、黑盒α散度、非高斯后验

一句话总结

本文提出流动变换隐式过程（FTIP），通过将正规化流引入隐式过程的有限样本代理系数空间，在保留前向采样可扩展性的同时诱导高度非高斯的函数空间后验，成功捕捉传统高斯近似无法建模的多峰与非对称不确定性。

研究背景与动机

贝叶斯神经网络（BNN）的传统推断大多建立在权重参数空间中，通过对权重施加先验并优化变分后验来逼近真实分布。然而，参数空间推断天然受制于神经网络复杂的参数几何结构，包括置换对称性、参数非可辨识性、参数间强相关性以及权重到预测函数的多对一冗余映射。由于实际预测仅取决于输入输出之间的映射函数，直接在函数空间开展变分推断成为更符合预测本质的方向。

然而，函数空间推断的核心计算瓶颈在于：对于由深度神经网络诱导的隐式随机过程先验，其函数空间边缘密度几乎无法写出闭式解析解。变分隐式过程（VIP）提出了一种基于采样的实用替代方案，即通过从隐式先验中前向采样有限个基函数构建经验均值与协方差代理，进而在低维组合系数空间上进行变分推断。但 VIP 将这组组合系数强行设定为多元高斯分布，这在本质上将函数后验限制在单峰、椭圆对称的几何族内；一旦真实函数的后验分布具有双峰分支、偏度或厚尾特性，高斯假设就会强行在峰值之间平滑平均或过度发散。另一方面，线性化 Laplace（TFSVI）等方法虽然提升了表达能力，却必须在测量点处计算高阶网络 Jacobian 矩阵，带来了沉重的算力负担。

针对既有方法在表达能力与计算开销之间的折角，需要一种既能保持基于样本的前向推断效率、又能突破高斯几何约束的非参数化函数后验表征形式。核心 idea：将可逆正规化流嵌套于隐式过程的有限先验样本系数空间中，利用保维度的非线性变换生成极富表达力的非高斯系数分布，从而在无需计算 Jacobian 的前提下诱导非对称与多峰的函数空间后验。

方法详解

整体框架

FTIP 的核心构想是解耦隐式先验采样与后验表达能力：利用隐式过程采样作为函数基底，同时利用正规化流作为系数分布的形状变形器。输入首先通过固定或适应性采样的先验网络集合提取先验基函数特征；随后从标准高斯基础分布中采样白噪声，经由仿射变换、有理二次样条耦合层和 LU 混合层构成的正规化流变换为非高斯系数；最后将系数与先验特征线性组合形成函数样本，并在黑盒 \(\alpha\) 目标函数下协同优化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入数据 x"] --> Feat["1. 有限先验样本特征提取<br/>前向采样 S 个隐式函数计算 m(x) 与基特征 φ(x)"]
    Noise["基础白噪声 ε ~ N(0, I)"] --> Flow["2. 嵌套正规化流变换<br/>仿射底座 + 样条耦合 + 1x1 LU 维度混合"]
    Flow --> Coeff["多峰/非对称系数 a = T_ψ(ε)"]
    Feat --> Combine["3. 函数空间预测合成<br/>F(x; a) = m(x) + Σ φ_s(x) a_s"]
    Coeff --> Combine
    Combine --> Obj["4. 目标函数优化与前向推断<br/>黑盒 α 目标函数 + 系数空间流 KL 散度"]

关键设计

1. 有限先验样本的低维函数空间代理

为了避免直接计算隐式先验函数空间的不可解析密度，FTIP 继承并泛化了变分隐式过程（VIP）的样本代理范式。具体而言，算法从隐式 BNN 权重先验 \(p(\bm{\theta})\) 中独立同分布抽取 \(S\) 组参数，生成 \(S\) 个确定性先验函数样本 \(f_s(\cdot) = g(\cdot, \bm{\theta}_s)\)。在此基础上，构造样本经验均值函数 \(m(\mathbf{x}) = \frac{1}{S}\sum_{s=1}^S f_s(\mathbf{x})\) 与无偏正规化基特征：

\[\phi_s(\mathbf{x}) = \frac{f_s(\mathbf{x}) - m(\mathbf{x})}{\sqrt{S - 1}}, \quad s = 1, \dots, S\]

整个函数空间的随机推断被投影到系数向量 \(\bm{a} = (a_1, \dots, a_S)^\top \in \mathbb{R}^S\) 的低维子空间上，代理预测函数写为 \(F(\mathbf{x}; \bm{a}) = m(\mathbf{x}) + \sum_{s=1}^S \phi_s(\mathbf{x}) a_s\)。该代理在系数先验 \(p(\bm{a}) = \mathcal{N}(\bm{0}, \bm{I}_S)\) 下严格匹配有限样本的均值与协方差。

2. 嵌套正规化流构造的非高斯系数后验

VIP 的核心瓶颈在于假设后验系数 \(q(\bm{a})\) 为高斯分布，从而抹杀了函数后验的多模态与非对称性。FTIP 将正规化流引入系数空间，令基础随机变量 \(\bm{\varepsilon} \sim \mathcal{N}(\bm{0}, \bm{I}_S)\) 经过可逆映射 \(\bm{a} = T_{\bm{\psi}}(\bm{\varepsilon})\) 得到系数后验。流结构设计为复合形式：

\[T_{\bm{\psi}} = T_{\mathrm{LU}}^{(L)} \circ T_{\mathrm{RQS}}^{(L)} \circ \dots \circ T_{\mathrm{LU}}^{(1)} \circ T_{\mathrm{RQS}}^{(1)} \circ T_{\mathrm{aff}}\]

最底层为可训练仿射变换 \(T_{\mathrm{aff}}(\bm{\varepsilon}) = \bm{M}\bm{\varepsilon} + \bm{b}\)；后续堆叠 \(L\) 层单调有理二次样条耦合层（RQS）与可逆 \(1\times 1\) LU 线性混合层。当非线性样条层初始化为恒等映射时，该模型退化为全协方差高斯 VIP，因此 FTIP 是对 VIP 的平滑嵌套扩展。通过变量代换公式，后验对数密度拥有封闭解析式：

\[\log q_{\bm{\psi}}(\bm{a}) = \log p_0(\bm{\varepsilon}) - \log|\det \bm{M}| - \sum_{\ell=1}^L \log\left|\det \frac{\partial T_{\mathrm{RQS}}^{(\ell)}}{\partial \bm{h}^{(2\ell-2)}}\right| - \sum_{\ell=1}^L \sum_{j=1}^S \log|(\bm{U}_\ell)_{jj}|\]

这赋予了后验系数分布任意扭曲、重尾甚至双峰分裂的能力，且 Jacobian 对数行列式在 \(\mathcal{O}(S)\) 复杂度内精确可算。

3. 黑盒 α 目标函数与两阶段热启动训练机制

为了在多峰预测下防止后验坍缩，FTIP 采用黑盒 \(\alpha\) 散度目标替代单一证据下界（ELBO）：

\[\mathcal{L}_\alpha = \frac{N_{\mathrm{data}}}{N_{\mathrm{batch}}}\sum_{n \in \mathcal{B}} \frac{1}{\alpha}\log\left[\frac{1}{K}\sum_{k=1}^K \exp\left(\alpha \log p(y_n \mid F(\mathbf{x}_n; \bm{a}^{(k)}))\right)\right] - \mathrm{KL}(q_{\bm{\psi}}(\bm{a}) \parallel p(\bm{a})) - R_{\mathrm{prior}}(\bm{\theta})\]

当 \(\alpha \to 0\) 时退化为标准 MC-ELBO，倾向于均值妥协解；而当 \(\alpha = 1\) 时转变为函数样本后验预测混合的对数似然，允许不同函数样本分别解释不同的条件分支。在优化稳定性上，模型采用热启动策略：前阶段仅训练仿射参数 \(\bm{M}, \bm{b}\)，令模型收敛至最佳高斯基底；第二阶段放开样条与 LU 层，在此基础上平滑形变为非高斯后验。同时引入对偶采样（Antithetic Sampling）将每次抽取的白噪声取反配对，使目标函数梯度方差显著降低 34%–88%。

损失函数 / 训练策略

整体损失包含三部分：基于 Log-Mean-Exp 的样本对数似然期望、变分系数后验与标准正态先验之间的 KL 散度 \(\mathbb{E}_{q_{\bm{\psi}}}[\log q_{\bm{\psi}}(\bm{a}) - \log p(\bm{a})]\)、以及从 VIP 继承的先验超参数边缘似然正则保护项 \(R_{\mathrm{prior}}(\bm{\theta}) \approx \log q_{\mathrm{GP}}(\bm{y}\mid \bm{X}, \bm{\theta})\)，防止隐式先验超参数在自适应微调时发生协方差过拟合。优化器采用 Adam 配合余弦退火策略。

实验关键数据

主实验

在基准评估中，作者在 UCI 回归、百万歌曲大规模回归数据集 YearPredictionMSD 以及真实行人轨迹长度预测任务上对比了 VIP、FTIP、FBNN、MFVI 和 TFSVI。评估指标包含均方根误差（RMSE）、负对数似然（NLL）、连续排位概率评分（CRPS）及中心化分位数指标（CQM）。

表 1：YearPredictionMSD 大规模回归基准测试结果（5 随机种子均值±标准差）

方法 RMSE \(\downarrow\) NLL \(\downarrow\) CRPS \(\downarrow\) CQM \(\downarrow\) 结构复杂度
VIP 10.23 ± 0.06 3.655 ± 0.015 5.326 ± 0.046 0.029 ± 0.003 高斯系数空间代理
FTIP (本文) 10.31 ± 0.05 3.407 ± 0.007 5.209 ± 0.028 0.010 ± 0.006 样条流非高斯代理
TFSVI 10.90 ± 0.12 3.720 ± 0.110 5.741 ± 0.085 0.042 ± 0.004 需计算 Jacobian
MFVI 10.80 ± 0.00 3.798 ± 0.000 5.778 ± 0.003 0.046 ± 0.000 参数空间对角高斯
FBNN 12.11 ± 1.52 3.888 ± 0.175 6.863 ± 1.054 0.075 ± 0.043 测量集测量误差

在合成双峰数据（Bimodal）和偏斜数据（Skewed）的精细诊断中，FTIP 展现出决定性优势：在双峰任务上，VIP 与 HMC/Deep Ensemble 均陷入单峰平均或虚宽分布（VIP NLL 为 3.49，CQM 为 0.222），而 FTIP 的 NLL 骤降至 2.10，CQM 降至 0.017，成功分离出两个完全独立的余弦/正弦函数分支。在 ETH/UCY 行人未来 12 步路径长度预测（LSTM 主干）中，FTIP 的 NLL 达到 1.26 ± 0.06（优于 VIP 的 1.42 及 MFVI 的 2.20）。

消融实验

作者在 YearPredictionMSD 数据集上对基函数数量 \(S\)、蒙特卡洛采样数 \(K\)、散度参数 \(\alpha\) 以及流深度（Coupling 深度）进行了单变量消融研究（5 种子中位数）。

表 2：YearPredictionMSD 核心超参数消融测试

消融维度 取值 RMSE \(\downarrow\) NLL \(\downarrow\) CRPS \(\downarrow\) CQM \(\downarrow\) 训练耗时 (min)
基特征数 \(S\) 5 10.76 3.46 5.50 0.010 17.7
基特征数 \(S\) 20 (默认) 10.30 3.41 5.22 0.007 17.9
基特征数 \(S\) 100 9.72 3.29 4.76 0.006 19.5
散度参数 \(\alpha\) 0.0 (ELBO) 10.23 3.74 5.49 0.050 18.5
散度参数 \(\alpha\) 0.5 10.19 3.74 5.47 0.051 17.7
散度参数 \(\alpha\) 1.0 (默认) 10.30 3.41 5.22 0.007 21.9
流深度 (Layers) 1 10.48 3.43 5.29 0.009 17.2
流深度 (Layers) 2 (默认) 10.30 3.41 5.22 0.007 17.9
流深度 (Layers) 6 10.30 3.40 5.21 0.006 18.6

关键发现

  • 非高斯表达力与点预测的权衡：FTIP 在 RMSE 上通常与 VIP 持平或略有微弱牺牲（如 10.23 对 10.31），但在刻画数据分布的 NLL 和 CQM 上获得跨数量级的显著改进。这是因为高斯均值预测天然利好点评估指标 RMSE，但在真实分布具有非对称或多模态结构时，点最优解并不代表概率质量的最佳分配。
  • 基函数容量是核心决定因素：基函数数量 \(S\) 对拟合精度的增益最为陡峭且单调，增大 \(S\) 能够直接拓宽函数空间子流形的基底容量；但受限于三次 Cholesky 分解计算，实践中以 \(S=20 \sim 100\) 为最佳平衡区间。
  • \(\alpha=1.0\) 是激活多峰功能的开关：当 \(\alpha \to 0\) 时，模型退化为纯单峰拟合，NLL 滞留在 3.74；当切换至 \(\alpha=1.0\) 时，对数期望混合似然机制充分激活了正规化流的多峰分配能力，NLL 骤降至 3.41。

亮点与洞察

  • 无缝嵌套的架构设计：FTIP 将最前置的变换形式显式设定为可训练仿射变换，使得初始化状态与标准 VIP 的全协方差高斯后验完全等价。这种 warm-start 训练策略既保证了优化的初级稳定性，又避免了冷启动阶段高维非线性流的模式崩塌。
  • 无 Jacobian 负担的函数空间推断：相比 TFSVI 等依赖一阶 Taylor 展开计算网络 Jacobian 的函数空间方法，FTIP 仅仅依赖隐式函数前向采样与低维正规化流变换，在反向传播中完全不涉及高阶求导，计算吞吐高出近 3 倍（17.0 ms/iter vs 49.5 ms/iter）。
  • 先验误设的自愈机制：在先验误设为单峰平稳高斯过程（GP）的情境下，FTIP 的后验流结构能够利用数据似然反向形变有限样本系数空间，诱导出分离的双峰后验，表现出对于错误基底先验的修正补偿弹性。

局限与展望

  • 计算与内存开销递增：虽然大幅轻于 Jacobian 方案，但相比经典 VIP，正规化流的多层样条参数和迭代前向采样仍带来了 30% 以上的训练时间上浮，在后验分布本质接近高斯的问题上收益有限。
  • 标量/低维输出约束：当前实验主要局限于单维回归、标量轨迹长度以及分类概率校准，尚未验证其在超高维多变量密集预测或大模型序列生成等复杂结构化输出下的有效性。
  • 基底容量上限：模型后验的表达能力受制于前向采样的 \(S\) 个离散基函数。若先验样本未能覆盖目标解所在的函数子空间，正规化流亦无法外推生成超越基向量张成空间的崭新模式。未来可探索自适应正交先验基生成策略以弥补基空间表达盲区。

相关工作与启发

  • vs Variational Implicit Processes (VIP)：VIP 是本文的直接基石，但其后验系数受困于单峰多元高斯形式。FTIP 在同一前向采样框架下串联神经样条流，打通了非高斯函数后验的路径。
  • vs TFSVI：TFSVI 基于权重的局部线性化构造函数分布，依赖测量点上下文和复杂的 Jacobian 计算；FTIP 则基于全局先验采样，以低维流转换完全规避了网络 Jacobian 计算瓶颈。
  • vs Transformed GPs：Transformed GP 将流应用于输入/输出空间的变量扭曲；FTIP 则是将流施加于随机过程的有限先验基组合系数上，作为变分推断的后验变形器，两者机制本质不同。

评分

  • 新颖性: ⭐⭐⭐⭐☆ [将正规化流创造性嵌入有限样本函数空间代理系数中，兼具数学优雅性与工程实用性]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖双峰/偏斜合成诊断、UCI、大规模 Year 数据、真实行人轨迹及先验误设分析]
  • 写作质量: ⭐⭐⭐⭐⭐ [框架逻辑清晰，公式推导详尽自洽，实验指标对比全面严谨]
  • 价值: ⭐⭐⭐⭐☆ [为贝叶斯深度学习中的非高斯函数空间推断提供了通用、可扩展的实现典范]

相关论文