Speculating Experts Accelerates Inference for Mixture-of-Experts

이 논문은 메모리 제약 환경에서 CPU-GPU 전송 병목 현상을 해결하기 위해 내부 모델 표현을 활용해 향후 전문가를 예측하여 사전 가져오기를 수행함으로써, MoE 모델의 추론 속도를 최대 14% 향상시키는 방법을 제안합니다.

핵심

🏪 비유: 거대한 슈퍼마켓과 바쁜 점원

1. 문제 상황: "무거운 짐을 들고 오느라 시간이 걸려요"
거대한 AI 모델 (예: Qwen3-30B) 은 마치 엄청나게 큰 슈퍼마켓과 같습니다. 이 슈퍼마켓에는 수천 개의 '전문가 코너 (Expert)'가 있는데, 각각은 수학, 코딩, 일반 상식 등 특정 분야의 전문가들입니다.

하지만 우리 집 컴퓨터 (GPU) 는 이 모든 전문가를 한 번에 다 모을 공간이 부족합니다. 그래서 대부분의 전문가들은 **창고 (CPU 메모리)**에 두고, 필요한 사람만 **매장 (GPU)**으로 데려와서 일하게 합니다.

• 기존 방식 (On-Demand Loading):
  고객이 "수학 문제를 풀어줘!"라고 하면, 점원은 창고로 달려가 "수학 전문가"를 데려와서 일을 시킵니다. 그다음 "코딩 문제를 풀어줘!"라고 하면 다시 창고로 가서 "코딩 전문가"를 데려옵니다.
  • 문제: 창고와 매장 사이를 오가는 이동 시간 (데이터 전송) 이 너무 길어서, 실제 일하는 시간보다 이동 시간이 훨씬 더 길어집니다. 마치 요리사가 재료를 다듬는 시간보다 마트에서 장보러 가는 시간이 더 긴 상황입니다.

2. 해결책: "내일의 메뉴를 미리 예상해서 준비하기"
이 논문은 **"다음에 누가 필요할지 미리 예측해서, 계산이 진행되는 동안 미리 데려오자"**는 아이디어를 제시합니다.

• 새로운 방식 (Expert Prefetching):
  점원이 지금 "수학 문제"를 풀고 있을 때, AI 는 "아, 다음 질문은 아마 코딩 문제일 거야"라고 미리 추측합니다.
  • 기적 같은 일: 점원이 지금 수학을 풀고 있는 동안, 다른 직원은 이미 창고에서 "코딩 전문가"를 데려와서 준비해 둡니다.
  • 결과: 계산이 끝나는 순간, 코딩 전문가가 바로 옆에 대기하고 있어서 이동 시간이 0 이 됩니다. 기다리는 시간이 사라진 것입니다.

🔍 이 기술의 핵심 3 가지

1. "내면의 신호"로 미래를 읽다 (Quasi-Hidden State)
AI 가 다음에 어떤 전문가를 부를지 어떻게 알까요?
AI 가 문제를 풀고 있는 중에도, 그 내부에서 흐르는 '생각의 흐름 (데이터)'을 분석합니다. 이 흐름을 보면 "다음 단계는 수학 전문가가 필요할 것 같다"는 신호가 이미 나옵니다.

• 비유: 요리사가 국을 끓이고 있는데, 국물 냄새를 맡아 "다음엔 소금이나 후추를 넣어야겠다"라고 미리 알아채는 것과 같습니다.

2. "틀려도 괜찮아" (Speculative Execution)
만약 AI 가 "다음은 코딩 문제야"라고 예측했는데, 실제로는 "역사 문제"였다면 어떨까요?
기존 방식은 "아, 틀렸네. 다시 창고로 가서 역사 전문가를 데려와야지"라고 해서 시간을 낭비합니다.
하지만 이 연구는 **"일단 코딩 전문가를 불러서 일하게 해보자"**고 합니다.

• 비유: 비가 올 것 같아서 우산을 챙겼는데, 안 왔다고 해서 우산을 다시 창고로 가져가는 게 아니라, 그냥 우산을 들고 있는 동안 다른 일을 하는 것입니다. 만약 틀렸다면, 그 전문가가 한 일은 버리고 다시 시작하면 되지만, 대부분의 경우 예측이 맞거나 큰 지장이 없기 때문에 전체 속도가 빨라집니다.

