Mixtral of Experts

Albert Q. Jiang; Alexandre Sablayrolles; Antoine Roux; Arthur Mensch; Blanche Savary; Chris Bamford; Devendra Singh Chaplot; Diego de las Casas; Emma Bou Hanna; Florian Bressand; Gianna Lengyel; Guillaume Bour; Guillaume Lample; Lélio Renard Lavaud; Lucile Saulnier; Marie-Anne Lachaux; Pierre Stock; Sandeep Subramanian; Sophia Yang; Szymon Antoniak; Teven Le Scao; Théophile Gervet; Thibaut Lavril; Thomas Wang; Timothée Lacroix; William El Sayed

논문 · arXiv 2401.04088

Mixtral of Experts

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour et al.

게시일 2024-01-08cs.LGcs.CL arXiv PDF

XYZNOT로 읽기 (전문 번역과 AI 해설)

세 줄 요약

🎯 과도한 추론 지연 시간과 연산 비용을 들이지 않고 트랜스포머의 용량과 성능을 확장합니다.
🔧 표준 피드포워드 블록을 8개의 SwiGLU 전문가로 대체하고, 학습된 라우터를 통해 토큰당 2개의 전문가를 선택하는 Sparse Mixture of Experts (SMoE) 디코더 온리 아키텍처입니다.
📈 토큰당 활성 파라미터는 13B에 불과하면서 총 47B 파라미터에 접근하여, MMLU(70.6% vs 69.9%) 및 GSM8K(74.4% vs 69.6%)에서 Llama 2 70B를 능가합니다.

초록

본 논문에서는 Sparse Mixture of Experts (SMoE) 언어 모델인 Mixtral 8x7B를 소개합니다. Mixtral은 Mistral 7B와 동일한 구조를 가지지만, 각 레이어가 8개의 피드포워드 블록(즉, 전문가)으로 구성된다는 점이 다릅니다. 모든 토큰에 대해 각 레이어에서 라우터 네트워크가 현재 상태를 처리할 두 개의 전문가를 선택하고 그 출력을 결합합니다. 각 토큰은 두 개의 전문가만 참조하지만, 선택되는 전문가는 타임스텝마다 달라질 수 있습니다. 그 결과 각 토큰은 총 47B 파라미터에 접근할 수 있으면서도 추론 시에는 13B의 활성 파라미터만 사용합니다. Mixtral은 32k 토큰의 컨텍스트 크기로 학습되었으며, 평가된 모든 벤치마크에서 Llama 2 70B 및 GPT-3.5와 대등하거나 이를 능가합니다. 특히 수학, 코드 생성, 다국어 벤치마크에서 Llama 2 70B를 크게 앞섭니다. 또한 명령어 준수(instruction-following)를 위해 파인튜닝된 Mixtral 8x7B - Instruct 모델을 제공하며, 인간 평가 벤치마크에서 GPT-3.5 Turbo, Claude-2.1, Gemini Pro, Llama 2 70B - chat 모델을 능가합니다. 베이스 모델과 인스트럭트 모델 모두 Apache 2.0 라이선스로 공개됩니다.

영어 원문 초록

We introduce Mixtral 8x7B, a Sparse Mixture of Experts (SMoE) language model. Mixtral has the same architecture as Mistral 7B, with the difference that each layer is composed of 8 feedforward blocks (i.e. experts). For every token, at each layer, a router network selects two experts to process the current state and combine their outputs. Even though each token only sees two experts, the selected experts can be different at each timestep. As a result, each token has access to 47B parameters, but only uses 13B active parameters during inference. Mixtral was trained with a context size of 32k tokens and it outperforms or matches Llama 2 70B and GPT-3.5 across all evaluated benchmarks. In particular, Mixtral vastly outperforms Llama 2 70B on mathematics, code generation, and multilingual benchmarks. We also provide a model fine-tuned to follow instructions, Mixtral 8x7B - Instruct, that surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B - chat model on human benchmarks. Both the base and instruct models are released under the Apache 2.0 license.

논문 해설

문제 및 동기

Dense 자기회귀 언어 모델은 처리되는 모든 토큰에 대해 총 파라미터 수에 직접적으로 비례하는 연산 비용이 발생합니다. 모델 크기가 70B 파라미터 이상으로 확장됨에 따라, 이를 낮은 지연 시간과 높은 처리량으로 서빙하는 데 드는 연산 비용이 매우 커집니다. 저자들은 토큰당 추론 비용을 훨씬 작은 모델 수준으로 유지하면서 모델 용량을 확장하고자 했으며, 이를 통해 Llama 2 70B 및 GPT-3.5와 같은 상용 베이스라인과 대등하거나 능가하는 오픈 가중치 모델을 구축하는 것을 목표로 합니다.

