技術要約：Transformer回路はクラスタリングアルゴリズムを実現できる

問題提起

クラスタリング、特に$k$-meansクラスタリングは、コンピュータサイエンスにおける基本的な組合せ最適化問題であるが、NP困難であることが知られている。ディープラーニングアーキテクチャ、特にTransformerは、生のデータから階層的な特徴を学習することに長けているが、これらが正確かつ離散的なアルゴリズム計算をどの程度実装できるのかについては不明なままである。本論文は、古典的な離散クラスタリングアルゴリズムと現代のニューラルアーキテクチャの間の溝を埋めるため、Transformerの回路メカニズムが、$k$-meansクラスタリングのためのLloydのアルゴリズムの反復ステップを正確にインスタンス化（実体化）できるかどうかを調査する。

手法

$k$-means Transformerアーキテクチャ

著者らは、Lloydのアルゴリズムをインコンテキストで実行するように設計された、特定のエンコーダ・デコーダTransformerアーキテクチャを提案している。このアーキテクチャは、標準的なTransformerコンポーネントである自己注意（self-attention）、クロス注意（cross-attention）、残差接続、およびフィードフォワードネットワーク（FFN）を利用しているが、特定のパラメータ化とアテンションスコアの定義を用いている。

1. 入力埋め込み（Input Embeddings）:

   • エンコーダ: $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$ はワンホット形式のクラスター割り当て（初期値はゼロ）を表す。
   • デコーダ: $k$ 個のクラスター中心 $C \in \mathbb{R}^{d \times k}$ を処理する。各中心 $c_j$ は、$(d+k)$ 次元のベクトル $[c_j^\top, e_j^\top]^\top$ に埋め込まれる。ここで $e_j$ は第 $j$ 特徴ベクトル（ワンホット・インデックス）である。

2. アテンション・メカニズム:

   • ユークリッド距離アテンション（Euclidean Attention）: エンコーダ（点割り当て）の更新において、本論文では負の二乗ユークリッド距離に基づくソフトマックス・アテンションの極限版 $A_{ji} = -\|Kz_j - Qx_i\|^2$ を利用する。温度パラメータ $\gamma \to \infty$ のとき、これはarg-max操作を近似し、点を最も近い中心に割り当てる。
   • ドット積アテンション（Dot-Product Attention）: デコーダ（中心更新）のステップでは、標準的なドット積アテンションを使用して、割り当てられた点の平均を計算する。
   • 残差接続（Residual Connections）: 現在の割り当てまたは中心をゼロにリセットするために使用され、これにより、Lloydのアルゴリズムにおける「割り当て」と「更新」のステップを単一のフォワードパスで実装する。

3. 理論的構成:

   • 定理 2.1: 特定の固定重み（式5で定義）と極限ソフトマックス・アテンション（$\gamma \to \infty$）を持つ$T$層のTransformerが、$T$回のLloydのアルゴリズムを正確に実行することを証明する。モデルサイズは $O(d+k)$ であり、入力サイズ $O(d(n+k))$ よりも小さい。
   • 定理 2.2: 標準的なドット積アテンションTransformerも、トークン埋め込みに二乗ノルムを付加し、中心を更新するために整流二次関数（rectified quadratic activation）を持つフィードフォワードネットワーク（FFN）を使用すれば、$k$-meansを正確に実行できることを示す。

学習実験

著者らは、Gaussian混合分布からのサンプリングを用いたクラスタリングタスクの分布に対し、$k$-means目的関数の平滑化された上界を最小化することで、単層Transformerを訓練した。訓練には、理論的な極限ケースではなく、標準的なソフトマックス（$\gamma=1$）を使用している。訓練されたモデルは、サンプルサイズ、次元数、および初期化戦略（ランダム vs $k$-means++）を変化させた未知のタスクに対して評価される。

主な貢献

1. Lloydのアルゴリズムの正確な実装: 本論文は、特定のパラメータ選択を持つTransformerアーキテクチャが、$k$-meansクラスタリングのためのLloydのアルゴリズムを正確に実装するという理論的証明と経験的検証を提供している。これは、離散的なアルゴリズムステップと連続的なニューラル回路メカニズムの間の溝を埋めるものである。
2. 学習可能なクラスタリングアルゴリズム: クラスタリングタスクの分布に対して訓練される際、このTransformerアーキテクチャは、標準的なLloydのアルゴリズムよりも優れたクラスタリング品質（より低い$k$-means目的関数）を持つ、汎用的なクラスタリングアルゴリズムを学習する。これは標準的なソフトマックス・アテンション（$\gamma=1$）を使用している場合でも成立する。
3. 多様なアルゴリズムのためのアーキテクチャ変種: 特定のTransformer回路メカニズム（アテンションの種類、正規化、活性化関数など）を修正することで、既存および新規の様々なクラスタリングアルゴリズムをインスタンス化できることを示している。
   • Soft $k$-means: 極限ソフトマックスを標準的なソフトマックスと線形アテンションに置き換えることで実現。
   • Spherical $k$-means: ドット積アテンションとトークンごとの正規化（例：RMSNorm）を使用することで実現。
   • Trimmed $k$-means: ソフトマックスをノルム最大（norm-max）アテンション（正則化されたarg-max）に置き換えることで実現。
   • Robust $k$-means: スパース・ソフトマックス・アテンションを使用することで実現。
   • Randomized $k$-medoids: Gumbel-softmaxアテンションを使用して、実際のデータ点を中心として選択することで実現。

結果

• 理論的等価性: 数値検証により、大きく有限な温度（$\gamma = 10^4$）において、Transformerのクラスター中心と最終的な割り当ての軌跡が、10回のイテレーションにわたってLloydのアルゴリズムと正確に一致することが確認された。
• 性能の優位性: 学習実験において、訓練された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のスケルトン（骨格）」として捉えることで、標準的なコンポーネント（アテンション、正規化）を修正することで、堅牢（robust）なバージョンやランダム化されたバージョンを含む、多様なクラスタリングの変種を生み出す枠組みを提供している。

著者らは制限事項についても述べている。具体的には、彼らの理論的結果はバックプロパゲーションが困難な極限ソフトマックス・アテンションに依存しているが、標準的なソフトマックスでもモデルが機能することを示している。また、モデルは $n$（点の数）の増大には汎化できるが、現在のところ、固定された $k$ で訓練された場合、より大きな $k$（クラスター数）への汎化はサポートしていない。本研究は、古典的な手法のニューラルな実装に関するさらなる探求のための基礎となるものである。