Decoupled Mixture-of-Experts for Parametric Knowledge Injection

기술 요약: 파라미터 지식 주입을 위한 디커플드 Mixture-of-Experts (DMoE)

문제 정의
대규모 언어 모델(LLM)은 사전 학습 후 정적인 파라미터 지식을 보유하며, 이로 인해 도메인 특화 또는 시의성이 중요한 질의에 대해 실패하여 환각 현상이나 오래된 응답을 생성하는 문제가 발생한다. 기존의 지식 주입 패러다임은 유연성과 통합 깊이 사이에서 근본적인 트레이드오프에 직면해 있다:

• 검색 증강 생성 (RAG): 지식을 외부에 유지하고 쉽게 업데이트할 수 있지만, 프롬프트 수준에서만 지식을 주입한다. 이는 얕은 통합을 제공하며, 반복적인 검색과 긴 컨텍스트 처리로 인해 추론 효율성이 저하된다.
• 사후 학습 방법 (예: SFT, LoRA): 지식을 모델 파라미터에 직접 인코딩하여 더 깊은 통합을 구현한다. 그러나 공유 파라미터를 업데이트하면 치명적 망각, 지식 충돌, 높은 재학습 비용의 위험이 있으며, 모듈성과 독립적인 업그레이드 능력이 부족하다.

식별된 핵심 구조적 병목 현상은 이질적인 지식을 격리하고, 특정 모듈으로 라우팅하며, 관련 없는 능력을 저해하거나 효율적인 자기 회귀적(auto-regressive) 추론을 깨뜨리지 않고 모델을 점진적으로 확장할 수 있는 메커니즘의 부재이다.

방법론: 디커플드 Mixture-of-Experts (DMoE)
저자들은 전문가 모듈과 라우터를 동결된 베이스 모델으로부터 분리하는 모듈형 아키텍처인 DMoE를 제안한다.

1. 아키텍처 및 전문가 구성:

   • 디커플링 (Decoupling): 라우터와 전문가가 공동 학습되어 백본 내에 임베딩되는 기존의 Mixture-of-Experts(MoE)와 달리, DMoE는 베이스 모델을 동결된 상태로 유지한다.
   • 전문가 (Experts): 외부 지식 코퍼스를 의미 단위로 분할한다. 각 단위에 대해, 해당 단위에서 파생된 인스트럭션 스타일 데이터로 독립적으로 학습된 경량 전문가 모듈(LoRA 어댑터로 구현)을 구축한다. 이러한 전문가들은 외부에 저장되며, 베이스 모델을 재학습시키지 않고도 추가, 제거 또는 업데이트할 수 있다.
   • 배치 (Placement): 효율적인 자기 회귀적 생성을 지원하기 위해, 전문가는 최종 레이어의 피드 포워드 네트워크(FFN)에만 부착된다. 전문가를 중간 레이어에 배치하면 후속 어텐션 블록으로 전달되는 은닉 상태(hidden states)를 변경하여 키-값(KV) 캐시를 무효화하고 재계산을 강제하게 된다. 전문가를 최종 레이어로 제한함으로써, DMoло는 이전 레이어로부터의 캐시된 어텐션 상태의 유효성을 보존한다.

2. 불확실성 인지 라우팅 (Uncertainty-Aware Routing):

   • 트리거링 (Triggering): 베이스 모델이 불충분한 지식을 보일 때만 작동하는 경량 라우터가 활성화된다. 이는 각 디코딩 단계에서 모델의 소프트맥스 분포 엔트로피로 계산되는 **토큰 불확실성 (Token Uncertainty, TU)**에 의해 결정된다. 만약 $TU > \tau$ (임계값)라면 라우팅이 트리거된다.
   • 선택 (Selection): 라우터는 각 전문가와 연관된 텍스트 대리물(surrogates)에 대해 BM25 기반 렉시컬 검색을 수행한다. 이를 통해 현재 생성 컨텍스트와 가장 잘 일치하는 문서를 가진 상위-$k$개의 전문가를 검색한다. 이 설계는 학습이 필요 없으며(training-free) 점진적으로 업데이트 가능하다.
   • 활성화 (Activation): 선택된 전문가들은 현재 디코딩 상태에 대해서만 베이스 모델 파라미터와 결합된다 ($\theta_{eff} = \theta + \sum \Delta\theta_i$). 중복적인 스왑을 피하기 위해 활성 세트는 다음 트리거링 이벤트 전까지 유지된다.

