EMO: 전문가 혼합(Mixture of Experts)을 이용한 창발적 모듈성 사전 학습

모델: https://huggingface.co/collections/allenai/emo | 📄

기술 보고서: https://allenai.org/papers/emo | 💻

코드: https://github.com/allenai/EMO | 📊

시각화: https://emovisualization.netlify.app/

오늘 저희는 EMO를 공개합니다. EMO는 인간이 정의한 사전 지식(priors)에 의존하지 않고, 모듈 구조 자체가 데이터로부터 직접 창발되도록 종단 간(end-to-end)으로 사전 학습된 새로운 전문가 혼합(Mixture-of-Experts, MoE) 모델입니다. EMO를 사용하면 주어진 작업에 대해 전체 전문가 중 작은 부분집합(단 12.5%)만을 사용하여 거의 전체 모델 수준의 성능을 유지할 수 있으며, 모든 전문가를 함께 사용할 경우 강력한 범용 모델로도 작동합니다.

대규모 언어 모델(Large language models)은 일반적으로 단일 시스템으로 학습되고 배포됩니다. 즉, 하나의 모델이 초기화되고, 사전 학습되며, 미세 조정되고, 하나의 통합된 개체로 서비스됩니다. 하지만 애플리케이션은 코드 생성, 수학적 추론 또는 도메인별 지식과 같은 특정 능력의 부분집합만을 필요로 하는 경우가 많습니다. 최첨단 언어 모델들이 루틴하게 수조 개의 파라미터에 도달함에 따라, 전체 모델을 사용하고 적응시키는 것은 대부분의 사용자에게 비실용적이며, 심지어 필요하지 않을 수도 있는 파라미터를 호스팅하는 데 불필요한 계산 비용과 메모리를 발생시킵니다.

Mixture-of-experts (MoE) 모델은 이러한 제약을 완화하는 자연스러운 방법처럼 보입니다. 각 레이어에서 하나의 거대한 피드포워드 네트워크를 사용하는 대신, MoE는 '전문가(expert)'라고 불리는 많은 작은 전문가들을 포함하며, 입력 토큰당 소수의 전문가만을 활성화합니다. 원칙적으로, 단일 기능만 필요한 작업은 관련 있는 전문가들만을 로드할 수 있습니다.

하지만 실제로는 기존의 MoE 모델들이 잘 작동하려면 여전히 전체 모델이 필요합니다. 심지어 단일 입력 내에서도 서로 다른 토큰들은 종종 서로 다른 전문가들을 활성화시키기 때문에, 하나의 작업이 생성되는 동안 모든 전문가를 사용할 수 있게 됩니다. 저희 논문에서 보여주듯이, 이는 표준 MoE의 전문가들이 높은 수준의 도메인이나 기능보다는 전치사나 구두점 같은 낮은 수준의 어휘 패턴에 특화되어 있기 때문에 부분적으로 발생합니다. 그 결과, 작은 하위 집합의 전문가들은 독립적으로 신뢰성 있게 사용하기 어렵습니다.

대신 저희는 선택적으로 사용하고 조합할 수 있도록 응집력 있는 그룹으로 조직된 전문가들을 가진 MoE 모델을 원합니다.

이러한 것을 사전 학습 중에 장려하는 한 가지 방법은 토큰을 수학(math), 생물학(biology) 또는 코드와 같은 미리 정의된 의미론적 도메인에 따라 전문가에게 라우팅하는 것입니다. BTX나 저희의 FlexOlmo 프로젝트와 같은 이전 연구들이 이를 시도했습니다. 하지만, 미리 정의된 도메인은 중요한 한계를 가집니다. 이들은 사전 학습 코퍼스 전반에 걸쳐 도메인 레이블을 요구하는데, 이는 모호하고 얻기 비용이 많이 들 수 있으며, 모델이 스스로 조직하는 방식에 너무 많은 인간의 편향을 주입할 수 있습니다. 더 중요하게는, 도메인을 미리 고정하면 모델의 모듈식 구조도 고정됩니다: 추론 시점에 새로운 도메인이나 능력이 나타나더라도 어떤 전문가를 사용해야 할지 명확하지 않습니다.

