ICLR-2024《P²OT: Progressive Partial Optimal Transport for Deep Imbalanced Clustering》

核心思想

这篇论文针对深度聚类（deep clustering）在现实世界中常见的不平衡数据分布问题，提出了一种新的问题设定——深度不平衡聚类（deep imbalanced clustering）。传统深度聚类方法多假设数据均匀分布，忽略了长尾分布的挑战，导致伪标签生成质量低下和确认偏差（confirmation bias）。核心思想是通过伪标签（pseudo-labeling）框架，引入渐进式部分最优传输（Progressive Partial Optimal Transport, P²OT）算法，将伪标签生成表述为一个单一的凸优化问题。该问题同时考虑类不平衡分布（通过KL散度约束）和样本置信度（通过总质量约束渐进选择高置信样本），从而生成不平衡感知的伪标签，并逐步从易样本过渡到难样本学习，避免手动阈值调参和退化解。

目标函数

给定模型对$N$个样本的预测$P\in {\mathbb{R}}_{+}^{N\times K}$（$K$为簇数）和伪标签$Q\in {\mathbb{R}}_{+}^{N\times K}$，初始交叉熵损失为：
$$\mathcal{L}=-\sum _{i=1}^N{Q}_i\log P_i=⟨Q,-\log P{⟩}_F,$$
其中$⟨\cdot ,\cdot {⟩}_F$为Frobenius内积。为避免退化解（所有样本赋给单一簇），引入KL散度约束假设先验均匀分布$\frac{1}{K}{1}_K$，并考虑不平衡分布，将伪标签生成表述为不平衡OT问题：
$$\underset{Q\in \Pi }{\min }⟨Q,-\log P{⟩}_F+\lambda \mathrm{KL}(Q^{\top }{1}_N,\frac{1}{K}{1}_K),$$
s.t. Π = { Q ∈ R + N × K ∣ Q 1 K = 1 N 1 N } \Pi = \{ Q \in \mathbb{R}^{N \times K}_+ \mid Q 1_K = \frac{1}{N} 1_N \} Π={Q∈R+N×K​∣Q1K​=N1​1N​}。

为处理初始表示差导致的噪声伪标签，引入课程学习思想：通过总质量约束$\rho$（渐增）选择高置信样本，并允许样本权重调整，避免阈值调参。最终P²OT目标函数为：
$$\underset{Q\in \Pi }{\min }⟨Q,-\log P{⟩}_F+\lambda \mathrm{KL}(Q^{\top }{1}_N,\frac{\rho }{K}{1}_K),$$
s.t. Π = { Q ∈ R + N × K ∣ Q 1 K ≤ 1 N 1 N ,   1 N ⊤ Q 1 K = ρ } \Pi = \{ Q \in \mathbb{R}^{N \times K}_+ \mid Q 1_K \leq \frac{1}{N} 1_N, \, 1_N^\top Q 1_K = \rho \} Π={Q∈R+N×K​∣Q1K​≤N1​1N​,1N⊤​Q1K​=ρ}，
其中$\rho \in [0,1]$为选中质量分数，渐增策略为$\rho ={\rho }_0+(1-{\rho }_0)\cdot e^{-5(1-t/T{)}^2}$（${\rho }_0=0.1$，$t$为当前迭代，$T$为总迭代），KL为非归一化散度以处理不平衡。

目标函数详细的优化过程

P²OT是非标准不平衡OT，为高效求解，重构为标准形式：

1. 引入虚拟簇吸收未选质量：扩展$Q$为$\hat{Q}=[Q,\xi ]\in {\mathbb{R}}_{+}^{N\times (K+1)}$，其中$\xi \in {\mathbb{R}}^{N\times 1}$为虚拟簇分配，满足$\hat{Q}{1}_{K+1}=\frac{1}{N}{1}_N$和$1_N^{\top }\xi =1-\rho$。成本矩阵$C=[-\log P,0_N]$，目标变为：
   $$\underset{\hat{Q}\in \Phi }{\min }⟨\hat{Q},C{⟩}_F+\lambda \mathrm{KL}({\hat{Q}}^{\top }{1}_N,\beta ),$$
   s.t. Φ = { Q ^ ∈ R + N × ( K + 1 ) ∣ Q ^ 1 K + 1 = 1 N 1 N } \Phi = \{ \hat{Q} \in \mathbb{R}^{N \times (K+1)}_+ \mid \hat{Q} 1_{K+1} = \frac{1}{N} 1_N \} Φ={Q^​∈R+N×(K+1)​∣Q^​1K+1​=N1​1N​}，$\beta =\left(\begin{array}{c}\frac{\rho }{K}{1}_K\\ 1-\rho \end{array}\right)$。

