Leaner Transformers Can Easily Learn to Cluster¶

会议: NeurIPS 2026
arXiv: 2610.09760
代码: https://github.com/rithram/kmeans-trf
领域: 优化/理论
关键词: 上下文学习, 聚类算法, Transformer表达力, 算法学习, 稀疏注意力平滑

一句话总结¶

本文提出了一种紧凑型二进制词元嵌入 Transformer 架构，将执行 Lloyd 聚类算法的嵌入维度从 $d+k$ 压缩至 $d+\lceil\log_2 k\rceil$，并系统性建立了端到端平滑替代目标下的优化收敛性与泛化理论。

研究背景与动机¶

近年来，利用注意力架构进行上下文学习（In-Context Learning, ICL）的研究逐步从经验性模式匹配走向形式化算法执行，证明了 Transformer 的前向传递过程能够模拟梯度下降、线性回归以及预条件优化等经典学习算法。在此背景下，Clarkson 等人（ICML 2026）首次确立了 Transformer 在欧氏空间离散聚类中的算法表达力，证明其能够精确模拟用于 $k$-means 聚类的 Lloyd 算法。然而，该基线方案依赖 one-hot（OH）聚类索引嵌入，强制要求模型嵌入维度达到 $d_{\mathsf{emb}} = d+k$，导致注意力投影矩阵参数量达 $\mathcal{O}((d+k)^2)$。当簇数 $k$ 较大时，模型参数规模急剧膨胀，计算吞吐与显存开销构成了严重的工程与理论瓶颈。

更为根本的矛盾在于离散聚类目标与端到端梯度学习之间的不可微鸿沟。传统的上下文学习大多聚焦于均方误差或交叉熵等天然光滑可微的目标函数，而 $k$-means 目标函数涉及离散硬分配（$\min$ 算子），天然不可微且属于 NP-hard 组合优化问题。要让 Transformer 在无人工显式编码参数的条件下，通过有限任务分布以随机梯度下降（SGD）自主学习聚类算法，必须引入平滑松弛；然而，松弛代理函数的凸正则化类型、平滑温度超参以及前述紧凑嵌入方式，如何联合制约优化曲面的 Lipschitz 常数、经验收敛速率乃至分布内泛化误差，此前缺乏系统的理论刻画与实证闭环。

针对上述瓶颈，本文从紧凑表征与平滑优化双重视角切入。核心 idea：利用 $\lceil\log_2 k\rceil$ 位二进制紧凑编码重构聚类索引嵌入以大幅压缩注意力参数量，并结合基于强凸正则化（Sparsemax / Softmax）的可微平滑代理目标，从理论与实证两端系统建立 Transformer 自主学习聚类算法的收敛速度与泛化界。

方法详解¶

整体框架¶

本文构建了一套用于单步聚类并可循环迭代执行的 Transformer 算法架构。对于包含 $d$ 维空间中 $n$ 个待聚类数据点 $\bm{X} \in \mathbb{R}^{d \times n}$ 以及 $k$ 个初始聚类中心 $\bm{C} \in \mathbb{R}^{d \times k}$ 的任务实例，模型首先通过特定的词元嵌入映射将输入升维至潜空间。随后，数据词元与中心词元通过交替的交叉注意力（Cross-Attention）与自注意力（Self-Attention）层进行特征更新，最后截取输出张量的前 $d$ 维作为更新后的聚类中心 $\bar{\bm{C}}(\Theta)$。整个流程在单步前向传播中完成单次 Lloyd 质心迭代更新，并在评估阶段支持重嵌入多步串联。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入点集 X 与初始中心 C"] --> Emb["二进制词元嵌入<br/>维度压缩至 d+ceil(log2 k)"]
    Emb --> Layer["注意力交互更新层<br/>交叉注意力与自注意力残差网络"]
    Layer --> Smooth["可微平滑代理目标<br/>L2/熵正则化 min 算子"]
    Smooth --> Out["输出精炼聚类中心 C_new"]

关键设计¶

