ICML-2024《Image Clustering with External Guidance》

推荐深蓝学院的《深度神经网络加速：cuDNN 与 TensorRT》，课程面向就业，细致讲解CUDA运算的理论支撑与实践，学完可以系统化掌握CUDA基础编程知识以及TensorRT实战，并且能够利用GPU开发高性能、高并发的软件系统，感兴趣可以直接看看链接：深蓝学院《深度神经网络加速：cuDNN 与 TensorRT》

论文分析

核心思想

论文提出了一个全新的图像聚类范式——外部引导聚类 (Externally Guided Clustering)，其核心思想是利用外部知识（external knowledge）作为监督信号来增强无监督图像聚类的性能。传统聚类方法（如 k-means、深度聚类、自监督聚类）主要依赖于从数据内部挖掘监督信号，受到数据本身信息量的限制。论文指出，外部知识（如文本语义）可以提供更丰富的语义信息，从而显著提升聚类效果。

具体而言，论文设计了一种名为 Text-Aided Clustering (TAC) 的方法，通过引入 WordNet 的文本语义来辅助图像聚类。TAC 利用预训练的视觉-语言模型（如 CLIP）挖掘图像和文本模态之间的语义关联，构造图像的文本对应物（text counterpart），并通过跨模态互蒸馏（cross-modal mutual distillation）协作图像和文本特征，最终实现高性能的图像聚类。

论文将聚类方法的演变分为四个阶段（见图 1）：

1. 经典聚类：基于数据分布假设（如紧密度、层次结构）。
2. 深度聚类：利用神经网络提取判别特征。
3. 自监督聚类：通过数据增强或动量策略构造监督信号。
4. 外部引导聚类（本文提出）：利用外部知识（如文本语义）指导聚类。

TAC 的核心创新在于从内部监督转向外部知识，突破了数据内部信息的限制，特别是在处理视觉上相似但语义不同的样本（如“柯基”和“柴犬”）时表现出色。

目标函数

TAC 的目标函数由三部分组成，综合考虑了跨模态互蒸馏、聚类置信度和聚类平衡性，形式如下：

$$L_{\text{TAC}}=L_{\text{Dis}}+L_{\text{Con}}-\alpha \cdot L_{\text{Bal}}$$

其中，$\alpha =5$ 为权重参数。以下详细分析各部分：

1. 跨模态互蒸馏损失 ($L_{\text{Dis}}$)：
   该损失旨在通过跨模态邻域信息蒸馏增强图像和文本模态之间的协作。定义两个聚类头：

   • $f:v\to p\in {\mathbb{R}}^K$：预测图像特征 $v$ 的软聚类分配 $p$。
   • $g:\bar{t}\to q\in {\mathbb{R}}^K$：预测文本对应物特征 $\bar{t}$ 的软聚类分配 $q$。

   对于图像 $v_i$ 和其文本对应物 ${\bar{t}}_i$，分别取随机最近邻 $\mathcal{N}(v_i)$ 和 $\mathcal{N}({\bar{t}}_i)$，计算软聚类分配矩阵 $P,P^{\mathcal{N}},Q,Q^{\mathcal{N}}$。蒸馏损失定义为：

   $$L_{\text{Dis}}=\sum _{i=1}^K\left(L_i^{v\to t}+L_i^{t\to v}\right)$$

   其中，图像到文本的蒸馏损失为：

   $$L_i^{v\to t}=-\log \frac{e^{\left(\text{sim}({\hat{q}}_i,{\hat{p}}_i^{\mathcal{N}})/\bar{\tau }\right)}}{{\sum }_k{e}^{\left(\text{sim}({\hat{q}}_i,{\hat{p}}_k^{\mathcal{N}})/\bar{\tau }\right)}+{\sum }_{k\ne i}{e}^{\left(\text{sim}({\hat{q}}_i,{\hat{q}}_k)/\bar{\tau }\right)}}$$

   文本到图像的蒸馏损失为：

   $$L_i^{t\to v}=-\log \frac{e^{\left(\text{sim}({\hat{p}}_i,{\hat{q}}_i^{\mathcal{N}})/\bar{\tau }\right)}}{{\sum }_k{e}^{\left(\text{sim}({\hat{p}}_i,{\hat{q}}_k^{\mathcal{N}})/\bar{\tau }\right)}+{\sum }_{k\ne i}{e}^{\left(\text{sim}({\hat{p}}_i,{\hat{p}}_k)/\bar{\tau }\right)}}$$

   其中，$\text{sim}(\cdot ,\cdot )$ 为余弦相似度，$\bar{\tau }$ 为 softmax 温度参数。$L_{\text{Dis}}$ 通过最小化类间相似度和鼓励模态间邻域一致性来增强聚类判别性。

