기술 요약: MoE 전문화에서의 기하학적 비대칭성

문제 제기

혼합 전문가 (Mixture-of-Experts, MoE) 아키텍처는 고정된 계산 비용으로 대규모 파라미터 수를 가능하게 함으로써 트랜스포머 모델 확장 표준 메커니즘이 되었으나, 전문가 전문화의 내부 기하학적 구조는 여전히 poorly understood(잘 이해되지 않음) 상태입니다. 기존 분석은 주로 라우팅 통계, 부하 균형, 또는 시스템 수준의 성능에 초점을 맞추고 있습니다. 전문가들이 수행하는 변환 (기능 공간) 과 라우팅된 토큰이 은닉 상태 내에서 어떻게 조직화되는지 (표현 공간) 를 이해하는 데 있어 중요한 격차가 존재합니다.

현재 직관은 MoE 레이어가 입력 분포를 각 전문가가 지역적 근사기로 작용하는 불연속 클러스터로 분할한다고 제안합니다. 그러나 실증적 증거는 전문가들이 종종 겹치는 토큰 영역에서 작동함을 시사하며, 이는 다음과 같은 질문을 제기합니다: 사전 학습된 MoE 트랜스포머 내부에서 라우팅이 유도하는 기하학적 구조는 무엇인가? 구체적으로, 전문가들은 서로 다른 함수를 구현하며, 그들의 표현은 서로 어떻게 관련되어 있는가?

방법론

저자들은 기능 공간과 표현 공간 모두에서 MoE 레이어를 분석하기 위한 통합된 야코비안-PCA-그라스마니안 프레임워크를 도입했습니다. 이 프레임워크는 사전 학습된 모델 (Mistral, Qwen) 과 처음부터 훈련된 통제된 3 층 트랜스포머 모델에 적용되었습니다.

1. 기능 기하학: 전문가 지역 야코비안

라우팅 역학과 독립적으로 개별 전문가의 기능을 특성화하기 위해, 저자들은 전문가 지역 야코비안 ($J_e(x) = \partial f_e(x) / \partial x$) 을 분석합니다.

• 지표: 그들은 데이터셋에 대해 각 전문가의 평균 야코비안 ($\bar{J}_e$) 을 계산하고 벡터화합니다.
• 정렬: 이러한 벡터화된 야코비안 간의 정규화된 프로베니우스 내적 (코사인 유사도) 을 통해 쌍별 기능 정렬을 측정합니다.
• 해석: 낮은 코사인 유사도는 기능적 비상관성을 나타내며, 이는 전문가들이 동일한 함수의 중복된 변형이 아니라 지역적으로 고유한 변환을 구현함을 의미합니다.

2. 표현 기하학: 라우팅된 PCA 와 그라스마니안 거리

라우팅된 표현의 기하학을 특성화하기 위해, 저자들은 라우팅 가중치로 가중치를 둔 은닉 상태 ($h_e(x) = g_e(x)h(x)$) 를 수집합니다.

• PCA: 라우팅된 은닉 상태에 표준 주성분 분석 (PCA) 을 수행하여 주성분 방향과 분산 스펙트럼을 식별합니다.
• 부분공간 거리: 서로 다른 전문가의 상위-$n$개 PCA 부분공간 간의 그라스마니안 측지선 거리를 사용하여 전문가 간의 기하학적 분리를 정량화합니다.
• 해석: 이 지표는 전체 부분공간 간의 다방향적 분리를 포착합니다. 거리가 0 에 가까우면 중복된 기능 역할을 의미하며, 이론적 최대값 ($\frac{\pi}{2}\sqrt{n}$) 에 가까워지면 직교성을 의미합니다.

3. 통제된 라우팅 실험

라우팅 메커니즘의 효과를 훈련 역학으로부터 분리하기 위해, 저자들은 고정된 전문가 파라미터를 가진 통제된 3 층 모델에서 두 가지 영역을 비교합니다:

• Top-k 라우팅: 희소하고 이산적인 선택.
• 완전 소프트 라우팅: 조밀하고 연속적인 혼합.

주요 결과

1. 구조적 비대칭성: 기능적 비상관성 대 표현 중첩

사전 학습된 MoE 트랜스포머 (Mistral-8x7B 및 Qwen1.5-MoE-A2.7B) 전반에 걸쳐 저자들은 일관된 구조적 비대칭성을 관찰합니다:

