기술 요약: DTOP-p MoE

문제 정의

희소 혼합 전문가(Sparse Mixture-of-Experts, MoE) 아키텍처는 파라미터 수와 연산 비용을 분리함으로써 파운데이션 모델을 효율적으로 확장하는 데 매우 중요합니다. 그러나 표준 Top-k 라우팅은 토큰의 난이도나 레이어별 요구 사항에 관계없이 고정된 수의 전문가를 선택하는 경직된 희소성 패턴을 강요합니다. 이는 입력 복잡성의 내재적 분산을 무시하고 적응형 자원 할당을 제한합니다.

반면 Top-p 라우팅(nucleus sampling)은 누적 확률이 임계값 $p$에 도달할 때까지 전문가를 선택함으로써 더 적응적인 대안을 제공하지만, 기존 구현체들은 두 가지 결정적인 한계점을 가집니다:

1. 고정된 글로벌 임계값: 나이브한(naive) Top-p 구현은 진화하는 훈련 역학에 적응하지 못하는 정적인 글로벌 확률 임계값을 사용합니다.
2. 제어되지 않는 희소성: 이러한 방식은 예측 불가능한 연산 비용을 초래합니다. 활성화되는 전문가의 수가 훈련 중에 크게 변동하므로, 엄격한 사전 학습 연산 예산을 준수하기 어렵습니다. 또한, 고정된 임계값은 하이퍼파라미터 선택에 매우 민감하여, 막대한 비용이 드는 광범리한 튜닝을 요구하는 경우가 많습니다.

방법론: DTOP-p MoE

저자들은 Top-p의 적응성과 대규모 사전 학습의 엄격한 계산 제약을 조화시키기 위해 설계된, 희소성 제어가 가능한 동적 라우팅 메커니즘인 DTOP-p를 제안합니다. 이 방법은 두 가지 주요 구성 요소로 이루어집니다:

1. 비례-적분(Proportional-Integral, PI) 컨트롤러

임계값 선택의 비미분성과 명시적 자원 제약의 부재를 해결하기 위해, DTOP-p는 고전 제어 이론의 피드백 제어 루프를 채택합니다.

• 메커니즘: 시스템은 목표 활성화 전문가 수($T$)를 설정값(setpoint)으로, 글로벌 확률 임계값($p$)을 제어 변수로 취급합니다.
• 업데이트 규칙: 컨트롤러는 현재 배치에서의 토큰당 평균 활성화 전문가 수($a_t$)를 모니터링합니다. 이는 희소성 오차 $e_t = (T - a_t)/N$를 계산하고, 이산 PI 제어 법칙을 사용하여 다음 배치를 위한 임계값($p_{t+1}$)을 업데이트합니다:
  $$p_{t+1} = p_0 + K_{pro} \cdot e_t + K_{int} \cdot \sum_{i=1}^{t} e_i$$
• 효과: 이는 양의 오차(불충분한 활성화)가 발생하면 $p$를 높여 더 많은 전문가를 선택하게 하고, 음의 오차가 발생하면 $p$를 낮추는 음의 피드백 루프를 생성합니다. 이를 통해 모델은 로짓 엔트로피의 변화와 관계없이 목표 희소성 수준으로 수렴합니다.

2. 동적 라우팅 정규화(Dynamic Routing Normalization, DRN)

모든 레이어에 단일 글로벌 임계값을 적용하는 것은 레이어마다 계산 요구 사항이 다르다는(예: 일반적 특징을 위한 얕은 레이어 vs 전문적 추론을 위한 깊은 레이어) 연구 결과와 상충하는, 균일한 전문가 신뢰도 분포를 가정하는 것입니다.

• 메커니즘: 로짓(logits) 상에서 직접 라우팅하는 대신, DTOP-p는 학습 가능한 스칼라 $\theta_l$을 사용하여 각 레이어 $l$ 내에서 로짓을 정규화합니다:
  $$P(x) = \text{softmax}\left(\theta_l \cdot \frac{z - \mu(z)}{\sigma(z)}\right)$$
  여기서 $z$는 원시 로짓이며, $\mu(z)$와 $\sigma(z)$는 각각의 평균과 표준편차입니다.
• 효과: 이는 글로벌 임계값을 국소적 통계적 특성으로부터 분리합니다. 더 큰 $\theta_l$은 분포를 날카롭게 만들어(더 적은 전문가 활성화) 적은 전문가를 선택하게 하고, 더 작은 $\theta_l$은 분포를 평탄하게 만들어(더 많은 전문가 활성화) 더 많은 전문가를 선택하게 합니다. 이를 통해 모델은 글로벌 계산 예산을 유지하면서도 레이어별로 구별되는 희소성 패턴을 학습할 수 있습니다.

