FlexMoE: MoE 언어 모델을 위한 One-for-All 중첩형 전문가 내부 프루닝 (Nested Intra-Expert Pruning)

Mixture-of-Experts (MoE) 언어 모델은 희소하게 활성화되는 전문가(experts)를 통해 모델의 능력을 확장하며, 이러한 구조는 현대 대규모 모델의 표준적인 레시피가 되었습니다. 그러나 희소 활성화(sparse activation)가 모든 전문가를 저장하고 서비스해야 하는 배포 부담을 제거해주지는 않으며, 가용한 배포 예산(deployment budget)은 장치, 사용자 및 워크로드에 따라 크게 달라질 수 있습니다. 기존의 MoE 압축 방법들은 여전히 주로 고정된 예산(fixed-budget) 방식이며, 일반적으로 선택된 각 목표 예산에 대해 하나의 압축된 엔드포인트를 최적화합니다. 우리는 다른 설정, 즉 대규모 사전 학습된 MoE LLM을 다양한 예산에 걸쳐 배포 가능한 중첩된 서브네트워크(subnetworks) 제품군으로 변환하는 것을 연구합니다. 우리의 방법은 먼저 전문가 FFN 채널을 중요도에 따라 순위를 매긴 다음, 각 전문가가 자신의 채널을 프루닝(pruning)하기 위한 이산적 행동(discrete action)을 학습하도록 합니다. 비용 압박을 점진적으로 증가시킴으로써, 단 한 번의 행동 학습(action-training) 실행을 통해 높은 예산에서 낮은 예산에 이르는 일련의 행동 마스크(action masks)를 내보내며, 각 마스크는 순위가 매겨진 기본 모델 내에 중첩된 신뢰할 수 있는 더 작은 서브네트워크를 식별합니다. 또한, 우리는 중간 프루닝 예산(40%)에서 단 한 번의 복구 미세 조정(recovery fine-tune)을 사용하여 저하된 모델 품질을 회복하고, 회복된 모델을 다른 보지 못한(unseen) 예산으로 전이합니다. 전반적으로, 우리의 프레임워크는 최근의 MoE 압축 베이스라인들을 능가합니다. 구체적으로 Qwen2-57B-A14B에서 우리의 방법은 미세 조정 없이도 라우팅된 전문가 파라미터의 50%를 프루닝하면서 기본 성능의 약 99.8%를 유지합니다. 배포 측면에서, 우리의 프루닝된 서브네트워크는 실제 메모리 감소와 처리량(throughput) 이득을 제공하며, 커널 수준의 공동 설계(kernel-level co-design)를 통해 실시간 온라인 예산 전환을 추가로 지원합니다.