💬 NLP

Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts

본 논문은 최적화된 과거 라우팅 결정의 메모리와 신뢰도 기반의 폴백 메커니즘을 활용하여 고정된 라우터보다 분포 변화를 더 효과적으로 처리함으로써 혼합 전문가 모델을 동적으로 개선하는 검색 증강 라우팅 프레임워크인 kNN-MoE 를 소개합니다.

원저자: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

게시일 2026-05-26

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신에게 방대하고 매우 지능적인 전문가 도서관 (즉, "혼합 전문가 (Mixture-of-Experts)" 모델) 이 있다고 상상해 보세요. 이 도서관에 질문을 던지면, 라우터 (Router) 가 사서처럼 행동합니다. 라우터의 역할은 당신의 질문을 살펴보고, 어떤 특정 전문가가 그 질문에 답하기에 가장 적합한지 빠르게 결정하는 것입니다.

일반적으로 이 사서는 한 번 훈련된 후 고정됩니다. 그들은 훈련 중에 배운 것을 기억에 의존합니다. 하지만 여기서 문제가 발생합니다. 사서가 본 적 없는 기괴하거나 어렵거나 완전히 새로운 유형의 질문을 던지면, 그들은 잘못 추측하여 당신을 잘못된 전문가에게 보낼 수 있습니다. 이는 사서가 단어들이 모호하게 비슷하다는 이유만으로 의학 질문을 역사 전문가에게 보내는 것과 같습니다.

이 논문은 과거의 성공 사례를 바탕으로 이 사서에게 "요약 자료 (cheat sheet)"를 제공하는 kNN-MoE라는 현명한 업그레이드를 소개합니다. 작동 원리는 다음과 같이 간단한 개념으로 나뉩니다:

1. "요약 자료" (기억 구성)

시스템이 실제 질문에 답하기 전에, 연구자들은 연습 질문 세트 (참조 데이터셋) 를 가져와 도서관을 통과시킵니다.

실험: 이러한 연습 질문의 모든 단어에 대해 다음과 같이 묻습니다. "만약 우리가 마법처럼 사서의 결정을 지금 당장 바꿀 수 있다면, 어떤 전문가가 절대적으로 최선의 답을 주었을 것인가?"
결과: 그들은 각 특정 순간에 대한 "완벽한" 전문가를 찾아 기록합니다. 그리고 이 쌍들을 저장합니다: "이 특정 질문 입력" + "완벽한 전문가 선택."
비유: 사서가 기말고사를 준비한다고 상상해 보세요. 규칙을 단순히 암기하는 대신, 그들은 거대한 인덱스 카드 덱을 만듭니다. 한쪽 면에는 까다로운 질문이 있고, 다른 쪽 면에는 과거에 그것을 완벽하게 해결한 정확한 전문가가 적혀 있습니다.

2. "지능적 조회" (추론)

이제 실제 사용자가 질문을 하면, 시스템은 두 가지 작업을 동시에 수행합니다:

고정된 사서: 원래 고정된 라우터가 훈련을 바탕으로 최선의 추측을 합니다.
요약 자료: 시스템은 인덱스 카드 덱에서 사용자의 질문을 조회하여 가장 유사한 과거 질문들을 찾습니다.

3. "신뢰도 투표" (적응형 혼합)

이 부분이 가장 중요합니다. 시스템은 요약 자료를 맹목적으로 신뢰하지 않습니다. 과거 사례들이 현재 질문과 얼마나 유사한지를 확인합니다.

높은 신뢰도: 과거 사례들이 현재 질문과 거의 동일하다면, 시스템은 "사서는 아마도 이 까다로운 질문에 대해 확신이 없을 것이지만, 우리의 요약 자료에 따르면 완벽한 전문가가 X 입니다"라고 말합니다. 그런 다음 사서의 추측과 요약 자료의 조언을 혼합하되, 요약 자료에 크게 의존합니다.
낮은 신뢰도: 현재 질문이 완전히 독특하고 요약 자료에 잘 맞는 것이 아무것도 없다면, 시스템은 "여기서는 요약 자료는 쓸모가 없다; 우리는 단순히 노이즈를 추가할지도 모른다"라고 말합니다. 이 경우 조회를 무시하고 원래 사서의 고정된 결정에 의존합니다.