3. "약한 예측을 보강하는 보조교사" (Lightweight Estimator)
어떤 모델은 처음 몇 단계에서 예측을 잘 못합니다. 이때는 아주 작고 간단한 '보조 교사 (Neural Estimator)'를 붙여줍니다. 이 교사는 다음 단계의 전문가를 더 정확하게 예측하도록 훈련되어, 예측 실패를 줄여줍니다.

🚀 실제 효과는?

이 기술을 적용한 결과, AI 가 한 마디를 내뱉는 데 걸리는 시간 (TPOT) 이 최대 14% 빨라졌습니다.

• 기존: 창고로 이동하는 시간이 전체 시간의 80% 이상을 차지했습니다.
• 이제: 이동 시간이 계산 시간과 겹쳐서 (중첩되어) 사라졌기 때문에, 훨씬 더 빠르게 응답합니다.

💡 요약

이 논문은 **"AI 가 다음에 무엇을 할지 미리 예측해서, 기다리는 시간을 계산하는 시간으로 바꿔버리는 기술"**입니다.

마치 요리사가 재료를 다듬는 동안, 도우미가 다음 재료를 미리 씻어놓는 것과 같습니다. 이렇게 하면 요리사 (GPU) 는 멈추지 않고 계속 일할 수 있어, 전체 요리 속도가 훨씬 빨라집니다. 덕분에 일반 가정용 컴퓨터에서도 무거운 AI 모델을 더 빠르고 부드럽게 사용할 수 있게 되었습니다.

기술 요약

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 용량을 확장하기 위해 혼합 전문가 (Mixture-of-Experts, MoE) 아키텍처가 널리 사용되고 있습니다. MoE 는 토큰당 계산량을 줄이면서 매개변수 수를 늘릴 수 있게 해줍니다. 그러나 메모리 제약이 있는 환경 (예: 소비자용 GPU) 에서 MoE 를 실행할 때 다음과 같은 병목 현상이 발생합니다.

• 메모리 부족: 모든 전문가 (Expert) 가 GPU 메모리에 상주할 수 없어, 대부분의 전문가 가중치를 CPU RAM 으로 오프로드 (offload) 해야 합니다.
• I/O 병목: 디코딩 (Decoding) 단계에서 활성화된 전문가를 CPU 에서 GPU 로 전송하는 시간이 전체 출력 토큰당 소요 시간 (TPOT, Time Per Output Token) 의 84~88% 를 차지합니다.
• 기존 방식의 한계: 기존 '온디맨드 (On-Demand)' 방식은 라우팅 (Routing) 후 필요한 전문가만 전송하므로, 전송 대기 시간이 계산과 병렬화되지 않아 성능이 저하됩니다.

2. 제안 방법 (Methodology)

이 논문은 예측 기반 전문가 프리페칭 (Expert Prefetching) 과 예측 실행 (Speculative Execution) 을 결합하여 CPU-GPU 간 데이터 전송과 GPU 계산을 겹치게 (Overlap) 하는 새로운 방식을 제안합니다.

A. 신호 기반 프리페칭 (Signals for Prefetching)

다음 레이어에서 어떤 전문가가 선택될지 예측하기 위해 모델 내부의 표현 (Representations) 을 활용합니다.

• 기본 벡터 (Default Vector, $d_l$): 각 전문가의 평균 활성화 값을 나타내는 오프라인 계산 벡터입니다.
• 준-잠재 상태 (Quasi-Hidden State, $q_l$): 현재 레이어의 정규화된 잔차 스트림 ($r_l$) 과 기본 벡터 ($d_l$) 를 결합하여 다음 레이어의 라우터 입력을 근사한 상태입니다.
  • $q_l = LN_{l+1}(d_l + r_l)$
• 효과: $q_l$ 은 실제 라우터 입력 ($s_{l+1}$) 과 더 높은 코사인 유사도를 보여, 다음 레이어의 전문가 선택을 더 정확하게 예측할 수 있음을 입증했습니다.

B. 예측 실행 (Speculative Execution)

기존의 프리페칭 방식은 예측이 틀리면 (Cache Miss) 다시 전송해야 했지만, 이 논문은 예측된 전문가를 그대로 실행합니다.

• 라우터가 선택한 '진짜' 전문가를 기다리지 않고, 프리페칭된 전문가를 즉시 계산에 사용합니다.
• 이는 전송 지연을 완전히 숨기고 (Hide), 계산과 전송을 최대한 겹치게 만듭니다.
• 정확도 유지: 하류 작업 (Downstream Task) 의 정확도를 유지하면서 예측된 전문가를 실행할 수 있음을 실험적으로 증명했습니다.

