End-to-end Differentiable Clustering with Associative Memories

Bishwajit Saha 1 Dmitry Krotov 2 Mohammed J. Zaki 1 Parikshit Ram 2 3

摘要

​  聚类是一种广泛使用的无监督学习技术，涉及到一个密集的离散优化问题。联想记忆模型或AMs是一种可微的神经网络，定义了一个递归动力系统，它已与各种深度学习架构集成。我们揭示了AM(Associative Memory)动力学和聚类中所需的固有离散分配之间的新联系，提出了一种新的无约束的连续松弛问题，使端到端可微聚类，称为ClAM(clustering with AM)。利用AMs的模式完成能力，我们进一步开发了一种新的自监督聚类损失。我们对不同数据集的评估表明，ClAM受益于自监督，并显著改进了传统的Lloyd‘s k-means算法和最近的连续聚类松弛（剪影系数高达60%）。

1. 引言

​  最近，传统的联想记忆或AM模型(Hopfield, 1982; 1984)被重新制定，以显著提高其记忆存储容量，并与现代深度学习技术集成(Krotov & Hopfield, 2016; Ramsauer et al., 2020; Krotov & Hopfield, 2021; Krotov, 2021)。这些新模型被称为密集联想记忆（Dense Associative Memories），是完全可微的系统，能够在突触权重中存储大量的多维向量，称为模式或“记忆”。它们可以使用反向传播算法在端到端完全可微的设置中进行训练，通常带有自监督损失。
​  我们认为，这种以端到端可微的方式学习AM的突触权值的能力，以及每个数据点恰好指向一个记忆的离散分配（关联），使AM唯一适合于可微聚类的任务。我们在图1中给出了一个简单的结果，其中标准的基于原型的聚类方案（离散的和可微的）无法找到正确的聚类，但我们提出的基于AM的方案是成功的。具体来说，我们做出了以下贡献：

  • 我们开发一个灵活的数学框架集群AM或ClAM，这是一个新颖的连续无约束放松的离散优化问题，允许集群端到端可微的方式，同时保持整个训练的离散集群分配，与线性时间集群分配。
  • 我们利用AMs的模式完成能力来开发一个可区分的自监督损失，以提高聚类质量。
  • 我们的经验证明，ClAM能够持续提高k-means高达60%，可以同时与基于光谱的聚类方法和与基于凝聚的聚类方法竞争，并产生深刻的解释。

2.相关工作

离散聚类算法

​  聚类是一个固有的硬组合优化问题。基于原型的聚类公式，如k-means聚类是np困难的，即使是对于k = 2（Dasgupta，2008）。然而，广泛使用的Lloyd’s算法或Voronoi迭代方法（Lloyd，1982）可以相当有效地找到局部最优解，这可以通过多次重启和仔细的迭代来改进(Vassilvitskii & Arthur, 2006)。它是一种直观的离散算法，在A（根据当前原型诱导的Voronoi分区（Aurenhammer，1991）为集群的分配点）和B（基于当前集群分配更新原型点）之间交替使用。然而，该方案的离散性质使其难以逃脱任何局部最小值。我们认为，基于SGD的原型聚类能够通过逃避局部最小值来提高集群的质量，同时也继承了SGD的计算效率。光谱聚类公式（Donath &霍夫曼，1973；Von Luxburg，2007）利用点的成对相似性或亲和性的特征谱（用图拉普拉斯式表示）将数据划分为“连接组件”，但没有明确地提取每个聚类的原型。这允许光谱聚类以Voronoi分区不允许的方式对数据进行分区。然而，它天真地需要二次时间（点数）来生成拉普拉斯矩阵，也需要三次时间来进行谱分解。层次聚类（Johnson，1967）基于先前建立的聚类找到连续的聚类，以自上而下的方式将它们划分，或以自下而上的方式进行组合。这些离散分层算法天真地在点的数量，尽管一些形式的自下而上的凝聚方案可以显示在二次（Sibson，1973）甚至次二次时间（3月等，2010）利用双树算法（Ram等人，2009；Curtin等人，2013；2015）。

密度聚类方案

​  如流行的DBSCAN（酯等，1996）和SNN（阿吉拉尔等，2001），不认为集群作为离散优化问题，但发现的问题模式数据分布，允许任意形状的集群鲁棒性噪声和异常值（桑德et al.，1998）。然而，多维密度估计是一项具有挑战性的任务，特别是使用非参数方法（西尔弗曼，1986；Ram & Gray，2011），导致使用参数模型，如高斯混合模型（雷诺兹，2009）（可以从期望最大化的数据中估计（Dempster et al.，1977））。

