MoE 모델의 채널 선택, 보정 데이터 없이 가능
원제 Channel Selection for Mixture-of-Experts Decoding Needs No Calibration Data
추천 22댓글 5
Key Point
현재 MoE 모델 최적화의 표준인 보정 기반 임계값이 예산 보장과 성능 향상을 모두 충족하지 못하며, 보정 없는 top-k와 경량 리더로 더 안정적이고 빠른 디코딩을 달성할 수 있음을 보여준다.
핵심 요약
- Mixture-of-Experts(MoE) 모델은 대화형 디코딩 중 라우팅된 전문가 가중치를 메모리에서 로드할 때 병목을 겪으며, 채널 선택기는 각 토큰에 실질적으로 영향을 주는 채널만 읽어 이를 완화한다.
- 기존 방식은 오프라인 데이터로 보정한 임계값을 사용하지만, 임계값을 통과하는 채널 수는 입력마다 달라져 실제 읽기 예산이 불안정하다: Prox는 18% 읽기 예산으로 보정되었으나 실제로는 평균 12.9%만 읽는다.
- Top-k 방식은 각 토큰의 k개 최고점수 채널을 선택하며 보정 데이터가 필요 없다.
- 13%와 18% 읽기 예산에서 top-k는 임계값 방식과 비교해 43개 비교 중 40개에서 동등하거나 우수하며, 생성 작업 23개는 모두 top-k가 낫다.
- Top-k는 보정된 임계값이 유발하는 Gemma 4의 반복 루프 붕괴(repetition-loop collapse) 문제도 피한다.
- 저자들이 소개한 리더는 모델 가중치만으로 구성되며, 삼진 스케치로 채널을 필터링하고 저랭크 맵으로 생략된 채널을 복구한다.
- 세 개의 사후학습 MoE 모델에서 리더는 라우팅된 전문가 읽기의 17.8%만 사용하면서 전체 모델 정확도의 98.1%를 유지한다.
- Triton 커널은 배치 1에서 최적화된 vLLM보다 MoE 블록을 1.60~1.90배 빠르게 실행하며, 더 큰 배치에서는 최대 2.69배 향상을 보인다.