[논문 리뷰] BASE Layers: Simplifying Training of Large, Sparse Models

Mike Lewis, Shruti Bhosale외 3명|arXiv (Cornell University)|2021. 03. 30.

Topic Modeling참고 문헌 29인용 수 63

한 줄 요약

BASE 레이어는 선형 배정 문제를 이용한 최적의 균형 토큰-전문가 할당을 도입하여 토큰당 한 명의 전문가가 있는 희소 라우팅을 가능하게 하고 추가 손실이나 하이퍼파라미터 없이 계산 효율을 개선합니다.

ABSTRACT

We introduce a new balanced assignment of experts (BASE) layer for large language models that greatly simplifies existing high capacity sparse layers. Sparse layers can dramatically improve the efficiency of training and inference by routing each token to specialized expert modules that contain only a small fraction of the model parameters. However, it can be difficult to learn balanced routing functions that make full use of the available experts; existing approaches typically use routing heuristics or auxiliary expert-balancing loss functions. In contrast, we formulate token-to-expert allocation as a linear assignment problem, allowing an optimal assignment in which each expert receives an equal number of tokens. This optimal assignment scheme improves efficiency by guaranteeing balanced compute loads, and also simplifies training by not requiring any new hyperparameters or auxiliary losses. Code is publicly released at https://github.com/pytorch/fairseq/

연구 동기 및 목표

• 훈련 비용을 줄이면서 언어 모델의 확장을 위해 희소 전문가 모델 사용을 촉진한다.
• 균형 잡힌 선형 할당 BASE 레이어를 제안하여 균형 손실과 용량 요인을 제거한다.
• 균형 할당으로 토큰당 단일 전문가 라우팅이 대형 모델에서도 강력한 계산 효율성을 낳는다는 것을 입증한다.

제안 방법

• 각 전문가가 동일한 수의 토큰을 처리하도록 토큰-전문가 할당을 선형 배정 문제로 공식화한다.
• 견고성을 위해 토큰 셔플링으로 엣지 케이스를 다루고, 작업자 간 병렬로 할당을 해결하기 위해 옥션 알고리즘을 사용한다.
• 각 토큰이 단일 전문가로 라우팅되는 BASE 레이어를 구현하고, 이어서 잔여 연결을 통한 전문가 출력의 소프트 혼합을 수행한다.
• 처리량 극대화를 위해 균형 토큰 할당으로 학습하고 추가 균형 손실을 피하며, 추론 시 탐욕적 전문가 선택으로 테스트한다.
• 고정된 GPU 런타임을 고려한 퍼플렉시티로 계산 효율을 측정하고, 밀집 및 이전의 희소 접근법과 비교한다.

실험 결과

연구 질문

• RQ1 선형 배정 기반 라우팅 방식이 추가 손실 항 없이 전문가 사용의 균형을 맞출 수 있는가?
• RQ2 단일 전문가 BASE 레이어가 대형 모델에서 기존의 밀집 및 희소 방법의 효율성을 상회하거나 일치하는가?
• RQ3 BASE 배치 위치와 크기가 성능 및 계산 효율에 어떤 영향을 미치는가?
• RQ4 훈련 중 균형 라우팅이 전문가 전문화 및 부하 분산에 어떤 영향을 미치는가?

주요 결과

• BASE 레이어는 균형 손실이나 용량 요인 없이 선형 배정을 통해 토큰-전문가 라우팅의 균형을 달성한다.
• 단일 BASE 레이어는 더 높은 컴퓨트 예산에서 밀집 데이터-병렬 및 모델-병렬 기준선보다 상당히 우수하게 성능을 발휘할 수 있다.
• BASE는 계산 효율성 면에서 Sparsely Gated MoE 및 Switch 트랜스포머와 비교해 우수하며, 유사 예산에서 때로는 이들의 성능에 비견하거나 이를 상회한다.
• 여러 BASE 레이어를 인터리빙하면 총 매개변수 수를 대략 일정하게 유지하면서 성능 이득이 증가한다.
• BASE 레이어는 다양한 배치 위치 및 BASE 서브레이어 수에서 강건한 성능을 보여주며, 아키텍처의 강건성을 시사한다.