C. 경량 신경 추정기 (Lightweight Neural Estimator)

일부 모델 (예: Qwen3-30B-A3B) 의 초기 레이어에서는 표현의 변화 (Drift) 가 커서 단순한 $q_l$ 기반 예측의 정확도가 떨어집니다. 이를 보완하기 위해:

• 신경 추정기: $q_l$ 에서 다음 레이어의 라우터 로짓 (Logits) 을 직접 예측하는 얕은 신경망을 학습시킵니다.
• 하이브리드 전략: 초기 레이어 등 예측 정확도가 낮은 구간에만 추정기를 적용하고, 나머지는 $q_l$ 기반 예측을 사용하는 하이브리드 방식을 도입했습니다.

3. 주요 기여 (Key Contributions)

1. 파라미터 없는 프리페칭: 추가 학습 없이 모델 내부 표현 ($q_l$) 만으로 미래의 전문가 라우팅을 예측할 수 있음을 발견했습니다.
2. 정확도 보존 예측 실행: 예측된 전문가를 재전송 없이 직접 실행하더라도 하류 작업의 정확도가 크게 저하되지 않음을 입증했습니다. 이는 전송 오버헤드를 줄이는 핵심 요소입니다.
3. 최적화된 추론 엔진 통합: 오픈소스 추론 엔진인 YALIS에 이 방식을 통합하여, 하드웨어 및 모델 구성에 관계없이 TPOT 5~14% 감소를 달성했습니다.
4. 경량 추정기 개발: 표현 변화가 큰 레이어에서 예측 정확도를 높여 성능 저하를 완화하는 경량 신경 추정기를 제안했습니다.

4. 실험 결과 (Results)

• 성능 향상 (TPOT 감소):
  • Qwen3-30B-A3B (A6000 GPU): 온디맨드 로딩 대비 9~14% 의 TPOT 단축. 특히 시퀀스 길이가 길어질수록 (1024 → 65536) 계산 시간이 증가하여 전송과 계산의 겹침 효과가 더 커짐.
  • GPT-OSS 모델: 다양한 모델에서 일관된 성능 향상을 보임.
  • 하드웨어 영향: A6000(상대적으로 약한 GPU) 에서 1214% 향상, A100/GH200(강력한 GPU) 에서는 58% 향상. GPU 의 계산 처리량이 높을수록 전송 시간이 상대적으로 더 큰 병목이 되어 프리페칭의 이점이 큼.
• 정확도 분석:
  • GPT-OSS: 단순 프리페칭 (Router-PF) 만으로도 정확도 유지.
  • Qwen3-30B-A3B: 초기 레이어의 높은 Drift 로 인해 Router-PF 는 정확도 저하 (StrategyQA, HumanEval 등) 발생. 하지만 Hybrid-PF(초기 레이어에 추정기 적용) 를 사용하면 기준선 (Baseline) 에 가까운 정확도로 회복됨 (GSM8k 에서 37% 의 정확도 격차 회복).
• 히트율 (Hit Rate):
  • 추정기를 사용하면 초기 레이어에서 예측 히트율을 최대 25% 까지 향상시킴.
  • 400 만 토큰 정도의 소량 학습 데이터로도 추정기가 90% 이상의 평균 히트율 달성.

5. 의의 및 결론 (Significance)

• 로컬 배포의 실용성: 이 연구는 소비자용 하드웨어 (단일 GPU) 에서도 대규모 MoE 모델을 효율적으로 실행할 수 있는 길을 열었습니다. CPU-GPU 간 전송 병목을 해결함으로써 로컬 LLM 추론의 속도를 크게 개선합니다.
• 아키텍처 혁신: 라우터 호출을 줄이거나 제거 (Router-free inference) 하여 추론 파이프라인을 단순화하고 병렬성을 높일 수 있는 가능성을 제시했습니다.
• 일반성: 다양한 MoE 아키텍처와 하드웨어 환경에서 적용 가능한 범용적인 솔루션을 제공합니다.

요약하자면, 이 논문은 MoE 모델의 메모리 오프로드 병목 현상을 해결하기 위해 모델 내부 표현을 활용한 미래 전문가 예측과 예측된 전문가의 직접 실행을 결합함으로써, 정확도를 희생하지 않으면서 추론 속도를 획기적으로 개선한 획기적인 연구입니다.