여기에 EMO가 등장합니다.

저희는 1조 개의 토큰으로 학습된 1B 활성, 14B 총 매개변수(8개 전문가 활성, 128개 전문가 총) MoE인 EMO가 선택적 전문가 사용을 지원한다는 것을 보여줍니다: 주어진 작업이나 도메인에 대해 전체 전문가의 작은 하위 집합(총 전문가의 단지 12.5%)만 사용하여 거의 전체 모델 성능을 유지할 수 있습니다. 동시에, 모든 전문가를 함께 사용할 때도 EMO는 강력한 범용 모델로 남아 있습니다. 이와 대조적으로, 동일한 데이터로 학습된 동등한 아키텍처의 표준 MoE는 그 전문가 하위 집합을 선택적으로 사용할 때 심각하게 성능이 저하됩니다.

EMO는 모듈성(modularity)을 일급 목표(first-class objective)로 학습된 MoE입니다. 특정 도메인(예: 수학, 코드, 생물의학)에 대해 사용자는 크기에 관계없이 적은 수의 전문가 하위 집합을 선택하여 거의 전체 모델 성능을 유지할 수 있습니다. 이는 단일 모델을 구성 가능한 아키텍처로 변환하며, 대규모 희소 MoE에서 향상된 메모리-정확도 트레이드오프를 가진 유연한 배포를 가능하게 합니다.

MoE에서는 라우터(router)라는 작은 네트워크가 각 토큰이 어떤 전문가들을 활성화할지 결정합니다. 우리는 라우터가 유사한 도메인의 토큰들이 유사한 전문가 하위 집합을 활성화하도록 학습하기를 원합니다. 우리의 핵심 관찰은 같은 문서의 토큰들은 보통 같은 도메인에서 온다는 것입니다. 따라서 우리는 문서 경계를 약한 지도 신호(weak supervisory signal)로 사용합니다: 훈련 중에, 한 문서 내의 모든 토큰은 공유 전문가 풀(shared expert pool)로부터 활성 전문가를 선택하도록 제한됩니다.

*(표준 MoE와 EMO 학습 비교 (k = 2, n = 10, 단순화를 위해 공유 전문가 생략). (왼쪽) 표준 MoE에서는 각 토큰이 독립적으로 상위-k개의 전문가를 선택합니다. 토큰 전반에 걸쳐 모든 전문가가 사용됩니다. (오른쪽) EMO에서는 라우터가 먼저 각 문서에 대한 전문가 하위 집합을 선택하고, 모든 토큰은 이 하위 집합 내에서 경로 지정(route)하도록 제약됩니다. 이는 문서 전반에 걸쳐 일관된 전문가 사용을 강제하여, 전문가 그룹들이 도메인 전문화(domain specialization)를 형성하도록 장려합니다.)

예를 들어, 총 10개의 전문가와 토큰당 2개의 활성 전문가를 가진 MoE에서, 문서 내의 모든 토큰은 위 그림에서 볼 수 있듯이 동일한 4개 전문가 풀(pool) 내에서 경로 지정(route)하도록 제한됩니다. 이 풀은 라우터 자체가 선택합니다. 즉, 문서 내 모든 토큰에 걸쳐 라우터의 전문가 선호도를 평균화한 다음, 가장 많이 사용된 전문가를 해당 문서의 공유 풀로 선택하는 것입니다. 서로 다른 문서는 서로 다른 풀을 사용할 수 있어, 반복적으로 나타나는 전문가 그룹이 훈련 데이터로부터 직접 출현할 수 있게 합니다.

시스템을 구현할 때 고려해야 할 몇 가지 사항들이 있습니다:

로드 밸런싱(Load balancing). 하나의 기술적 과제는 로드 밸런싱입니다. 표준 MoE 훈련에서는 모델이 소수의 전문가에게만 집중되는 것을 방지하기 위해 로드-밸런싱 목적 함수(objective)가 사용됩니다. 언뜻 보기에는 이것이 EMO의 훈련 목표와 상충하는 것처럼 보입니다. 우리는 각 문서를 명시적으로 전문가 하위 집합만을 사용하도록 제한하고 있기 때문입니다.

