[논문 리뷰] Sparsely Activated Mixture-of-Experts are Robust Multi-Task Learners

Shashank Gupta, Subhabrata Mukherjee외 5명|arXiv (Cornell University)|2022. 04. 16.

Domain Adaptation and Few-Shot Learning인용 수 8

한 줄 요약

이 논문은 자연어처리(NLP)에서 다중 작업 학습을 향상시키기 위해 희소로 활성화되는 전문가의 혼합(MoE) 모델을 위한 작업 인식 게이팅 메커니즘을 제안한다. 서로 다른 작업의 토큰을 작업별 전문가에게 라우팅하면서 일부 가중치를 공유함으로써, 더 높은 파라미터 효율성과 강건성을 달성한다. 이는 밀도 높은 모델에 비해 자원이 제한된 작업으로의 전이 성능이 뛰어나며, 새로운 작업에 대해 샘플 효율적으로 일반화하고, 치명적인 잊음에 대한 저항성도 보여준다.

ABSTRACT

Traditional multi-task learning (MTL) methods use dense networks that use the same set of shared weights across several different tasks. This often creates interference where two or more tasks compete to pull model parameters in different directions. In this work, we study whether sparsely activated Mixture-of-Experts (MoE) improve multi-task learning by specializing some weights for learning shared representations and using the others for learning task-specific information. To this end, we devise task-aware gating functions to route examples from different tasks to specialized experts which share subsets of network weights conditioned on the task. This results in a sparsely activated multi-task model with a large number of parameters, but with the same computational cost as that of a dense model. We demonstrate such sparse networks to improve multi-task learning along three key dimensions: (i) transfer to low-resource tasks from related tasks in the training mixture; (ii) sample-efficient generalization to tasks not seen during training by making use of task-aware routing from seen related tasks; (iii) robustness to the addition of unrelated tasks by avoiding catastrophic forgetting of existing tasks.

연구 동기 및 목표

• 경쟁하는 작업들이 모델 가중치를 서로 다른 방향으로 이끄는 간섭을 해결하기 위해.
• 모든 입력에 대해 모든 가중치를 사용하는 밀도 높은 모델의 한계를 극복하여, 가중치 간 충돌과 최적 성능 이하의 성능을 방지하기 위해.
• 희소 MoE 모델이 다중 작업 학습에 적합하도록, 입력을 작업별 전문가에게 라우팅하는 작업 인식 게이팅 메커니즘을 설계함으로써.
• 세 가지 핵심 차원에서 희소 MoE 모델의 강건성을 평가하기 위해: 자원이 제한된 작업으로의 전이, 새로운 작업으로의 일반화, 치명적인 잊음에 대한 저항.
• 희소 MoE 모델에 작업 인식 라우팅을 적용할 경우, 다중 작업 밀도 모델에 비해 전이 가능성, 샘플 효율성, 작업 추가 시 안정성에서 뛰어난 성능을 보임을 입증하기 위해.

제안 방법

• 다양한 작업별 전문가와 작업 인식 게이팅 네트워크를 포함하는 Mixture-of-Experts(MoE) 레이어를 갖춘 트랜스포머 기반 아키텍처를 채택한다.
• 작업 정체성에 기반해 입력 토큰을 적절한 전문가에게 라우팅하는 작업별 게이팅을 사용하여, 작업 간 간섭을 줄인다.
• 강력한 기초 표현에서 효과적인 다중 작업 적응을 가능하게 하기 위해, 사전 학습된 BERT 가중치를 사용해 MoE 모델을 초기화한다.
• 계산 효율성을 유지하면서 조건부 계산과 높은 파라미터 용량을 가능하게 하기 위해, 최상위 1개 라우팅과 로드 밸런싱을 사용한다.
• 공유 인코더와 작업별 전문가 헤드를 사용하여 다양한 NLU 작업의 혼합으로 모델을 동시에 훈련시키며, 토큰 단위로 작업 정체성에 기반한 라우팅 결정을 내린다.
• 각 작업별로 라우팅 결정을 분리함으로써 공유 게이팅에 대한 경쟁을 피하고, 간섭 없이 전문화를 가능하게 하는 라우팅 전략을 구현한다.

실험 결과

연구 질문

• RQ1 희소 MoE 모델에서 작업 인식 게이팅이 공유 게이팅 MoE 또는 밀도 높은 모델에 비해 다중 작업 NLP 학습에서 작업 간 간섭을 줄일 수 있는가?
• RQ2 훈련 혼합에서 관련 작업을 활용할 경우, 희소 MoE 모델이 자원이 제한된 작업으로 지식을 전이할 수 있는 정도는 어느 정도인가?
• RQ3 희소 MoE 모델이 관련된 본 적 있는 작업들로부터의 라우팅을 활용해, 소수의 예시만으로도 새로운 작업으로의 일반화 성능이 얼마나 잘 되는가?
• RQ4 특히 이전에 학습한 작업의 기억을 상실하는 치명적인 잊음을 방지하기 위해, 관련 없는 작업의 추가에 대해 희소 MoE 모델은 얼마나 강건한가?
• RQ5 MoE 모델에서 작업 인식 라우팅을 사용할 경우, 다양한 NLU 벤치마크에서 밀도 높은 다중 작업 모델에 비해 성능과 안정성 향상이 이루어지는가?

주요 결과

• 제안된 MT-TaG 모델은 GLUE++의 Small Tasks 서브셋에서 단일 작업 기준선 대비 2.28점의 절대적 향상을 보이며, 83.56점(비교 기준선 81.28점)의 성능을 달성한다.
• 전체 All Tasks 벤치마크에서 MT-TaG는 평균 점수 86.46점을 기록하여, 단일 작업 기준선(85.00점)을 초월하며 다양한 작업으로의 확장성은 뛰어나다.
• 훈련 혼합에서 관련 작업을 활용함으로써, RTE, MRPC, WiC와 같은 자원이 제한된 작업으로의 전이 성능이 뛰어나다.
• MT-TaG는 관련된 본 적 있는 작업들로부터의 입력을 적절한 전문가에게 효과적으로 라우팅함으로써, 큰 피팅 데이터가 필요 없이 샘플 효율적으로 새로운 작업으로 일반화한다.
• 관련 없는 작업을 추가한 후에도 이전에 학습한 작업의 성능을 유지함으로써, 치명적인 잊음에 대한 강건성을 입증한다.
• 실험 결과는 작업 인식 MoE 라우팅이 공유 게이팅 MoE 및 밀도 높은 다중 작업 모델에 비해 모든 세 가지 핵심 강건성 차원(전이, 일반화, 잊음 저항)에서 뚜렷한 성능 향상을 보임을 확인한다.