• 강한 기능적 비상관성: 전문가 간 야코비안 정렬은 일관되게 0 에 가깝습니다 (평균 $\approx$ 0.000–0.068). 이는 전문가들이 중복되지 않으며 고유한 지역적 변환을 구현함을 나타냅니다.
• 부분적 표현 중첩: 기능적 비상관성에도 불구하고, 라우팅된 표현은 고유하지만 부분적으로 중첩된 부분공간을 차지합니다. 그라스마니안 거리는 상당합니다 (분리를 나타냄) 하지만 완전한 직교성을 위한 이론적 최대값보다 현저히 낮습니다.
• 결론: 기능적 비상관성과 표현 중첩이 공존합니다. 전문가들은 불연속 클러스터로 입력 공간을 분할하는 대신, 공유 매니폴드 위의 중첩된 영역에서 지역적으로 비상관된 연산자로 작용합니다.

2. 기하학적 조절 장치로서의 라우팅 희소성

통제된 실험은 라우팅의 날카로움이 이 기하학을 직접 형성함을 보여줍니다:

• Top-k 라우팅: 더 날카로운 기능적 분리 (낮은 야코비안 정렬) 와 더 큰 그라스마니안 거리 (더 큰 부분공간 발산) 를 유도합니다.
• 완전 소프트 라우팅: 높은 야코비안 정렬 (중복 매핑으로의 붕괴를 나타냄) 과 축소된 그라스마니안 거리 (높은 부분공간 중첩) 를 가진 얽힌 전문가 구조를 생성합니다.
• 시사점: 라우팅 희소성은 전문화를 강제하는 핵심 요소입니다.

3. 조밀한 레이어와의 비교

MoE 전문가들은 조밀한 피드포워드 레이어에 비해 분산 스펙트럼에서 더 가파른 고유값 감쇠를 보입니다. MoE 전문가의 첫 번째 주성분은 분산의 훨씬 더 높은 비율을 포착하며, 이는 라우팅이 토큰을 더 낮은 차원의 더 집중된 부분공간으로 집중시킨다는 것을 시사합니다.

4. 모델 간 일관성

관측된 기하학적 패턴 (낮은 야코비안 정렬 + 높지만 최대값 미만인 그라스마니안 분리) 은 서로 다른 모델 (Mistral, Qwen), 레이어 (초기, 중간, 후기), 그리고 데이터셋 전반에 걸쳐 일관되게 나타나며, 이는 특정 훈련 실행의 산물이 아닌 MoE 아키텍처의 견고한 속성임을 시사합니다.

의의 및 주장

본 논문은 현대 트랜스포머 아키텍처의 조건부 계산을 연구하기 위한 일반적인 진단 프레임워크를 제공한다고 주장합니다. 주요 기여점은 다음과 같습니다:

1. 통합 진단 도구: 단순한 라우팅 통계나 활성화 패턴을 넘어 전문가 전문화를 분석하기 위한 체계적인 방법을 제공하며, 기능 공간과 표현 공간을 공동으로 특성화합니다.
2. MoE 의 기하학적 해석: 결과는 MoE 레이어가 공유 표현 매니폴드의 소프트 분할을 구현한다는 관점을 지지합니다. 전문가들은 기능적으로 고유하지만 중첩된 부분 매니폴드에서 작동합니다.
3. 기작적 통찰: 이 연구는 라우팅 날카로움 (희소성) 을 전문화 정도에 대한 직접적인 조절 메커니즘으로 식별하여 아키텍처 설계 선택을 측정 가능한 기하학적 구조와 연결합니다.
4. 모델 설계에 대한 시사점: 저자들은 표현 유사성 (예: 활성화 중첩) 만을 기반으로 한 전문가 병합 또는 가지치기 전략은 결함이 있을 수 있다고 제안합니다. 전문가들은 고유한 기능을 수행하면서도 중첩된 표현을 가질 수 있으므로, 진정한 중복성을 올바르게 식별하기 위해 기능적 지표(야코비안 정렬 등) 가 필요할 수 있습니다.

본 논문은 분석이 국소적 (1 차 야코비안) 이며 선형 (PCA) 이고, 훈련 중의 전역 비선형 행동이나 기하학의 진화를 포착하지는 않는다는 점을 인정하며 겸손한 범위를 유지합니다. 그러나 이는 MoE 모델이 계산을 어떻게 조직화하는지에 대한 기초적인 기하학적 관점을 확립합니다.