1. 二进制词元嵌入：压缩嵌入维度并降低参数依赖 为消除 Clarkson 等人（2026）方案对簇数 $k$ 的二次方参数依赖 $\Omega(k^2)$，本文设计了基于二进制编码的嵌入方案（BN）。对于 $k$ 个簇中心，不再使用 $k$ 维独热向量 $\bm{I}_k$，而是采用包含 $\lceil\log_2 k\rceil$ 个比特位的二进制矩阵 $\bm{B}_k \in \{0, 1\}^{\lceil\log_2 k\rceil \times k}$，其中第 $i$ 列 $\bm{b}_i$ 为整数 $(i-1)$ 的二进制表征。初始质心嵌入定义为 $\bar{\bm{C}}^{(0)} = [\bm{C}; \bm{B}_k]$，数据点嵌入定义为 $\bar{\bm{X}}^{(0)} = [\bm{X}; \mathbf{0}_{\lceil\log_2 k\rceil \times n}]$。这一改进将整体嵌入维度从 $d_{\mathsf{emb}} = d+k$ 削减至 $d_{\mathsf{emb}} = d+\lceil\log_2 k\rceil$，使注意力的查询、键、值投影矩阵参数量直接降低至 $\mathcal{O}((d+\log_2 k)^2)$，在 $k \approx d$ 时参数量削减高达 75%。

