기술 요약: MobileMoE: 온디바이스 혼합 전문가 확장

문제 제기

혼합 전문가 (MoE) 아키텍처는 클라우드 서버에서 수십억 개의 파라미터를 가진 대규모 언어 모델 (LLM) 을 확장하는 표준으로 자리 잡았으나, 온디바이스 배포를 위한 수십억 미만 규모에서의 잠재력은 여전히 largely 탐구되지 않은 상태입니다. 기존 온디바이스 LLM(예: MobileLLM, MobileLLM-Pro) 은 압도적으로 밀집 (dense) 구조로, 품질 향상을 위해 활성 파라미터 수와 추론 연산량을 증가시켜야 합니다. 이는 메모리 (DRAM) 와 연산 (배터리 수명, 지연 시간) 에 대한 엄격한 제약을 받는 에지 AI 에 병목 현상을 초래합니다. 최근 스마트폰 DRAM 용량이 증가했음 (예: 12~16GB) 에도 불구하고, 이러한 특정 모바일 제약 하에서 총 모델 크기, 활성 파라미터, 추론 비용 간의 균형을 맞추는 MoE 아키텍처 확장을 위한 원칙 있는 방법론은 아직 확립되지 않았습니다.

방법론

1. 온디바이스 MoE 확장 법칙

저자들은 메모리 및 연산 제약 하에서 아키텍처를 공동 최적화하기 위한 일반화된 온디바이스 MoE 확장 법칙을 수립했습니다. 손실 함수 $L$은 활성 파라미터 ($N_{act}$), 학습 데이터 크기 ($D$), 전문가 수의 단조 변환 ($\hat{E}$), 그리고 아키텍처 선택 ($x$, 전문가 세분성 $g$ 및 공유 전문가 $s$ 포함) 의 함수로 모델링됩니다.

최적화 목표는 다음 제약 조건 하에서 손실을 최소화하는 것입니다:

• 연산 제약: 학습 FLOPs ($F_{train} = 6N_{act}D$) 및 토큰당 추론 FLOPs ($F_{inf} = 2N_{act}$).
• 메모리 제약: 양자화된 가중치 ($b_w$-비트) 와 컨텍스트 길이 $T$에서의 KV 캐시 ($b_{kv}$-비트) 를 포함하는 총 메모리 풋프린트 $M(N_{total}, T)$.

2. 분할 정복을 통한 아키텍처 설계

확장 법칙을 활용하여 저자들은 세 가지 통제된 제거 실험을 통해 온디바이스 MoE 를 위한 최적의 "적정 지점"을 식별했습니다:

• 전문가 수 ($E$): 적절한 희소성이 최적입니다. $E$를 늘리면 손실이 감소하지만 $E=8$을 초과하면 체감 효과가 줄어듭니다. 메모리 제약이 있는 기기를 위해 $E=8$을 실용적인 적정 지점으로 선정했습니다.
• 전문가 세분성 ($g$): 세분화된 전문가 (각 전문가를 $g$개의 서브 전문가로 분할) 는 더 다양한 라우팅 경로를 가능하게 합니다. 연구 결과, $g=8$이 과도한 오버헤드 없이 거친 세분성 전문가보다 상당한 손실 감소를 제공하는 연산 최적의 세분성으로 밝혀졌습니다.
• 공유 전문가 ($s$): 항상 활성화된 공유 전문가 (밀집 경로) 를 포함하면 라우팅된 전문가들을 보완합니다. 고정된 연산 하에서 공유 전문가가 있는 구성 ($s=\checkmark$) 은 라우팅 전용 구성보다 더 낮은 손실을 달성합니다.

결과 아키텍처: 도출된 MobileMoE 계열은 활성 파라미터가 0.3B, 0.5B, 0.9B 이고 총 파라미터가 각각 1.3B, 2.8B, 5.3B 인 세 가지 규모 (S, M, L) 로 구성됩니다. 이 설계는 $E=8$, $g=8$(64 개의 세분화된 라우팅 전문가), 그리고 상위 4 개 라우팅을 사용하는 공유 전문가를 적용합니다.

3. 4 단계 학습 레시피

아키텍처적 이점을 실현하기 위해 오픈소스 데이터셋에서 포괄적인 학습 파이프라인이 적용되었습니다:

1. 사전 학습: 광범위한 언어적 커버리지와 전문가 전문화를 보장하기 위해 웹 중심 혼합으로 약 6T 토큰을 학습합니다. 보조 손실 없는 균형 조정, 라우터 z-loss, 효율성을 위한 그룹화 MLP(GMM) 커널 등의 기법을 포함합니다.
2. 중간 학습: 컨텍스트를 8k 로 확장하고 전문가 전문화를 강화하기 위해 수학, 코드, 지식 등 고품질 도메인별 데이터에 초점을 맞춘 약 500B 토큰을 학습합니다.
3. 지시 미세 조정 (SFT): 구조화된 지시 신호를 보존하기 위해 드롭리스 토큰 디스패칭을 사용하여 다양한 도메인의 약 80M 샘플로 학습합니다.
4. 양자화 인식 학습 (QAT): INT4 가중치 양자화 (그룹 크기 32) 와 INT8 동적 활성화가 적용됩니다. 특히 라우팅 안정성을 유지하기 위해 라우터 가중치는 FP32 로 유지되며, 이는 미미한 메모리 오버헤드만 추가합니다.

