ConMoE: Expert-Pool Consolidation via Prototype Reassignment for MoE Compression

ConMoE: MoE 압축을 위한 프로토타입 재할당을 통한 전문가 풀 통합에 대한 기술적 요약

1. 문제 제기

전문가 혼합 (Mixture-of-Experts, MoE) 언어 모델은 토큰당 소수의 전문가만 활성화하여 토큰당 계산량을 줄임으로써 확장성을 달성합니다. 그러나 전체 라우팅된 전문가 풀은 여전히 저장되고 서비스되어야 하므로, 배포 시 상당한 메모리 병목 현상을 초래합니다. 기존 사후 학습 (post-training) 압축 방법들은 주로 전문가 가지치기(사용도나 중요도에 따라 전문가 제거) 또는 전문가 병합(여러 전문가를 새로운 융합 모듈로 결합) 을 통해 이를 해결합니다.

저자들은 이러한 접근법들이 종종 두 가지 명확히 구분된 결정을 혼동한다고 주장합니다:

1. 어떤 전문가 매개변수를 유지해야 하는가?
2. 압축 후 라우터의 원래 전문가 참조는 어떻게 표현되어야 하는가?

가지치기는 매개변수를 삭제하는 반면, 병합은 새로운 가중치를 구성하며, 종종 압축 후 미세 조정 (fine-tuning) 이나 복잡한 융합 연산자를 필요로 합니다. 본 논문은 **전문가 풀 통합 (Expert-Pool Consolidation)**이라는 보완적인 관점을 제시합니다. 이 접근법은 재사용 가능한 프로토타입으로 더 작은 세트의 사전 학습된 전문가를 유지하고, 각 원래 전문가 참조를 선택된 프로토타입 중 하나로 결정론적으로 재매핑합니다. 이는 축소된 풀의 저장과 재사용 구조를 분리하여, 가중치 업데이트 없이 원래 라우터 인터페이스를 변경 없이 유지할 수 있게 합니다.

2. 방법론: ConMoE

제안된 프레임워크인 ConMoE는 프로토타입 선택과 결정론적 재할당을 통해 작동하는 학습이 필요 없는 원샷 (one-shot) 압축 방법입니다.

A. 문제 정의
이 방법은 하나 이상의 인접 MoE 레이어를 포함하는 로컬 범위 $G$를 정의합니다. 축소 비율 $\rho$가 주어졌을 때, 목표는 원래 전문가 풀 $E_G$에서 크기 $K$의 축소된 프로토타입 풀 $P_G$를 구성하는 것입니다. 각 원래 전문가 $e$가 선택된 프로토타입 $p$로 매핑되는 재할당 맵 $m_G: E_G \to P_G$가 설정됩니다.
이상적인 목표는 재할당 비용 $L_G(P)$를 최소화하여 유지된 전문가의 중요성과 다른 전문가를 재매핑함으로써 발생하는 오류 사이의 균형을 맞추는 것입니다.

B. 프로토타입 점수 매기기 및 선택
결합 최적화 문제를 직접 해결하는 대신, ConMoE 는 상위 $K$개의 전문가를 프로토타입으로 선택하기 위해 휴리스틱 점수 규칙을 사용합니다. 각 전문가 $e$에 대해 보정 데이터를 사용하여 두 가지 점수가 계산됩니다:

1. 라우팅 조건부 기여도 ($a_e$): 라우터에 의해 선택될 때 전문가의 평균 기여도를 라우팅 확률과 출력 노름으로 가중치를 두어 측정합니다.
2. 대체 가능성 ($b_e$): 범위 내 가장 가까운 이웃까지의 거리를 측정합니다. 거리가 클수록 해당 전문가를 대체하기 어렵다는 것을 의미합니다.

최종 점수는 정규화된 기여도와 대체 가능성 점수의 곱입니다 ($s_e = \bar{a}_e \bar{b}_e$). 상위 $K$개의 전문가가 프로토타입 풀 $P_G$로 선택됩니다.

C. 결정론적 재할당
프로토타입이 선택되면, 모든 원래 전문가가 $P_G$ 내의 가장 가까운 프로토타입에 할당됩니다.

• 가중치 업데이트 없음: 선택된 프로토타입은 사전 학습된 전문가로 직접 재사용됩니다. 기본 파이프라인에서는 가중치 융합이나 미세 조정이 발생하지 않습니다.
• 라우터 보존: 원래 라우터는 변경되지 않습니다. 라우터가 토큰에 대한 전문가 세트를 선택할 때, 통합된 연산자는 이러한 호출을 할당된 프로토타입으로 리디렉션합니다. 여러 라우팅된 전문가가 동일한 프로토타입으로 매핑되는 경우, 그들의 라우팅 가중치는 단일 계수로 집계됩니다.

D. 범위 및 레이어 간 재사용
ConMoE 는 로컬 레이어 간 공유를 지원합니다. 범위는 여러 인접 레이어를 포함할 수 있으며, 중복성이 있는 경우 프로토타입이 이들 간에 공유될 수 있습니다. 그러나 범위는 먼 레이어에서 발견되는 깊이별 변환의 불일치를 방지하기 위해 제한됩니다.