핵심 아이디어

Sparse Mixture of Experts (SMoE): 32개의 트랜스포머 레이어 전체에 걸쳐 dense 피드포워드 네트워크(FFN) 블록을 8개의 서로 다른 전문가 네트워크로 대체하고, 각 토큰을 상위 2개의 전문가(top-2)로만 라우팅합니다.
분리된 파라미터 수(Decoupled Parameter Counts): 이 아키텍처는 레이어 전체에 걸쳐 총 47B 파라미터를 보유하지만, 토큰 순전파 시에는 13B의 활성 파라미터만 실행하므로 중간 규모 dense 모델 수준의 연산량만 필요로 합니다.
온전한 32k 컨텍스트 윈도우: passkey 벤치마크에서 100% 검색 정확도를 보이며 32,768 토큰 시퀀스 길이를 기본적으로 지원합니다.
주제 기반이 아닌 구문적 라우팅: 토큰 라우팅을 실증적으로 분석한 결과, 전문가 선택은 고수준의 주제 영역보다는 주로 구문적(syntactic) 기능과 시간적 국소성에 의해 결정되는 것으로 나타났습니다.

방법론 (핵심 수식 포함)

Mixtral은 Mistral 7B 아키텍처( $32$ 레이어, 은닉 차원 $4096$ , $32$ 어텐션 헤드, $8$ 키-값 헤드)를 수정하여, 모든 FFN 서브 블록을 중간 차원이 $14336$ 인 SwiGLU 네트워크로 구현된 $n = 8$ 개의 전문가( $E_{0}, \dots, E_{n - 1}$ )로 이루어진 MoE 모듈로 대체합니다.

