Expected Reward Prediction, with Applications to Model Routing

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 응답 품질을 평가하기 위해 **리워드 모델 (Reward Model)**이 널리 사용되지만, 기존 리워드 모델은 주로 고정된 프롬프트에 대해 여러 개의 응답 (Response) 을 생성한 후 그중 가장 좋은 것을 선택하는 **응답 수준 (Response-level)**의 점수 매기기에 초점을 맞추고 있습니다.

이 연구는 다음과 같은 핵심 문제를 제기합니다:

• 모델 적합성 예측의 부재: 실제 응답을 생성하기 전에, 특정 프롬프트에 대해 특정 LLM 이 얼마나 좋은 응답을 생성할지 (즉, 해당 모델의 적합성) 를 사전에 예측할 수 있는가?
• 예측 가능성의 불확실성: LLM 의 응답 분포가 프롬프트에 따라 어떻게 변하는지, 그리고 리워드 모델이 학습된 상대적 선호도 (Pairwise preference) 데이터의 특성상 절대적 점수가 정의되지 않는다는 점 때문에, 프롬프트 기반의 **기대 리워드 (Expected Reward)**를 예측하는 것이 이론적으로나 실용적으로 가능한지 불분명했습니다.

2. 방법론 (Methodology)

가. 기대 리워드 예측 (Expected Reward Prediction, ERP)

저자들은 프롬프트 $x$와 모델 $\pi$, 리워드 함수 $r$이 주어졌을 때, 모델이 생성하는 무작위 응답 $Y$에 대한 기대 리워드 $E_{Y \sim \pi(Y|x)}[r(x, Y)]$를 프롬프트만으로부터 예측하는 방법을 제안합니다.

• 데이터 생성: 각 프롬프트에 대해 특정 모델 (예: Llama3.1, Gemma 등) 로부터 $K=32$개의 응답을 샘플링하고, 리워드 모델로 점수를 매겨 평균을 내어 'Ground Truth' 기대 리워드를 계산합니다.
• 간단한 선형 모델 (Linear Probe): 프롬프트의 임베딩 벡터 (Embedding, 예: gte-large-en-v1.5) 를 입력으로 받아 기대 리워드를 예측하는 리지 회귀 (Ridge Regression) 선형 모델을 학습시킵니다.
  • 입력: 프롬프트 임베딩 $v(x)$
  • 출력: 예측된 기대 리워드 $\hat{ER}(x)$
  • 목적 함수: $\text{argmin}_{\theta} E[(\theta^\top v(X) - \text{Empirical } ER)^2] + \beta \|\theta\|^2$

나. 모델 라우팅 (Model Routing) 적용

예측된 기대 리워드를 활용하여, 주어진 프롬프트를 가장 비용 효율적으로 처리할 수 있는 모델을 선택하는 모델 라우팅 프로토콜을 설계합니다.

• 목표: 계산 비용 (Cost) 을 제한 조건으로 두면서 기대 리워드를 최대화하는 라우팅 정책 $\rho$를 학습합니다.
  • $\max_\rho E[r(x, Y)]$ subject to $E[c(\rho(x))] \le C$
• 비용 조정 정책: 비용과 리워드의 상대적 중요도를 조절하는 파라미터 $\lambda$를 도입하여, 각 프롬프트에 대해 $\hat{ER}(x) - \lambda \cdot \text{Cost}$를 최대화하는 모델을 선택합니다.
• 확장성: 새로운 모델이 풀 (Pool) 에 추가될 때, 기존 모델 간의 쌍별 (Pairwise) 비교 데이터가 필요 없이 해당 모델만의 예측 모델만 학습하면 되므로 확장성이 매우 뛰어납니다.

3. 주요 기여 (Key Contributions)