2. 置信度损失 ($L_{\text{Con}}$)：
   该损失鼓励模型产生更自信的聚类分配，定义为：

   $$L_{\text{Con}}=-\log \sum _{i=1}^n{p}_i^{\top }{q}_i$$

   当 $p_i$ 和 $q_i$ 接近 one-hot 向量时，$L_{\text{Con}}$ 达到最小值，从而提高聚类分配的确定性。

3. 平衡损失 ($L_{\text{Bal}}$)：
   为防止样本过度集中于少数簇，引入平衡损失，定义为：

   $$L_{\text{Bal}}=-\sum _{i=1}^K\left({\bar{p}}_i\log {\bar{p}}_i+{\bar{q}}_i\log {\bar{q}}_i\right)$$

   其中，$\bar{p}=\frac{1}{n}{\sum }_{i=1}^n{p}_i\in {\mathbb{R}}^K$ 和 $\bar{q}=\frac{1}{n}{\sum }_{i=1}^n{q}_i\in {\mathbb{R}}^K$ 分别为图像和文本模态的聚类分配分布。$L_{\text{Bal}}$ 通过最大化熵来鼓励均匀的簇分布。

目标函数的优化过程

TAC 的优化过程分为两个主要阶段：文本对应物构造 和 跨模态互蒸馏，具体优化步骤如下：

1. 文本对应物构造：

   • 图像语义中心计算：
     使用 k-means 在图像嵌入 $v_i$ 上聚类，估计簇数 $k=N/\tilde{N}$，其中 $\tilde{N}=300$ 为期望的簇大小。计算语义中心：

     $$s_l=\sum _{i=1}^N{1}_{v_i\in l}{v}_i,l\in [1,k]$$

     其中，$1_{v_i\in l}$ 为指示函数，表示图像 $v_i$ 是否属于第 $l$ 个簇。

   • 判别性名词选择：
     利用 CLIP 的零-shot 分类范式，将 WordNet 的所有名词反向分类到图像语义中心，计算名词 ${\mathbf{t}}_i$ 属于第 $l$ 个中心的概率：

     $$p(y=l|{\mathbf{t}}_i)=\frac{\exp \left(\text{sim}(t_i,s_l)\right)}{{\sum }_{j=1}^k\exp \left(\text{sim}(t_i,s_j)\right)}$$

     选择每个中心的前 $\gamma =5$ 个高置信度名词，组成文本空间 { t ˉ i } i = 1 M \{\bar{\mathbf{t}}_i\}_{i=1}^M {tˉi​}i=1M​。

   • 文本对应物检索：
     为每个图像 $v_i$ 检索文本对应物 ${\bar{t}}_i$，通过软检索计算：

     $${\bar{t}}_i=\sum _{j=1}^{M}p({\bar{t}}_j|v_i){\bar{t}}_j$$

     其中，

     $$p({\bar{t}}_j|v_i)=\frac{\exp \left(\text{sim}(v_i,{\bar{t}}_j)/\bar{\tau }\right)}{{\sum }_{k=1}^M\exp \left(\text{sim}(v_i,{\bar{t}}_k)/\bar{\tau }\right)}$$

     $\bar{\tau }=0.005$ 控制检索的软性，避免对应物坍缩到单一节点。

   • 基线聚类：
     通过拼接图像和文本特征 $[{\bar{t}}_i,v_i]$，直接应用 k-means 得到一个无需训练的基线聚类结果。