이것이 중요한 이유

이 논문은 이 방법이 두 가지 극단 사이의 "황금 지점"이라고 주장합니다:

아무것도 하지 않음 (Zero-Shot): 고정된 사서만 사용하는 것. 이는 빠르지만 어렵거나 새로운 질문에서는 실패할 수 있습니다.
재훈련 (지도 미세 조정): 사서에게 처음부터 새로운 규칙을 가르치는 것. 이는 잘 작동하지만 매우 느리고 비용이 많이 들며, 모든 새로운 작업마다 전체 과정을 다시 수행해야 합니다.

kNN-MoE는 무거운 비용 없이 재훈련의 성능 향상을 얻습니다. 이는 사서에게 다시 학교에 다니게 하는 대신, 그들의 가장 훌륭한 과거 행동에 대한 동적이고 검색 가능한 기억을 제공하는 것과 같습니다.

논문에서 도출된 주요 발견 사항

어려운 질문에 효과적: 시스템은 원래 사서가 혼란스러워할 때 (높은 "퍼플렉시티") 가장 큰 도움을 줍니다. 사서가 이미 확신할 때는 시스템을 물러나게 하여 상황을 악화시키는 것을 방지합니다.
적은 것이 더 많다: 놀랍게도, 단 하나의 과거 예시 (가장 가까운 단일 일치) 를 보는 것이 많은 예시를 평균내는 것보다 더 잘 작동했습니다. 논문은 전문가 라우팅의 경우 과거의 "의견"이 너무 많으면 실제로 신호가 희석된다고 제안합니다.
속도 대 정확도: 전체 모델을 재훈련하는 것보다 이 "요약 자료"를 만드는 것 (오프라인) 이 훨씬 빠릅니다. 실제 답변 단계에서는 아주 작은 지연 (약 3-4% 느려짐) 을 추가하지만, 의학 시험 및 코딩과 같은 어려운 작업에서 정확도를 크게 향상시킵니다.

주의점 (한계점)

이 논문은 이 시스템이 요약 자료를 구축하기 위해 레이블이 지정된 데이터의 "참조 세트"가 필요하다고 지적합니다. 만약 학습할 수 있는 유사한 과거 사례가 전혀 없는 상황에 처해 있다면, 이 방법은 도움이 될 수 없습니다. 이 방법은 "과거에 작동했던 것"이 "지금 작동할 것"에 대한 좋은 지침이라는 가정에 의존합니다.

요약하자면, kNN-MoE는 AI 전문가들이 완벽한 결정의 역사를 엿볼 수 있게 함으로써 더 똑똑하게 만드는 방법이지만, 오직 그렇게 하는 것이 안전할 때만 가능합니다.

기술 요약: kNN-MoE

문제 정의

혼합 전문가 (MoE) 아키텍처는 토큰을 희소 하위 집합의 전문가로 분배하는 매개변수 기반 "라우터"를 사용하여 대규모 언어 모델 (LLM) 을 효율적으로 확장합니다. 그러나 표준 라우터는 한 번 훈련된 후 고정되는 경량 분류기입니다. 분포 내 데이터에서는 효과적이지만, 분포 변화 (예: 전문 도메인 또는 분포 외 작업) 에 직면할 때 추론 중 유연성이 부족합니다. 이러한 시나리오에서 라우터는 고정된 사전 지식에서 외삽해야 하므로 종종 비최적의 전문가 선택과 높은 퍼플렉시티를 초래합니다. C3PO 와 같은 기존 테스트 시간 적응 방법은 온라인으로 라우팅 로짓을 미세 조정하려 시도하지만, 상당한 계산 오버헤드 (기반 대비 약 5 배) 를 초래합니다. 다른 접근 방식은 지도 미세 조정 (SFT) 을 통한 영구적인 매개변수 업데이트를 요구하며, 이는 계산 집약적이고 추론 효율성이 낮습니다.