이러한 충돌은 로드 밸런싱이 일반적으로 적용되는 규모에서 발생합니다. 많은 MoE 구현에서, 로드 밸런싱은 종종 소수의 문서만 포함하는 마이크로 배치(micro-batch) 내에서 국소적으로 계산됩니다. 이 국지적 목적 함수는 동일한 문서 내의 토큰들이 여러 전문가에 걸쳐 퍼지도록 밀어낼 수 있으며, 이는 문서 내에서 전문가 사용을 일관되게 유지하려는 EMO의 목표와 정면으로 배치됩니다.

이를 해결하기 위해 우리는 여러 문서에 걸쳐 전역적으로 부하 균형(load balancing)을 적용합니다. 이 더 큰 규모에서는 두 가지 목적 함수가 상호 보완적이 됩니다. EMO는 동일한 문서 내의 토큰들이 일관된 전문가 풀을 사용하도록 장려하는 반면, 전역적 부하 균형은 서로 다른 문서들이 모든 전문가를 집단적으로 커버하도록 장려합니다. 실제로 우리는 전역적 부하 균형이 안정적인 훈련에 중요하다는 것을 발견했습니다.

문서 풀 크기(Document pool size). 문서 풀 크기는 모듈성 제약 조건이 얼마나 제한적인지를 제어합니다. 더 작은 풀은 동일한 문서 내의 토큰들이 더 타이트한 전문가 세트를 공유하도록 강제하여 더 강력한 모듈성을 장려하는 반면, 더 큰 풀은 모델에 더 많은 유연성을 제공하지만 제약을 약화시킵니다.

특정 풀 크기를 고정하기보다는 훈련 중에 무작위로 샘플링합니다. 이는 EMO가 단일 하위 집합 크기에 과적합되는 것을 방지하고 추론 시간에도 다양한 전문가 하위 집합 크기를 지원할 수 있게 합니다.

범용 벤치마크에서 EMO는 표준 MoE 모델의 성능과 일치하며, 모듈성 목적 함수가 전체 모델 성능을 희생시키지 않음을 보여줍니다. 하지만 더 중요한 질문은 우리가 전문가의 하위 집합만 유지했을 때도 모델이 여전히 작동할 수 있는지 여부입니다. 이 설정에서는 소량의 작업 검증 데이터에 따라 전문가를 순위 매겨, 가장 많이 사용된 전문가를 유지하고 나머지는 폐기함으로써 작업별 전문가 하위 집합을 구성합니다.

아래 그림은 EMO가 선택적 전문가 사용 하에서도 견고하게 작동함을 보여줍니다. 전문가의 25%만 유지할 때(32개 전문가 하위 집합), EMO는 모든 벤치마크에서 절대 성능이 약 1%만 감소합니다. 심지어 전문가의 12.5%만 유지할 때(16개 전문가 하위 집합)에도 전체적인 하락 폭은 단지 약 3%에 불과합니다. 이는 미세 조정 전후 모두 동일하게 나타납니다. 이와 대조적으로, 매칭 표준 MoE는 전문가 하위 집합이 작아질수록 급격히 성능이 저하되며, 가장 작은 전문가 하위 집합 설정에서는 종종 무작위 성능 수준에 가깝거나 그 이하로 떨어집니다.

게다가, 태스크에 적합한 전문가를 선택하는 것이 놀라울 정도로 비용 효율적이라는 것을 보여줍니다. 소수의 예시만 있는 퓨샷(few-shot) 시연으로도 전체 검증 세트를 사용하여 선택한 것과 동등하게 성능을 발휘하는 모듈을 식별하기에 충분합니다. 그리고 EMO는 특정 선택 방식에 국한되지 않습니다: Easy-EP와 같은 기존의 전문가 가지치기(expert-pruning) 접근 방식과 잘 작동하며, 두 방법은 상호 보완적입니다.

