기술 요약: 트랜스포머 회로가 클러스터링 알고리즘을 구현할 수 있음

문제 정의

클러스터링, 특히 $k$-means 클러스터링은 컴퓨터 과학에서 근본적인 조합 최적화 문제이지만, NP-난해(NP-hard)인 것으로 알려져 있다. 딥러닝 아키텍처, 특히 트랜스포머는 원시 데이터로부터 계층적 특징을 학습하는 데 탁월하지만, 이들이 정확하고 이산적인 알고리즘 계산을 어느 정도까지 구현할 수 있는지는 불분나하다. 본 논문은 고전적인 이산 클러스터링 알고리즘과 현대적인 신경망 아키텍처 사이의 간극을 메우기 위해, 트랜스포머 회로 메커니즘이 $k$-means 클러스터링을 위한 로이드(Lloyd's) 알고리즘의 반복 단계를 정확하게 인스턴스화할 수 있는지 조사한다.

방법론

$k$-means 트랜스포머 아키텍처

저자들은 로이드 알고리즘을 인컨텍스트(in-context)로 실행하도록 설계된 특정 인코더-디코더 트랜스포머 아키텍처를 제안한다. 이 아키텍처는 표준 트랜스포머 구성 요소인 셀프 어텐션(self-attention), 크로스 어텐션(cross-attention), 잔차 연결(residual connections), 피드포워드 네트워크(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$는 원-핫 클러스터 할당(초기값은 0)을 나타낸다.
   • 디코더: $k$개의 클러스터 중심 $C \in \mathbb{R}^{d \times k}$를 처리한다. 각 중심 $c_j$는 $(d+k)$ 차원 벡터 $[c_j^\top, e_j^\top]^\top$로 임베딩된다. 여기서 $e_j$는 $j$번째 특성 벡터(원-핫 인덱스)이다.

2. 어텐션 메커니즘 (Attention Mechanisms):

   • 유클리드 어텐션 (Euclidean Attention): 인코더(포인트 할당) 업데이트를 위해, 본 논문은 음의 제곱 유클리드 거리 기반의 소프트맥스 어텐션의 극한 버전을 활용한다: $A_{ji} = -\|Kz_j - Qx_i\|^2$. 온도 파라미터 $\gamma \to \infty$로 갈 때, 이는 arg-max 연산을 근사하여 포인트를 가장 가까운 중심에 할당한다.
   • 닷 프로덕트 어텐션 (Dot-Product Attention): 디코더(중심 업데이트) 단계를 위해, 할당된 포인트들의 평균을 계산하는 표준 닷 프로덕트 어텐션이 사용된다.
   • 잔차 연결 (Residual Connections): 크로스 어텐션 업데이트 전에 현재의 할당이나 중심을 0으로 리셋하는 데 사용되며, 이는 로이드 알고리즘의 "할당" 및 "업데이트" 단계를 단일 순전파 과정에서 효과적으로 구현한다.

3. 이론적 구성 (Theoretical Construction):

   • 정리 2.1: 특정 고정 가중치(식 5에 정의됨)와 극한 소프트맥스 어텐션($\gamma \to \infty$)을 갖춘 $T$-레이어 트랜스포머가 $T$번의 로이드 알고리즘 반복을 정확히 수행함을 증명한다. 모델 크기는 $O(d+k)$로, 입력 크기 $O(d(n+k))$보다 작다.
   • 정리 2.2: 표준 닷 프로덕트 어텐션 트랜스포머 또한 토큰 임베딩이 제곱 노름(squared norms)으로 보강되고, 중심을 업데이트하기 위해 멱함수 이차 활성화 함수를 가진 피드포워드 네트워크(FFN)가 사용될 경우 로이드 알고리즘을 정확히 실행할 수 있음을 보여준다.

학습 실험

저자들은 가우시안 혼합 모델(Gaussian mixtures)로부터 샘플링된 클러스터링 작업 분포에 대해, $k$-means 목적 함수의 평활화된 상한(smoothed upper bound)을 최소화하며 단일 레이어 트랜스포머를 학습시킨다. 학습에는 이론적 극한 사례($\gamma \to \infty$) 대신 표준 소프트맥스($\gamma=1$)를 사용한다. 학습된 모델은 샘플 크기, 차원, 초기화 전략(랜덤 vs $k$-means++)이 변하는 미학습 태스크에 대해 평가된다.

주요 기여

1. 로이드 알고리즘의 정확한 구현: 본 논문은 특정 파라미터 선택을 갖춘 트랜스포머 아키텍처가 $k$-means를 위한 로이드 알고리즘을 정확하게 구현한다는 이론적 증명과 경험적 검증을 제공한다. 이는 이산적 알고리즘 단계와 연속적인 신경 회로 메커니즘 사이의 간극을 메운다.
2. 학습 가능한 클러스터링 알고리즘: 클러스터링 작업 분포에 대해 학습될 때, 트랜스포머 아키텍처는 표준 로이드 알고리즘보다 더 나은 클러스터링 품질(더 낮은 $k$-means 목적 함수)을 보이는 일반화된 클러스터링 알고리즘을 학습한다. 이는 표준 소프트맥스 어텐션($\gamma=1$)을 사용하는 경우에도 마찬가지이다.
3. 다양한 알고리즘을 위한 아키텍처 변형: 특정 트랜스포머 회로 메커니즘(예: 어텐션 유형, 정규화, 활성화 함수)을 수정함으로써, 저자들은 다음과 같은 다양한 기존 및 새로운 클러스터링 알고리즘을 인스턴스화할 수 있음을 보여준다:
   • Soft $k$-means: 극한 소프트맥스를 표준 소프트맥스와 선형 어텐션으로 교체하여 달성.
   • Spherical $k$-means: 닷 프로덕트 어텐션과 토큰 단위 정규화(예: RMSNorm)를 사용하여 달성.
   • Trimmed $k$-means: 소프트맥스를 노름-맥스(norm-max) 어텐션(정규화된 arg-max)으로 교체하여 달성.
   • Robust $k$-means: 희소(sparse) 소프트맥스 어텐션을 사용하여 달성.
   • Randomized $k$-medoids: 실제 데이터 포인트를 중심으로 선택하기 위해 Gumbel-softmax 어텐션을 사용하여 달성.

결과

• 이론적 등가성: 수치적 검증을 통해, 크지만 유한한 온도($\gamma = 10^4$)에서 트랜스포머의 클러스터 중심 궤적과 최종 할당이 10회의 반복 동안 로이드 알고리즘과 정확히 일치함을 확인하였다.
• 성능 우위: 학습 실험에서, 학습된 트랜스포머는 동일한 태스크에 대해 로이드 알고리즘($\sim 5.25$)보다 더 낮은 평균 로그-$k$-means 목적 함수($\sim 5.0$)로 수렴한다. 이러한 개선은 트랜스포머가 오직 랜덤 초기화로만 학습되었음에도 불구하고, 샘플 크기(일반화) 및 초기화 방식(예: $k$-means++)이 다른 태스크에서 테스트될 때도 유지된다.
• 임베딩 차원: 절제 연구(Ablation studies)는 토큰 임베딩의 추가적인 $k$ 차원(총 $d+k$)이 매우 중요하다는 것을 보여준다. $d$ 차원만 가진 모델은 로이드의 성능을 따라가지 못하며, 더 높은 목적 함수($\sim 5.35$)로 수렴한다.
• 가중치 분석: 학습된 가중치는 이론적 구성(예: $Q^\top K$ 행렬의 지배적인 대각 성분 형태)과 구조적 유사성을 보이지만, 동시에 뚜렷한 차이점도 보여준다. 이는 학습된 알고리즘이 표준 로이드 알고리즘과 특정 방식(예: 중심 업데이트 방식)에서 다르게 동작함을 시사한다.

의의 및 주장

본 논문은 트랜스포머 회로 메커니즘이 클러스터링을 위한 정확한 알고리즘 루틴을 인스턴스화할 수 있으며, 이는 "해석 가능한 아키텍처 수정을 통해 대안적이고 새로운 클러스터링 알고리즘을 개발하기 위한 일반적인 접근법을 제공한다"고 주장한다.

• 알고리즘 투명성: 본 연구는 트랜스포머가 단순한 블랙박스 함수 근사기가 아니라, 로이드 알고리즘과 같은 명시적이고 잘 정의된 알고리즘을 수행하도록 명시적으로 설계될 수 있음을 입증한다.
• 학습 가능성: 결정적으로, 본 논문은 이러한 수기로 지정된 아키텍처가 이론적으로 가능할 뿐만 아니라, 클래식한 베이스라인을 능가하는 알고리즘을 생성하도록 데이터로부터 학습 가능하다는 것을 보여준다.
• 혁신을 위한 프레임워크: 클러스터링 알고리즘을 "트랜스포머 스켈레톤"으로 간주함으로써, 저자들은 표준 구성 요소(어텐션, 정규화)를 수정하여 견고한(robust) 버전 및 무작위(randomized) 버전을 포함한 다양한 클러스터링 변형을 얻을 수 있는 프레임워크를 제공한다.

저자들은 한계점도 언급하였는데, 구체적으로 그들의 이론적 결과는 역전파(backpropagation)가 어려운 극한 소프트맥스 어텐션에 의존하지만, 모델이 표준 소프트맥스로도 작동함을 보여주었다는 점이다. 또한, 모델이 더 큰 $n$(포인트 수)에는 일반화될 수 있지만, 고정된 $k$로 학습되었을 경우 더 큰 $k$(클러스터 수)로의 일반화는 현재 지원하지 않는다고 밝혔다. 이 연구는 고전적 방법의 신경망 구현에 대한 추가 탐구를 위한 기초로서 자리매김하고 있다.