SlimWise, 디코딩 단계에서만 MoE 전문가를 줄여 추론 효율 개선

배경

Mixture-of-Experts(MoE) 모델은 토큰마다 일부 전문가만 활성화해 대규모 모델의 계산량을 줄입니다. 그러나 실제 배치 기반 자동회귀 디코딩에서는 이 구조의 장점이 곧바로 서비스 효율로 이어지지 않을 수 있습니다. 여러 요청이 서로 다른 전문가로 라우팅되면 하나의 배치가 전문가 풀 대부분에 접근하게 되고, 전문가 가중치를 반복적으로 읽고 이동하는 비용이 커집니다. 이때 연산량보다 가중치 트래픽이 더 중요한 병목이 될 수 있습니다.

SlimWise는 추론 전체에 동일한 가지치기 정책을 적용하지 않습니다. 입력 문맥을 처리하는 프리필 단계와 토큰을 순차적으로 생성하는 디코딩 단계의 특성이 다르다는 점에 주목해, 프리필은 전체 모델로 수행하고 디코딩만 축소하는 구조를 택합니다.

핵심 구성

• 프리필에서는 전체 전문가 유지: 프리필은 입력 문맥을 한 번에 처리하는 단계이며 계산 중심의 성격이 강합니다. 이 단계에서 전문가를 제거해도 처리량 이득이 제한적일 수 있으므로, SlimWise는 불필요한 품질 저하를 피합니다.
• 디코딩에서 전문가 풀 축소: 디코딩은 토큰마다 전문가 접근을 반복합니다. 전문가 수를 줄이면 생성 루프에서 발생하는 가중치 이동과 메모리 트래픽을 낮출 수 있습니다.
• KV 캐시 직접 재사용: 전체 모델이 프리필에서 생성한 KV 캐시를 가지치기된 디코더가 추가 변환 없이 사용합니다. 이 학습 없는 캐시 전달은 단순히 디코더를 줄였을 때 발생하는 품질 손실을 줄이는 핵심 요소입니다.
• 선택적 경량 증류: 남아 있는 정확도 저하와 출력 길이 변화를 완화하기 위해, 디코더가 전체 모델의 KV 캐시에서 생성 continuation을 이어가도록 학습할 수 있습니다. 업데이트 대상은 소수의 파라미터로 제한됩니다.

저자들은 일반적인 벤치마크 정확도만으로는 가지치기의 실제 영향을 놓칠 수 있다고 설명합니다. 모델의 점수가 비슷하더라도 생성 길이가 달라지면 실제 계산량과 서비스 비용이 달라질 수 있기 때문입니다. SlimWise는 vLLM에 구현됐으며, 프리필과 디코딩을 서로 다른 작업자로 분리하는 PD disaggregation과 한 환경에 함께 배치하는 방식 모두를 지원합니다. 두 가지 MoE 백본과 세 가지 가지치기 기준을 평가했으며, Qwen3.6-35B-A3B에서는 전문가를 50% 줄였을 때 디코딩 처리량이 최대 1.81배 향상됐고 정확도 손실은 작게 보고됐습니다.

의미와 한계

SlimWise의 중요한 시사점은 모델 압축도 추론 단계별 병목에 맞춰야 한다는 것입니다. 문맥 처리에서는 전체 모델을 사용해 표현 품질을 지키고, 전문가 가중치 접근이 반복되는 디코딩 루프만 줄이면 품질과 효율 사이의 균형점을 만들 수 있습니다.

다만 효과가 모든 환경에서 동일하게 나타난다고 보기는 어렵습니다. 배치 크기, 입력 및 출력 길이, 전문가 라우팅 분포, 프리필과 디코딩의 배치 방식에 따라 실제 이득이 달라질 수 있습니다. 증류를 선택하면 추가 학습과 운영 비용도 필요합니다. 따라서 SlimWise는 무조건적인 최대 가지치기라기보다, 전체 모델과 공격적인 압축 사이를 서비스 특성에 맞게 조정하는 설계로 보는 편이 적절합니다.

출처: Hugging Face Daily Papers