Router-R1: Teaching LLMs Multi-Round Routing and Aggregation via Reinforcement Learning

하임·2026년 1월 9일

Routing

목록 보기
31/44

한 줄 요약

Router-R1은 “라우터도 LLM으로 만들자”는 발상으로, 내부 추론()과 외부 LLM 호출(→)을 여러 라운드로 엮어가며 답을 조립하도록 RL로 학습한 멀티-LLM 라우팅 프레임워크입니다. 단발(one-shot) 라우팅보다 복잡 질의에서 더 좋은 정확도를 내고, 비용(모델 크기·토큰 수) 도 보상에 넣어 정확도-비용 균형을 스스로 학습합니다.


1) 핵심 아이디어

  • 멀티-라운드 라우팅을 RL로 정식화: 라우터(정책 LLM)가 라우팅 풀의 후보 LLM들을 여러 번 선택·질의하고, 받은 정보를 컨텍스트에 누적하며 추론을 이어가 최종 답을 냅니다.
  • 라우터도 LLM: 라우터 자체가 LLM이라 추론()과 라우팅()을 교대(interleave) 하면서 동적으로 결정을 내립니다.
  • 포맷/정답/비용 보상을 결합하고, 계층형 보상으로 포맷 위반 시 다른 보상을 무효화해 안정적으로 학습합니다.
  • 프롬프트 템플릿로 내부 추론→필요 시 외부 LLM 호출→정보 통합→최종 답()의 루프를 강제합니다.
  • 일반화 용이: 후보 LLM의 “설명문”을 프롬프트에 넣는 구조라 새 LLM을 추가해도 재학습 없이 라우팅 풀을 확장할 수 있습니다.

2) 방법: 수식과 보상 설계

2-1. RL 최적화 목적(식 1)

KL 정규화를 포함한 일반적 정책 최적화 형태로, 정책 π가 라우팅 풀 PP을 활용해 생성한 출력 yy에 대해

maxπExD,yπ(x;P)[rϕ(x,y)βlogπ(yx;P)πref(yx;P)]\max_\pi \mathbb{E}_{x\sim D,\,y\sim \pi(\cdot|x;P)}\big[r_\phi(x,y) - \beta \log \tfrac{\pi(y|x;P)}{\pi_{\text{ref}}(y|x;P)}\big]

를 최적화합니다(β는 참조 정책과의 이탈 억제). 이 틀은 PPO/GRPO/ KL-제약 계열로 구현 가능합니다.

2-2. 보상 구성

  • 포맷 보상(식 2): 요구한 출력 포맷을 지키면 0, 아니면 −1. 학습 안정화 목적.
  • 최종 정답 보상(식 3): Exact Match(EM) 정확 일치만 가산. 간명하고 효과적.
  • 비용 보상(식 4): 선택 LLM의 파라미터 수출력 토큰 수에 비례해 불이익(∝ −m(PLLM)·Tout) → 큰 모델·장문일수록 보상 감소.
  • 종합 보상(식 5): rϕ=Rformat+(1α)Routcome+αRcost.r_\phi = R_{\text{format}} + (1-\alpha) R_{\text{outcome}} + \alpha R_{\text{cost}}.α 로 정확-비용 균형 제어.
  • 계층형 결합: 포맷이 틀리면 다른 보상은 0으로 무효화(보상 해킹 억제·안정화).

3) 프롬프트 & 상호작용 프로토콜

  • 템플릿: <think>…</think>로 내부 사고 → 필요하면 <search> Candidate LLM: Query </search>로 외부 LLM 호출 → 응답은 <info>…</info> 로 들어오고 → 마지막에 <answer>…</answer>.
  • 학습 시 동작: 생성 중 <search> 토큰이 나오면 해당 LLM을 호출, 응답을 시퀀스에 삽입하고 계속 추론(안정화를 위해 <info>는 손실 계산에서 제외). 간단한 질의는 외부 호출 없이 자체 지식으로 답변.

4) 실험 설정

  • 데이터·지표: 7개 QA(단일/멀티-홉) — NQ, TriviaQA, PopQA, HotpotQA, 2WikiMultiHopQA, Musique, Bamboogle; EM과 F1로 평가.
  • 베이스/학습: Base LLM은 Qwen2.5-3B-Instruct, LLaMA-3.2-3B-Instruct. 최대 4회 라우팅 스텝, veRL+PPO로 학습, 배치 64, 225 스텝, 기본 α=0.0\alpha=0.0. 학습 데이터는 NQ 7k + Hotpot 7k(총 14k).
  • 라우팅 풀(후보 6종): Qwen2.5-7B, LLaMA-3.1-8B/70B, Mistral-7B, Mixtral-8×22B, Gemma-2-27B.

5) 결과 요약

  • 메인 결과: 7개 벤치마크에서 Direct/CoT/SFT/RAG 및 Search-R1 등 광범위 베이스라인을 일관되게 상회. Llama 베이스로 평균 EM 0.409, Qwen 베이스로 0.416 기록.
    • 특히, RAG보다 큰 폭으로 우세(정적 문서 검색 대신 전문 LLM을 동적으로 호출). Search-R1보다도 전반적으로 우세.
    • 단발(single-round) 라우터(KNN/MLP/BERT/RouterDC/GraphRouter) 대비 멀티-라운드 + 추론-교대 설계가 성능 격차의 원인으로 제시.
  • 일반화: 학습은 NQ/Hotpot(인-도메인) 만으로 했지만, 나머지 5개 OOD 데이터에서도 강건.
  • 비용 제어(α 분석): α\alpha↑(비용 가중↑)이면 정확도는 완만히↓, 비용 보상은↑. 작은 모델부터 질의하고 필요시 큰 모델로 승급하는 “자연 발현” 라우팅 전략을 관찰.

6) 왜 잘 되나? (설계 인사이트)

  1. 포맷→정답→비용의 계층 보상이 학습을 안정화하고, 실제 운영 제약(가격·토큰)을 정책이 내재화하게 만듭니다.
  2. LLM-as-Router 구조가 질의별로 추론과 라우팅을 섞어가며, 단발 라우팅의 정보 손실을 줄입니다.
  3. 모델 설명문 기반 프롬프트로 라우팅 풀 확장이 쉬워, 새 모델이 등장해도 유연하게 적응합니다.

profile
NLP 공부합니당

0개의 댓글