跳转至

Supervised Distributional Reduction via Optimal Transport and Dependence Maximization

会议: NeurIPS 2026
arXiv: 2605.27619
OpenReview: https://openreview.net/forum?id=3bqTlAQf5B
代码: https://github.com/huguesva/Distributional-Reduction
领域: 优化/理论
关键词: 分布约简、最优传输、熔合Gromov-Wasserstein、依赖性最大化、高斯过程非平稳核

一句话总结

本文提出有监督分布约简(SDR)框架,结合熔合 Gromov-Wasserstein 最优传输与基于 CKA 的显式依赖性最大化以解决监督梯度瓶颈,并通过再生核希尔伯特空间(RKHS)正则化投影实现样本外泛化与非平稳自适应高斯过程核构建。

研究背景与动机

学习既能保留输入数据内在几何结构、又能紧致保留下游预测相关特征的低维表示,是现代机器学习的核心挑战之一。传统的降维算法(如 t-SNE、UMAP)与聚类算法(如 k-means、谱聚类)大多孤立运行,而基于最优传输(Optimal Transport, OT)的分布约简方法(如 DistR)通过求解输入度量结构与低维代表点之间的 Gromov-Wasserstein(GW)问题,成功将样本点聚类压缩与特征维度缩减统一在严格的分布匹配数学框架内。然而,现有的分布约简技术基本局限于无监督范式,所学紧致分布完全由无标签的输入空间成对几何距离主导,无法有意识地保留任务强相关的判别信号;同时,OT 固有的离散耦合计算难以对未见测试样本产生显式泛化函数映射,限制了其在下游预测流水线中的直接应用。

若简单将经典的熔合 Gromov-Wasserstein(Fused Gromov-Wasserstein, FGW)直接引入有监督分布约简,会遭遇严重的结构性优化瓶颈。在标准 FGW 框架中,特征层级监督与成对度量结构共享同一个传输耦合矩阵。当通过 Bregman 散度对偶/原始性质消去未知的原型标签后,监督损失退化为仅依赖于传输耦合矩阵的泛函,对于低维表示点本身的直接梯度恒为零。此时表示向量的更新仅能依赖最优耦合对表示坐标的二阶敏感性(隐式链式法则传递)。一旦最优耦合矩阵被严格的双边边际或强结构几何约束主导,传输矩阵对表示微调局部不敏感,反向传播至表示空间的监督梯度就会急剧衰减,导致低维原型点退化为纯几何聚类而丧失判别能力。

为了打破这种由耦合矩阵中介所导致的“监督瓶颈”,必须在最优传输的全局保结构能力与直接的表示层级判别力之间建立双重解耦与协同机制。核心 idea:将重心监督型熔合 Gromov-Wasserstein 最优传输与基于中心化核对齐(CKA)的全局依赖性最大化项统一在一个变分目标中,利用块坐标下降解耦传输几何对齐与表示判别更新,并通过 RKHS 岭回归显式正则化样本外投影映射以诱导高斯过程非平稳核。

方法详解

整体框架

SDR 接收具有 \(n\) 个样本的输入数据 \(X \in \mathbb{R}^{n \times d}\) 及其对应的回归或分类监督目标 \(Y\),目标是学习包含 \(m \le n\) 个代表性原型点、特征维度降至 \(p \le d\) 的紧致低维分布 \(Z \in \mathbb{R}^{m \times p}\) 以及对应的非均匀权重向量 \(h_Z \in \Sigma_m\)。整个方法由三大模块紧密协同构成:第一阶段基于重心消解技术将未知原型标签解析消除,建立半松弛重心有监督 FGW 目标并用条件梯度求解传输耦合 \(T\);第二阶段引入输运加权的 CKA 统计依赖性损失,在表示空间为 \(Z\) 注入强监督判别梯度;第三阶段通过最小化投影一致性正则项拟合 RKHS 变换算子 \(L\),提供闭式样本外泛化映射并将几何扭曲映射为高斯过程的自适应非平稳协方差。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据 X 与监督标签 Y"] --> B["重心有监督 FGW 对齐<br/>解析消除原型标签并解耦求解耦合 T"]
    B --> C["显式 CKA 依赖性最大化<br/>输运加权核矩阵直接驱动低维表示 Z"]
    C --> D["RKHS 样本外连续延拓<br/>闭式核岭回归拟合投影算子 L"]
    D --> E["非平稳自适应高斯过程建模<br/>以低维 SDR 映射构造空间非平稳核"]

