🤖 machine learning

Transformer Circuits Can Realize Clustering Algorithms

本文证明了名为 $k$ -means transformer 的 Transformer 架构在理论和实证上均能利用标准电路机制实现用于 $k$ -means 聚类的 Lloyd 算法，同时还能学习超越传统算法的性能，并通过架构修改自然地泛化到多种不同的聚类变体。

原作者： Kenneth L. Clarkson, Lior Horesh, Takuya Ito, Charlotte Park, Parikshit Ram

发布于 2026-08-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Kenneth L. Clarkson, Lior Horesh, Takuya Ito, Charlotte Park, Parikshit Ram

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一个这样的世界：计算机不再仅仅是猜测模式，而是真正学会了像人类数学家那样遵循精确的、循序渐进的指令。这就是机器学习的领域，特别是被称为“Transformer”的一个分支。你可能知道 Transformer 是聊天机器人和图像生成器背后的超级智能引擎，但在其核心，它们只是庞大的数学网络，通过观察数据并试图寻找其中的联系。通常，我们教这些网络擅长预测句子中的下一个词，或者识别照片中的猫。但科学家们一直在问一个大问题：这些灵活的学习机器是否真的能执行精确、僵化的数学问题，比如将一堆乱七八糟的物体分类成整齐的组？这不仅仅是为了做出更好的聊天机器人；这是为了理解这些数字大脑是真的能像计算机程序一样“思考”，还是仅仅在非常擅长模仿一个程序。

你即将探索的这篇论文深入探讨了这个谜团，它解决了一个经典的难题，叫做 k-means 聚类。把这想象成一场分拣弹珠的游戏。想象你有一个巨大的袋子，里面装满了各种颜色和大小混杂在一起的弹珠。你的目标是将它们分成 $k$ 个组（比如 5 组），使得同一组内的弹珠看起来尽可能相似。几十年来，完成这项任务的标准方法是一个被称为 Lloyd 算法 的算法。这是一个非常具体、僵化的配方：随机选取 5 个位置作为“中心”，将每颗弹珠移动到离其最近的中心，然后将中心移动到其新弹珠的平均位置，如此循环往复，直到分组不再发生变化。这是一场完美的、数学化的舞蹈，但很难教会学习机器精确地执行它，因为机器通常倾向于“猜测”而不是遵循严格的规则。

来自 IBM 研究院和麻省理工学院的研究人员提出了一个大胆的问题：我们能否构建一个 Transformer，它不仅是猜测如何分拣弹珠，而是真正执行 Lloyd 算法的精确步骤？更酷的是，我们能否教它做得比原始配方更好？

他们构建了一种特殊的 Transformer，称之为 “k-means Transformer”。他们并没有让机器通过试错来学习如何分拣，而是设计了机器的内部齿轮（其注意力机制和连接方式）来物理性地模拟 Lloyd 算法的数学逻辑。他们从数学上证明了，如果将机器的权重设置得当，这个 Transformer 的一层就正好执行 Lloyd 算法的一个步骤。如果堆叠十层，它就会执行十个步骤，完美地复制了经典算法。这就像建造一个机器人，它不仅仅是学习走路，而是你直接用机械结构锁定了它的腿部，使其步幅与人类完全一致。

但故事并未止步于仅仅复制旧的配方。团队随后让这个机器从零开始学习，向它展示了数千个不同的分拣谜题。他们发现了一个令人惊讶的现象：训练后的 Transformer 不仅仅是复制了 Lloyd 算法，它还学习了一种更聪明的新分拣方式。在对未见过的全新数据进行测试时，这个经过训练的机器创建的组别比经典的 Lloyd 算法更紧凑、更准确。这仿佛这个机器人已经把舞步练到了极致，以至于它发明了一套更好的编舞。

