[논문 리뷰] One Student Knows All Experts Know: From Sparse to Dense

Fuzhao Xue, Xiaoxin He외 3명|arXiv (Cornell University)|2022. 01. 26.

COVID-19 diagnosis using AI인용 수 5

한 줄 요약

이 논문은 두 단계의 훈련 과정을 통해 희박한 Mixture-of-Experts (MoE) 모델에서 지식을 정제하여 고밀도 학생 네트워크(OneS)에 통합하는 새로운 프레임워크인 지식 통합을 제안한다: SVD-KG 또는 Top-KG를 통한 지식 수집, 다음으로 지식 정제. OneS는 오직 1500만 파라미터로 ImageNet에서 상위-1 정확도 78.4%를 달성하며, MoE의 61.7%에 해당하는 성능 이점을 유지하고, MoE 대비 3.7배 빠른 추론 속도를 제공하며, NLP 과제에서 최대 51.7% 높은 성능을 기록해 기존 기준을 초월한다.

ABSTRACT

Human education system trains one student by multiple experts. Mixture-of-experts (MoE) is a powerful sparse architecture including multiple experts. However, sparse MoE model is easy to overfit, hard to deploy, and not hardware-friendly for practitioners. In this work, inspired by the human education model, we propose a novel task, knowledge integration, to obtain a dense student model (OneS) as knowledgeable as one sparse MoE. We investigate this task by proposing a general training framework including knowledge gathering and knowledge distillation. Specifically, to gather key knowledge from different pre-trained experts, we first investigate four different possible knowledge gathering methods, \ie summation, averaging, Top-K Knowledge Gathering (Top-KG), and Singular Value Decomposition Knowledge Gathering (SVD-KG) proposed in this paper. We then refine the dense student model by knowledge distillation to offset the noise from gathering. On ImageNet, our OneS preserves $61.7\%$ benefits from MoE and achieves $78.4\%$ top-1 accuracy ImageNet with only $15$M parameters. On four natural language processing datasets, OneS obtains $88.2\%$ MoE benefits and outperforms the best baseline by $51.7\%$ using the same architecture and training data. In addition, compared with the MoE counterpart, OneS can achieve $3.7 imes$ inference speedup due to less computation and hardware-friendly architecture.

연구 동기 및 목표

• 희박한 MoE 모델의 실용적 한계(과적합, 어려운 구현, 열악한 하드웨어 효율성)를 해결하기 위해, 그 지식을 고밀도로 구현 가능한 학생 모델로 이전하는 것.
• 단일 고밀도 모델이 다수의 사전 훈련된 전문가로부터 학습하는 새로운 과제인 지식 통합을 제안하여, 인간 교육 모델에서 여러 교사로부터 배우는 방식을 모방하는 것.
• 사전 훈련된 MoE에서 효과적으로 초기화하고 정교화하는 고밀도 학생 모델을 위한 일반적인 훈련 프레임워크를 개발하는 것. 이 프레임워크는 지식 수집과 지식 정제를 통합한다.
• 결과적으로 도출된 고밀도 모델(OneS)이 MoE 성능을 따라하거나 초월하면서도 훨씬 더 효율적이고 구현이 용이하다는 것을 입증하는 것.

제안 방법

• 지식 수집은 네 가지 방법을 사용한다: 합산, 평균화, 상위-k 지식 수집(Top-KG), 그리고 새로운 특이값 분해 지식 수집(SVD-KG)으로, MoE 전문가들로부터 핵심 지식을 추출하고 통합한다.
• SVD-KG는 전문가 가중치를 특이값 분해(SVD)를 통해 가장 중요한 특이 성분을 유지하여, 고밀도 학생 모델의 피드포워드 네트워크(FFN) 레이어를 초기화하는 데 사용한다.
• Top-KG는 전문가들 간에 가장 활성화된 상위-k 뉴런 또는 가중치 성분을 선택하여 학생 모델을 초기화함으로써 고영향도 지식에 집중한다.
• 지식 정제는 전체 MoE 교사 모델의 소프트 레이블을 사용하여 수행되며, 지식 수집 단계에서 발생할 수 있는 노이즈를 줄이고 일반화 성능을 향상시킨다.
• 이 프레임워크는 두 단계 훈련 과정을 사용한다: 먼저 지식 수집을 통해 학생 모델을 초기화하고, 다음으로 소프트 레이블과 교차 엔트로피 손실을 사용하여 지식 정제를 통해 미세조정한다.
• 학습자 모델 아키텍처는 표준 고밀도 모델과 동일하여 하드웨어 호환성과 효율적인 추론을 보장한다.

실험 결과

연구 질문

• RQ1 사전 훈련된 MoE에서 유래한 지식을 사용해 고밀도 학생 모델을 효과적으로 초기화할 수 있는가? 이는 성능을 유지하면서도 구현 가능성을 향상시키는가?
• RQ2 합산, 평균화, Top-KG, SVD-KG 중 어떤 지식 수집 방법이 MoE 전문가들로부터 고밀도 학생 모델을 초기화할 때 가장 높은 성능을 낼 수 있는가?
• RQ3 지식 정제가 지식 수집 단계에서 발생한 노이즈를 줄이고 초기화 과정의 성능을 향상시키는 데 얼마나 기여하는가?
• RQ4 정확도, 파라미터 효율성, 추론 속도 측면에서 최종 고밀도 학생 모델은 MoE와 어떻게 비교되는가?
• RQ5 지식 통합 고밀도 모델은 시각 및 NLP 벤치마크에서 기존의 정제 및 희박 모델 기준을 초월하는가?

주요 결과

• ImageNet에서 OneS는 오직 1500만 개의 파라미터로 상위-1 정확도 78.4%를 달성하며, MoE의 성능 이점의 61.7%를 유지한다.
• 네 가지 NLP 벤치마크에서 OneS는 파라미터의 46%만 사용하면서 MoE의 성능 이점의 88.2%를 달성하며, 최고의 기준 모델 대비 최대 51.7% 높은 성능을 기록한다.
• 계산량 감소와 하드웨어 우수한 고밀도 아키텍처 덕분에 OneS는 MoE 대비 3.7배 더 빠른 추론 속도를 제공한다.
• 절단 실험 결과, 지식 수집과 지식 정제 모두 필수적임을 확인: 둘 중 하나를 제거하면 성능이 크게 하락하며, 특히 지식 수집이 초반 성능 향상에 가장 기여한다.
• 유사한 훈련 에포크 수를 기준으로 OneS는 기준 모델을 능가하며, 더 많은 훈련을 거치면서도 지속적으로 성능 향상을 보이며, MoE 기반 모델인 WideNet과 달리 성능이 정체되는 경향을 보이지 않는다.
• SQuAD1.1 및 SST-2와 같은 소규모 데이터셋에서도 OneS는 MoE 교사 모델을 초월하는 성능을 기록하여, 고밀도로 일반화 가능한 학생 아키텍처 덕분에 과적합이 감소했음을 시사한다.