2. 跨模态互蒸馏：

   • 初始化聚类头：
     定义两个两层 MLP 聚类头 $f$ 和 $g$，分别处理图像和文本特征，维度为 512-512-$K$，$K$ 为目标簇数。

   • 邻域信息计算：
     在训练前，计算每个图像 $v_i$ 和文本对应物 ${\bar{t}}_i$ 的 $\bar{N}=50$ 个最近邻 $\mathcal{N}(v_i)$ 和 $\mathcal{N}({\bar{t}}_i)$。

   • 优化目标函数：
     使用 Adam 优化器（初始学习率 $1e-3$，批量大小 512，训练 20 轮）最小化 $L_{\text{TAC}}$。具体步骤：

     • 计算 $L_{\text{Dis}}$，通过图像和文本模态的邻域一致性优化聚类分配。
     • 计算 $L_{\text{Con}}$，增强聚类分配的置信度。
     • 计算 $L_{\text{Bal}}$，平衡簇分布。
     • 综合加权求和，更新聚类头参数。

   • 特殊设置：
     对于大规模数据集（如 UCF-101 和 ImageNet-1K），调整 $\bar{\tau }=5.0$，批量大小为 8192，训练轮数为 100，以适应大簇数场景。

3. 聚类结果生成：
   优化完成后，使用聚类头 $f$ 或 $g$ 的软分配（或其平均值）进行最终聚类，通常通过 argmax 得到硬分配。

主要贡献点

1. 提出外部引导聚类范式：
   首次提出利用外部知识（如文本语义）指导图像聚类，突破了传统内部监督信号的局限性，开启了聚类研究的新方向。

2. 设计 TAC 方法：
   提出了一种简单而有效的实现方法 TAC，通过 WordNet 名词和 CLIP 模型协作图像和文本模态，显著提升聚类性能。

3. 文本对应物构造策略：
   创新性地通过反向分类 WordNet 名词到图像语义中心，构造判别性文本空间，并在无类名先验的情况下实现高性能聚类。

4. 跨模态互蒸馏机制：
   提出了一种新颖的跨模态互蒸馏策略，通过邻域信息共享增强图像和文本模态的协作，显著提升聚类判别性。

5. 广泛实验验证：
   在八个数据集（包括 ImageNet-1K）上验证了 TAC 的优越性，超越了零-shot CLIP 和其他最先进的深度聚类方法。

实验结果

实验在八个数据集上进行，包括五个经典数据集（STL-10、CIFAR-10、CIFAR-20、ImageNet-10、ImageNet-Dogs）和三个更具挑战性的数据集（DTD、UCF-101、ImageNet-1K）。评估指标包括 归一化互信息 (NMI)、准确率 (ACC) 和 调整兰德指数 (ARI)。

1. 经典数据集性能（表 2）：

   • TAC 在所有五个数据集上均达到最先进性能，平均性能（AVG）为 82.2，显著优于零-shot CLIP（78.7）和 SIC（79.9）。
   • 无训练版本（TAC no train）通过拼接图像和文本特征已显著优于 k-means，例如在 ImageNet-Dogs 上 ARI 提升 43.5%。
   • 跨模态互蒸馏进一步提升性能，例如在 ImageNet-Dogs 上，TAC 的 NMI、ACC、ARI 分别为 83.3、83.6、83.0，远超零-shot CLIP（73.8、73.8、58.2）。

2. 挑战性数据集性能（表 3）：

   • TAC 在 DTD 和 UCF-101 上超越零-shot CLIP，例如在 UCF-101 上，TAC 的 NMI、ACC、ARI 分别为 82.3、68.7、60.1，优于 CLIP（74.8、61.7、51.9）。
   • 在 ImageNet-1K 上，TAC 无训练版本提升 k-means 约 10% 的 ACC，训练后性能接近零-shot CLIP，表明其在无类名先验下的竞争力。

3. 消融研究（表 4、5）：

   • 损失组合分析表明，$L_{\text{Dis}}$、$L_{\text{Con}}$ 和 $L_{\text{Bal}}$ 共同作用最佳，去掉任一项均导致性能下降。
   • 蒸馏方向实验显示，双向蒸馏优于单向蒸馏，文本到图像蒸馏优于图像到文本，验证了文本语义的判别性。