레이어별 변형(Layer-wise Variant)

엄격한 레이어별 예측 가능성이 필요한 배포 시나리오를 위해, 저자들은 글로벌 교차 레이어 예산 공유를 대신하여 근사적으로 결정론적인 레이어별 추적을 수행하는, 각 레이어가 자체 PI 컨트롤러와 임계값을 유지하는 레이어별 변형 모델도 제안합니다.

주요 기여

1. 고정 임계값 Top-p 분석: 본 논문은 고정 임계값 Top-p MoE가 Top-k에 비해 미미한 이점만을 제공하며, 임계값 $p$에 대한 과도한 민감도와 불안정한 희소성 문제를 겪어 제어 메커니즘 없이는 대규모 사전 학습에 부적합함을 입증합니다.
2. DTOP-p 프레임워크: 희소성 제어를 위한 PI 컨트롤러와 레이어별 적응성을 위한 동적 라우팅 정규화를 통합한 동적 Top-p MoE를 제안합니다.
3. 포괄적인 실증적 조사: 자연어 처리(NLP) 및 컴퓨터 비전(CV) 도메인 전반에 걸친 광범위한 검증을 통해, DTOP-p가 Top-k의 평균 FLOPs를 충족하면서도 Top-k 및 고정 Top-p 베이스라인보다 우수한 성능을 보임을 입증합니다.

실험 결과

저자들은 DCLM-Baseline 데이터셋(NLP)과 고품질 내부 이미지 데이터셋(CV)을 사용하여 LLM 및 확산 트랜스포머(DiT)에서 DTOP-p를 평가했습니다.

• 성능: NLP 벤치마크(MMLU, HellaSwag, ARC를 포함한 13개 데이터셋)에서 DTOP-p는 목표 활성화 전문가 수(예: 8개)를 엄격히 준수하면서 Top-k MoE 베이스라인 대비 평균 **1.9%**의 성능 향상을 달달성했습니다.
• 안정성: 활성화되는 전문가 수의 변동이 심한 고정 Top-p와 달리, DTOP-p는 목표 희소성에 빠르게 수렴하며(첫 1B 토큰 이내), 낮은 표준 편차($\approx 1$)를 유지합니다.
• 스케일링 특성: DTOP-p는 다음과 같은 영역에서 견고한 스케일링을 보여주었습니다:
  • 전문가 세분화(Expert Granularity): 세분화 정도가 높아질수록(예: 32개에서 128개 전문가로) 성능 이득이 확대되었습니다.
  • 전문가 용량(Expert Capacity): 총 전문가 용량이 증가함에 따라 성능이 단조롭게 향상되었습니다.
  • 모델 및 데이터셋 크기: 모델 크기(활성 파라미터 최대 2.4B) 및 데이터셋 크기(최대 300B 토큰)가 증가함에 따라 베이스라인 대비 우위를 유지했습니다.
• 레이어별 분석: 활성화 분포 분석을 통해 DTOP-p가 계층적 희소성 패턴을 학습함을 확인했습니다. 즉, 얕은 레이어에서는 더 적은 전문가를 활성화하고 깊은 레이어에서는 더 많은 전문가를 활성화하여, 서로 다른 깊이에서 처리되는 특징의 가변적인 복잡도와 일치하는 모습을 보였습니다.

의의 및 주장

본 논문은 DTOP-p가 파운데이션 모델 사전 학습을 위한 강력하고 효율적인 MoE 프레임워크를 제공한다고 주장합니다. 이 연구의 주요 의의는 동적 적응성과 엄격한 계산 제약을 성공적으로 조화시킨 데 있습니다. 확률 임계값을 제어 이론을 통해 효과적으로 학습 가능하게 만들고 레이어별 유연성을 부여함으로써, DTOP-p는 정의된 FLOP 예산을 초과하지 않으면서 토큰의 난이도와 레이어 깊이에 따라 계산 자원을 할당할 수 있게 합니다.

저자들은 이 연구를 **지속 가능한 그린 AI(Green AI)**를 향한 단계로 포지셔닝하며, 희소 모델의 계산 효율성을 개선합니다. 또한, 이 방법이 효율성을 개선하긴 하지만 훈련 데이터의 내재된 편향을 완화하는 것은 아니므로 엄격한 안전 정렬(safety alignment)과 병행되어야 함을 명시합니다. 본 연구는 고정된 Top-k는 너무 경직되어 있고 제어되지 않은 Top-p는 너무 예측 불가능한 상황에서 파운데이션 모델을 확장하기 위한 실질적인 솔루션으로 제시됩니다.