기술 요약: MESH – Mixture-of-Experts 학습을 위한 메모리 효율적 Sinkhorn 최적화

1. 문제 정의

본 논문은 메모리 효율적인 Mixture-of-Experts (MoE) 언어 모델 최적화 문제를 다룹다. AdamW는 시간적 1차 모멘트 평활화(temporal first-moment smoothing)와 좌표별 적응형 스케일링(coordinate-wise adaptive scaling)의 결적인 조합 덕분에 사전 학습(pretraining)의 표준으로 사용되지만, 파라미터당 두 개의 전체 상태 텐서를 저장해야 하므로 높은 메모리 비용을 발생시킨다.

SAGE/Sinkhorn 하이브리드와 같은 Sinkhorn 스타일의 행렬 최적화 기법은 밀집 행렬(dense matrices)에서 모멘트 상태를 제거함으로써 이러한 비용을 줄이는 메모리 효율적인 대안이다. 그러나 저자는 이러한 상태가 없는(stateless) 방식들을 라우팅된 MoE 전문가(routed MoE experts)에 직접 적용하는 것은 실패한다는 것을 입증했다. 제어된 110M 파라미터 규모의 "nanowhale" DeepSeek 스타일 MoE 설정에서, 표준 SAGE/Sinkhorn 하이브리드는 옵티마이저 상태를 크게 줄였으나(0.883GB에서 0.331GB로), 평가 손실(evaluation loss)에서 심각한 성능 저하를 보였다 (AdamW 베이스라인인 3.58–3.64 대비 3.8265).

핵심적인 실패 원인은 라우팅된 전문가 특유의 **순서 문제(ordering problem)**로 식별되었다:

• 라우팅된 그래디언트는 조건부적이다: 밀집 MLP 그래디언트와 달리, 라우팅된 전문가 그래디언트는 라우터의 결정과 게이트 가중치에 기반한 조건부 샘플이며, 이는 시간적으로 변한다.
• Sinkhorn의 비선형성: Sinkhorn 정규화는 비선형 연산이다. 순간적이고 변동성이 큰 라우팅 유도 그래디언트($Sinkhorn(G_t)$)에 이를 적용하는 것은 시간적으로 평균화된 신호($Sinkhorn(E[G_t])$)에 적용하는 것과 다르게 동작한다.
• 노이즈 증폭: 상태가 없는(stateless) Sinkhorn은 매 단계마다 라우팅으로 인한 노이즈를 정규화하지만, 상태가 없는 방식의 실패 원인은 행렬 정규화가 일어나기 전의 시간적 평활화(temporal smoothing)가 결여되어 있다는 점에 있다.

2. 방법론: MESH 및 MESH-B

저자는 MoE 아키텍처를 위해 특별히 설계된 하이브리드 옵티마이저인 MESH(Memory-Efficient Sinkhorn for Experts)를 제안한다.

핵심 메커니즘: 숨겨진 모멘텀 (Hidden Momentum)

MESH는 MoE 전문가를 위한 stateless Sinkhorn 업데이트를 숨겨진 모멘텀 접근 방식으로 대체한다. 명시적인 1차 모멘트 텐서를 옵티마이저 상태에 저장하는 대신(이는 메모리 절감 목표를 저해함), MESH는 그래디언트 버퍼 라이프사이클(gradient-buffer lifecycle) 내에서 1차 모멘트를 유지한다:

1. 누적 (Accumulation): 정규화되지 않은 전문가 모멘텀 누적기가 유지된다: $H_{e,t} = \beta_1 H_{e,t-1} + G_{e,t}$.
2. 버퍼 라이프사이클: 누적기 $H_{e,t}$는 옵티마이저 상태 딕셔너리가 아닌 그래디언트 버퍼에 저장된다.
3. 업데이트 사이클: 역전파(backward pass) 전에, 버퍼가 그래디언트로 재부착된다 ($p.grad \leftarrow \beta_1 H_{e,t}$). 역전파는 현재의 그래디언트를 더하며, 결과적으로 다음 단계를 위한 누적기를 업데이트한다.
4. 정규화: Sinkhorn 업데이트는 평활화된 신호에 대해 계산된다: $U_{e,t} = Sinkhorn(H_{e,t})$.

이 접근 방식은 옵티마이저를 위한 명시적인 텐서를 저장하지 않고도 라우팅된 전문가에 필요한 시간적 1차 모멘트 신호를 복구한다.

선택적 확장: MESH-B

또한 본 논문은 Sinkhorn 방향에 거친 블록/뉴런 역-RMS 승수(coarse block/neuron inverse-RMS multiplier)를 추가하는 MESH-B를 평가한다. 이는 블록 수준의 프리컨디셔너(preconditioner) 역할을 한다. 구현의 단순성을 유지하기 위해 블록 통계는 별도의 로우(raw) 그래디언트 스트림이 아닌 평활화된 소스($H_{e,t}$)로부터 계산된다.

역할 할당 전략 (Role Assignment Strategy)

MESH는 서로 다른 파라미터 클래스에 대해 하이브리드 전략을 유지한다:

• 임베딩 / LM Head: SAGE (경량 상태 업데이트)가 처리한다.
• 비전문가 밀집 행렬 (Non-Expert Dense Matrices): stateless Sinkhorn이 처리한다.
• 라우팅된/공유 전문가 (Routed/Shared Experts): MESH (숨겨진 모멘텀 Sinkhorn)가 처리한다.
• Norms / Biases: SAGE가 처리한다.

3. 주요 기여

1. 실패 지점 식별: 본 논문은 라우팅된 MoE 전문가 행렬이 직접적인 SAGE/Sinkhorn 최적화를 사전 학습에 적용할 때 발생하는 특정 실패 지점임을 식별하였으며, 이를 밀집 행렬이나 임베딩과 구별하였다.
2. MESH 알고리즘: 그래디언트 버퍼 라이프사이클을 활용하여 명시적인 옵티마이저 상태를 추가하지 않고도 전문가를 위한 1차 모멘트 신호를 복구하는 숨겨진 모멘텀 Sinkhorn 업데이트를 도입하였다.
3. 절제 실험 분석 (Ablation Analysis): 2차 모멘트 전용, 부호/모멘텀 전용, 신뢰도 전용, 스케일 전용 개입이 성능 격차를 설명하지 못함을 보여주는 체계적인 부정적 절제 실험을 수행하였다. 주요 인과적 요소는 행렬 정규화 전의 시간적 평활화로 확인되었다.
4. MESH-B 평가: 블록 프리컨디셔닝된 변형 모델을 평가하여, 특정 시드에서는 메모리-품질 프런티어를 개선할 수 있지만, MESH가 제공하는 시간적 평활화와 비교했을 때 보편적으로 필수적인 요소는 아님을 보여주었다.

4. 실험 결과

실험은 FineWeb-Edu 데이터셋에서 5,000 스텝 동안 훈련된 110M 파라미터 규모의 DeepSeek 스타일 MoE 모델을 대상으로 수행되었다.

• 메모리 효율성:
  • AdamW: 0.883GB 옵티마이저 상태.
  • MESH / MESH-B: 0.331GB 옵티마이저 상태 (AdamW 대비 62.5% 감소).
  • 피크 CUDA 할당 (Peak CUDA Allocation): 약 12.6% 감소 (~4.37GB에서 ~3.82GB로).
• 성능 (평가 손실):
  • AdamW 베이스라인: ~3.5877 (시드 100/512 평균).
  • Stateless SAGE/Sinkhorn: 3.8265 (심각한 저하).
  • MESH: 3.6430 (AdamW에 근접하지만 약간 높음).
  • MESH-B: 3.6365 (MESH보다 약간 우수하지만 여전히 AdamW에 뒤처짐).
• 진단 결과 (Seed 42):
  • 명시적으로 1차 모멘트를 저장하는 풀 스테이트 변형 모델(MESH-FULL, MESH-B-FULL)은 AdamW와 유사한 품질(예: V(H)가 3.6102에 도달한 MESH-B-FULL)을 회복하면서도 AdamW의 약 절반 수준의 옵티마이저 상태를 가졌다.
  • 이는 시간적 평활화가 핵심적인 결핍 요소였음을 확인시켜 주며, 숨겨진 모멘텀 구현이 명시적인 상태 비용 없이 이를 성공적으로 근사했음을 입증한다.

5. 의의 및 주장

본 논문은 MESH가 MoE 훈련을 위한 메모리-품질 프런티어에서 실질적으로 더 낮은 상태 비용 지점을 제공한다고 겸허히 주장한다.

• 주요 결론: MoE 전문가는 Sinkhorn 스타일의 옵티마로를 효과적으로 작동시키기 위해 행렬 정규화 전의 시간적 평활화가 필요하다. 이 평활화가 AdamW와의 성능 격차를 좁히는 지배적인 인과적 요소이다.
• 부차적 결론: 블록/뉴런 프리컨디셔닝(MESH-B)은 특정 설정에서 결과를 개선할 수 있지만, AdamW와 비교했을 때 보편적으로 필수적인 구성 요소로 확립되지는 않았다.
• 한계점: 저자는 MESH/MESH-B와 AdamW 사이의 잔여 성능 격차를 인정한다. 저자는 이것이 비전문가 역할 할당(특히 어휘 행렬 및 비전문가 밀집 행렬) 때문이라고 가설을 세웠는데, 이들은 본 설정에서 SAGE 또는 stateless Sinkhorn을 사용하고 있기 때문이다.
• 범위: 본 연구는 110M 아키텍처와 5,000 스텝 실행에 초점을 맞춘 의도적으로 좁은 범위의 연구이다. 저자는 MESH가 모든 지표에서 AdamW를 능가한다고 주장하는 것이 아니라, stateless 방식에서 손실된 대부분의 품질을 회로하면서 메모리 오버헤드를 획기적으로 줄일 수 있음을 보여주고자 한다.

요약하자면, MESH는 MoE 훈련에서 메모리 효율적인 옵티마이저의 실패가 MoE 구조 자체의 내재적 결함이 아니라, 라우팅된 그래디언트에 대한 시간적 평활화의 부재 때문임을 입증한다. "숨겨진 모멘텀"을 구현함으로써 MESH는 이 순서 문제를 해결하고, 메모리 효율적인 MoE 사전 학습을 위한 실행 가능한 경로를 제시한다.