4. 参数分析（图 5）：

   • TAC 对期望簇大小 $\tilde{N}$、判别名词数 $\gamma$ 和邻域数 $\bar{N}$ 表现出鲁棒性，表明方法对超参数不敏感。
   • 适当的 $\tilde{N}$（如 300）平衡了语义粒度，过大或过小会导致性能下降。

5. 可视化（图 4）：

   • TAC 学习的特征在 ImageNet-Dogs 上表现出更清晰的簇分离，ARI 从 CLIP 图像嵌入的 38.1 提升到 72.2。

算法实现过程

TAC 的实现分为以下步骤，结合 CLIP 模型（ViT-B/32 图像编码器和 Transformer 文本编码器）完成图像聚类：

1. 数据准备：

   • 输入图像数据集，提取图像嵌入 $v_i$（通过 CLIP 图像编码器）。
   • 加载 WordNet 名词，结合提示模板（如 “A photo of [CLASS]”）生成文本嵌入 $t_i$（通过 CLIP 文本编码器）。

2. 文本对应物构造：

   • 估计簇数：根据数据集大小 $N$，设置 $k=N/300$，对于小数据集（如平均簇大小 < 300），设 $k=3K$（$K$ 为目标簇数）。
   • k-means 聚类：对图像嵌入 $v_i$ 应用 k-means，计算语义中心 $s_l$。
   • 名词分类：计算每个名词 ${\mathbf{t}}_i$ 属于语义中心 $s_l$ 的概率 $p(y=l|{\mathbf{t}}_i)$，选择前 $\gamma =5$ 个高置信度名词，组成文本空间 { t ˉ i } i = 1 M \{\bar{\mathbf{t}}_i\}_{i=1}^M {tˉi​}i=1M​。
   • 软检索：为每个图像 $v_i$ 计算文本对应物 ${\bar{t}}_i$，使用软检索公式避免坍缩。
   • 基线聚类：拼接 $[{\bar{t}}_i,v_i]$，应用 k-means 得到无训练基线结果。

3. 跨模态互蒸馏：

   • 初始化聚类头：定义 $f$ 和 $g$ 为两层 MLP，输入维度为图像/文本嵌入维度，输出维度为 $K$。
   • 邻域计算：预计算 $\bar{N}=50$ 个最近邻 $\mathcal{N}(v_i)$ 和 $\mathcal{N}({\bar{t}}_i)$。
   • 训练：
     • 使用 Adam 优化器，设置学习率 $1e-3$，批量大小 512，训练 20 轮（UCF-101 和 ImageNet-1K 调整为 8192 和 100 轮）。
     • 每轮计算 $L_{\text{Dis}}$（跨模态蒸馏）、$L_{\text{Con}}$（置信度）、$L_{\text{Bal}}$（平衡），加权求和更新 $f$ 和 $g$。
   • 超参数设置：$\tau =5e-3$，$\bar{\tau }=0.5$（UCF-101 和 ImageNet-1K 为 5.0），$\alpha =5.0$。

4. 聚类输出：

   • 使用优化后的 $f$ 或 $g$ 的软分配，通过 argmax 得到硬聚类标签。
   • 评估 NMI、ACC、ARI 指标，比较基线和训练后结果。

5. 实现细节：

   • 实验在单个 Nvidia RTX 3090 GPU 上运行，CIFAR-10 训练仅需 1 分钟。
   • 代码基于 PyTorch，利用预训练 CLIP 模型和 WordNet 数据库。

总结

TAC 通过引入外部文本语义，突破了传统聚类方法的局限性，提出了一种高效的外部引导聚类范式。其目标函数巧妙结合跨模态蒸馏、置信度和平衡性，通过简单的构造和优化流程实现卓越性能。实验结果验证了 TAC 在多种数据集上的优越性，尤其是在无类名先验的情况下超越零-shot CLIP。算法实现清晰，兼具理论创新和实用价值，为未来聚类研究提供了新方向。