可微深度聚类

​  可微性在深度聚类中至关重要，我们希望同时学习点的潜在表示，并在该潜在空间中执行聚类(Ren et al., 2022;Zhou et al., 2022)。现有的可微聚类公式，如范数和（Panahi et al.，2017）或非负矩阵分解（Kim & Park，2008），本质上是解决约束优化，不直接适合端到端可微深度学习管道。因此，各种方案利用了某种形式的软聚类公式，如模糊k-均值（Bezdek et al.，1984）。Xie等人（2016）提出了一种新的概率k-means公式，启发于t-SNE（Minton，2008），对预先训练的自编码器（AE）的表示进行可微聚类（DEC），Guo et al. (2017a) （IDEC）显示了从联合学习表征和聚类中获得的收益。对于图像深度聚类中的表示学习，Song等人（2013）、Xie等人（2016）和Yang等人（2017）使用了堆叠自动编码器，Guo等人（2017b）（DCEC）使用卷积自动编码器（CAE）进行了进一步的改进。Chazan等人（2019年）使用了多个AEs，每个聚类一个，但（软）聚类分配是使用模糊k-means进行的。Cai等人（2022）在DEC放松中增加了一个“焦点损失”，以惩罚非离散任务，同时也增强了表征学习。深度子空间聚类（Peng等人，2016；Ji等人，2017；Zhang等人，2019）放宽了组合光谱聚类问题（同样采用部分聚类分配），并利用“self-expressive”层以可微的方式同时学习必要的拉普拉斯矩阵（相似度矩阵）和聚类。因此，深度聚类通常使用一些概率部分分配的点给多个簇，偏离了原始k-means的离散性质。此外，这些方案通常利用通过Lloyd（1982）进行的预先训练过的AEs和/或k-means来初始化网络参数。我们提出的ClAM保持了聚类的离散分配性质，并且在没有任何特殊播种的情况下工作良好。

联想记忆（AM）

​  这是一个神经网络，它可以存储一组多维向量-记忆-作为一个循环动态系统的定点吸引子状态。它被设计为将初始状态（呈现给系统）与固定点（内存）上的最终状态关联起来，从而定义吸引力的不相交盆地或数据域的分区，因此可以表示聚类。该网络在数学上被形式化为经典的Hopfield网络（霍普菲尔德，1982），但已知其内存容量有限，只能在d维数据域中存储≈0.14d随机内存（Amit等人，1985；McEliece等人，1987）。对于相关数据，容量甚至更小，并且通常会导致一个不动点将整个数据集吸引到单个盆地中。对于相关数据，容量甚至更小，并且通常会导致一个不动点将整个数据集吸引到单个盆地中。这种行为对于集群来说是有问题的，因为集群的数量——稳定的固定点的数量——应该与数据维数解耦。Krotov & Hopfield（2016）提出了现代霍普菲尔德网络或密集AM，通过在动态系统中引入快速增长的非线性激活函数，允许更密集的内存排列，以及超线性(d)内存容量。对于某些激活函数，致密AMs具有幂律甚至指数容量（克罗托夫&霍普菲尔德，2016；德米西吉尔等人，2017；拉姆索尔等人，2020年；卢西贝罗和梅扎德，2023年）。Ramsauer等人（2020）证明了变压器的注意机制（Vaswani等人，2017）是具有软max激活的密集AMs的一种特殊极限情况。密集的AMs也被用于描述整个变压器块（Hoover等人，2023年），以及集成在复杂的基于能量的神经结构中（Hoover等人，2022年）。对这些结果的回顾见Krotov（2023）。密集的AMs最近也被研究与缝线连接（Burns & Fukai，2023），并应用于异质结合环境（Liang et al.，2022）。在我们的工作中，我们将证明密集算法的循环动态和大内存容量使它们唯一适合聚类。

3.集群中的联想记忆

​  在本节中，我们提出了(i)基于AMs的ClAM的数学框架，（ii）激发了其对聚类的适用性，以及（iii）提出了一种学习记忆的方法，以进行良好的聚类。我们将所有这些放在我们的新的基于原型的聚类算法ClAM中。

3.1.联想记忆：数学框架

