[논문 리뷰] On the Representation Collapse of Sparse Mixture of Experts

Zewen Chi, Dong Li외 10명|arXiv (Cornell University)|2022. 04. 20.

Expert finding and Q&A systems인용 수 16

한 줄 요약

이 논문은 희소 믹스처 오브 엑스퍼트(SMoE) 모델에서 토큰 표현이 전문가 중심점 주변으로 군집되어 성능이 저하되는 표현 붕괴 현상을 규명한다. 이를 바탕으로 X-MoE를 제안하며, 은닉 표현을 저차원 초구면으로 투영하고, L2 정규화된 라우팅 점수와 학습 가능한 온도 게이트를 사용함으로써, 일곱 개인 다국어 벤치마크에서 일관된 성능 향상을 달성하고, 라우팅 안정성을 향상시킨다.

ABSTRACT

Sparse mixture of experts provides larger model capacity while requiring a constant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such a routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse. In this work, we propose to estimate the routing scores between tokens and experts on a low-dimensional hypersphere. We conduct extensive experiments on cross-lingual language model pre-training and fine-tuning on downstream tasks. Experimental results across seven multilingual benchmarks show that our method achieves consistent gains. We also present a comprehensive analysis on the representation and routing behaviors of our models. Our method alleviates the representation collapse issue and achieves more consistent routing than the baseline mixture-of-experts methods.

연구 동기 및 목표

• 희소 믹스처 오브 엑스퍼트(SMoE) 모델에서 다루지 않은 표현 붕괴 문제를 규명하고 분석하는 것.
• 라우팅 점수 계산 방식을 재정의하여 SMoE에서 라우팅 안정성과 표현 다양성을 향상시키는 것.
• 보다 견고한 라우팅 알고리즘을 통해 다국어 사전 훈련 및 미세조정 성능을 향상시키는 것.
• 여러 훈련 실행 간 일관된 라우팅 행동을 달성하여 모델 행동의 변동성을 줄이는 것.

제안 방법

• 라우팅 점수를 계산하기 전에 토큰 은닉 표현과 전문가 임베딩을 저차원 공간으로 투영하는 것.
• 라우팅 점수를 초구면에서 측정하기 위해 토큰 표현과 전문가 임베딩 양쪽에 L2 정규화를 적용하는 것.
• 소프트 전문가 게이트에서 학습 가능한 온도 파rameter를 사용하여 전문가 활성화를 제어하고, 라우팅의 유연성을 향상시키는 것.
• 표준 비정규화된 내적 곱을 대체하여, 정규화된 표현 간의 내적 유사도로 라우팅을 공식화하는 것.
• Transformer 기반 SMoE 모델에 새로운 라우팅 메커니즘을 통합하고, top-1 라우팅 및 소프트맥스와 시그모이드 게이팅 함수를 모두 사용하는 것.
• 다국어 사전 훈련 및 하류 미세조정 작업에서 모델을 훈련하고 평가하여 성능와 라우팅 일관성의 영향을 평가하는 것.

실험 결과

연구 질문

• RQ1 희소 믹스처 오브 엑스퍼트 모델에서 사전 훈련 및 미세조정 과정에서 표현 붕괴가 어느 정도 나타나는가?
• RQ2 저차원 초구면에 표현을 투영하는 것이 표현 붕괴를 완화하고 모델 성능을 향상시키는 데 기여하는가?
• RQ3 제안된 라우팅 메커니즘이 미세조정 과정에서 다양한 난수 시드 간 라우팅 일관성에 어떤 영향을 미치는가?
• RQ4 L2 정규화된 라우팅 점수와 학습 가능한 온도 게이트의 사용이 더 안정적이고 효과적인 전문가 할당을 이끌어내는가?
• RQ5 제안된 방법은 다국어 벤치마크 성능과 표현 다양성 측면에서 기준 SMoE 모델과 비교해 어떻게 다를까?

주요 결과

• 제안된 X-MoE 방법은 사전 훈련 및 미세조정 모두에서 일곱 개인 다국어 벤치마크에서 일관된 성능 향상을 달성한다.
• SMoE 모델에서 표현 붕괴가 경험적으로 관찰되며, 토큰 표현이 전문가 중심점 주변으로 빽빽이 군집되어 있음을 확인했다.
• X-MoE는 초구면 라우팅과 L2 정규화를 통해 다양성을 강제함으로써 표현 붕괴를 크게 감소시켰다.
• X-MoE의 라우팅 행동은 높은 실행 간 일관성을 보이며, 난수 시드에 관계없이 유사한 라우팅 패턴으로 수렴하는 반면, SMoE 기준 모델은 그렇지 않다.
• 이 방법은 언어 모델링 성능과 하류 작업 성능을 모두 향상시키며, 더 안정적이고 다양한 전문가 라우팅 덕분에 성능 향상이 발생한다.
• 소프트 게이트에서 학습 가능한 온도의 사용은 라우팅 적응성을 향상시키고, 성능 향상과 안정성 향상에 기여한다.