기술 요약: 혼합 전문가 (MoE) 계층을 활용한 하이퍼파라미터 전이

1. 문제 제기

혼합 전문가 (Mixture-of-Experts, MoE) 계층은 순전파 (forward pass) 동안 학습 가능한 매개변수의 총수와 활성 매개변수를 분리함으로써 현대 신경망을 확장하는 데 중요한 메커니즘이 되었습니다. 그러나 희소 MoE 모델을 학습시키는 것은 하이퍼파라미터 (HP) 선택과 관련하여 상당한 복잡성을 야기합니다. 밀집 (dense) 모델과 달리 MoE 아키텍처는 신중한 조정이 필요한 새로운 학습 가능한 매개변수 (라우터 가중치) 와 새로운 아키텍처 차원 (전문가 수, 전문가 크기) 을 도입합니다.

대규모에서 하이퍼파라미터 (학습률, 초기화 스케일, 가중치 감쇠 등) 를 직접 조정하는 것은 계산적으로 불가능합니다. 작은 모델에서 찾은 최적의 하이퍼파라미터를 더 큰 모델로 외삽할 수 있게 해주는 HP 전이 기법이 밀집 트랜스포머에는 존재하지만, 이러한 방법들은 희소 MoE 계층의 특정 확장 역학에 대해 엄격하게 적용되지 않았습니다. 핵심적인 과제는 매번 규모를 변경할 때마다 재조정하지 않으면서도 학습 안정성과 성능을 보장한 채, 모델 너비, 깊이, 전문가 수, 전문가 크기를 동시에 증가시킬 때 하이퍼파라미터를 어떻게 확장할지 결정하는 것입니다.

2. 방법론

2.1. 제안된 파라미터화

저자들은 MoE 계층을 가진 트랜스포머 모델에 대한 새로운 파라미터화를 제안하며, 이전에 밀집 트랜스포머를 위해 개발된 CompleteP 파라미터화를 MoE 특화 확장 규칙을 포함하도록 확장합니다. 목표는 모델 차원이 확장됨에 따라 원시 하이퍼파라미터 값 (초기화 표준 편차 $\sigma$ 및 학습률 $\eta$) 이 어떻게 변해야 하는지 예측하는 규칙 세트를 정의하여 학습 역학이 일관되게 유지되도록 하는 것입니다.

이 유도 과정은 네트워크 구성 요소 (사전 활성화 및 잔차 업데이트) 가 초기화 시 $O(1)$이 되고 학습 단계당 $\Theta(1)$의 업데이트를 받아야 한다는 최대 업데이트 파라미터화 ($\mu$P) 원리에 의존합니다. 저자들은 이를 MoE 로 확장하여 최대 업데이트 조건이 계층 출력뿐만 아니라 개별 전문가 구성 요소 (믹싱 계수와 전문가 출력) 에 대해서도 성립해야 한다고 요구합니다.

MoE 모듈을 위해 유도된 주요 확장 규칙은 다음과 같습니다:

• 라우터 가중치: 학습률은 $\eta \propto n_{embd}^{-1}$로 스케일링됩니다. 초기화는 $n_{embd}^{-\gamma}$ ($\gamma \geq 0.5$) 로 스케일링됩니다.
• 전문가 편향: 희소성이 고정된 경우, 0 으로 초기화되며 일정한 학습률 $\eta \propto 1$을 가집니다 (전문가 수와 무관).
• 전문가 MLP 가중치:
  • 업-프로젝션 ($W_{up}$): $\sigma_{init} \propto n_{embd}^{-1/2}$, $\eta \propto n_{embd}^{-1}$.
  • 다운-프로젝션 ($W_{down}$): $\sigma_{init} \propto \alpha_{ffn}^{-1} n_{embd}^{-1/2}$, $\eta \propto \alpha_{ffn}^{-1} n_{embd}^{-1}$.
  • 여기서 $\alpha_{ffn}$은 임베딩 차원에 대한 전문가 은닉 크기 배수를 나타냅니다.

중요하게도, 저자들은 총 전문가 수 ($n_{exp}$) 와 활성 전문가 수 ($n_{act}$) 를 확장하는 동안 희소성 비율 $\kappa = n_{act}/n_{exp}$ (활성 전문가의 비율) 을 상수로 고정합니다. 이는 전체 풀을 증가시키는 동안 활성 전문가 수를 고정하는 접근법과 대조됩니다.