4. 효율적인 온디바이스 추론

배포 간극을 해소하기 위해 저자들은 ExecuTorch 에서 사용자 정의 퓨즈드 MoE 커널을 구현했습니다. 이 커널은 다음을 수행합니다:

• 토큰 재배열을 통해 희소 전문가 디스패치를 밀집 그룹화 GEMM 으로 변환합니다.
• 커널 시작 및 양자화 오버헤드를 분산시키기 위해 모든 서브 연산 (라우팅, 디스패치, 프로젝션, 활성화, 스캐터) 을 단일 연산자로 퓨즈합니다.
• 상용 스마트폰 CPU(Snapdragon 8 Elite, Apple A18 Pro) 및 GPU 에서 효율적인 추론을 가능하게 합니다.

주요 결과

벤치마크 성능

14 가지 기초 벤치마크 (상식, 지식, 과학, 독해, 추론) 에서 MobileMoE 는 새로운 파레토 프론티어를 확립했습니다:

• 밀집 모델 대비 효율성: MobileMoE-S/M 은 비교 가능한 메모리 풋프린트에서 2~4 배 적은 추론 FLOPs를 사용하여 선도적인 밀집 온디바이스 LLM(예: MobileLLM-Pro, Llama 3.2 1B) 과 맞먹거나 능가합니다.
• MoE 모델 대비 효율성: MobileMoE-M 은 약 60% 적은 활성 및 총 파라미터로 최첨단 OLMoE-1B-7B 와 정확도를 맞춥니다. MobileMoE-L 은 30% 적은 활성 파라미터와 23% 작은 메모리 풋프린트로 OLMoE-1B-7B 를 능가합니다.
• 토큰 효율성: MobileMoE-L 은 약 6T 의 사전 학습 토큰으로 우수한 성능을 달성하여, 1.5~2 배 더 많은 토큰으로 학습된 밀집 베이스라인 (예: 9T 의 Llama 3.2 1B) 보다 우수한 성능을 보입니다.

온디바이스 런타임 프로파일링

INT4 가중치로 Samsung Galaxy S25 및 iPhone 16 Pro 에 배포된 결과:

• 지연 시간: 비교 가능한 INT4 가중치 메모리에서 MobileMoE-S 는 밀집 MobileLLM-Pro 베이스라인 대비 프리필 1.8~3.8 배, 디코딩 2.2~3.4 배 더 빠릅니다.
• 메모리: 8k 컨텍스트에서 MobileMoE-S 는 밀집 베이스라인 대비 최대 22% 적은 피크 RAM을 소비합니다. MobileMoE-L 은 8k 컨텍스트에서 5GB 이내의 DRAM 에 적합하여 현대 스마트폰에서의 실현 가능성을 확인했습니다.
• 정확도: MobileMoE-L 은 중간 정도의 런타임 비용으로 (컨텍스트 길이가 증가함에 따라 감소) MobileLLM-Pro(평균 45.5) 보다 훨씬 높은 정확도 (평균 57.8) 를 제공합니다.

중요성 및 주장

본 논문은 에지 배포를 위해 특별히 최적화된 최초의 수십억 미만 활성 MoE 계열을 제시하며, 이는 새로운 온디바이스 확장 법칙에서 도출된 것입니다. 그 중요성은 세 가지 주요 기여에 있습니다:

1. 이론적 기반: 메모리 및 연산 제약 하의 MoE 설계에 대한 원칙 있는 프레임워크를 제공하며, 모바일 기기에 최적화된 구성으로 적절한 희소성 ($E=8$), 세분화된 전문가 ($g=8$), 그리고 공유 전문가를 식별합니다.
2. 학습 효율성: 밀집 모델에 비해 훨씬 적은 학습 토큰(~6T) 으로 최첨단 온디바이스 성능을 달성할 수 있음을 입증하여, 수십억 미만 규모에서 MoE 의 학습 효율성을 검증합니다.
3. 실용적 배포: 상용 스마트폰에서의 첫 번째 효율적인 MoE 추론을 제공함으로써 모바일 배포의 "마지막 마일"을 연결합니다. 결과는 MoE 가 클라우드 규모 아키텍처일 뿐만 아니라 차세대 에지 AI 를 위한 실용적이고 효율적인 경로임을 확인시켜 주며, 스마트폰 및 웨어러블 기기에서 저지연, 개인정보 보호, 비용 효율적인 애플리케이션을 가능하게 합니다.

저자들은 기존 온디바이스 모델이 여전히 주로 밀집 구조임을 인정하면서도, MobileMoE 가 MoE 가 에지를 위한 우월한 대안임을 입증하여 정확도, 메모리, 지연 시간을 균형 있게 조화시키는 새로운 파레토 프론티어를 확립했다고 결론지었습니다.