예상 보상 예측과 모델 라우팅 적용

초록

본 논문은 프롬프트만으로 언어 모델이 생성할 응답의 기대 보상을 정확히 예측할 수 있음을 보이고, 이를 이용해 비용‑효율적인 모델 라우팅 정책을 설계한다. 간단한 선형 프로브와 사전 학습된 프롬프트 임베딩만으로도 높은 R²를 달성했으며, Llama 3.1, Gemma 2, Gemma 1 시리즈를 포함한 다중 모델 풀에서 기대 보상 기반 라우팅(ERP)이 기존 고정 라우팅 및 복합 라우팅 기법을 능가한다.

상세 분석

이 연구는 두 가지 핵심 질문에 답한다. 첫째, “프롬프트만으로 특정 LLM이 보상 모델에 의해 부여받을 기대 보상 ERπ(x)를 예측할 수 있는가?” 둘째, “그 예측값을 실제 추론 단계에서 모델 선택에 활용할 수 있는가?”
예측 가능성 검증을 위해 저자들은 (1) 다양한 도메인(일반 대화, 명령 수행, 수학·코드 추론)에서 4,000개의 프롬프트를 포함하는 open‑perfectblend 데이터셋을 사용했고, (2) Llama 3.1‑IT(8B/70B), Gemma 2‑IT(9B/27B), Gemma 1‑IT(7B) 등 6개의 생성 모델을 대상으로 각각 32개의 샘플을 생성해 보상 모델(OpenAssistant‑RM2, GRM‑2B‑RM, InternLM‑RM)으로 평균 보상을 계산했다.
예측 모델은 gte‑large‑en‑v1.5 임베딩(1024 차원)을 입력으로 하는 Ridge‑regularized 선형 회귀이며, β=1로 정규화하였다. 결과는 놀라울 정도로 높은 R²(전체 평균 0.450.59, 카테고리별 0.190.59)로, 단순 선형 프로브가 기대 보상의 대부분 변동성을 설명함을 보여준다. 특히 “코딩”·“수학” 같은 고난이도 카테고리에서도 Llama 3.1‑70B는 0.54 이상의 예측력을 유지했다.
두 번째 실험에서는 예측된 기대 보상을 라우팅 기준으로 사용한다. 각 프롬프트 x에 대해 ERP는
 argmax_π