방법론: kNN-MoE

저자들은 모델의 매개변수를 수정하지 않고 전문가 할당 결정을 정제하는 검색 증강 라우팅 프레임워크인 kNN-MoE를 제안합니다. 이 방법은 두 가지 명확한 단계로 작동합니다:

1. 메모리 구축 (오프라인)

시스템은 각 MoE 레이어에 대해 키 - 값 쌍 $(k_t, v_t)$ 으로 구성된 비매개변수 메모리 저장소 $\mathcal{M}$ 을 구축합니다:

키 ( $k_t$ ): 참조 데이터셋 $D_{ref}$ 에서 가져온 라우터 입력 (은닉 상태) $x_t$ .
값 ( $v_t$ ): 고정된 네트워크 매개변수를 기반으로 타겟 토큰 $y_t$ $y_{t}$ 의 가능도를 최대화하도록 라우팅 로짓을 최적화하여 유도된 "국소 최적" 전문가 할당.
- 고정된 라우터의 출력을 저장하는 대신, 이 방법은 $D_{ref}$ 의 각 토큰에 대해 최적화 문제를 해결합니다: 고정된 네트워크 매개변수가 주어졌을 때 타겟 토큰의 음의 로그 가능도를 최소화하는 로짓 벡터 $r^*_t$ 를 찾는 것입니다.
- 이 최적화는 원래 매개변수 로짓에서 시작하여 경사 하강법 (일반적으로 $S=1$ ) 을 몇 단계 수행하여 이루어집니다.
- 이렇게 생성된 최적화된 로짓 $r^{(S)}_t$ 가 메모리 내 값 $v_t$ 로 사용됩니다.

2. 신뢰도 인식 적응적 혼합 (추론)

추론 중 주어진 라우터 입력 $x$ 에 대해:

검색: 시스템은 $x$ 와 저장된 키 간의 유사성을 기반으로 메모리 $\mathcal{M}$ 에서 $K$ 개의 최근접 이웃 $\mathcal{N}(x)$ 를 검색합니다.
비매개변수 제안: 검색된 최적 로짓을 입력과의 유사도로 가중치를 두어 집계하여 메모리 기반 로짓 제안 $r_{mem}(x)$ 를 계산합니다.
신뢰도 기반 혼합: 최종 라우팅 로짓은 원래 매개변수 라우터 로짓 $r(x)$ $r (x)$ 와 메모리 기반 제안 $r_{mem}(x)$ $r_{m e m} (x)$ 사이의 선형 보간입니다:
$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$
혼합 계수 $\lambda(x)$ $λ (x)$ 는 검색된 이웃의 평균 유사도로 정의됩니다. 이는 신뢰도 점수로 작용합니다:
- 높은 유사성 ( $\lambda \approx 1$ ): 모델이 검색된 국소 최적 할당을 신뢰합니다.
- 낮은 유사성 ( $\lambda \approx 0$ ): 모델이 관련 없는 사례의 노이즈를 피하기 위해 원래 고정된 라우터로 돌아갑니다.
출력: 최종 로짓은 희소 전문가 할당을 생성하기 위해 Top-K Softmax 를 통과하며, 모델의 원래 희소성 제약을 유지합니다.

주요 기여

