跳转至

IConE: Batch Independent Collapse Prevention for Self-Supervised Representation Learning

会议: ECCV 2026
论文: ECCV 原文
代码: https://media.eventhosts.cc/Conferences/ECCV2026/pdfs/5862.pdf
领域: 医学图像
关键词: 自监督学习, 联合嵌入架构, 表征崩溃, 批大小无关, 生物医学成像

一句话总结

IConE 将自监督学习中的防崩溃机制从动态 Mini-batch 统计解耦至全局可学习的实例嵌入表,利用全数据集级多样性正则化提供稳定的几何排斥力,使模型在批大小低至 1 及极端类别不平衡场景下依然保持高内在维度与出色的表征性能。

研究背景与动机

自监督学习(SSL)中的联合嵌入架构(JEAs)在通用视觉预训练中取得了显著成功,其核心机制在于通过拉近同一样本不同增强视角的距离来学习语义不变性,同时引入防崩溃(Anti-collapse)机制防止所有表示退化至单点。然而,现有的防崩溃手段普遍严重依赖于当前训练批次(Mini-batch)内部的统计交互——对比学习方法(如 SimCLR)需要大批次提供足够丰富的负样本分布,而非对比方法(如 VICReg、Barlow Twins)以及聚类/自蒸馏方法(如 SwAV、DINO)则依赖批内样本方差、协方差或批均值中心化来阻止坍塌。

这种对大批次统计的紧耦合,在高维度科学计算与生物医学成像领域造成了严重的实践困境。在三维计算机断层扫描(CT)、核磁共振(MRI)以及高分辨率显微镜图像中,单个 3D 体数据即可占满 GPU 显存,迫使预训练只能在极小的批大小(甚至单样本批次 \(B=1\) 或 \(2\))下进行。在小批次下,样本方差极不稳定,负样本池极度匮乏,且批内极易缺失少数类别,导致现有方法发生严重的表征崩溃与特征退化。此外,高维流形本身的稀疏性理论上要求更强的正则化,而显存约束又限制了批大小,形成了“高维数据需要强正则、但大显存开销剥夺了批级正则可靠性”的固有矛盾。

本文的核心切入点在于:表征学习所必需的几何均匀性(Uniformity)与多视角对齐性(Alignment)在数学优化上是完全可以解耦的。核心 idea:引入一组独立于动态批次的全局持久可学习实例锚点表,将抗崩溃排斥约束转移至全数据集维度的参数空间并施加几何正交正则化,使编码器仅通过纯吸引力目标与稳定几何锚点对齐,实现彻底脱离批大小依赖的无崩溃自监督表征学习。

方法详解

整体框架

IConE 摒弃了在编码器特征空间直接施加批内负样本对比或方差正则的传统范式,构建了一个双重几何优化系统:在静态参数空间维护一个全数据集尺寸的实例嵌入表,通过全局正交惩罚维持特征在单位超球面的均匀散开;在动态表征空间,主干编码器仅负责提取输入的多视角表征,并通过纯吸引力损失对齐各视角及对应的实例几何锚点。

输入批次由 \(B\) 个独立样本构成,每个样本经由随机数据增强生成 \(V\) 个视角。编码器将输入直接映射至单位超球面 \(\mathcal{S}^{d-1}\),无需额外的投影头(Projector Head)。与当前批次样本对应的实例嵌入从持久表中索引取出并做行归一化。训练过程中,编码器参数仅受视角间吸引与视角-实例吸引力驱动,而实例嵌入表则同时吸收与编码器表征的对齐梯度以及全局实例间的排斥梯度。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入样本 x_i (批大小 B)"] --> B["多视角增强生成 (V 个视图)"]
    B --> C["主干编码器与超球投影<br/>z_i^(v) = f(x_i^(v)) / ||f(x_i^(v))||"]
    D["全局实例嵌入表 E (N x d)<br/>持久参数存储全数据几何"] --> E["实例锚点归一化 e_i"]
    C --> F["视角间一致性学习<br/>多视角特征余弦相似度拉近"]
    C & E --> G["视角-实例对齐更新<br/>编码器表征向实例锚点单向吸引"]
    D --> H["数据集级几何多样性正则<br/>Gram 矩阵非对角平方铰链排斥"]

关键设计

1. 全局持久实例锚点表:解耦瞬态批次统计与表征防崩溃 针对小批次下样本多样性与统计量急剧退化导致的表征坍塌痛点,IConE 在模型架构之外引入一个形状为 \(N \times d\) 的可学习参数矩阵 \(E\)(其中 \(N\) 为训练集样本总数,\(d\) 为表征维度),初始参数采样自高斯分布 \(\mathcal{N}(0, 0.02)\)。每一行 \(e_i\) 充当样本 \(x_i\) 的专属几何锚点,在每次计算前进行单位归一化 \(\tilde{e}_i = e_i / \|e_i\|_2\)。该设计将原本瞬态绑定在 Mini-batch 前向传播中的防崩溃几何约束永久固化在模型参数中;在小批次训练甚至单样本批次下,当前样本均有确定的全局坐标参照系,使模型彻底免疫批内负样本缺失的冲击。