입력 토큰 표현 $x$ 에 대해, 라우터 네트워크는 상위 $K$ 개 로짓( $K = 2$ )을 사용하여 전문가에 대한 희소 게이팅 분포를 계산합니다: $G (x) := Softmax (TopK (x \cdot W_{g}))$ 여기서 Top-K 연산자는 로짓 $ℓ = x \cdot W_{g} \in R^{n}$ 의 상위 $K$ 개 값을 제외한 모든 좌표에 $- \infty$ 를 할당합니다: $(TopK (ℓ))_{i} := {ℓ_{i} - \infty ℓ_{i} 가 ℓ 의 상위 K 개 좌표에 속하는 경우, 그렇지 않은 경우 .$

레이어의 최종 출력 $y$ 는 전문가 출력에 게이팅 가중치를 적용한 합으로 계산됩니다: $y = i = 0 \sum n - 1 G (x)_{i} \cdot SwiGLU_{i} (x)$ 효율적인 연산 수행을 위해 MoE 연산을 희소 행렬 곱셈으로 표현하는 MegaBlocks와 같은 커널을 활용하고, 여러 GPU에 걸친 Expert Parallelism을 적용합니다.

결과 (주요 수치 중심)

일반 및 추론 벤치마크: MMLU (5-shot)에서 Mixtral 8x7B는 70.6%를 기록하여 Llama 2 70B의 69.9% 및 GPT-3.5의 70.0%를 능가합니다. GSM8K (8-shot, maj@8)에서는 Llama 2 70B의 69.6% 대비 74.4%를 달성합니다. MATH 벤치마크(4-shot, maj@4)에서는 Llama 2 70B(13.8%)보다 훨씬 높은 28.4%를 기록합니다.
코드 생성: Mixtral은 HumanEval(40.2% vs 29.3%)과 MBPP(60.7% vs 49.8%)에서 Llama 2 70B를 크게 앞섭니다.
다국어 작업: 다국어 데이터 업샘플링을 통해 ARC-Challenge, HellaSwag, MMLU에서 프랑스어, 독일어, 스페인어, 이탈리아어에 걸쳐 Llama 2 70B를 능가합니다 (예: 프랑스어 MMLU 70.9% vs 64.3%).
명령어 파인튜닝 및 Arena Elo: SFT와 Direct Preference Optimization으로 학습된 Mixtral 8x7B – Instruct는 MT-Bench에서 8.30점을 기록하고 LMSys Chatbot Arena Elo에서 1121점을 달성하여 Claude-2.1 (1117), GPT-3.5-Turbo (1117), Llama-2-70b-chat (1077)을 앞섰습니다.
컨텍스트 검색: 전체 32k 컨텍스트 윈도우에 걸친 합성 passkey 검색 작업에서 100%의 검색 정확도를 보입니다.

한계점 및 열린 질문

메모리 사용량 대 연산량: 연산 비용은 13B 활성 파라미터 수준이지만, 서빙 시 47B 파라미터 전체를 유지하기 위한 충분한 VRAM이 필요하므로 다중 GPU 환경이나 오프로딩 기법이 요구됩니다.
Expert Parallelism 병목 현상: 본 논문에서는 높은 시간적 국소성을 발견했습니다. 연속된 토큰이 무작위 확률보다 훨씬 자주 동일한 전문가에 할당됩니다 (레이어 15에서 무작위 선택 12.5% 대비 첫 번째 선택 반복률 최대 28.4%). 이는 Expert Parallelism 환경에서 로드 불균형과 특정 GPU 과부하를 초래할 수 있습니다.
주제별 전문화의 부재: 라우터의 선택이 개념적 도메인(예: 생물학 vs 철학)에 따라 자연스럽게 나뉘지 않으므로, 도메인별 전문가 특화가 명시적으로 필요한 경우 게이팅을 어떻게 구성해야 할지에 대한 연구가 열린 질문으로 남아 있습니다.

핵심 용어

Sparse Mixture of Experts (SMoE): 모든 파라미터를 활성화하는 대신 조건부 연산을 통해 각 입력 토큰을 피드포워드 서브 네트워크(전문가)의 희소한 하위 집합으로 라우팅하는 신경망 아키텍처입니다.
Active Parameter Count: 단일 토큰 추론 시 실제로 로드되어 연산에 참여하는 파라미터 수로, Mixtral의 경우 총 47B의 희소 파라미터 중 13B입니다.
SwiGLU: 각 전문가 피드포워드 블록의 내부 구조로 사용되는, Swish와 원소별 곱셈을 결합한 게이트 선형 유닛(GLU) 변형 활성화 함수입니다.
TopK Gating: 토큰 표현에 선형 투영을 적용하고 상위 K개의 가장 큰 값만 유지한 채 나머지를 음의 무한대로 마스킹한 뒤, 소프트맥스를 통해 정규화하는 라우팅 메커니즘입니다.
Expert Parallelism (EP): 서로 다른 전문가를 개별 GPU에 매핑하고, 디바이스 간에 토큰 활성화를 할당된 전문가로 라우팅하는 분산 병렬화 전략입니다.
Direct Preference Optimization (DPO): 명시적인 강화학습 보상 모델을 학습시키지 않고, 선호도 쌍 데이터를 기반으로 언어 모델 파라미터를 직접 최적화하는 정렬(alignment) 알고리즘입니다.
Passkey Retrieval: 최대 32k 토큰 길이의 긴 컨텍스트 프롬프트 내에 무작위로 삽입된 짧은 패스키를 찾아 반환하는 모델의 능력을 측정하는 합성 평가 테스트입니다.
Positional Locality: 시퀀스 내에서 연속적인 토큰들이 균등 무작위 확률을 크게 웃도는 빈도로 동일한 전문가에게 반복 할당되는 관찰된 경향성입니다.

자주 묻는 질문

Mixtral은 파라미터가 56B가 아니라 47B인데 왜 8x7B라고 불리나요?

레이어당 8개의 전문가가 있고 아키텍처가 Mistral 7B에서 파생되었지만, 셀프 어텐션 가중치와 임베딩 레이어는 모든 전문가가 공유합니다. 피드포워드 서브 블록만 8개로 복제되기 때문에 8 × 7B = 56B가 아닌 총 47B 파라미터를 갖게 됩니다.

추론 중 각 토큰에 대해 몇 개의 전문가가 활성화되나요?

라우터는 각 레이어에서 사용 가능한 8개의 전문가 중 정확히 2개( $K = 2$ )의 전문가를 선택하므로, 순전파 패스당 13B 활성 파라미터가 사용됩니다.

전문가마다 서로 다른 학문 분야나 도메인 주제를 전문화하여 학습하나요?

명확한 도메인 전문화는 발견되지 않았습니다. The Pile 데이터셋을 통한 분석에 따르면 ArXiv, PubMed, PhilPapers와 같은 주제 전반에 걸쳐 전문가 할당 분포가 거의 동일했습니다. 라우팅 결정은 주제보다는 구문, 토큰 정체성 및 위치적 국소성과 깊은 연관이 있습니다.

Mixtral 8x7B 모델은 어떤 소프트웨어 라이선스로 공개되나요?

Mixtral 8x7B 베이스 모델과 Mixtral 8x7B – Instruct 모델 모두 학술 및 상업적 용도로 사용 가능한 관대한 Apache 2.0 라이선스로 배포됩니다.

이 페이지의 해설은 AI가 논문을 바탕으로 작성했으며 오류가 있을 수 있습니다. 자세한 내용은 원문을 확인하세요.

XYZNOT로 읽기 (전문 번역과 AI 해설)