검색 증강 라우팅: 과거 유사 사례의 메모리에서 오프라인으로 계산된 국소 최적 전문가 할당을 재사용하여 라우팅 결정을 개선하는 프레임워크를 도입합니다.
비매개변수 적응: SFT 나 온라인 적응 방법과 달리 kNN-MoE 는 모델 매개변수를 업데이트하지 않습니다. 계산 부하를 오프라인 메모리 구축 단계로 이전하여 효율적인 추론을 가능하게 합니다.
신뢰도 기반 폴백: 검색된 이웃의 평균 유사도를 동적 혼합 계수로 활용하여 관련 역사적 사례가 발견되지 않을 때 고정된 라우터로 복귀함으로써 견고성을 보장합니다.
효율성: 이 방법은 메모리 구축 중 토큰당 단일 경사 하강 단계만 필요로 하며, 온라인 최적화나 전체 모델 미세 조정의 무거운 계산 비용을 피합니다.

실험 결과

저자들은 세 가지 MoE 모델 (OLMoE, GPT-OSS, Qwen3) 에 대해 여섯 가지 벤치마크 (GPQA, SuperGPQA, MMLU, USMLE, MedMCQA, MBPP) 에서 kNN-MoE 를 평가했습니다.

기반 대비 성능: kNN-MoE 는 모든 모델과 데이터셋에서 제로샷 기반 및 "SFT(라우터 전용)" 기반을 일관되게 능가했습니다.
SFT 와의 경쟁력: 이 방법은 전체 지도 미세 조정 (SFT) 과 경쟁력 있는 성능을 달성했으며, 여러 데이터 부족 설정 (예: GPQA, USMLE) 에서는 SFT 를 능가했습니다. SFT 는 작은 참조 세트에서 과적합을 겪는 경우가 있었습니다.
안정성: 입력 길이 문제로 인해 특정 모델 (GPT-OSS, OLMoE) 의 성능을 저하시킨 5-샷 프롬프팅과 달리, kNN-MoE 는 안정적인 개선을 보여주었습니다.
절차 연구:
- 이웃 수 ( $K$ ): 일반적인 kNN-LM 결과와 반대로 성능은 $K=1$ 에서 정점을 찍었습니다. 더 많은 이웃을 집계 ( $K>1$ ) 하면 노이즈가 발생하여, 실제로 유익한 라우팅 신호는 희귀하고 구체적임을 시사합니다.
- 최적화 단계 ( $S$ ): $S=1$ 을 초과하여 경사 하강 단계수를 늘리는 것은 구성 비용을 증가시키면서 정확도 향상은 미미했습니다.
- 참조 세트 크기: 성능은 참조 세트 크기 $|D_{ref}|$ 에 따라 단조롭게 향상되었으며, 0.5k~1k 예제를 초과하면 체감 수익이 감소했습니다.
효율성: 메모리 구축은 1k 예제에 대해 0.46 시간이 소요되었습니다 (SFT 는 약 1.34 시간). IVFPQ(역색인 파일 + 제품 양자화) 변형은 기반 대비 추론 지연 오버헤드를 약 3.67% 로, VRAM 사용량 증가를 약 9.43% 로 줄였습니다.

중요성 및 주장

이 논문은 kNN-MoE 가 정적 매개변수 라우팅과 동적 추론 요구 사항 사이의 간극을 메운다고 주장합니다. 그 주요 중요성은 미세 조정의 계산적 부담 없이 MoE 적응성을 향상시키는 확장 가능한 경로를 제공한다는 점에 있습니다.

저자들은 이 방법이 라우터가 불확실할 때 (높은 퍼플렉시티와 상관관계가 있음) 특히 효과적이라고 강조합니다. 검색 메커니즘이 유사하고 고품질인 역사적 사례에서 교정 신호를 제공하기 때문입니다. 이 방법은 테스트 데이터와 유사한 분포를 가진 레이블이 지정된 참조 세트의 가용성에 의존하지만, 비용이 많이 드는 적응 기술에 대한 실용적이고 비매개변수적인 대안을 제공하며, 매개변수 업데이트 없이 준비 시간을 크게 단축하고 SFT 와 경쟁력 있는 결과를 달성합니다.