2. 数据集级几何多样性正则:Gram 矩阵正相关惩罚 为保证实例锚点在单位超球面均匀分布而不缩聚到局部区域,IConE 直接在参数矩阵 \(\tilde{E}\) 上施加显式排斥正则化。计算整个嵌入表的 Gram 矩阵 \(G = \tilde{E}\tilde{E}^\top\),并通过平方铰链损失约束非对角元素: $\(\mathcal{L}_{div} = \frac{1}{N(N-1)} \sum_{i \ne j} \left( \max(0, G_{ij}) \right)^2\)$ 由于大多数现实数据规模满足 \(N > d\),严格全正交无法实现,因此该目标放宽限制,仅对点积大于 0 的正相关对进行平方惩罚。这一机制赋予相似锚点强烈的排斥梯度,而使近似正交的对保持稳定。整个计算仅涉及 \(d\) 维轻量向量矩阵相乘,内存占用仅约数兆字节,成功将全局排斥压力转化为完全独立于训练 Batch Size 的确定性几何先验。

3. 双向视角-实例对齐与单向排斥梯度隔离 传统对比学习中编码器直接接收来自批内所有负样本的推开梯度,导致小批次梯度的方差极大。IConE 通过视角-实例对齐损失 \(\mathcal{L}_{vi}\) 建立了编码器表征 \(z_i^{(v)}\) 与持久锚点 \(\tilde{e}_i\) 的联系: $\(\mathcal{L}_{vi} = \frac{1}{B \cdot V} \sum_{i \in B} \sum_{v=1}^V \langle z_i^{(v)}, \tilde{e}_i \rangle\)$ 该损失驱动编码器输出向当前实例的静态坐标聚集,同时拉动锚点 \(\tilde{e}_i\) 缓慢更新至各视角表征的中心。至关重要的是梯度隔离机制:主干网络编码器 \(\theta\) 的梯度 \(\nabla_\theta \mathcal{L} = \nabla_\theta \mathcal{L}_{vi} + \nabla_\theta \mathcal{L}_{vv}\) 仅包含吸引分量,编码器表征之间从未显式互相推斥;所有排斥压力全由嵌入表更新 \(\nabla_{e_i} \mathcal{L} = \nabla_{e_i} \mathcal{L}_{vi} + \nabla_{e_i} \mathcal{L}_{div}\) 承担。这种设计从数学根源上阻断了批大小对编码器梯度的干扰。

损失函数 / 训练策略

IConE 的总损失函数由三项直接求和构成,完全摒弃了对比学习中极为敏感的温度超参数(Temperature): $\(\mathcal{L} = \mathcal{L}_{vi} + \mathcal{L}_{vv} + \mathcal{L}_{div}\)$ 其中视角间一致性损失 \(\mathcal{L}_{vv}\) 定义为同一样本所有增强视角间的余弦距离均值: $\(\mathcal{L}_{vv} = \frac{1}{B} \sum_{i \in B} \frac{2}{V(V-1)} \sum_{1 \le m < n \le V} \langle z_i^{(m)}, z_i^{(n)} \rangle\)$ 即使在极端批大小 \(B=1\) 下,只要保证视角数 \(V \ge 2\),\(\mathcal{L}_{vv}\) 与 \(\mathcal{L}_{vi}\) 均具有明确良定义的梯度。在优化算法中,编码器参数 \(\theta\) 与嵌入表 \(E\) 均通过标准优化器(AdamW)同步以学习率 \(\eta\) 进行端到端梯度下降更新。

实验关键数据

主实验

在涵盖 6 个 3D 医学影像数据集(Organ3D、Nodule3D、Adrenal3D、Fracture3D、Synapse3D、Vessel3D)的 MedMNIST3D 基准上,评估不同批大小(\(B \in \{4, 16, 64\}\))下各自监督方法学得表征的非参数化 k-NN 分类平衡准确率(Balanced Accuracy, %):

方法 B=4 (k=1) B=4 (k=5) B=4 (k=20) B=16 (k=1) B=16 (k=5) B=16 (k=20) B=64 (k=1) B=64 (k=5) B=64 (k=20)
BYOL 54.6 53.7 49.6 58.1 56.3 49.7 58.1 55.3 53.0
DIET 55.6 55.6 54.7 55.3 55.3 55.3 57.1 57.0 54.4
DINO 55.8 54.6 50.5 57.7 54.9 51.9 57.4 56.6 54.6
MoCo-v2 50.8 51.1 49.7 56.7 56.2 55.4 57.2 56.3 53.3
SimCLR 49.9 48.4 47.7 52.0 51.6 51.8 53.1 54.5 53.9
SwAV 48.0 46.5 44.2 50.8 47.8 46.6 50.6 48.9 45.8
VICReg 50.8 47.3 45.2 55.9 56.7 55.7 59.1 57.8 57.0
IConE (本文) 63.6 62.2 60.6 63.3 62.8 62.4 62.5 60.6 58.8

在监督学习扩展中,IConE 的参数锚点思想同样显著优于标准交叉熵(Cross-Entropy)与有监督对比学习(SupCon),特别是在 \(B=1\) 的极值场景下:

方法 2D (B=1) 2D (B=4) 2D (B=16) 2D (B=64) 3D (B=1) 3D (B=4) 3D (B=16) 3D (B=64)
Cross-Entropy 76.20 78.17 79.87 80.32 66.61 72.70 69.65 72.85
SupCon 72.20 74.08 79.70 79.73 59.49 66.50 72.59 71.55
IConE-Instance 76.81 78.43 79.53 79.63 68.77 75.16 76.56 75.31
IConE-Class 78.71 82.42 82.16 82.31 69.59 74.25 77.14 76.21

消融实验

在 2D 与 3D 数据集上测试将默认的 Gram 矩阵正交损失替换为其他无批次限制的统计正则器(VCReg 与 SIGReg)时的 Top-1 平衡准确率(%):

正则项配置 2D (B=1) 2D (B=4) 2D (B=16) 2D (B=32) 3D (B=1) 3D (B=4) 3D (B=16) 3D (B=32)
IConE (默认正交铰链) 59.6 60.6 62.0 60.1 62.1 69.7 68.8 65.4
IConE-VCReg 59.3 60.8 60.4 61.2 63.0 68.8 69.1 67.9
IConE-SIGReg 57.3 58.3 58.5 58.4 62.9 68.4 67.5 68.0

关键发现

  • 小批次稳定性极强:在 3D 医学影像基准中,当批大小从 64 缩减至 4 时,VICReg 从 59.1% 断崖式下滑至 50.8%,SimCLR 从 53.1% 下滑至 49.9%,而 IConE 从 62.5% 反向提升至 63.6%,实现了极低的批大小依赖性与完全平坦的性能曲线。
  • 高内在维度与防坍塌:几何特征分析(RankMe、LiDAR 及 Effective Rank)显示,基线方法在小批次下有效秩急剧下降;而 IConE 在全批次范围内保持显著更高的奇异值谱熵与有效秩,说明其特征空间未发生维数坍塌(Dimensional Collapse)。
  • 非平衡类别隐式记忆:在阶梯衰减(Step-k)类别不平衡实验中,当不平衡比例达到 20-25 时,基线方法由于小批次经常无法采样到少数类样本,少数类准确率跌至 10-20%;而 IConE 通过常驻参数表持续施加多样性排斥,少数类准确率依然维持在 40% 以上。

亮点与洞察

  • 解耦吸引与排斥的载体:传统 JEAs 强行让编码器在同一个潜在空间同时学习吸引(对齐)和排斥(抗崩溃),IConE 将排斥压力完全剥离给全局参数表,使编码器只用做最稳定的流形平滑回归。
  • 无需温度系数超参数:对比学习对温度系数 \(\tau\) 极度敏感,微小的缩放不当即可引发训练崩溃;IConE 采用直接的平方铰链正交损失替代 Softmax 概率分布竞争,完全省去了该调参开销。
  • 显存与算力友好:去掉了高算力开销的 Projection Head,直接在 Backbone 特征上操作;嵌入表占用仅在数兆字节级别,为高分辨率 3D 大体积骨干网络留出了宝贵的显存空间。

局限与展望

  • 海量样本下的计算复杂度:当训练集规模激增至数百万甚至十亿级别时,默认的 \(N \times N\) Gram 矩阵计算将面临 \(O(N^2)\) 的时间瓶颈,必须平滑切换为 VCReg 或 SIGReg 等线性复杂度正则。
  • 动态新增样本的增量拓展:每个实例均对应一个显式参数行,在持续学习或在线流式数据新增时,如何为新样本动态扩容与初始化几何锚点尚需深入探索。

相关工作与启发

  • vs SimCLR / MoCo: 对比学习强依赖批内负样本池或动量内存队列(Memory Queue),后者在训练早期编码器剧烈变化时容易积累过期特征;IConE 直接通过梯度下降端到端同步更新静态锚点,且无双倍模型参数开销。
  • vs VICReg / Barlow Twins: 非对比方法在特征协方差矩阵上施加统计正则,但在 \(B < 16\) 时样本方差有偏且不可逆;IConE 将统计约束迁移到全数据集 \(N\) 的维度上计算,批大小即使降到 1 统计量依旧饱满。
  • vs DIET: 参数实例判别(PID)方法通过分类交叉熵隐式促成正交,对多视角对齐缺乏显式拉近约束;IConE 显式引入视角-视角一致性并用直观的超球面距离约束取代 Softmax 竞争。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙解耦自监督多视角对齐与抗崩溃机制,将排斥力完全托管给全局参数表,思想精炼深刻。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D/3D 医学影像基准、自然图像验证、理论几何分析(RankMe/LiDAR/UMAP)及不平衡与监督学习拓展。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨,问题动机直击 3D 科学计算显存痛点,数学推导与实验论证高度闭环。
  • 价值: ⭐⭐⭐⭐⭐ 为高维大体积数据(3D 医疗、遥感、显微断层扫描)在单卡/小批次下的自监督学习开辟了切实可行的通用新路径。