UniMoMo: 추천 MoE를 전문가 병합으로 빠르게 배포하기

도입

희소 mixture-of-experts(MoE)는 추천 모델의 용량을 키우는 데 유용하다. 입력마다 일부 전문가만 활성화하므로 계산을 조건부로 수행하면서도 더 다양한 사용자와 아이템 패턴을 담을 수 있다. 하지만 배포 단계에서는 문제가 남는다. 학습된 checkpoint는 여전히 전체 전문가 뱅크를 저장하고, 라우터도 원래의 모든 전문가를 대상으로 동작한다. 대규모 추천 서비스에서는 메모리, 지연 시간, 운영 비용 부담으로 이어질 수 있다.

UniMoMo는 바로 이 배포 문제를 다룬다. 이미 학습된 추천 MoE를 주어진 전문가 예산에 맞춰 더 작은 표준 MoE로 변환하되, 추론 시 별도의 압축 전용 모듈을 추가하지 않는 것이 목표다.

핵심 포인트

• 파라미터 유사성보다 기능적 유사성: 전문가의 가중치가 비슷하다고 해서 실제 입력에 같은 방식으로 반응한다고 볼 수는 없다. UniMoMo는 무라벨 보정 데이터를 사용해 동일한 추천 상태에 대한 전문가 출력이 얼마나 비슷한지 측정하고, 기능적으로 가까운 전문가를 묶는다.

• 제약이 있는 그래프 조대화 문제로 정식화: 전문가 병합을 명시적인 전문가 수 예산 아래에서의 constrained graph coarsening 문제로 본다. 단순히 전문가를 제거하거나 평균내는 방식보다, 어떤 전문가를 같은 그룹으로 합칠지 체계적으로 결정할 수 있다.

• 트래픽이 많은 전문가 보호: 실제 MoE 추천 모델에서는 전문가 사용량이 균등하지 않다. 자주 라우팅되는 전문가는 중요한 추천 경로를 담당할 가능성이 높다. UniMoMo는 라우팅 노출도를 기반으로 계층 적응형 보호 메커니즘을 적용해 고트래픽 전문가가 무리하게 병합되는 것을 제한한다.

• 압축 후에도 표준 MoE 유지: 변환된 checkpoint는 일반적인 top-k MoE로 동작한다. 추론 서버에 특별한 모듈을 추가하지 않아도 되므로, 운영 복잡도와 통합 리스크를 줄일 수 있다.

• 보정 기반 재구성: SwiGLU 전문가를 병합한 뒤에는 최소제곱 보정을 사용해 원래 모델의 중간 활성값을 더 잘 복원한다. 이는 병합으로 인한 표현 변화와 성능 저하를 줄이는 장치다.

결과와 의미

Amazon Beauty, KuaiRec, TenRec 데이터셋과 2개, 4개, 6개 MoE block 설정에서, UniMoMo의 4전문가 checkpoint는 원본 대비 5회 평균 NDCG@10 비율 99.92%〜102.30%를 기록했다. A100에서 측정한 속도 향상은 1.28배〜1.63배다. 더 공격적인 2전문가, top-1 설정에서는 NDCG@10 비율 98.36%〜104.24%, 속도 향상 1.47배〜2.21배가 보고됐다.

이 결과는 학습된 추천 MoE 내부에 기능적으로 병합 가능한 중복성이 있을 수 있음을 보여준다. 또한 UniMoMo의 장점은 이를 실제 서빙 예산에 맞춘 변환 워크플로로 만든 데 있다. 서비스 팀은 품질과 지연 시간 요구에 따라 4전문가 또는 2전문가 버전을 선택할 수 있고, 모델 구조는 표준 MoE로 유지된다.

물론 보정 데이터가 실제 온라인 분포를 얼마나 잘 대표하는지는 중요하다. 트래픽 패턴이 크게 바뀌면 보호해야 할 전문가도 달라질 수 있다. 그럼에도 UniMoMo는 추천 MoE 압축을 단순한 모델 축소가 아니라, 라우팅 행동과 추론 효율, 배포 단순성을 함께 고려하는 문제로 다룬다는 점에서 의미가 크다.

출처: Hugging Face Daily Papers