研究人员还展示了这种“机器即算法”的思想具有极高的灵活性。通过微调 Transformer 的内部组件——比如改变它对数据的关注方式或如何对数字进行归一化处理——他们可以瞬间将他们的机器转化为不同类型的分拣算法。他们可以让它处理“软”分拣（即一颗弹珠可以部分属于两个组）、“球面”分拣（针对存在于球形空间的数据），甚至是“修剪”分拣（忽略那些无法融入任何地方的异常值弹珠）。

简而言之，这篇论文证明了 Transformer 不仅仅是模糊的猜测者；它们足够强大，可以被构建为复杂数学问题的精确、循序渐进的计算器。更重要的是，它表明当我们让这些计算器进行学习时，它们可以发现解决这些问题的全新且改进的方法，从而架起了僵化的计算机科学与灵活的人工智能之间的桥梁。

技术摘要：Transformer 电路可以实现聚类算法

问题陈述

聚类，特别是 $k$ -means 聚类，是计算机科学中一个基本的组合优化问题，但已知它是 NP 难的。虽然深度学习架构（特别是 Transformer）擅长从原始数据中学习层次化特征，但它们在多大程度上能够实现精确的离散算法计算仍不明确。本文通过研究 Transformer 电路机制是否可以精确实例化 Lloyd 算法中 $k$ -means 聚类的迭代步骤，来填补经典离散聚类算法与现代神经架构之间的空白。

方法论

$k$ -means Transformer 架构

作者提出了一种特定的编码器-解码器 Transformer 架构，旨在上下文（in-context）中执行 Lloyd 算法。该架构利用了标准的 Transformer 组件：自注意力、交叉注意力、残差连接和前馈网络（FFN），但采用了特定的参数化方式和注意力得分定义。

输入嵌入（Input Embeddings）：
- 编码器（Encoder）： 处理 $n$ 个数据点集合 $X \in \mathbb{R}^{d \times n}$ 。每个点 $x_i$ 被嵌入为一个 $(d+k)$ 维向量 $[x_i^\top, y_i^\top]^\top$ ，其中 $y_i \in \{0,1\}^k$ 代表独热（one-hot）聚类分配（初始为零）。
- 解码器（Decoder）： 处理 $k$ 个聚类中心 $C \in \mathbb{R}^{d \times k}$ 。每个中心 $c_j$ 被嵌入为一个 $(d+k)$ 维向量 $[c_j^\top, e_j^\top]^\top$ ，其中 $e_j$ 是第 $j$ 个特征向量（独热索引）。
注意力机制（Attention Mechanisms）：
- 欧几里得注意力（Euclidean Attention）： 对于编码器（点分配）更新，论文使用了一种基于负平方欧几里得距离的限制性 Soft-max 注意力： $A_{ji} = -\|Kz_j - Qx_i\|^2$ 。当温度参数 $\gamma \to \infty$ 时，它近似于 arg-max 操作，将点分配给最近的中心。
- 点积注意力（Dot-Product Attention）： 对于解码器（中心更新）步骤，使用标准的点积注意力来计算分配点的均值。
- 残差连接（Residual Connections）： 用于在交叉注意力更新之前重置当前的分配或中心为零，从而有效地在单次前向传递中实现 Lloyd 算法的“分配”和“更新”步骤。
理论构建（Theoretical Construction）：
- 定理 2.1： 证明了一个具有特定固定权重（定义在等式 5 中）和限制性 Soft-max 注意力（ $\gamma \to \infty$ ）的 $T$ 层 Transformer 可以精确执行 $T$ 次 Lloyd 算法迭代。该模型大小为 $O(d+k)$ ，小于输入大小 $O(d(n+k))$ 。
- 定理 2.2： 表明如果通过增加平方范数并使用带有整流二次激活函数的反馈网络（FFN）来增强标记嵌入，标准的点积注意力 Transformer 也可以精确执行 Lloyd 算法。

学习实验

作者在一个分布式的聚类任务（来自高斯混合模型的样本）上训练了一个单层 Transformer，通过最小化 $k$ -means 目标的平滑上界进行训练。训练使用的是标准 Soft-max ( $\gamma=1$ ) 而非理论上的极限情况。随后，该训练好的模型在未见的任务上进行评估，这些任务在样本量、维度和初始化策略（随机 vs. $k$ -means++）方面各不相同。

