跳转至

T-REGS: Minimum Spanning Tree Regularization for Self-Supervised Learning

会议: NeurIPS 2025
arXiv: 2510.23484
代码: GitHub (available)
领域: 自监督学习 / 表示学习 / 正则化
关键词: 自监督学习, 最小生成树, 维度坍缩, 均匀性, 正则化

一句话总结

提出 T-REGS——一种基于最小生成树(MST)长度最大化的自监督学习正则化框架，理论证明可同时防止维度坍缩并促进表示分布均匀性，在紧致黎曼流形上成立，实验在标准 JE-SSL 基准上验证了有效性。

研究背景与动机

  1. 领域现状：联合嵌入自监督学习（JE-SSL）通过让同一图像不同视图的嵌入相似来学习表示。核心挑战是防止表示坍缩（所有输入映射到相同向量）和维度坍缩（表示仅占低维子空间）。

  2. 现有痛点：

  3. 对比方法（SimCLR、MoCo）需要大量负样本和大批量，计算开销大
  4. 冗余缩减方法（BarlowTwins、VICReg）仅利用协方差矩阵（二阶矩），对分布的集中点等高阶信息视而不见，无法保证收敛到均匀分布
  5. 非对称方法（BYOL、DINO）缺乏理论解释为何非对称架构能防止坍缩

  6. Fang et al. 提出的最优传输正则化计算昂贵，且闭式加速公式仅在球面上成立、存在数值稳定性问题

  7. 核心矛盾：好的 SSL 正则化需要同时满足四项性质（实例置换不变、实例克隆、特征克隆、特征约束），现有方法要么无法全部满足，要么计算成本高。

  8. 本文要解决什么：设计一种概念简单、计算高效、理论上可证明防止维度坍缩和促进均匀性的 SSL 正则化方法。

  9. 切入角度：MST 长度是分布熵的经典估计量（Steele 定理），最大化 MST 长度等价于最大化表示分布的 Rényi 熵，自然促进均匀性。同时 MST 长度对点云的维度敏感，可自然检测和避免维度坍缩。

  10. 核心idea一句话：通过最大化嵌入空间中最小生成树的长度（加球面约束），一个简单的正则项同时解决维度坍缩和非均匀性问题。

方法详解

整体框架

T-REGS 框架将 T-REG 正则化应用于 JE-SSL 的两个分支。每个分支独立计算 MST 长度并最大化。整体损失为视图不变性损失（如 MSE 或已有 SSL 方法的损失）+ T-REG 正则化。

关键设计

  1. MST 长度最大化损失 \(\mathcal{L}_E\)：
  2. 做什么：最大化嵌入点云 \(Z\) 的最小生成树长度
  3. 核心思路：\(\mathcal{L}_E(Z) = -\frac{1}{n} E(\text{MST}(Z))\)。MST 中每条边连接的两个点会产生排斥力（梯度公式简洁：\(\nabla_x E = \sum_{(x,z) \in \text{MST}} \frac{x-z}{\|x-z\|_2}\)）。可利用 GPU 并行的 MST 算法高效计算

  4. 设计动机：MST 长度与 Rényi \(\frac{d-1}{d}\)-熵相关（Steele 定理），最大化 MST 长度 → 最大化熵 → 趋向均匀分布

  5. 球面约束 \(\mathcal{L}_S\)：

  6. 做什么：软约束嵌入点到单位球面
  7. 核心思路：\(\mathcal{L}_S(Z) = \frac{1}{n} \sum_i (\|z_i\|_2 - 1)^2\)，防止单纯最大化 MST 长度导致点无限发散

  8. 组合为 T-REG：\(\mathcal{L}_{\text{T-REG}}(Z) = \gamma \mathcal{L}_E(Z) + \lambda \mathcal{L}_S(Z)\)

  9. 理论保证：

  10. 小样本行为（定理4.1）：当 \(n \leq d+1\) 时，MST 长度在球面上取最大值于正则单纯形顶点——即点自然扩散到球面上最大间隔配置
  11. 大样本渐近（定理4.4 + 命题4.5）：当 \(n \to \infty\)，在紧致黎曼流形上，最大化 MST 长度等价于最大化对 Rényi 熵的估计量，其唯一最大值在均匀分布处取得。并且更高维流形上的均匀分布有更大的上界，自然鼓励全维度使用

使用方式

T-REGS 可作为：(a) 独立正则化——仅用 MSE + T-REG 替代全部 SSL 方法；(b) 辅助损失——叠加到已有方法（VICReg、BarlowTwins 等）上增强性能。

实验关键数据

主实验

在 CIFAR-10、CIFAR-100、STL-10、Tiny-ImageNet 上使用 ResNet-18/50 进行标准 SSL 评估（线性探测、KNN）。

方法 关键指标 说明
T-REGS (standalone) 与 VICReg/BarlowTwins 可比 仅用 MSE + T-REG，无需负样本或协方差正则化
VICReg + T-REG 优于 VICReg 作为辅助损失提升现有方法
BarlowTwins + T-REG 优于 BarlowTwins 同上

消融实验

配置 结果
仅 \(\mathcal{L}_E\)（无球面约束） 点发散到无穷，无法收敛
仅 \(\mathcal{L}_S\) 无正则化效果
T-REG 完整 稳定收敛到球面均匀分布
高维（256维）合成数据 成功收敛到正则单纯形

关键发现

  • MST 正则化确实防止维度坍缩：嵌入的有效维度（rank）显著高于未使用 T-REG 的基线
  • 均匀性提升：T-REG 使表示在球面上更均匀分布，优于仅做协方差正则化
  • 作为辅助损失普遍有效：叠加 T-REG 到多种已有方法上均带来一致提升
  • 满足 Fang et al. 的四项性质：理论证明 T-REG 满足 instance permutation、instance cloning、feature cloning、feature baby 四项性质

亮点与洞察

  • 从拓扑数据分析到 SSL 的优美连接：MST 长度通过 Steele 定理和 Rényi 熵与均匀性直接相关，这个连接既自然又强大
  • 概念最简的 SSL 正则化：相比 VICReg 的多项损失或对比学习的负样本机制，"最大化 MST 长度 + 球面约束"极其简洁
  • 在黎曼流形上的通用理论：不限于欧式空间或球面，理论对任意紧致黎曼流形成立

局限性 / 可改进方向

  • MST 计算成本：尽管有 GPU 并行算法，MST 计算仍至少 \(O(n^2)\)，对非常大的 batch size 可能成为瓶颈
  • 超参数 \(\gamma, \lambda\) 的调节：两个权重参数控制扩散与约束的平衡，需要经验调参
  • 大规模实验有限：主要实验在 CIFAR/STL/Tiny-ImageNet 上，缺乏 ImageNet-1K 级别的验证
  • 与 DINO/DINOv2 等 SOTA 方法的对比：未与当前最强的自监督方法直接对比

相关工作与启发

  • vs VICReg/BarlowTwins：这些方法用协方差正则化（二阶矩），T-REG 通过 MST 隐式利用更高阶分布信息
  • vs Fang et al. OT 方法：同样满足四项性质，但 T-REG 计算更简单，不需要 SVD 和根号运算
  • vs 对比学习：T-REG 不需要负样本，通过 MST 边的排斥力自然实现均匀性

评分

  • 新颖性: ⭐⭐⭐⭐ MST 正则化用于 SSL 是新颖的跨领域迁移，理论elegant
  • 实验充分度: ⭐⭐⭐ 合成数据和中等规模验证充分，但缺乏大规模实验