1. 예측 가능성의 입증: 복잡한 LLM 의 응답 분포 특성을 프롬프트 임베딩을 통해 단순한 선형 모델로도 높은 정밀도로 예측할 수 있음을 실증적으로 증명했습니다. 이는 리워드 모델이 프롬프트별 기대 성능을 포착할 수 있음을 의미합니다.
2. 모델 라우팅 성능 향상: 예측된 기대 리워드를 기반으로 한 라우팅 (ERP) 이, 카테고리별 평균 성능이 가장 좋은 모델을 고정적으로 선택하는 베이스라인이나, 더 복잡한 라우팅 방법 (Zooter 등) 보다 우수한 성능을 보임을 입증했습니다.
3. 확장성과 효율성: 기존 선호도 기반 (Preference-based) 라우팅이 모델 쌍 ($N^2$) 에 비례하는 데이터와 학습 비용이 드는 반면, ERP 방식은 모델 수 ($N$) 에 비례하여 데이터와 학습 비용이 증가하므로, 모델 풀이 커질수록 훨씬 효율적입니다.
4. 이론적 통찰: 기대 리워드의 차이가 모델 간 승률 (Win rate) 과 직접적으로 연관됨을 Proposition 1 을 통해 수학적으로 설명했습니다.

4. 실험 결과 (Results)

• 데이터셋 및 모델: open-perfectblend 데이터셋 (4000 개의 프롬프트) 을 사용했으며, Llama3.1 (8B/70B), Gemma2 (9B/27B), Gemma1 (7B) 등 다양한 크기의 오픈 웨이트 모델과 OpenAssistant-RM, GRM-2B-RM, InternLM-RM 등 3 가지 리워드 모델을 평가했습니다.
• 예측 정확도 ($R^2$):
  • 프롬프트 카테고리 (코딩, 수학, 일반 대화 등) 내에서도 기대 리워드 예측이 유의미하게 수행되었습니다.
  • 대부분의 경우 $R^2$가 0.4~0.6 사이로, 프롬프트 임베딩이 모델의 세부적인 능력 차이를 잘 설명함을 보여줍니다.
• 라우팅 성능:
  • Regret-Cost Tradeoff: ERP 기반 라우팅은 Zooter(복잡한 신경망 기반) 와 유사하거나 더 나은 Pareto 프론티어를 형성했습니다.
  • 베이스라인 비교: 고정된 모델 사용이나 카테고리별 최적 모델 선택 (Oracle) 보다 낮은 Regret(기회 손실) 을 달성했습니다. 특히 GRM-2B-RM 의 경우 일부 카테고리에서 한 모델이 압도적으로 잘하는 경향이 있어, ERP 가 이를 잘 포착하여 효율적인 라우팅을 수행했습니다.
  • 쌍별 비교: ERP 를 기반으로 한 이진 분류 (어떤 모델이 이길까?) 의 AUROC 는 직접적인 쌍별 선호도 학습 (Logistic Regression) 과 거의 동일한 성능을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 과 리워드 모델의 중요한 특성을 발견했습니다. 즉, 응답 수준의 리워드를 모델 수준의 기대 리워드로 '들어 올릴 (Lift)' 수 있으며, 이를 프롬프트 임베딩만으로 간단히 예측할 수 있다는 점입니다.

• 실용적 가치: 추론 시간 (Inference time) 에 모델 라우팅을 통해 비용을 절감하면서도 품질을 극대화할 수 있는 실용적인 솔루션을 제공합니다.
• 시스템 설계의 유연성: 새로운 모델이 추가될 때마다 전체 시스템을 재학습할 필요 없이 해당 모델에 대한 예측기만 추가하면 되므로, 동적 모델 풀 관리에 매우 유리합니다.
• 미래 전망: 모델 라우팅 외에도 프롬프트 수정, 시스템 프롬프트 핫스왑 (Hot-swapping) 등 다양한 추론 시간 최적화 작업에 적용 가능한 잠재력을 가지고 있습니다.

요약하자면, 이 연구는 복잡한 모델 비교를 단순한 기대값 예측으로 대체함으로써 LLM 시스템의 효율성과 확장성을 획기적으로 개선할 수 있음을 보여준 선구적인 작업입니다.