关键设计

1. 重心有监督 FGW 对齐:解析消除未知原型标签并避免伪标签优化

传统 FGW 依赖源域与目标域均具有明确特征向量,但在分布约简任务中,\(m\) 个低维代表原型点 \(Z\) 的目标属性 \(G = (g_1, \dots, g_m)\) 初始完全未知,若作为自由变量联合迭代极易陷入不良局部极小。本文揭示出原型目标在监督损失项中具备可分性结构:

\[g_j^*(T) = \arg\min_{g \in \mathcal{G}} \sum_{i=1}^n L_s(y_i, g) T_{ij}\]

当特征损失 \(L_s\) 属于由凸函数 \(\phi\) 生成的广义 Bregman 散度族(如用于回归的平方 \(L_2\) 损失及用于分类的广义交叉熵损失)时,最优原型标签具有显式的一阶极值闭式解:

\[g_j^*(T) = \frac{1}{\pi_j} \sum_{i=1}^n T_{ij} y_i, \quad \text{其中 } \pi_j = \sum_{i=1}^n T_{ij}\]

这意味着原型的伪标签本质上是源样本标签关于输运质量的加权重心平均。将该闭式解代回原目标,便消除了额外的参数搜索,得到仅关于耦合矩阵 \(T\) 与低维表示 \(Z\) 的重心有监督 FGW(BSFGW)紧致变分形式,并采用半松弛条件梯度算子高效求解 \(T\) 与原型边缘概率 \(h_Z = T^\top \mathbf{1}_n\)。

2. 显式 CKA 依赖性最大化:打破耦合中介瓶颈并直接注入判别梯度

理论分析证明,固定耦合矩阵 \(T\) 时,重心监督项对表示变量的梯度严格为零:\(\nabla_Z S(T) = 0\)。监督信号若仅能经由 \(T^*(Z)\) 的局部扰动敏感性反传,梯度幅值将被算子范数 \(\|\frac{\partial T^*(Z)}{\partial Z}\|\) 显著压制。针对这一监督瓶颈,SDR 引入了基于归一化 Hilbert-Schmidt 独立性准则的中心化核对齐(Centered Kernel Alignment, CKA)直接约束低维表示:

\[\text{CKA}(Z, \tilde{Y}) = \frac{\langle K_Z, K_{\tilde{Y}} \rangle_{\mathcal{F}}}{\|K_Z\|_{\mathcal{F}} \|K_{\tilde{Y}}\|_{\mathcal{F}}}\]

在 \(m < n\) 的聚类降维设定下,原型空间的标签成对相似度核矩阵通过传输诱导投影精确压缩计算:\(K_{\tilde{Y}\tilde{Y}} = T^\top K_{YY} T\)。利用 CKA 具备 \([0, 1]\) 规范化区间的优异性质,SDR 在更新表示 \(Z\) 时直接施加 \(-\eta \, \text{CKA}(Z, \tilde{Y})\) 目标,以一阶动量优化器(Adam)直接调整原型坐标,确保低维簇间在保持流形测地结构的同时产生优异的语义类间间距。

3. RKHS 样本外连续延拓:正则化表示流形并提供稳定闭式预测映射

离散最优传输求解仅产生训练集内部点到代表原型的静态映射,无法评估测试集新样本。SDR 将样本外变换规约为再生核希尔伯特空间(RKHS)上的向量值核岭回归问题,将 \(Z\) 约束在核映射的像空间附近:

\[\mathcal{J}_{\text{SDR-OOS}}(Z, T, L) = \mathcal{J}_{\text{SDR}}(Z, T) + \frac{\mu}{2} \|T Z - K L\|_F^2 + \frac{\lambda_L}{2} \text{Tr}(L^\top K L)\]