130B 토큰 규모의 작은 설정. 다양한 메모리 예산 하에서 16개 MMLU 카테고리에 걸친 평균 성능. EMO 전문가 하위 집합은 메모리-정확도 트레이드오프에서 파레토 프론티어(Pareto frontier)를 밀어내며, 표준 MoE나 처음부터 학습된 고정 예산 모델보다 우수한 성능을 보입니다.

EMO가 훈련 후 실제로 무엇을 학습했는지 확인하기 위해, 12K개의 사전 훈련 문서에 걸친 첫 100개 토큰의 라우터 활성화(router activations)를 클러스터링했습니다. 이는 표준 MoE와는 확연한 차이를 보입니다.

EMO의 토큰 클러스터는 건강, 의학 및 웰빙, 뉴스 보도, 미국 정치 및 선거, 그리고 영화 및 음악과 같은 것들과 일치합니다. 표준 MoE는 전치사(Prepositions), 고유 명사(Proper Names), 연결 동사(Copula Verbs) 또는 *정관사(Definite Articles)*와 같은 클러스터를 생성합니다. EMO에서는 주어진 문서의 토큰들이 대부분 동일한 클러스터에 속하는 반면, 표준 MoE에서는 여러 클러스터에 흩어지게 됩니다.

이 대비는 단일 예시로 가장 쉽게 볼 수 있습니다. 건강 관련 기사를 예로 들어보겠습니다: EMO에서는 거의 모든 토큰이 건강, 의학 및 웰빙 클러스터로 라우팅될 것입니다. 표준 MoE의 경우, 최상위 클러스터는 *소유격 및 정관사(Possessives & Definite Articles)*입니다. 이 모델은 해당 텍스트가 무엇에 관한 것인지와 관계없이 단어 the나 your를 사용하는 모든 다른 텍스트와 함께 기사를 그룹화할 것입니다.

1조 개의 토큰으로 훈련된 MoE의 사전 학습 데이터 토큰 클러스터. EMO 클러스터는 의미론적으로 의미 있는 도메인에 해당하며, 동일 문서의 토큰들이 크게 그룹 지어집니다. 표준 MoE 훈련은 표면 수준 또는 구문적 특징의 클러스터를 생성하며, 문서 토큰들은 여러 클러스터에 분산됩니다.

EMO가 표면적 특징이 아닌 의미론적 도메인에 매핑되는 모듈을 형성하기 때문에, 작은 전문가(expert) 하위 집합만 선택해도 여전히 기능하는 모델을 가질 수 있습니다. 즉, 그 그룹은 실제 역량에 해당합니다.

저희의 인터랙티브 시각화에서 직접 클러스터링 결과를 가지고 놀아보실 수 있습니다.

저희는 전체 EMO 학습 모델, 동일한 데이터로 훈련된 매칭 표준-MoE 기준선(baseline) 모델, 그리고 훈련 코드를 공개합니다. 저희는 이러한 결과물들이 MoE에서 창발적 모듈성(emergent modularity)을 연구하는 다른 그룹들에게 유용하기를 바랍니다.

아직 더 많은 작업이 남아 있습니다. EMO는 대규모 희소 모델(large sparse models)을 더욱 모듈화하는 방향으로 나아가기 위한 초기 단계일 뿐이지만, 여전히 많은 질문들이 남아 있습니다. 즉, 전문가 하위 집합(expert subsets)을 어떻게 더 잘 선택하고 구성할지, 전체 모델을 방해하지 않으면서 모듈을 어떻게 업데이트할지, 그리고 해석 가능성(interpretability)과 제어(control)를 개선하기 위해 모듈 구조를 어떻게 사용할 것인지 등입니다. 이러한 모델들을 공개하는 것이 커뮤니티가 이 질문들을 연구하고 배포(deploy), 적응(adapt), 검사(inspect), 구성(compose)하기 더 쉬운 모듈형 언어 모델을 구축하는 데 도움이 되기를 바랍니다.