주요 기여

• 모듈형 파라미터 주입: DMoE는 지식 전문가와 라우터를 베이스 모델로부터 완전히 분리하여, 독립적인 업데이트와 모듈형 통합을 가능하게 하는 프레임워크를 도입한다.
• 캐시 안전 설계 (Cache-Safe Design): 전문가를 오직 최종 레이어 FFN에만 부착함으로써, DMoE는 KV-캐시 재사용을 보존하면서 파라미터 수준의 지식 증강을 가능하게 하며, 이는 자기 회귀적 효율성의 핵심 요소이다.
• 불확실성 기반 활성화: 시스템은 토큰 엔트로피를 활용하여 전문가가 필요할 때만 선택적으로 활성화함으로써, 지속적인 검색이나 파라미터 수정의 오버헤드를 방지한다.

실험 결과
저자들은 Llama-3.2-1B 및 Qwen2.5-1.5B를 베이스 모델로 사용하여 네 가지 지식 집약적 벤치마크(HotpotQA, ComplexWebQuestions (CWQ), Quasar-T, StrategyQA)에서 DMoE를 평가하였다.

• 효과성: DMoE는 덴스(dense) 베이스라인보다 일관되-게 우수한 성능을 보였으며, RAG(Basic-RAG, FLARE) 및 어댑터 기반 방법(PRAG, SFT-LoRA)과 비교하여 경쟁력 있거나 더 뛰어난 성능을 달성했다. 보고된 14개 지표 중 11개에서 최고점 또는 공동 최고점을 기록했다.
• 효율성:
  • DMoE는 FLARE와 같은 동적 RAG 방식에 비해 추론 지연 시간을 약 3배 감소시켰다.
  • FLARE 대비 GPU 메모리 사용량을 1.6배~1.9배 줄였다.
  • 절제 연구(Ablation studies)를 통해 중간 레이어에 전문가를 부착하는 것(KV-캐시 재사용을 깨뜨림)이 성능과 효율성을 크게 저하시킨다는 것을 확인하였으며, 이를 통해 최종 레이어 배치 전략의 타당성을 입증했다.
• 결합형 MoE와의 비교: 튜닝된 결합형 MoE 백본(OLMoE)과 비교했을 때, DMoE는 훨씬 낮은 추론 시간(2.73.9초 vs 20.0초)과 메모리 점유율(7.28.3 GB vs 26.1 GB)로 더 높은 효과성을 달성하였으며, 이는 이점이 단순한 파라미터 스케일링이 아닌 아키텍처의 디커플링에서 기인함을 보여준다.

의의 및 주장
본 논문은 DMoE가 LLM의 확장 가능하고 업데이트 가능한 지식 주입을 위한 유망한 방향임을 제시한다. 전문가를 모델이 불확실할 때만 호출되는 외부의 경량 모듈로 취급함으로써, DMoE는 표준 MoE 아키텍처나 전통적인 RAG보다 엄격하게 개선된 효과-효율성 트레이드오프를 달ر성한다. 이 접근 방식은 공유 파라미터 업데이트(망각/충돌)와 프롬프트 수준 주입(얕은 통합)의 한계를 해결하는 동시에, KV-캐시 보존을 통해 자기 회귀적 생성의 계산적 이점을 유지한다. 저자들은 이 시스템이 하이퍼파라미터 선택(예: top-$k$ 전문가, 트리거링 임계값)에 견고하며, 백본을 재학습시키지 않고도 점진적인 업데이트를 지원한다는 점을 강조한다.