DLLG: Dynamic Logit-Level Gating of LLM Experts

기술 요약: DLLG: LLM 전문가를 위한 동적 로짓 레벨 게이팅(Dynamic Logit-Level Gating of LLM Experts)

문제 정의
좁은 도메인(예: 수학적 추론, 코드 생성)에 최적화된 특화된 대규모 언어 모델(LLM)의 급증은 다음과 같은 과제를 제기한다: 어떻게 하면 독립적으로 학습된 "블랙박스" 전문가들을 막대한 재학습 비용 없이 그들의 상호 보완적인 강점을 활용하는 하나의 통합된 시스템으로 결합할 것인가? 기존 방식들은 적응성, 견고성, 그리고 실용성 사이의 지속적인 긴장 상태를 겪고 있다:

• **라우팅 기반 방식(Routing-based methods)**은 하드한 초기 전문가 선택에 의존한다. 이는 조기 확정(premature commitment)을 초래한다. 만약 선택된 전문가가 단일 응답 내의 특정 하위 작업에서 실패할 경우, 그 오류는 복구 불가능하다. 또한, 많은 라우터는 광범어적인 벤치마크별 캘리브레이션을 필요로 한다.
• 휴리스틱 앙상블(Heuristic Ensembling)(예: 퍼플렉시티 가중치, 토큰 투표)은 추론 시점의 프록시 신호를 사용한다. 이러한 지표들은 종종 경험적이며 실제 작업의 정답성과 일치하지 않아, 최적이 아닌 출력을 낳는다.
• 파라미터 병합(Parameter Merging)(예: 모델 수핑, 태스크 아리스메틱)은 가중치를 정적으로 융합한다. 이는 모듈성을 희생하며, 서로 충돌하는 그래디언트 방향을 가진 전문가들을 결합할 때 파괴적 간섭을 유도하여 성능 저하를 일으킬 수 있다.
• 혼합 전문가(Mixture-of-Experts, MoE) 아키텍처는 일반적으로 게이트와 전문가를 엔드 투 엔드로 공동 학습해야 하므로, 이미 만들어진(off-the-shelf) 사전 학습된 전문가들을 활용하는 데 한계가 있다.

방법론: DLLG
저자들은 희소한 응답 수준의 감독(sparse, response-level supervision)을 통해 학습된 경량 게이팅 모델을 사용하여 동적인 토큰 레벨 로짓 융합을 수행하는 프레임워크인 **DLLG (Dynamic Logit-Level Gating)**를 제안한다.

1. 아키텍처:

   • 동결된 전문가(Frozen Experts): $K$개의 특화된 LLM은 동결된 상태를 유지한다. 각 디코딩 단계 $t$에서, 프롬프트 $x$와 접두사 $y_{<t}$가 주어지면, 각 전문가 $E_i$는 은닉 상태 $h^{(i)}_t$와 다음 토큰 로짓 $\ell^{(i)}_t$를 생성한다.
   • 게이팅 모델: 경량 모델($g_\phi$)은 모든 전문가의 연결된 은닉 상태를 처리한다. 이 모델은 저차원 투영(low-rank projection)을 사용하여 이 표현들을 공유 임베딩 공간으로 투영하고, 자기회귀 스템(예: Qwen2.5-0.5B-Instruct 및 LoRA 어댑터)을 통해 처리한다.
   • 가중치 예측: 예측 헤드는 시그모이드 변환을 통해 토큰별 융합 가중치 $w_{t,i}$를 출력한다. 이 가중치들은 전문가들의 로짓에 적용된다: $\ell_t = \sum_{i=1}^K w_{t,i} \ell^{(i)}_t$. 가중치는 합이 1로 제한되지 않으므로, 부드럽고 유연한 기여를 허용한다.