2. 替换为加权KL确保严格约束：原KL无法严格保证$1_N^{\top }\xi =1-\rho$，引入加权KL：
   $$\widehat{\mathrm{KL}}({\hat{Q}}^{\top }{1}_N,\beta ,\lambda )=\sum _{i=1}^{K+1}{\lambda }_i[{\hat{Q}}^{\top }{1}_N{]}_i\log \frac{[{\hat{Q}}^{\top }{1}_N{]}_i}{{\beta }_i},$$
   并设${\lambda }_{K+1}\to +\infty$（实际用大值$\iota$），以强制虚拟簇大小为$1-\rho$。目标为：
   $$\underset{\hat{Q}\in \Phi }{\min }⟨\hat{Q},C{⟩}_F+\widehat{\mathrm{KL}}({\hat{Q}}^{\top }{1}_N,\beta ,\lambda ).$$

3. 理论保证：命题1证明${\hat{Q}}^{\star }=[Q^{\star },{\xi }^{\star }]$，其中$Q^{\star }$为原P²OT最优解。

4. 熵正则化与缩放求解：添加$-ϵH(\hat{Q})$（$ϵ=0.1$），问题可由Sinkhorn缩放算法求解。令$M=\exp (-C/ϵ)$，$f=\frac{\lambda }{\lambda +ϵ}$（逐元素），$\alpha =\frac{1}{N}{1}_N$：
   $${\hat{Q}}^{\star }=\mathrm{diag}(a)M\mathrm{diag}(b),$$
   其中$a,b$由迭代更新：
   $$a\leftarrow \frac{\alpha }{Mb},b\leftarrow \left(\frac{\beta }{M^{\top }a}\right)\circ f,$$
   直到$b$收敛（变化<1e-6或迭代达1000）。提取$Q=\hat{Q}[:,:K]$。证明见附录C。

该过程将总质量约束融入边际，KL松弛为不平衡，优于通用不平衡OT求解器（快2倍）。

主要贡献点

1. 问题与基准：首次形式化深度不平衡聚类，建立新基准（长尾CIFAR100、ImageNet-R、iNaturalist子集），填补现实差距。
2. 框架与算法：提出渐进PL框架，用P²OT统一建模不平衡分布与置信选择，避免退化与偏差。
3. 优化创新：重构P²OT为不平衡OT（理论保证），用高效缩放算法求解，支持mini-batch与记忆缓冲稳定。
4. 实验验证：在基准上SOTA（如CIFAR100 ACC提升0.9%，ImageNet-R提升2.4%），证明鲁棒性。

算法实现过程

算法实现为Algorithm 1（缩放算法）：

• 输入：成本$-\log P$，$ϵ,\lambda ,\rho ,N,K$，大值$\iota$。
• 预处理：
  • $C\leftarrow [-\log P,0_N]$（扩展成本）。
  • $\lambda \leftarrow [\lambda ,\dots ,\lambda ,\iota {]}^{\top }$（K+1维，加权）。
  • $\beta \leftarrow [\frac{\rho }{K}{1}_K^{\top },1-\rho {]}^{\top }$。
  • $\alpha \leftarrow \frac{1}{N}{1}_N$。
  • 初始化$b\leftarrow 1_{K+1}$，$M\leftarrow \exp (-C/ϵ)$，$f\leftarrow \frac{\lambda }{\lambda +ϵ}$。
• 迭代：
  • while $b$未收敛：
    • $a\leftarrow \frac{\alpha }{Mb}$。
    • $b\leftarrow \left(\frac{\beta }{M^{\top }a}\right)\circ f$。
• 输出：$Q\leftarrow \mathrm{diag}(a)M\mathrm{diag}(b)[:,:K]$（前K列为伪标签）。

实际中，用mini-batch处理大数据集，记忆缓冲存储历史预测稳定优化。整体框架交替：用当前表示计算$P$→P²OT得$Q$→用$Q$更新表示→重复至收敛。超参：$\lambda =1$，$ϵ=0.1$，${\rho }_0=0.1$。

论文局限性总结与分析

该论文在深度不平衡聚类上创新性强，但存在以下局限：

1. 依赖预训练模型：实验基于DINO预训练ViT-B16，未探讨从零训练场景。现实中，若无高质量预训练，初始表示差可能放大确认偏差，影响P²OT收敛。
2. $\rho$策略固定：渐增用sigmoid函数，依赖迭代数$t/T$，非自适应。论文承认未来需基于学习进度动态调整，否则早期$\rho$过小可能欠拟合，晚期过快引入噪声。
3. 计算效率与规模：虽缩放算法高效（矩阵-向量乘），但对超大规模数据集（如全iNaturalist），mini-batch与缓冲增加内存开销。未比较与投影梯度下降等方法的权衡。
4. 不平衡假设：假设先验均匀，仅用KL松弛处理长尾，未显式建模极端不平衡（如R>>100）或多模态分布。噪声标签或分布外样本可能导致虚拟簇滥用。
5. 评估局限：基准虽新，但ImageNet-R“分布外”挑战依赖预训练分布；未测试开放集聚类或动态簇数$K$。指标（如ACC）对尾类敏感，但ARI不适不平衡（仅附录）。
6. 理论深度不足：虽有命题保证等价性，但缺乏全局收敛或泛化界；熵$ϵ$调参敏感，可能影响不平衡捕捉。

总体，该工作推进PL在不平衡场景的应用，但未来可探索自监督初始化、自适应约束和理论强化，以提升通用性。