[논문 리뷰] Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training

Zexuan Zhong, Mengzhou Xia외 2명|arXiv (Cornell University)|2024. 05. 06.

Topic Modeling인용 수 6

한 줄 요약

Lory는 인과적 세그먼트 라우팅과 유사성 기반 데이터 배치화를 사용하여 자동회귀 언어 모델의 사전 학습을 위한 완전 미분 가능 MoE 프레임워크를 도입하고, 세그먼트 수준의 전문가 병합으로 강력한 perplexity와 하위 작업 이득을 달성합니다.

ABSTRACT

Mixture-of-experts (MoE) models facilitate efficient scaling; however, training the router network introduces the challenge of optimizing a non-differentiable, discrete objective. Recently, a fully-differentiable MoE architecture, SMEAR, was proposed (Muqeeth et al., 2023), which softly merges experts in the parameter space; nevertheless, its effectiveness was only demonstrated in downstream fine-tuning on classification tasks. In this paper, we present Lory, the first approach that scales such architectures to autoregressive language model pre-training. Lory introduces two key techniques: (1) a causal segment routing strategy that achieves high efficiency for expert merging operations while preserving the autoregressive nature of language models; (2) a similarity-based data batching method that encourages expert specialization by grouping similar documents in training instances. We pre-train a series of Lory models on 150B tokens from scratch, with up to 32 experts and 30B (1.5B active) parameters. Experimental results show significant performance gains over parameter-matched dense models on both perplexity (+13.9%) and a variety of downstream tasks (+1.5%-11.1%). Despite segment-level routing, Lory models achieve competitive performance compared to state-of-the-art MoE models with token-level routing. We further demonstrate that the trained experts in Lory capture domain-level specialization without supervision. Our work highlights the potential of fully-differentiable MoE architectures for language model pre-training and advocates future research in this area.

연구 동기 및 목표

• 인코더/분류 작업을 넘어, 자동회귀 LMs용 완전 미분 가능 MoE 아키텍처를 동기부여하고 가능하게 한다.
• 계산량을 줄이면서 자동회귀 생성을 보존하기 위해 인과적 세그먼트 라우팅을 통한 효율적인 전문가 병합을 제안한다.
• 사전 학습 중 도메인별 전문가 전문화를 유도하기 위해 유사성 기반 데이터 배치를 도입한다.
• 150B 토큰에서 최대 32개의 전문가로 최대 30B 전체 매개변수의 Lory 모델을 사전 학습하여 확장성을 입증한다.

제안 방법

• Softmax 라우터에서 얻은 e_i로 합쳐진 theta_i를 사용하여 o_x = FFN(h_x; sum_i e_i * theta_i)와 같은 전문가 병합을 도입한다.
• 인과적 세그먼트 라우팅을 사용한다: 입력을 길이 T의 세그먼트로 나누고, 앞선 세그먼트에서 병합된 theta를 계산하여 현재 세그먼트에 적용해 자동회귀 인과성을 보존한다.
• 학습 중에는 이전 세그먼트의 평균 은닉 표현에 기반해 라우팅 가중치를 계산하고, 누설을 방지하기 위해 stop-gradient를 적용한다.
• 의미적으로 유사한 문서를 연결하도록 유사성 기반 데이터 배치를 적용하여 전문가 전문화를 촉진한다.
• 추론 시에는 프롬프트 전용 라우팅을 사용해 전체 생성에 대해 각 계층당 하나의 병합된 FFN을 결정한다.
• Dense 모델 초기화 및 워밍업 일정에 따라 150B 토큰, 0.3B 및 1.5B 활성 매개변수, 계층당 8/16/32명의 전문가를 사용하여 학습한다.

실험 결과

연구 질문

• RQ1 완전 미분 가능 MoE 아키텍처를 자동회귀 언어 모델 사전 학습에 효과적으로 확장할 수 있는가?
• RQ2 인과적 세그먼트 라우팅이 세그먼트 수준에서 전문가 병합을 가능하게 하면서 자동회귀 생성을 보존하는가?
• RQ3 유사성 기반 데이터 배치가 도메인 수준의 전문가 전문화를 촉진하고 다운스트림 작업 성능을 향상시키는가?
• RQ4 완전 미분 가능 MoE Lory 모델은 perplexity와 다운스트림 작업에서 전통적 토큰 수준 MoE 및 Dense 베이스라인과 어떻게 비교되는가?

주요 결과

• Lory는 평가된 도메인 전반에서 매개변수 매칭 밀집(Dense) 베이스라인 대비 유의한 perplexity 이득을 달성했으며, 0.3B/32E 모델에서 Books에 대해 상대적 +13.9% 개선 및 arXiv, Books, Wikipedia, C4, Python 등에서 더 넓은 perplexity 감소를 보였다.
• 선정된 구성에서 MoE를 사용한 다운스트림 작업은 일관된 개선을 보이며: +3.7% (상식 추론), +3.3% (독해), +1.5% (무지문 QA), +11.1% (텍스트 분류).
• 인과 병합이 포함된 세그먼트 수준 라우팅은 최첨단 토큰 수준 MoE 모델에 비해 경쟁적 성능을 보이면서도 완전한 미분 가능 학습을 가능하게 한다.
• 전문가들은 감독 없이도 도메인 수준의 전문화를 보여주며, 중간/상위 계층에서 도메인 신호가 더 강하게 나타난다(예: arXiv, Python), 로컬 토큰 패턴을 넘어선 효과적인 전문화를 시사한다.
• 유사성 기반 배치는 무작위 배치에 비해 MoE 성능을 크게 향상시키며, 효과적인 전문가 학습과 활용에 필수적이다.