Mixture-of-Experts as Soft Clustering: A Dual Jacobian-PCA Spectral Geometry Perspective

논문 개요

이 논문은 Mixture-of-Experts (MoE) 아키텍처가 단순히 계산 효율성을 위한 도구를 넘어, 학습된 함수와 내부 표현의 기하학적 구조 (Geometry) 에 어떤 영향을 미치는지 분석합니다. 저자는 MoE 의 라우팅 (Routing) 메커니즘을 입력 도메인의 소프트 클러스터링 (Soft Clustering) 으로 해석하며, 이를 정량화하기 위해 이중 자코비안 (Dual Jacobian) 과 주성분 분석 (PCA) 기반의 스펙트럼 기하학 프로브를 제안합니다.

1. 연구 문제 (Problem)

• 배경: MoE 는 대규모 언어 모델 (LLM) 의 확장성을 위해 널리 사용되지만, 라우팅 메커니즘이 학습된 함수의 국소적 민감도 (Local Sensitivity) 와 표현 공간의 기하학적 구조를 어떻게 변화시키는지는 명확히 규명되지 않았습니다.
• 한계: 기존 연구는 주로 확장 법칙 (Scaling Laws) 이나 효율성에 집중했으며, 라우팅이 함수의 곡률 (Curvature), 국소적 민감도, 그리고 표현의 차원성 (Dimensionality) 에 미치는 기하학적 영향을 체계적으로 분석한 사례는 부족합니다.
• 목표: 작업 성능이나 데이터 규모와 무관하게, 라우팅 메커니즘 그 자체가 함수 공간과 표현 공간의 기하학을 어떻게 재구성하는지 메커니즘을 규명하는 것입니다.

2. 방법론 (Methodology)

저자는 혼란 변수를 제거하기 위해 제어된 실험 환경 (Synthetic Setting) 을 구축하고 두 가지 주요 기하학적 프로브를 결합했습니다.

• 실험 설정:

  • 모델: 완전 연결 신경망 (MLP) 기반의 MoE 아키텍처를 사용하여 정밀한 자코비안 (Jacobian) 계산과 공분산 추정이 가능하도록 함.
  • 비교 대상: 밀집형 (Dense), Top-k 라우팅, 완전 소프트 (Fully-soft) 라우팅 아키텍처를 용량 (Capacity) 을 동일하게 맞춘 상태에서 비교.
  • 데이터: 등방성 가우시안 분포를 가진 합성 데이터 (Synthetic Inputs) 를 사용하여 통계적 안정성 확보.

• 주요 분석 도구:

  1. 자코비안 스펙트럼 프로브 (Jacobian Spectral Probe):
     • 각 전문가 (Expert) 의 국소 자코비안 $J_e(x)$ 에 대한 특이값 (Singular Values) 스펙트럼을 분석.
     • 목적: 함수의 국소적 민감도 (Curvature) 와 이방성 (Anisotropy) 정량화.
     • 지표: 최대 특이값 ($\sigma_1$), 스펙트럼 감쇠율, 누적 에너지 분포.
  2. 가중치 PCA 프로브 (Weighted PCA Probe):
     • 라우팅 가중치 $g_e(x)$ 를 샘플 가중치로 사용하여 숨겨진 상태 (Hidden States) 에 대한 PCA 수행.
     • 목적: 각 전문가가 실제로 처리하는 입력 분포에 기반한 표현 공간의 분산 구조 분석.
     • 지표: 설명 분산 비율, 유효 차원 (Effective Rank), 누적 분산 곡선.
  3. 교차 전문가 정렬 분석:
     • 평균 자코비안 간의 코사인 유사도를 계산하여 전문가 간의 기능적 중첩 (Functional Overlap) 정도를 측정.

3. 주요 결과 (Key Results)

가. 함수 기하학 (Function Geometry): 자코비안 분석

• 민감도 감소: MoE 전문가들의 자코비안은 밀집형 (Dense) 모델보다 더 작은 주특이값과 더 빠른 스펙트럼 감쇠를 보임. 이는 MoE 가 국소적 곡률을 평탄화 (Flattening) 하여 입력에 대한 민감도를 낮춘다는 것을 의미.
• 낮은 정렬 (Low Alignment): 전문가 간 자코비안의 코사인 유사도는 0 에 가깝게 분포. 이는 각 전문가가 공유된 매핑의 스케일 변형이 아닌, 중첩이 적은 고유한 변환 (Low-overlap Transformations) 을 학습함을 시사.
• 에너지 집중: MoE 는 밀집형 모델보다 적은 수의 주된 방향에 민감도 에너지를 집중시킴.