当 \(\mu=1\) 且在纯降维情形(\(m=n\),\(T\) 为排列矩阵)下,映射算子矩阵 \(L \in \mathbb{R}^{n \times p}\) 具有闭式极小解 \(L^* = (K + \lambda_L I)^{-1} Z\)。为了平衡表示流形的几何灵活性与函数可延拓性,SDR 在每次外层迭代引入显式松弛步更新表示:

\[Z' = (1 - \beta) Z + \beta K L^*\]

这一设计兼具两重价值:既能防止离散表示过拟合孤立训练点,又显式导出了任意测试样本 \(x^*\) 在低维空间的解析投射 \(z(x^*) = K(x^*, X) L\)。

4. 基于分布约简的非平稳高斯过程核构建:解耦端到端网络联合优化的稳健替代方案

标准高斯过程(GP)采用平稳协方差核(如 RBF),难以拟合具有空间非均匀平滑性或突变边界的复杂函数。深度核学习(DKL)依靠深度网络端到端联合反向传播学习输入翘曲,但小样本下极易对超参数敏感并面临非凸局部极值;深度高斯过程(DGP)则依赖变分推断且收敛缓慢。SDR 将低维目标敏感映射嵌入平稳基核 \(\psi\),诱导出完全数据依赖的非平稳核函数:

\[k_{\text{SDR-GP}}(x, x') = \psi(z(x), z(x')) = \psi(K(x, X)L, K(x', X)L)\]

由于表示学习过程通过最优传输与依赖最大化预先完成,高斯过程仅需在该低维翘曲流形上拟合标准超参数,从而彻底解耦了非平稳几何发现与贝叶斯推断。

损失函数 / 训练策略

SDR 的总体联合变分优化目标定义为:

\[\mathcal{J}_{\text{SDR}}(Z, T, h_Z) = (1 - \alpha) \sum_{i=1}^n \sum_{j=1}^m L_s(y_i, g_j^*(T)) T_{ij} + \alpha \, \text{GW}(Z; T) - \eta \, \text{CKA}(Z, \tilde{Y})\]

在考虑样本外泛化(OOS)扩展时,优化过程采用代理交替块坐标下降(Block Coordinate Descent, BCD)策略迭代求解: - \(T\) 更新步:暂时忽略高阶非线性 CKA 项,利用条件梯度(Frank-Wolfe 变体)算法针对半松弛重心 FGW 目标求解传输方案 \(T\),并由列和自适应恢复原型权重 \(h_Z = T^\top \mathbf{1}_n\); - \(Z\) 更新步:保持 \(T\) 固定,使用 Adam 优化器针对 \(\alpha \, \text{GW}(Z; T) - \eta \, \text{CKA}(Z, \tilde{Y})\) 执行至多 1000 步梯度下降; - \(L\) 与松弛更新步:利用核正则化闭式解求解 \(L^*\),随后以步长 \(\beta\) 对表示执行流形向 RKHS 投影的软收缩 \(Z \leftarrow Z'\)。在全量基准测试中,超参数缺省配置为 \(\alpha=0.2\)、\(\eta=1000\)、\(\lambda_L=10^{-2}\)、\(\beta=0.5\)。

实验关键数据

主实验

论文在真实 UCI 回归基准与图像分类基准上对 SDR-GP 及主流非平稳核/特征提取高斯过程进行了系统评测。评估指标包括均值对数似然(MLL,越大越好)、均方误差(MSE,越小越好)、均值对数预测概率(MLP,越大越好)以及分类准确率(ACC,越大越好)。各指标均基于 5 个随机种子计算平均值与标准差。

数据集 任务类型 指标 SDR-GP (本文) DKL (联合训练) UMAP-GP KSPCA-GP NCA-GP
Boston Housing 回归 (\(d=13, n=506\)) MLL (↑) -0.32 ± 0.18 -0.42 ± 0.21 -1.32 ± 0.43 -0.99 ± 0.27 -0.76 ± 0.31
Boston Housing 回归 MSE (↓) 0.14 ± 0.08 0.13 ± 0.07 0.39 ± 0.11 0.38 ± 0.13 0.25 ± 0.11
Energy Efficiency 回归 (\(d=8, n=768\)) MLL (↑) -0.37 ± 0.02 -0.45 ± 0.03 -0.53 ± 0.08 -0.67 ± 0.05 -0.81 ± 0.38
Energy Efficiency 回归 MSE (↓) 0.05 ± 0.01 0.09 ± 0.02 0.13 ± 0.06 0.22 ± 0.03 0.29 ± 0.38
Concrete Strength 回归 (\(d=8, n=1030\)) MLL (↑) -0.40 ± 0.07 -0.65 ± 0.11 -1.38 ± 0.19 -1.28 ± 0.16 -1.01 ± 0.16
Concrete Strength 回归 MSE (↓) 0.13 ± 0.02 0.21 ± 0.05 0.67 ± 0.18 0.79 ± 0.23 0.45 ± 0.12
MNIST-2k 图像分类 (\(n=2000\)) MLP (↑) -0.23 ± 0.02 -0.30 ± 0.05 -1.05 ± 0.06 -0.46 ± 0.05 -0.31 ± 0.04
MNIST-2k 图像分类 ACC (↑) 0.95 ± 0.00 0.90 ± 0.02 0.82 ± 0.00 0.88 ± 0.02 0.91 ± 0.02
COIL-20 图像分类 (\(n=1440\)) MLP (↑) -0.07 ± 0.02 -0.09 ± 0.02 -0.90 ± 0.22 -0.13 ± 0.02 -0.15 ± 0.02
COIL-20 图像分类 ACC (↑) 0.99 ± 0.00 0.99 ± 0.00 0.83 ± 0.04 0.97 ± 0.01 0.98 ± 0.01

消融实验

1. 监督瓶颈消融:纯 FGW vs 显式 CKA 依赖项

在聚类分布约简设定下(\(m=200\) 个代表点),对比纯 FGW 架构(通过传输中介传递监督)与包含显式 CKA 项的 SDR 框架在特征空间的 CKA 保持度与聚类轮廓系数(Silhouette Score):

实验配置 数据集 原型空间 CKA 分数 (↑) 轮廓系数 Silhouette (↑) 说明
原始数据基准 COIL-20 0.48 — 原始高维特征与目标依赖性
纯 FGW 变体 COIL-20 0.31 ± 0.03 26.96 ± 6.19 依赖度反常下降,证明监督信号严重受阻
SDR (完整方法) COIL-20 0.56 ± 0.02 58.93 ± 2.70 CKA 大幅提升,簇间分离度提高一倍以上
原始数据基准 SNAREseq 0.45 — 单细胞多组学染色质数据基准
纯 FGW 变体 SNAREseq 0.38 ± 0.04 6.19 ± 4.08 稀疏几何约束下结构项压制标签传递
SDR (完整方法) SNAREseq 0.91 ± 0.02 64.13 ± 9.72 形成高度致密且可判别的细胞亚群簇

2. 投影连续性正则化强度 \(\beta\) 消融分析

在 UCI 回归数据集上针对 RKHS 投影更新参数 \(\beta \in [0.0, 1.0]\) 进行均匀采样测试(3 个随机种子): - \(\beta \to 0\)(无 RKHS 正则约束):训练误差最低,但由于低维表示完全自由拟合离散点几何,核回归外推映射对未见测试样本产生严重过拟合,测试 MSE 明显升高; - \(\beta \to 1\)(强行退化为解析核投影):测试外推方差降低,但过度拘束了最优传输流形对标签分布的自适应扭曲,导致训练损失与拟合能力受损; - \(\beta = 0.5\)(中间权衡折衷):在流形表达灵活性与样本外泛化鲁棒性之间达到最佳平衡,取得最低测试误差与最佳对数似然。

关键发现

  • 显式依赖项对于消除监督稀释是不可或缺的:仅依赖 FGW 传输矩阵更新表示时,原型空间的 CKA 甚至低于原始数据(如 COIL-20 从 0.48 骤降至 0.31),这充分印证了 \(\nabla_Z S(T) = 0\) 导致的“梯度遮蔽效应”。SDR 引入 CKA 后成功将得分拉升至 0.56 与 0.91,使原型呈现清晰的类别几何隔离。
  • 解耦特征学习在不确定性校准上具有显著优势:如图 4 与图 11 的校准分析所示,DKL 虽然预测精度较高,但在非凸联合鞍点处容易过度自信(Over-confident);SDR-GP 在实现最低 MLL/MSE 的同时,其平均绝对校准误差(MACE)紧贴对角理想线,兼顾了贝叶斯后验的不确定性可靠度。
  • 与预训练特征微调相比仍具结构性优势:附录针对 DKL 使用监督预训练网络初始化的实验表明,DKL-pretrained 在 Concrete 数据集上的 MLL 仅为 -0.82,而 SDR-GP 达到 -0.40,证明 SDR 的提升根植于最优传输几何保真机制,而非单纯网络容量差异。

亮点与洞察

  • Bregman 散度原型消解技巧:巧妙利用 Bregman 散度的凸对偶与重心平移性质,在无监督原型标签缺失的情形下解析消去未知变量,避免了将原型标签作为可训练参数带来的优化灾难。
  • 最优传输与连续核机械的无缝衔接:突破了传统离散最优传输不可外推的理论宿命,通过引入 RKHS 向量值核岭回归作为内生正则化项,赋予离散低维分布严格连续的样本外映射函数。
  • 非参数化非平稳高斯过程构建范式:开辟了一条不同于端到端 DKL 和重型 DGP 的全新路径,利用度量配准直接构建数据自适应输入扭曲,对低数据、高质量不确定性量化场景极具借鉴意义。

局限与展望

  • 复杂度随样本量立方级增长:SDR 在纯降维设定下计算复杂度为 \(\mathcal{O}(n^3)\),分布约简设定下为 \(\mathcal{O}(nm^2 + mn^2)\)。当训练样本量突破数千点时,内存与成对距离矩阵运算将成为严重瓶颈,未来需探索基于低秩张量分解或局部切片近似的线性化 GW 加速算法。
  • 映射表达能力受限于预选再生核:样本外泛化算子依赖预设的正定核(如 RBF),当真实输入流形具备高度非线性复杂拓扑或多模态结构时,线性系数矩阵 \(L\) 拟合的 RKHS 映射可能欠拟合。引入深度核或神经网络隐式参数化投影函数是自然的扩展方向。
  • 超参数调节对非凸平衡较敏感:权衡几何保持与标签对齐的权重 \(\eta\) 需根据度量尺度跨越数个数量级调优(通常设为 \(10^3\)),尽管 CKA 本身有界,但 GW 损失量级随点对测度动态变化,自动自适应加权机制有待深入。

相关工作与启发

  • vs DistR (Assel et al., 2025): DistR 首次利用 Gromov-Wasserstein 实现了无监督聚类与降维的优雅统一,但缺乏监督目标融合且无法对测试集进行样本外预测。SDR 在其基础上拓展了重心监督 FGW、引入 CKA 破除监督瓶颈,并通过核回归补齐了 OOS 映射能力。
  • vs Deep Kernel Learning (DKL, Wilson et al., 2016): DKL 利用前馈神经网络联合优化 GP 超参数,具有强大的拟合能力但在小样本非平稳问题上极其敏感易过拟合,且易破坏不确定性校准。SDR-GP 采用几何解耦优化路线,在样本量有限(\(n \le 2000\))的场景下显著超越 DKL 的似然表现并给出更精准的置信区间。
  • vs Supervised UMAP (McInnes et al., 2020) & KSPCA (Barshan et al., 2011): S-UMAP 与 KSPCA 属于离散启发式或核主成分有监督投影。实验证明这些传统降维方法直接接入高斯过程时,往往扭曲原始度量几何导致后验协方差严重畸变;而 SDR 凭借最优传输的全局保距性质保证了平滑非平稳变换。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次打通了有监督分布约简、最优传输原型消解与非平稳高斯过程核构建的理论桥梁。
  • 实验充分度: ⭐⭐⭐⭐☆ 包含标准聚类降维基准、UCI 回归、图像分类、消融对比及不确定性校准全套评估,小数据集表现惊艳;缺少万级以上大规模验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密流畅,从监督瓶颈分析到引理构造层层递进,动机非常清晰。
  • 价值: ⭐⭐⭐⭐☆ 为高斯过程非平稳建模和度量流形压缩提供了极具原创性的理论工具与可落地的解耦解法。