RouterDC: Query-Based Router by Dual Contrastive Learning for Assembling Large Language Models

https://arxiv.org/pdf/2409.19886

한 줄 요약: RouterDC는 “한 쿼리에 여러 LLM이 다 잘하는 상황”에서 기존 라우팅 학습(softmax 점수 정렬)이 약해지는 문제를 해결하기 위해, (1) 쿼리–LLM 간 대조학습으로 “잘하는 LLM들 쪽으로 끌어당기고 못하는 LLM들은 밀어내며”, (2) 쿼리–쿼리 대조학습으로 임베딩 공간을 안정화하는 Dual Contrastive 라우터를 제안한다.

---

1) 문제의식 / 동기

라우팅이 왜 필요한가

• 여러 오픈소스 LLM(일반형/특화형)을 조합(assemble) 하면 단일 모델보다 성능이 좋아질 수 있다. 하지만 모든 LLM을 매번 호출하는 앙상블(예: Voting) 은 추론 비용이 매우 크다.
• 그래서 쿼리마다 “가장 적합한 LLM 하나만” 고르는 라우팅이 효율적이다.

기존 라우팅(점수 분포 정렬)의 핵심 한계

• 대표 접근(예: ZOOTER)은 후보 LLM들의 점수(보상/평가)를 softmax 정규화한 분포를 정답으로 두고, 라우터 출력 분포를 그에 맞추는 형태(KL 최소화)를 쓴다.
• 그런데 실제로는 상위 LLM들 간 점수 차이가 매우 작고(top-2 gap이 작음), softmax 후 분포가 너무 평평해져(supervision이 약해져) 라우터 학습이 비효율적이라는 관찰을 제시한다.

---

2) 핵심 아이디어: “분포를 맞추지 말고, Top/Bottom을 대조하자”

RouterDC는 “정답 분포를 그대로 맞추는” 대신, 각 쿼리에서 (a) 잘하는 LLM 집합(Top-K+) 과 (b) 못하는 LLM 집합(Bottom-K−) 을 골라 대조학습(contrastive learning) 으로 라우터를 학습한다.

• 이렇게 하면 Top 성능 LLM이 여러 개 인 경우에도, 라우터는 그 중 하나를 ‘동등하게’ 선택할 수 있어 “상위가 여러 개인 상황”에서 학습 신호가 약해지는 문제를 줄인다.

---

3) 방법론 상세

3.1 라우터 구조 (Query → Embedding → LLM 선택)

• RouterDC는

  1. 쿼리를 임베딩으로 바꾸는 Encoder (E(x; w)) 와

  2. 각 후보 LLM을 대표하는 학습 가능한 LLM 임베딩 (k_t)

     로 구성된다.

• 선택 확률은 cosine similarity + softmax 로 계산한다.
• 중요 포인트: 후보 LLM들은 동결(frozen) 이고, 학습은 Encoder + LLM 임베딩만 업데이트한다(훈련 효율/비용 측면 장점).

---

3.2 LLM “점수 산정(Scoring)” — 라벨을 만들기 위한 단계

학습을 위해 각 쿼리 ($x_i$) 에 대해 LLM별 성능 점수 ($s_i^{(t)}$) 를 만든다.

• Open-ended generation(예: GSM8K)
  • 샘플링/빔 기반으로 여러 번 생성한 뒤, 정답 여부를 평균내어 점수를 만든다.
• Multiple-choice(예: MMLU)
  • 정답 옵션을 맞췄을 때 그 옵션 확률을 (옵션합으로) 정규화해 점수로 사용하고, 오답이면 0으로 둔다.

이 점수는 “분포를 직접 맞추기”가 아니라 Top/Bottom을 고르는 데만 사용된다.

---

3.3 Loss 1: Sample–LLM Contrastive Loss (핵심 라우팅 학습)

• 쿼리 (x_i) 에 대해
  • Positive LLM set ($I_i^+$): 점수 상위 Top-K+ LLM들
  • Negative LLM set ($I_i^-$): 점수 하위 Bottom-K− 중 점수가 낮은 것들(조건 포함) 로 구성한다.
    
• 그리고 쿼리 임베딩이 positive LLM 임베딩에 가깝고, negative LLM 임베딩에서 멀어지도록 하는 대조 손실을 최소화한다. (식 (4))

직관: “이 쿼리는 수학 강한 모델군으로 보내고, 약한 모델군으로는 보내지 마라”를 분포 정렬 없이 직접 강제.

---

3.4 Loss 2: Sample–Sample Contrastive Loss (학습 안정화)

저자들은 Sample–LLM loss만 쓰면 학습이 불안정할 수 있다고 보고, 이를 보완하는 쿼리–쿼리 대조학습을 추가한다.

• 사전 임베딩(mDeBERTaV3-base) → t-SNE → k-means로 쿼리를 N개 군집으로 묶는다.
• 같은 군집(in-group) 쿼리끼리는 가깝게, 다른 군집(out-group) 쿼리들과는 멀어지게 하는 loss를 둔다. (식 (5))

---

3.5 최종 목적식 / 학습·추론 절차

• 최종 loss는 두 대조 손실의 합: [ $\sum L_{\text{sample-LLM}} + \lambda , L_{\text{sample-sample}}$ ] 로 구성된다.
  
• 추론 시에는 쿼리 임베딩과 LLM 임베딩의 similarity가 최대인 LLM을 골라 그 모델 1개만 호출한다.
• 알고리즘 요약(Algorithm 1): 점수 산정 → 군집화 → 미니배치 학습(두 loss) → argmax 선택 추론.

---

4) 실험 설정 (무엇을, 어떻게 평가했나)

후보 LLM (7개)

• Mistral 계열 5개 + Llama-3 계열 2개로 구성.

데이터셋

• In-distribution (학습에 포함되는 태스크): MMLU, GSM8K, CMMLU, ARC-C, HumanEval
• OOD (학습에 포함되지 않는 태스크): PreAlgebra, MBPP, C-EVAL

주요 하이퍼파라미터(대표 설정)

• Encoder: mDeBERTaV3-base (86M)
• ($K^+=3, K^-=3, H=3, \lambda=1, N=5$) 등.

---

5) 결과 요약 (핵심만)

In-distribution

• Table 1 기준, RouterDC가 평균 정확도에서 최고 성능(Avg 58.54)이며, 라우팅 베이스라인(ZOOTER 등) 대비 일관적으로 우수한 결과를 보고한다.

Out-of-distribution

• Table 2 기준, OOD 평균에서도 RouterDC가 최고 성능(Avg 45.85)을 보고한다.

“Top-1만 positive로 두면 안 된다”는 근거

• ($K^+=1$) 구성(사실상 일반 다중분류)에선 성능이 떨어지고, 여러 top LLM을 positive로 두는 설계가 중요하다는 감도분석을 제시한다.

비용-효율 관점 (추가 분석)

• GSM8K/MBPP에서 비용을 점수에 포함시키는 변형((s_i^{(t)} + c_i^{(t)}))을 통해 cost-effectiveness 비교를 수행했고, RouterDC가 더 비용 효율적이라는 결과를 제시한다.

---