2.2. 이론적 근거: 동적 평균장 이론 (DMFT)

이러한 휴리스틱 확장 규칙을 정당화하기 위해 저자들은 동적 평균장 이론 (DMFT) 을 사용합니다. 그들은 고정된 활성화 희소성 $\kappa$를 유지하면서 무한한 너비 ($n_{embd}$), 깊이 ($L$), 전문가 크기 ($n_{hid}$), 그리고 전문가 수 ($n_{exp}$) 의 동시 극한에서 MoE 계층을 가진 잔차 네트워크의 학습 역학을 분석합니다.

분석은 새로운 3 단계 평균장 계층 구조를 드러냅니다:

1. 잔차 스트림 표현: 전문가 출력에 대한 평균장.
2. 전문가 출력: 개별 전문가 뉴런에 대한 평균장.
3. 개별 뉴런: 각 전문가 내부의 평균장.

DMFT 분석은 제안된 파라미터화 하에서 극한 학습 역학이 다음과 같음을 보여줍니다:

• FFN 비율 ($\alpha_{ffn}$) 에 독립적: 결합 확장 극한이 취해진다면, 역학은 임베딩 차원에 대한 전문가의 특정 크기에 의존하지 않습니다.
• 희소성 $\kappa$에만 의존적: 희소성 비율이 고정되어 있다면 역학은 모든 확장 매개변수에서 일관됩니다.
• 스케일 불변성: 네트워크 요약 통계량 (예: 계층별 특징 커널) 의 진화는 규모에 걸쳐 일관되므로, 이론적으로 하이퍼파라미터 전이가 신뢰할 수 있음을 보장합니다.

2.3. 실험 설정

저자들은 FineWeb 및 C4 데이터셋의 디코더 전용 트랜스포머 언어 모델을 사용하여 경험적으로 파라미터화를 검증합니다.

• 베이스 모델: 약 38M 개의 활성화된 매개변수를 가진 작은 모델에서 조정됨.
• 확장: 너비, 깊이, 전문가 수, 전문가 크기를 변화시키며 총 2B 개의 매개변수를 가진 모델까지 확장됨.
• 토큰 예산: 초기 학습 역학을 분리하기 위해 고정된 1B 토큰 (2000 스텝) 의 토큰 예산으로 실험을 수행했으며, 더 긴 시간 범위 (최대 7.5B 토큰) 에 대해서도 실험을 수행했습니다.
• 옵티마이저: 표준 Adam 옵티마이저.
• 로드 밸런싱: 손실 함수에 정규화 항을 추가하지 않고 전문가 편향을 직접 업데이트하여 로드 밸런싱을 장려하는 보조 손실 없는 전략을 사용했습니다.

3. 주요 기여

1. MoE 파라미터화: 이 논문은 희소 MoE 모델로 CompleteP 파라미터화를 확장하여 너비, 깊이, 전문가 수, 전문가 크기에 걸쳐 라우터 가중치, 전문가 편향, 전문가 MLP 가중치에 대한 명시적인 확장 규칙을 제공합니다.
2. DMFT 를 통한 이론적 정당화: 저자들은 DMFT 를 사용하여 파라미터화에 대한 엄격한 이론적 기초를 제공합니다. 그들은 무한한 너비/깊이 극한에서 학습 역학에 대한 명시적인 설명을 유도하여, 역학이 전문가 차원의 특정 확장에 의존하지 않고 희소성에만 의존하는 안정적이고 스케일 불변인 시스템으로 수렴함을 증명합니다.
3. HP 전이의 경험적 검증: 이 연구는 작은 베이스 모델 (38M 활성 매개변수) 에서 식별된 최적의 하이퍼파라미터 (학습률 및 초기화 스케일) 가 다양한 아키텍처 차원에 걸쳐 훨씬 더 큰 모델 (최대 2B 총 매개변수) 로 신뢰할 수 있게 전이됨을 보여줍니다.
4. 아키텍처 통찰: 저자들은 경험적으로 해당 파라미터화 하에서 총 매개변수를 고정하면서 전문가 수를 증가시키는 것이 개별 전문가의 크기를 증가시키는 것보다 더 나은 성능을 낸다는 것을 확인했습니다. 이 발견은 최근 문헌과 일치하지만, 각 규모에서 비싼 하이퍼파라미터 스윕이 필요하지 않게 달성되었습니다.