3. 주요 기여

1. 정식화: 본 논문은 원샷 MoE 압축을 전문가 풀 통합으로 공식화하여, 어떤 매개변수를 저장할지 결정하는 것과 원래 라우터 슬롯을 어떻게 표현할지 결정하는 것을 명시적으로 분리합니다.
2. 프레임워크 (ConMoE): 원래 라우터 인터페이스를 보존하면서 논리적 라우팅 전문가 풀을 축소하는 학습이 필요 없는 재매핑 방법입니다. 이는 그래디언트 업데이트 없이 프로토타입을 선택하기 위해 보정 기반 기여도 및 대체 가능성 신호에 의존합니다.
3. 실증적 검증: 저자들은 재매핑 기반 통합이 Qwen3, DeepSeekMoE, OLMoE 등 여러 사전 학습된 MoE 모델에서 가지치기 및 병합에 대한 실현 가능한 대안임을 입증했습니다.
4. 애블레이션 통찰:
   • 결정론적 재할당은 프레임워크의 가장 안정적인 구성 요소로 확인되었습니다.
   • 레이어 간 공유는 유익하지만 엄격하게 로컬입니다. 범위를 너무 광범위하게 확장 (예: 8~16 레이어) 하면 깊이 불일치로 인해 성능이 저하됩니다.
   • 사후 가중치 융합(예: Arcee, 가중 평균) 은 이 설정에서 직접적인 프로토타입 재사용보다 덜 안정적인 것으로 나타났습니다.

4. 실험 결과

실험은 Qwen3-30B-A3B, deepseek-moe-16b-base, OLMoE-1B-7B-0125 세 가지 모델을 대상으로 수행되었습니다. 압축은 6 가지 다중 선택 벤치마크 (WinoGrande, ARC-C, ARC-E, BoolQ, HellaSwag, PIQA) 를 사용하여 25% 및 50% 논리적 라우팅 전문가 축소 비율로 평가되었습니다.

• 성능: ConMoE 는 강력한 가지치기 (Frequency, REAP) 및 병합 (M-SMoE, HC-SMoE) 베이스라인과 일치하거나 능가합니다.
  • DeepSeekMoE에서 ConMoE 는 25% 및 50% 축소 모두에서 최고의 평균 점수를 달성했습니다.
  • OLMoE에서는 25% 축소 시 가장 좋은 성능을 보였으며 50% 축소 시에도 경쟁력을 유지했습니다.
  • Qwen3에서는 가지치기가 강력했습니다 (특히 50% 에서), 하지만 ConMoE 는 경쟁력을 유지하며 WinoGrande 및 BoolQ 와 같은 특정 작업에서 강력한 성능을 유지했습니다.
• 애블레이션:
  • 범위 크기: 작은 범위 (14 레이어) 는 동등하거나 최상의 결과를 산출했습니다. 더 큰 범위 (816 레이어) 는 전문가 대체 가능성이 로컬임을 확인시켜 주는 상당한 성능 저하를 초래했습니다.
  • 선택 정책: 라우팅 조건부 기여도가 지배적인 신호였습니다. 거리만 기반한 선택은 성능이 낮았으며, 대체 가능성은 용량 할당을 위한 정규화제로서 가장 잘 작동했습니다.
  • 융합: 사후 가중치 융합을 적용한 방법보다 직접적인 프로토타입 재매핑이 더 우수했으며, 이는 이득이 가중치 병합이 아니라 재매핑 구조 자체에서 비롯됨을 시사합니다.

5. 중요성 및 주장

본 논문은 결정론적 프로토타입 재할당이 전통적인 가지치기 및 병합에 대한 견고하고 학습이 필요 없는 대안을 제공한다고 주장합니다. 압축을 재사용 가능한 프로토타입으로의 전문가 풀 통합으로 간주함으로써, ConMoE 는 원래 모델의 라우터 인터페이스를 유지하고 가중치 융합이나 미세 조정과 관련된 불안정성을 피하면서 상당한 메모리 절감 (논리적 축소) 을 달성합니다.

저자들은 그들의 접근 방식이 범위 측면에서 겸손하다고 강조합니다:

• 이는 학습이 필요 없으나 라우팅 통계를 위해 보정 데이터에 의존합니다.
• 레이어 간 재사용은 로컬입니다. 모델 전체 깊이 전반에 걸친 전문가의 글로벌 상호 교환성을 가정하지 않습니다.
• 보고된 절감액은 논리적입니다. 배포 시 물리적 메모리 절감을 실현하려면 단순히 구체화된 체크포인트가 아닌 공유 프로토타입 런타임 또는 체크포인트 형식이 필요합니다.

이 연구는 "재사용 구조"가 삭제나 융합의 부산물이 아니라 압축된 모델의 명시적인 부분이 될 수 있음을 보여줌으로써 MoE 압축을 위한 안정적인 기준선을 제공한다는 점에서 차별화됩니다.