2. 负欧氏距离注意力与质心更新机制：精确复刻 Lloyd 迭代 模型层内交互采用基于负平方欧氏距离的 Softmax 注意力操作 $\mathcal{A}_2^{\gamma}$： $$\mathcal{A}_2^{\gamma}(\bm{Z}, \bm{Z}'; (\bm{Q}, \bm{K}, \bm{V})) = \bm{V} \bm{Z}' \, \mathrm{softmax}\left(-\frac{\gamma}{\sqrt{d_{\mathsf{emb}}}} \, \mathrm{sqdist}(\bm{Q}\bm{Z}, \bm{K}\bm{Z}')\right)$$ 在更新数据点分配时，点词元通过查询与中心的距离更新其二进制分配段 $\bm{Y}$；在更新质心时，中心词元通过查询匹配属于该簇的所有数据点并执行平均池化。理论证明当逆温度 $\gamma \to \infty$ 时，该结构完全退化为硬分配与质心重计算，等价于精确执行一次标准 Lloyd 算法（定理 2.1）。在实际可微分训练中，固定采用有限逆温度 $\gamma = 1$。

3. 正则化平滑代理目标：打通不可微聚类的端到端梯度优化 标准离散 $k$-means 目标函数 $\mathcal{L}(\bm{X}, \bar{\bm{C}}(\Theta)) = \sum_{i=1}^n \min_{j \in [k]} \|\bm{X}_{:,i} - \bar{\bm{C}}(\Theta)_{:,j}\|^2$ 包含不可微的硬 $\min$ 算子。为此，引入单纯形概率分布 $\bm{\lambda} \in \Delta^k$ 上的强凸正则化项 $\Omega(\bm{\lambda})$ 及惩罚系数 $\tau > 0$，构建凸正则化变分下界代理 $\tilde{L}_{\Omega}^{\tau}(\bm{x}, \bm{C}) = \langle P_{\Omega}(\bm{r}/\tau), \bm{r} \rangle$（其中 $r_j = \|\bm{x} - \bm{C}_{:,j}\|^2$）。对比负香农熵正则化（NE，对应 Softmax 权重，上界间隙随 $\sqrt{2\log k}$ 增长），本文重点引入基于 $L_2$ 正则化的 Sparsemax 算子： $$P_{\text{L2}}(\bm{r}/\tau)_j = [-(r_j/\tau) - \kappa(-\bm{r}/\tau)]_+$$ Sparsemax 能够诱导严格稀疏的分配权重，并在类间间距较大时达到零代理间隙（$C_{\Omega}^{\tau} = 0$），为优化提供更紧凑的上界与更高保真度的梯度流。

理论分析与设计权衡¶

定理 3.1 建立了平滑目标损失 $\tilde{\mathcal{L}}_{\Omega}^{\tau}$ 关于模型参数 $\Theta$ 的 Lipschitz 连续性上界： $$\alpha \sim \mathcal{O}\left(\frac{k n \omega_{\tau}}{\tau} \cdot \max\{1, \gamma^2\} \cdot \sqrt{d_E} \cdot \max\left\{1, \left(\frac{d_E}{\sqrt{d_{\mathsf{emb}}}}\right)^2\right\} \cdot \exp\left(\gamma \frac{d_E}{\sqrt{d_{\mathsf{emb}}}}\right)\right)$$ 其中 $d_E$ 为中心嵌入向量的最大二范数平方。在 OH 编码下 $d_E = d+1, d_{\mathsf{emb}} = d+k$；而在本文的 BN 编码下 $d_E = d_{\mathsf{emb}} = d + \lceil\log_2 k\rceil$。

该定理揭示了深刻的理论与工程权衡： 1. 收敛速度 vs 参数容量：在数据维度 $d$ 较小而 $k$ 较大时，BN 的 $d_E/\sqrt{d_{\mathsf{emb}}}$ 项显著大于 OH，导致其 Lipschitz 常数 $\alpha$ 更大，理论上需要更多优化步数才能达到同等平滑收敛界。 2. 渐近等价性：随着数据自身维度 $d$ 的增大（例如 $d \ge 32$），$d_E$ 与 $d_{\mathsf{emb}}$ 的相对差距迅速被 $d$ 主导，BN 与 OH 之间的收敛速率差距彻底消失。 3. 平滑度 vs 泛化紧度：降低平滑温度 $\tau$ 可以缩小代理间隙 $C_{\Omega}^{\tau}$ 并收紧测试集泛化界，但会导致 Lipschitz 常数以 $1/\tau$ 比例发散并减缓训练收敛。

实验关键数据¶

主实验：计算开销与运行效率对比¶

在固定数据维度 $d=32$、数据点 $n=1024$ 的设定下，比较独热编码（OH）与本文二进制编码（BN）在不同簇数 $k$ 时的计算延迟与显存占用（数据源自论文 Table 1）：

簇数 $k$	增维 $(d_{\mathsf{emb}}-d)$ (OH / BN)	前向耗时 (ms) (OH / BN [加速比])	反向耗时 (ms) (OH / BN [加速比])	峰值显存 (MB) (OH / BN [显存降幅])
6	6 / 3	4.25 / 4.83 [0.88×]	9.53 / 10.66 [0.89×]	419.93 / 416.29 [0.9%]
25	25 / 5	5.24 / 4.99 [1.05×]	11.21 / 11.01 [1.02×]	441.08 / 423.23 [4.0%]
100	100 / 7	9.76 / 4.52 [2.16×]	20.21 / 9.93 [2.03×]	532.96 / 447.04 [16.1%]
400	400 / 9	38.68 / 7.89 [4.90×]	66.13 / 16.47 [4.02×]	1095.60 / 531.02 [51.5%]
1000	1000 / 10	193.07 / 15.90 [12.14×]	308.28 / 32.33 [9.54×]	3930.34 / 1207.75 [69.3%]

消融与分析实验：训练优化配置与泛化行为¶

在由高斯混合分布生成的聚类任务上进行广泛消融（包含不同温度 $\tau$、嵌入模式、正则化方法与多步执行模式）：

实验维度 / 配置	核心设置	验证集相对 Lloyd 损失	关键现象与结论说明
平滑正则项对比 ($d=4, \tau=1$)	NE (Softmax) vs L2 (Sparsemax)	相对损失 $>1.0$ vs $<1.0$	L2 正则化提供更紧凑的上界，能在较大平滑温度下仍实现正向算法学习
温度惩罚敏感度 ($d=4, \tau=0.1$)	$\tau=0.1$ 紧凑松弛	相对损失 $\ll 1.0$	显著超越单步 Lloyd 基线；但更小 $\tau$ 会使得优化初期的梯度平滑度下降
嵌入消融	OH / BN vs NA (无辅助编码)	$<0.95$ vs $>1.05$	缺失聚类标识的 NA 架构完全无法超越 Lloyd 基线，证实额外嵌入维度的必要性
多步执行策略 ($T=20$)	朴素递归迭代 vs 重嵌入串联迭代	目标函数严重发散 vs 单调下降且大幅超越 Lloyd	单步训练模型无法直接端到端多层循环，必须在每步对中间质心重新嵌入
分布族外泛化	Normal 训练 $\to$ Gumbel/Laplace/Lognormal	保持泛化优势，超越 Lloyd	算法具备跨同类轻尾分布的泛化能力；但在厚尾 Cauchy 分布上失效

关键发现¶

计算收益随簇数急剧提升：当 $k \ge 100$ 时，BN 编码展现出压倒性的工程优势；当 $k=1000$ 时，前向推理实现 12.14 倍加速，训练反向传播实现 9.54 倍加速，并节省近 70% 的 GPU 显存。
多步迭代必须引入重新嵌入：单步训练的 Transformer 不能直接通过深层堆叠进行多步前向递归，否则分配段特征会迅速累积漂移；而将每一步视为独立的单步并重新注入二进制嵌入，模型能够在 20 步迭代中持续优化并稳定超越多步 Lloyd 算法。
厚尾分布失效之谜：模型在 Cauchy 分布任务上训练与泛化完全失效，核心原因在于厚尾离群点使得平方欧氏距离极大，破坏了正则化平滑上界的紧致性，导致代理间隙比正态分布扩大整整一个数量级。

亮点与洞察¶

信息编码紧凑化带来结构性减负：利用 $\lceil\log_2 k\rceil$ 位二进制替代 $k$ 维独热向量，打破了以往注意力架构模拟聚类算法时对簇数的平方级参数依赖，证明算法表达力无需依赖过度参数化。
平滑代理与优化几何的理论闭环：不仅提出了工程加速架构，更严谨给出了可微代理优化的 Lipschitz 常数与泛化误差界，清晰阐明了平滑度、嵌入维度与收敛速率之间的内在张力。
跨步与长度泛化潜力：单层训练模型在测试期通过“单步重嵌入串联”方式展现出优异的多步迭代收敛性，并能直接零样本外推至点数规模翻倍（从 512 到 2048 点）的聚类任务。

局限与展望¶

固定维度与簇数限制：当前架构要求为特定的数据维度 $d$ 与簇数 $k$ 单独训练模型，无法直接动态泛化至任意未知簇数 $k$。
厚尾离群点鲁棒性缺失：在 Cauchy 等重尾或高噪声分布下平滑代理界失效，未来需探索截断机制或鲁棒中心目标（如 $k$-medoids / Trimmed $k$-means）的端到端学习。
训练范式局限：模型仅针对单步更新监督训练，尚未探索多步端到端强化学习或循环解卷（unrolling）训练策略。

评分¶

新颖性: ⭐⭐⭐⭐☆ 首次将紧凑二进制编码引入上下文聚类算法执行，并在理论上解决离散目标的平滑学习。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3000 次独立模型训练、多分布泛化探查、多步外推、超参数敏感性与详细显存耗时基准。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰严谨，从构造性表达力证明到优化收敛理论层层递进。
价值: ⭐⭐⭐⭐☆ 为神经架构学习非凸非光滑经典算法提供了极具启发性的理论框架与工程范例。

实验维度 / 配置	核心设置	验证集相对 Lloyd 损失	关键现象与结论说明
平滑正则项对比 (\(d=4, \tau=1\))	NE (Softmax) vs L2 (Sparsemax)	相对损失 \(>1.0\) vs \(<1.0\)	L2 正则化提供更紧凑的上界，能在较大平滑温度下仍实现正向算法学习
温度惩罚敏感度 (\(d=4, \tau=0.1\))	\(\tau=0.1\) 紧凑松弛	相对损失 \(\ll 1.0\)	显著超越单步 Lloyd 基线；但更小 \(\tau\) 会使得优化初期的梯度平滑度下降
嵌入消融	OH / BN vs NA (无辅助编码)	\(<0.95\) vs \(>1.05\)	缺失聚类标识的 NA 架构完全无法超越 Lloyd 基线，证实额外嵌入维度的必要性
多步执行策略 (\(T=20\))	朴素递归迭代 vs 重嵌入串联迭代	目标函数严重发散 vs 单调下降且大幅超越 Lloyd	单步训练模型无法直接端到端多层循环，必须在每步对中间质心重新嵌入
分布族外泛化	Normal 训练 \(\to\) Gumbel/Laplace/Lognormal	保持泛化优势，超越 Lloyd	算法具备跨同类轻尾分布的泛化能力；但在厚尾 Cauchy 分布上失效