[논문 리뷰] Parameter-Efficient Mixture-of-Experts Architecture for Pre-trained Language Models

Zefeng Gao, Peiyu Liu외 3명|arXiv (Cornell University)|2022. 03. 02.

Topic Modeling인용 수 11

한 줄 요약

이 논문은 매트릭스 곱 연산자(MPO) 분해를 통해 전역 중심 텐서를 전문가들 간에 공유함으로써, 사전 훈련된 언어 모델을 위한 파라미터 효율적인 믹스처 오브 응용(MoE) 아키텍처인 MPOE를 제안한다. 전문가의 가중치 행렬을 공유되는 중심 텐서와 전문가별로 고유한 보조 텐서로 분해함으로써, MPOE는 Switch Transformers 대비 파라미터를 27.2배 감소시키면서 T5 및 GPT-2 모델에서 성능을 유지하거나 향상시킨다. 또한, 균형 잡히지 않은 최적화 문제를 해결하기 위해 기울기 마스크 전략을 도입한다.

ABSTRACT

Recently, Mixture-of-Experts (short as MoE) architecture has achieved remarkable success in increasing the model capacity of large-scale language models. However, MoE requires incorporating significantly more parameters than the base model being extended. In this paper, we propose building a parameter-efficient MoE architecture by sharing information among experts. We adopt the matrix product operator (MPO, a tensor decomposition from quantum many-body physics) to reconstruct the parameter matrix in the expert layer and increase model capacity for pre-trained language models by sharing parameters of the central tensor (containing the core information) among different experts while enabling the specificity through the auxiliary tensors (complementing the central tensor) of different experts. To address the unbalanced optimization issue, we further design the gradient mask strategy for the MPO-based MoE architecture. Extensive experiments based on T5 and GPT-2 show improved performance and efficiency of the pre-trained language model (27.2x reduction in total parameters for the superior model performance, compared with the Switch Transformers). Our code is publicly available at https://github.com/RUCAIBox/MPOE.

연구 동기 및 목표

• 대규모 사전 훈련된 언어 모델(PLM)에서 믹스처 오브 응용(MoE)의 높은 파라미터 비용이 모델 확장성에 제한을 줌에도 불구하고, 모델 용량이 증가하는 문제를 해결하기 위해.
• 전문가 간의 가중치 행렬에 존재하는 공통된 구조적 정보를 식별함으로써 전문가 간의 파라미터 중복을 줄이기 위해.
• 특히 자원이 제한된 환경이나 효율적인 훈련 시나리오에서 모델 용량이나 성능을 희생시키지 않고도 파라미터 효율적인 MoE를 가능하게 하기 위해.
• 공유된 중심 텐서와 전문가별 보조 텐서 간의 최적화 불균형을 완화할 수 있는 훈련 전략을 설계하기 위해.
• 구조적 행렬 분해(MPO)가 효과적으로 고용량의 컴팩트한 MoE 아키텍처를 구축하는 데 활용될 수 있음을 입증하기 위해.

제안 방법

• MPOE는 매트릭스 곱 연산자(MPO) 분해를 사용하여 각 전문가의 가중치 행렬을 전역적으로 공유되는 중심 텐서와 다수의 전문가별 보조 텐서의 곱으로 분해한다.
• 다수의 파라미터를 포함하는 중심 텐서는 모든 전문가 간에 공유되며, 이로 인해 총 파라미터 수가 크게 감소한다.
• 전문가별 보조 텐서는 작업 또는 샘플에 특화된 변동성을 포착하도록 훈련되어, 파라미터 효율성을 유지하면서도 용량을 확장할 수 있도록 한다.
• 공유된 중심 텐서의 기울기 전파를 억제함으로써 훈련 중에 중심 텐서가 업데이트에 지배적이지 않게 하기 위해 기울기 마스크 전략을 도입한다.
• 모델은 T5 및 GPT-2에 적용되며, 희소 활성화와 효율적인 추론을 가능하게 하기 위해 MoE 라우팅 메커니즘을 유지한다.
• MPO 분해를 통해 보조 텐서의 변화가 중심 텐서로 전파되어, 파라미터 공유에도 불구하고 모델의 표현력을 유지한다.

실험 결과

연구 질문

• RQ1 MoE 아키텍처에서 전문가 간의 파라미터 중복을 모델 용량을 훼손하지 않고 체계적으로 줄일 수 있는가?
• RQ2 MPO 분해에서 공유되는 중심 텐서가 사전 훈련된 언어 모델의 다수의 전문가 간 핵심 지식을 효과적으로 표현할 수 있는가?
• RQ3 제안된 기울기 마스크 전략이 공유 중심 텐서와 전문가별 보조 텐서 간의 최적화 불균형을 효과적으로 완화하는가?
• RQ4 MPO 기반 MoE는 기존 표준 MoE 방법에 비해 얼마나 많은 파라미터 수를 줄일 수 있으며, 성능을 유지하거나 향상시킬 수 있는가?
• RQ5 MPOE는 Switch Transformers와 같은 기존 MoE 강화 PLM보다 더 뛰어난 파라미터 효율성을 달성할 수 있는가?

주요 결과

• MPOE는 Switch Transformers 대비 총 파라미터 수를 27.2배 감소시키면서 T5 및 GPT-2에서 성능을 동일하거나 초월한다.
• GLUE 및 WikiText-2 벤치마크에서 강력한 성능을 유지함으로써, 공유 중심 텐서가 필수적인 용량을 유지하고 있음을 입증한다.
• 기울기 마스크 전략은 훈련 안정성을 크게 향상시키고, 미세조정 중 공유 중심 텐서의 과적합을 방지한다.
• 실험 결과, 전문가 간 중심 텐서가 매우 유사함을 확인하여, MPO 분해를 통한 파라미터 공유의 타당성을 검증한다.
• MPOE 아키텍처는 최소한의 파라미터 오버헤드로 효율적이고 고용량의 MoE 추론을 가능하게 하여, 자원이 제한된 환경이나 배포 제약 조건에 적합하다.
• 기존의 경쟁 MoE 강화 PLM에 비해 파라미터 효율성에서 뛰어나며, 하류 작업 정확도에 하락이 없음을 입증한다.