核心贡献

Lloyd 算法的精确实现： 本文提供了理论证明和实证验证，表明具有特定参数选择的 Transformer 架构可以精确实现 $k$ -means 聚类的 Lloyd 算法。这架起了离散算法步骤与连续神经电路机制之间的桥梁。
可学习的聚类算法： 当在聚类任务分布上进行训练时，Transformer 架构学习到了一种通用的聚类算法，其聚类质量（更低的 $k$ -means 目标值）优于标准的 Lloyd 算法，即使在使用标准 Soft-max 注意力（ $\gamma=1$ ）的情况下也是如此。
用于多样化算法的架构变体： 通过修改特定的 Transformer 电路机制（例如注意力类型、归一化、激活函数），作者展示了如何实例化各种现有及新型聚类算法：
- 软 $k$ -means（Soft $k$ -means）： 通过将限制性 Soft-max 替换为标准 Soft-max 和线性注意力来实现。
- 球面 $k$ -means（Spherical $k$ -means）： 通过使用点积注意力和标记归一化（如 RMSNorm）来实现。
- 修剪 $k$ -means（Trimmed $k$ -means）： 通过将 Soft-max 替换为范数最大化注意力（norm-max attention，正则化 arg-max）来实现。
- 鲁棒 $k$ -means（Robust $k$ -means）： 通过使用稀疏 Soft-max 注意力来实现。
- 随机化 $k$ -medoids： 通过使用 Gumbel-softmax 注意力来选择实际的数据点作为中心。

结果

理论等价性： 数值验证确认，在具有较大但有限的温度（ $\gamma = 10^4$ ）时，Transformer 的聚类中心轨迹和最终分配与 Lloyd 算法完全匹配，涵盖了 10 次迭代。
性能优越性： 在学习实验中，训练好的 Transformer 收敛到的平均 log- $k$ -means 目标约为 $\sim 5.0$ ，优于在相同任务上的 Lloyd 算法（ $\sim 5.25$ ）。即使在测试不同样本量（泛化性）和不同初始化方案（例如 $k$ -means++）的任务时，这种改进依然成立，尽管其仅在随机初始化上进行了训练。
嵌入维度： 消融研究表明，标记嵌入中额外的 $k$ 个维度（总计 $d+k$ ）至关重要。仅具有 $d$ 维度的模型无法匹配 Lloyd 的性能，其收敛到的目标值较高（ $\sim 5.35$ ）。
权重分析： 学习到的权重表现出与理论构建相似的结构特征（例如 $Q^\top K$ 矩阵中的主对角线形式），但也显示出明显的差异，这表明学习到的算法在特定方面（例如中心更新方式）偏离了标准 Lloyd 算法。

意义与主张

本文声称，Transformer 电路机制可以实例化聚类的精确算法程序，从而提供一种“通过可解释的架构修改来开发替代性和新颖聚类算法的通用方法”。

算法透明度： 这项工作证明了 Transformer 不仅仅是黑盒函数近似器，还可以被明确设计为执行离散且定义良好的算法（如 Lloyd 算法）。
可学习性： 至关重要的是，本文表明这种人工指定的架构不仅在理论上可行，而且可以通过数据学习来产生超越经典基准的算法。
创新框架： 通过将聚类算法视为“Transformer 骨架”，作者提供了一个框架，通过修改标准组件（注意力、归一化）即可产生多样化的聚类变体，包括鲁棒型和随机型版本。

作者指出了局限性，特别是他们的理论结果依赖于难以进行反向传播的限制性 Soft-max 注意力，尽管他们展示了模型使用标准 Soft-max 也能正常工作。他们还指出，虽然模型可以泛化到更大的 $n$ （点数），但如果是在固定的 $k$ 上训练，目前尚不支持向更大 $k$ （聚类中心数）的泛化。这项工作为进一步探索经典方法的神经实现奠定了基础。