나. 표현 기하학 (Representation Geometry): PCA 분석

• 분산 재분배: 합성 데이터 (Synthetic Data) 환경에서 MoE 전문가들의 표현은 밀집형 모델보다 더 높은 유효 차원 (Effective Rank) 을 가짐. 즉, 분산이 소수의 주성분으로 집중되지 않고 더 많은 방향으로 고르게 분포됨.
• 라우팅 강도에 따른 차이:
  • Top-k (Sharp): 입력이 좁은 영역으로 세분화되어 분산이 더 집중되고 유효 차원이 낮아짐.
  • Fully-soft (Soft): 더 이질적인 입력을 처리하여 분산이 넓게 퍼지고 유효 차원이 높게 유지됨.

다. 자연어 데이터에서의 검증 (WikiText-2)

• 곡률 감소의 보편성: 3 레이어 트랜스포머 모델에서도 자코비안 곡률 감소 경향은 유지됨 (Dense: $\sigma_1 \approx 2.95$ vs MoE: $\sigma_1 \approx 0.44$).
• PCA 스펙트럼의 반전 (Inversion): 합성 데이터와 달리, 자연어 데이터에서는 Dense 모델이 높은 유효 차원을, MoE 전문가가 낮은 유효 차원을 보임.
  • 해석: 자연어는 저차원 곡면 (Low-dimensional Manifold) 에 존재하므로, MoE 라우팅이 토큰을 더 단순한 하위 다양체 (Sub-manifold) 로 분할하여 각 전문가가 국소적으로 더 낮은 차원의 지도 (Chart) 를 학습하게 됨. 이는 소프트 클러스터링 가설을 강력히 지지함.

4. 주요 기여 (Key Contributions)

1. MoE 의 기하학적 프레임워크 정립: MoE 를 함수 공간의 소프트 파티셔닝 (Soft Partitioning) 으로 해석하고, 이를 '전문가 국소 지도 (Expert-local Charts)'의 집합으로 정의.
2. 스펙트럼 비교 분석: 용량이 동일한 조건에서 Dense, Top-k, Fully-soft 라우팅의 기하학적 특성을 체계적으로 비교.
3. 새로운 프로브 제안: 자코비안 특이값 스펙트럼과 가중치 PCA 를 결합하여 라우팅이 국소 민감도와 표현 차원성에 미치는 영향을 정량화하는 방법론 제시.
4. 실제 모델에 대한 가설 도출: 대규모 모델 (70B 클래스) 에 적용 가능한 3 가지 검증 가능한 가설 제시.

5. 의의 및 시사점 (Significance & Implications)

이 연구는 MoE 의 작동 원리를 효율성 관점이 아닌 기하학적 관점에서 재해석하며, 다음과 같은 실용적 통찰을 제공합니다.

• 기하학적 팔꿈치 (Geometric Elbow): 전문가 수 증가에 따른 분산 집중의 '팔꿈치' 지점을 찾아 최적의 전문가 수를 결정할 수 있는 기준 마련.
• 할루시네이션 억제: 자코비안 노름의 감소 (평탄한 매핑) 는 미세한 입력 변화에 대한 민감도를 낮추어, 장기 컨텍스트 생성 시 발생할 수 있는 오류 누적 (할루시네이션) 을 줄일 가능성을 시사.
• 앙상블 신뢰성: Top-k 라우팅이 Fully-soft 보다 전문가 간 자코비안 정렬이 낮아 (직교성 높음) 기능적 다양성이 더 크므로, 앙상블로서의 신뢰성이 높을 것으로 예측.
• 디자인 가이드라인: 라우팅의 날카로움 (Sharpness) 과 표현의 차원성 (Rank) 사이의 균형을 맞추는 것이 대규모 트랜스포머 설계에 중요함을 강조.

결론적으로, 이 논문은 MoE 가 단순히 파라미터를 효율적으로 사용하는 것을 넘어, 함수 공간을 평탄화하고 표현 공간을 적절히 분할 (Soft Clustering) 함으로써 학습된 모델의 기하학적 구조를 근본적으로 변화시킨다는 메커니즘을 규명했습니다.