技术摘要：T-REGS（用于自监督学习的最小生成树正则化）

问题陈述

自监督学习（SSL）已成为在没有标签数据的情况下学习数据表示的主导范式，尤其是在视觉领域。然而，联合嵌入自监督学习（JE-SSL）面临的一个关键挑战是表示崩溃（representation collapse），即网络无论输入为何都输出相同或无信息的向量。近期的研究发现了一种更微妙的失效形式，称为维度崩溃（dimensional collapse），即学习到的特征仅占据低维子空间，导致特征相关性高且多样性降低。此外，有效的表示需要均匀性（uniformity），确保嵌入在表示空间中均匀分布，以保留最大信息并辅助下游泛化。

现有解决方案面临以下局限性：

• 对比学习方法（如 SimCLR、MoCo）需要较大的批次大小和大量的负样本，这在计算上非常昂贵。
• 冗余减少方法（如 Barlow Twins、VICReg）强制协方差矩阵接近单位矩阵，但仅依赖于二阶矩，可能无法捕捉到防止收敛至均匀分布的密度集中点。
• 非对称方法（如 BYOL）在如何防止崩溃方面缺乏理论依据。
• 基于最优传输的方法（如 Fang 等人）满足必要的理论属性，但计算成本高昂，并可能因 SVD 计算而产生数值不稳定性。

方法论：T-REG 与 T-REGS

作者提出了 T-REG（最小生成树正则化），这是一个旨在同时缓解维度崩溃并促进分布均匀性的框架。

核心概念：MST 长度最大化

其核心思想是最大化学习到的嵌入的**最小生成树（Minimum Spanning Tree, MST）**的长度。理论分析将 MST 长度与统计维度估计和熵最大化联系起来。

• MST 长度 ($E(G)$)： 定义为树中连接节点之间的欧几里得距离之和：$E(G) = \sum_{(z,z') \in E} \|z - z'\|^2$。
• 排斥力： 在优化过程中，构成 MST 边的点对彼此施加排斥力，自然地将点推开，以增加总树长。

T-REG 损失函数

为了防止点向无穷远处发散的平凡解（因为单纯最大化 MST 长度会鼓励扩张），T-REG 将嵌入约束在一个紧致流形上，具体是一个超球面。总损失由两部分组成：

1. MST 长度最大化 ($L_E$)： $L_E(Z) = -\frac{1}{n} E(\text{MST}(Z))$。该项鼓励扩散。
2. 软球面约束 ($L_S$)： $L_S(Z) = \frac{1}{n} \sum_i (\|z_i\|^2 - 1)^2$。该项惩罚偏离单位球面的行为。

组合损失为：
$$L_{\text{T-REG}}(Z) = \gamma L_E(Z) + \lambda L_S(Z)$$
其中 $\gamma$ 和 $\lambda$ 是平衡扩散与维持球面约束之间权衡的超参数。

T-REGS 的扩展

T-REGS 将此正则化应用于 JE-SSL 框架。给定两个增强视图 $X$ 和 $X'$，模型生成嵌入 $Z$ 和 $Z'$。T-REGS 作为独立应用于每个分支的正则化项：
$$L_{\text{T-REGS}}(Z, Z') = L_{\text{T-REG}}(Z) + L_{\text{T-REGS}}(Z')$$
它可以以两种方式使用：

1. 独立使用： 直接与视图不变性项（如均方误差）结合。
2. 辅助使用： 添加到现有的 SSL 目标函数（如 BYOL、Barlow Twins）中以增强其性能。

理论分析

论文为 T-REG 提供了严格的理论保证：

• 小样本情况 ($n \le d+1$)： 定理 4.1 证明，在球面约束下，最大化 MST 长度会迫使点落在**正规单纯形（regular simplex）**的顶点上。这种配置最大化了成对距离，通过利用完整的表示维度有效地防止了维度崩溃。
• 大样本情况 ($n \to \infty$)： 分析将 MST 长度与内在 Rényi 熵联系起来。定理 4.4 和推论 4.6 表明，最大化 MST 长度渐进地驱动嵌入分布趋向于紧致黎曼流形（球面）上的均匀分布。
• 原则属性： T-REG 自然满足 Fang 等人 [23] 提出的四个属性（实例置换、实例克隆、特征克隆和特征婴儿约束），这些属性对于防止维度崩溃和确保均匀性是必要的。

实验结果

作者在合成数据和标准 SSL 基准测试上评估了 T-REGS。

合成数据验证

• 维度崩溃： 在具有受控崩溃水平的 1024 维高斯数据集上，T-REG 的损失敏感度与 W2 度量（来自 Fang 等人）一致，证实了其惩罚崩溃的能力。
• 均匀性： 点云优化表明，T-REG 成功地将退化或聚集的数据转化为均匀分布在球面上的点。在高维（256 维）情况下，余弦相似度分布收敛到 0 附近的 Dirac $\delta$ 函数，表明呈现出正规单纯形配置。

基准性能

• CIFAR-10/100： T-REGS 取得了具有竞争力的结果，达到或略微超过了 W2 正则化方法。当作为 BYOL 和 Barlow Twins 的辅助损失时，它始终提高了相对于使用其他正则化器（如 Lu、W2）的基线和变体的 Top-1 准确率。
• ImageNet-100/1k： T-REGS 展示了可扩展性。作为 BYOL 和 Barlow Twins 的辅助损失，它提升了在两个数据集上的性能。独立的 T-REGS（结合 LMSE）在 ImageNet-1k 上实现了 68.8% 的 Top-1 准确率，与 INTL 等最先进方法相比具有竞争力。
• 多模态应用（图像-文本检索）： 应用于微调 CLIP 时，T-REGS 提高了 Flickr30k 和 MS-COCO 上的检索性能（R@1 和 R@5）。这表明 T-REGS 通过在图像和文本分支中强制执行均匀性，有助于缓解“模态间隙（modality gap）”。

计算效率

• 复杂度： MST 计算使用 Kruskal 算法，最坏情况复杂度为 $O(B^2(D \cdot \log B))$，其中 $B$ 是批次大小，$D$ 是维度。
• 运行时间： 尽管存在理论复杂度，但通过使用 torchph 的 GPU 并行化实现，在实践中 T-REGS 的每步实际运行时间可以与 SimCLR 和 VICReg 持平。

意义与主张

论文声称 T-REGS 提供了一个概念简单、易于实现且计算高效的现有正则化策略的替代方案。其主要意义在于：

1. 理论完备性： 它提供了防止维度崩溃并促进紧致流形上均匀分布的证明保证，其基础是拓扑数据分析和熵最大化。
2. 实际有效性： 它在多种架构（对比式和非对比式）和多种模态（单模态和多模态）中提高了表示质量，且不需要大型负样本批次或复杂的数值运算（如 SVD）。
3. 通用性： 它既可以作为独立的优化目标，也可以作为模块化的辅助损失，使其能够适应各种 SSL 流水线。

作者总结道，T-REGS 成功地弥合了良好 SSL 表示的理论要求与实际、可扩展实现之间的差距。