​  在d维欧几里得空间中，考虑M内存\(\boldsymbol{\rho}_\mu\in\mathbb{R}^d,\mu\in[M]\triangleq\{1,\ldots,M\}\)（我们稍后将讨论记忆是如何学习的）。这个数学框架的关键方面是能量函数和attractor dynamics(Krotov & Hopfield, 2021; Millidge et al., 2022)。一个适合聚类的能量应该是一个点（或一个粒子）\(v\in\mathbb{R}^d\)的连续函数。此外，能量应该有M个局部最小值，对应于每个内存。最后，当一个粒子逐渐接近一个记忆时，它的能量应该主要由单个记忆决定，而剩余的M−1记忆的贡献应该很小。一个满足这些要求的能量函数为： \[E(\boldsymbol{v})=-\frac{1}{\beta}\log\left(\sum_{\mu\in[M]}\exp(-\beta\|\boldsymbol{\rho}_{\mu}-\boldsymbol{v}\|^{2})\right)\] ​  用一个标量的\(\beta > 0\)被解释为一个逆的“温度”。随着\(\beta\)的增长，方程1中的exp（·）确保只有前导项仍然显著，而其余的M−1项被抑制。因此，整个能量函数将用最近记忆周围的抛物线来描述。该空间将被划分为每个记忆周围的M个吸引盆地，每个记忆周围的能量函数的形状将主要由最近的记忆来定义，并对其他记忆进行了小的修正。
​  attractor dynamics通过dv/dt控制v在时间内空间中的移动，同时确保能量减少。这就是dE (v)/dt < 0，它确保了一个粒子收敛到一个局部最小值——动力学的一个不动点。attractor dynamics可以通过能量景观上的梯度下降来描述： \[\tau\frac{d\boldsymbol{v}}{dt}=-\frac{1}{2}\nabla_{\boldsymbol{v}}E=\sum_{\mu\in[M]}(\rho_{\mu}-\boldsymbol{v}) \boldsymbol{\sigma}(-\beta\|\rho_{\mu}-\boldsymbol{v}\|^{2})\] ​  其中\(\tau > 0\)是一个特征时间常数，描述了粒子在能量景观上移动的速度，\(\sigma(\cdot)\)是对记忆缩放距离的softmax函数，定义为\(\sigma(z_{\mu})=\exp(z_{\mu})\Big/\sum_{m=1}^{M}\exp(z_{m})\)的任何\(\mu \in[M]\)。这一点如图2a所示。这就保证会减少能量 \[\frac{dE(\boldsymbol{v})}{dt}\stackrel{(\mathbf{a})}{=}\nabla_{\boldsymbol{v}}E(\boldsymbol{v})\cdot\frac{d\boldsymbol{v}}{dt}\stackrel{(\mathbf{b})}{=}-2\tau\Big\|\frac{d\boldsymbol{v}}{dt}\Big\|^2\stackrel{(\mathbf{c})}{\leq}0\] ​  其中(a)为链式规则，(b)遵循方程2，(c)中的等式表示局部平稳性。在一个离散的时间步长t+1下，对于点v从状态\(v_t\)到\(v^{t+1}=v^t+\delta^{t+1}\)的有效更新\(\delta^{t+1}\)是通过有限差分： \[\delta^{t+1}=\frac{dt}{\tau}\sum_{\mu\in[M]}(\rho_\mu-v^t) \sigma(-\beta\|\rho_\mu-v^t\|^2)\] ​  给定一个点S的数据集，对于每个点v∈S，我们可以用一些噪声破坏它，从而产生一个扭曲的点v˜。这可以作为AM网络v0←v˜的初始状态。AM动态是由可学习的权重ρµ，µ=1，…，M，它也对应于动态（记忆）的固定点。根据方程4，T递归的网络随时间演化，其中选择T以确保充分收敛到一个固定点（见图2b中的例子）。将最终状态vT与未损坏的v进行比较来定义损失函数 \[\mathcal{L}=\sum_{\boldsymbol{v}\in S}\|\boldsymbol{v}-\boldsymbol{v}^{T}\|^{2},\mathrm{~where~}\boldsymbol{v}^{0}\leftarrow\tilde{\boldsymbol{v}}\] ​  相对于AM参数ρµ，它的反向传播最小。在集群的背景下，AM模型自然诱导一个分区的数据空间不重叠盆地的吸引力，隐式定义了一个硬集群分配定点在每个盆地，并通过完全连续的和可微的动态，允许学习与标准深度学习框架，我们接下来讨论。