2. 학습 (응답 수준의 감독):

   • 프레임워크는 토큰 레벨의 주석(annotation) 부족 문제를 자동 검증기(예: 코드의 유닛 테스트, 수학의 정답 확인)로부터 도출된 응답 수준의 정답 신호($s_i \in \{0, 1\}$)를 사용하여 해결한다.
   • 교사 강요(Teacher Forcing): 학습 중, 전문가들은 정답 접두사 $y_{<t}$에 조건화된 은닉 상태를 생성한다. 게이팅 모델은 각 토큰 위치에 대한 가중치 $w_{t,i}$를 예측한다.
   • 목적 함수: 모델은 예측된 가중치와 브로드캐스트된 응답 수준의 정답 신호 사이의 평균 제곱 오차(MSE)를 최소화하도록 학습된다: $\mathcal{L} = \frac{1}{T} \sum_{t=1}^T \sum_{i=1}^K (w_{t,i} - s_i)^2$. 이를 통해 게이트가 토큰 레벨의 융합을 전체 궤적의 성공과 일치시키도록 가르친다.

3. 추론:

   • DLLG는 자기회귀적으로 작동한다. 각 단계에서 게이팅 모델은 현재 전문가의 은닉 상태를 기반으로 융합 가중치를 예측하고, 로짓이 결합되며, 다음 토큰이 선택된다. 이를 통해 **복구 가능한 전문가 활용(recoverable expert utilization)**이 가능하다. 만약 전문가가 초기에 오류를 범하더라도, 게이팅 모델은 후속 단계에서 더 유능한 전문가로 가중치를 동적으로 전환할 수 있다.

주요 기여

1. 프레임워크 제안: 토큰 레벨의 주석이나 온라인 강화 학습 없이도 미세한 전문가 활용을 가능하게 하며, 토큰 레벨 생성과 희소한 응답 수준의 감독 사이를 잇는 동적 로짓 레벨 앙상블 프레임워크인 DLLG를 소개한다.
2. 메커니즘 분석: 학습된 게이팅 메커니즘이 롤아웃(rollout) 중에 전문가의 기여를 동적으로 조정하며, 진화하는 문맥(예: 단일 응답 내에서 수학적 추론에서 코드 생성으로 전환)에 따라 가중치를 조절함을 입증한다.
3. 실증적 검증: 다양한 추론 및 코드 벤치마크(GSM8K, MATH, HumanEval 등)에서 평가를 수행하여, 라우팅, 휴리스틱 앙상블 및 파라미터 병합 베이스라인 대비 일관된 성능 향상을 보여준다.

결과

• 성능: 0.5B 및 1.5B 모델 규모 모두에서, DLLG는 RouterDC, EmbedLLM, 다양한 로짓 앙상블 방법(GaC, UniTe) 및 파라미터 병합 기술(Linear, SLERP, Task Arithmetic)을 포함한 강력한 베이스라인들과 비교하여 일관되게 가장 높은 평균 성능을 달면한다.
• 견고성: 개별 전문가가 약한 저용량(0.5B) 환경에서, DLLG는 취약한 라우팅 및 휴리스틱 방식보다 우수한 성능을 보이며, 소프트 융합(soft fusion)을 통해 초기 실수를 완화하는 능력을 입증한다.
• 일반화: 방법론은 도메인 외 벤치마크(예: HumanEval, MBPP, BBH)에서도 강력한 성능을 보여주며, 이는 학습된 융합 규칙이 특정 학습 분포를 넘어 일반화됨을 나타낸다.
• 절제 연구(Ablation): 실험을 통해 효과적인 토큰 레벨 융합을 위해서는 충분히 표현력이 있는 게이팅 아키텍처(보조 LLM 및 저차원 투영 포함)가 필수적이며, 이것이 단순한 MLP 헤드보다 우수함을 확인하였다.

의의 및 주장
본 논문은 DLLG가 특화된 전문가들을 통합하기 위한 견고하고 확장 가능한 패러다임을 제공한다고 주장한다. 약한 감독(weak supervision)으로부터 로짓 레벨 융합을 학습함으로써, 전문가를 재학습시키는 비용이나 파라미터 간섭의 위험 없이 MoE와 같은 적응성을 달 Achieve한다. 저자들은 DLLG가 전문가의 모듈성을 보존하며, 미세하고 복구 가능한 전문가 활용을 가능하게 하여, 라우팅의 조기 확정 문제나 파라미터 병합의 정적인 특성을 극복한다고 강조한다. 이 연구는 미세한 전문가 활용을 약한 감독으로부터 학습하는 것이 점점 더 전문화되는 언어 모델을 활용하기 위한 실질적인 경로임을 시사한다.