4. 결과

• 신뢰할 수 있는 HP 전이: 제안된 확장 규칙 하에서 최적의 베이스 학습률 및 초기화 표준 편차는 51M 에서 2B 총 매개변수에 이르는 모델 간에 효과적으로 전이됩니다. 확장된 모델의 손실 곡선은 초기 반복에서 베이스 모델과 수렴한 후 발산합니다 (더 큰 모델이 더 낮은 손실을 달성함).
• 안정성: 파라미터화는 전문가 수를 확장할 때도 균일한 전문가 로드 밸런싱을 포함한 안정적인 학습 역학을 보장합니다. 저자들은 MoE 사전 학습이 상수 스케일 하이퍼파라미터 ($\Theta(1)$로 취급되는 승수) 에 특히 민감하며, 이를 조정하는 것이 안정성에 중요하다고 지적합니다.
• 성능: 작은 모델에서 전이된 제로샷 하이퍼파라미터로 학습된 모델은 활성 매개변수 수를 맞췄을 때 밀집 베이스라인 (예: GPT-2 small/medium) 과 경쟁력 있는 성능을 달성합니다.
• 전문가 수 vs 크기: 실험은 고정된 매개변수 수에서 전문가 크기를 증가시키는 것보다 전문가 수를 증가시키는 것이 더 매개변수 효율적임을 확인합니다. 이 이점은 더 긴 학습 시간 범위 (5B 토큰) 에서도 지속됩니다.
• 희소성 민감성: 이 연구는 하이퍼파라미터 전이가 희소성 비율 $\kappa$가 고정된 경우에만 유효함을 강조합니다. 활성 전문가 수를 고정하면서 전문가 수를 확장하는 것 ($\kappa \to 0$으로 보냄) 은 최적 하이퍼파라미터의 전이성을 깨뜨립니다.

5. 의미 및 주장

이 논문은 MoE 모델을 확장하기 위한 실용적이면서도 엄격한 프레임워크를 제공한다고 주장합니다. 휴리스틱 파라미터화와 DMFT 분석을 결합함으로써 저자들은 다음과 같은 방법을 제시합니다:

• 학습 비용 절감: 작고 저렴한 베이스 모델만 조정하여 대규모 MoE 모델에 대한 최적 하이퍼파라미터를 선택할 수 있게 합니다.
• 안정성 보장: 희소 MoE 학습에서 흔히 발생하는 학습 불안정성 (예: 전문가 붕괴 또는 발산) 을 방지하는 규칙을 제공합니다.
• 아키텍처 설계 지침: 고정된 희소성 하에서 전문가 크기를 증가시키는 것보다 전문가 수를 증가시키는 것이 바람직하다는 이론적 및 경험적 증거를 제공하여 효율적인 대규모 모델 설계에 도움을 줍니다.

저자들은 한계를 인정하며, 현재 작업은 학습률과 초기화에 초점을 맞추고 배치 크기, 가중치 감쇠, LR 스케줄과 같은 다른 하이퍼파라미터는 향후 연구에 맡긴다고 명시합니다. 또한 DMFT 분석이 무한한 너비 극한을 지원하지만, 작은 너비 전이의 행동은 추가적인 이론적 연구가 필요하다고 지적합니다. 더 나아가, 이 논문은 MoE 에 대한 "컴퓨팅 최적" 확장 법칙 (예: Chinchilla 지수) 을 해결한다고 주장하지는 않습니다. 희소성으로 인한 하드웨어 제약으로 인해 MoE 의 FLOP 대 성능 트레이드오프는 밀집 모델과 크게 다르기 때문입니다.

요약하자면, 이 작업은 학습 역학에 대한 새로운 3 단계 평균장 이론에 기반하여 작은 MoE 모델에서 큰 MoE 모델로 하이퍼파라미터를 신뢰할 수 있게 외삽할 수 있는 기초적인 파라미터화를 확립합니다.