3.2.AM作为一个可微离散arg min求解器

3.3.理解由AMs引起的分区

3.4.ClAM：聚类 与AMs和自我监督

​  AM框架允许对离散的k-means问题（方程6）进行一种新的端到端可微的无约束连续松弛（方程8）。

方程6: \(\min_{R}\sum_{\boldsymbol{x}\in S}\|\boldsymbol{x}-\boldsymbol{\rho}_{\mu_{\boldsymbol{x}}^{\star}}\|^{2},\text{s.t.}\mu_{\boldsymbol{x}}^{\star}=\arg\min_{\mu\in[k]}\|\boldsymbol{x}-\boldsymbol{\rho}_{\mu}\|^{2}\) \[\min_{R}\sum_{\boldsymbol{x}\in S}\|\boldsymbol{x}-\boldsymbol{\rho}_{\mu_{\boldsymbol{x}}^{\star}}\|^{2}\\\text{s.t.}\mu_{\boldsymbol{x}}^{\star}=\arg\min_{\mu\in[k]}\|\boldsymbol{x}-\boldsymbol{\rho}_{\mu}\|^{2}\] 方程8: \(\min_R\sum_{x\in S}\left\|x-x_R^T\right\|^2,\mathrm{~where~}x^0\leftarrow x\)

​  接下来，我们希望利用上述AMs（3.1）的强大模式完成能力。我们使用标准掩膜自我监督学习——我们应用（随机）掩膜\(m\in\{0,1\}^ d\)到点\(x\in S\subset\mathbb{R}^{\bar{d}}\)和利用上e递归完成部分模式\(x^0 = m\odot x\)和ground-truth之间的失真掩盖值和完成模式我们的损失如下: \[\mathcal{L}=\sum_{x\in S}\mathbb{E}_{m\sim\mathcal{M}}\|\bar{m}\odot(x-x_{R}^{T})\|^{2}, x^{0}\leftarrow m\odot x\] ​  其中\(x^T_R\)由\(x^0\)演变而来，方程4的T递归。原型R是通过SGD最小化自监督模式完成损失（方程9）来学习的。精确的学习算法在算法1的TrainClAM子程序中给出，并如图5所示。

image-20241007095524047

图5：算法1。对于x∈S，我们首先应用一个掩码（紫色的）m到x来得到AM递归的初始迭代x0。对于T递归，我们有一个完整的版本x T R。原型R用自监督损失L上的梯度∇RL进行更新（公式9）。

​  在随机播种原型（第2行）之后，我们对数据（第3行）（数据批B中的每个x），我们应用随机掩码（第8行），使用当前原型（第9-12行）完成AM递归的模式，累积公式9（第13行），并使用批梯度（第14行）更新原型。

image-20241006230200990

​  命题3.1：算法1中的TrainClAM子例程占用\(O(dkTN|S|)\)时间，其中\(|S|\)是S的基数，收敛到一个\(O(N^{-1/2})\)-平稳点。
​  \(N^{-1/2}\)项来自于光滑非凸问题的标准SGD的收敛速度（内斯特罗夫，2003），可以通过动量改进到\(N^{-3/2}\)（Fang等人，2018），考虑到我们的提出的端到端可微性，这是很简单的。算法1的InferClAM子程序使用学习到的原型R分配集群，并一次通过S。这在聚类结束时调用一次。
​  命题3.2：算法1中的InferClAM子例程需要占用\(O(dkTN|S|)\)时间，其中\(|S|\)是S的基数。
​  从命题3.1和3.2中，我们可以看到，ClAM训练和推理在数据集S中的点\(|S|\)数、维数d和簇数k上是线性的。对于数据集S上的N个时期，ClAM训练取\(O(dkTN|S|)\)，其中T为AM递归深度，而k-means训练（采用Lloyd算法）取\(O(dkN|S|)\)。因此，k-means和ClAM都与点数成线性关系，尽管ClAM运行时有一个乘法系数T，其中T通常小于10。所以k-means直觉上比ClAM快。需要注意的是，k-means和ClAM明显高于谱聚类和凝聚聚类，凝聚聚类的样本数量通常在二次和三次之间。注意，对于相同数量的时代，我们可以为基于SGD的ClAM建立收敛保证（这可以通过使用基于动量的SGD来改进），而k-means没有类似的收敛保证。

（wei'wan）