Reward-Based Online LLM Routing via NeuralUCB

1. 연구 배경 및 문제 정의 (Problem Definition)

• 배경: 대형 언어 모델 (LLM) 은 작업에 따라 성능 (품질) 과 추론 비용 (Inference Cost) 이 크게 달라집니다. 고성능 모델은 비용이 비싸고, 저렴한 모델은 복잡한 추론이나 코딩 작업에서 실패할 수 있습니다.
• 문제: 각 사용자 쿼리에 대해 적절한 LLM 을 동적으로 선택하여 답변의 품질과 추론 비용 사이의 균형을 최적화하는 것입니다.
• 제약 조건:
  • 부분 피드백 (Partial Feedback): 배포 환경에서는 각 쿼리에 대해 하나의 모델만 실행되므로, 선택된 모델의 결과만 관찰하고 다른 모델들의 결과는 알 수 없습니다.
  • 비선형성: 쿼리와 최적의 모델 선택 간의 매핑은 매우 비선형적일 수 있습니다.
• 수식적 정의: 이를 컨텍스트 밴딧 (Contextual Bandit) 문제로 정의합니다. 시간 $t$에서 쿼리 $x_t$가 주어졌을 때, 후보 모델 집합 $A$에서 모델 $a_t$를 선택하여 장기적인 기대 유틸리티 (Utility) 를 최대화하는 정책을 학습합니다.
  • 목표: $\max_{\pi} E [\sum_{t=1}^{T} r(x_t, a_t)]$

2. 제안 방법론 (Methodology)

2.1 유틸리티 함수 (Utility Function)

품질 ($q$) 과 비용 ($c$) 의 트레이드오프를 고려한 보상 함수를 정의합니다.

• 보상 함수: $r(x, a) = q(x, a) \exp(-\lambda \tilde{c}(x, a))$
  • $\tilde{c}$: 로그 정규화를 통해 [0, 1] 범위로 변환된 비용.
  • $\lambda$: 비용 패널티의 강도를 조절하는 하이퍼파라미터.
• 이 함수는 높은 품질을 유지하면서 비용을 최소화하는 행동을 장려합니다.

2.2 UtilityNet (신경망 구조)

라우팅 정책의 핵심인 유틸리티 예측기입니다.

• 입력: 텍스트 임베딩 ($x_{emb}$), 보조 특징 ($x_{feat}$), 도메인 ID ($d$), 행동 ID ($a$).
• 구조:
  1. 인코딩: 텍스트, 특징, 도메인, 행동 (Action) 을 각각 임베딩 및 MLP 를 통해 인코딩합니다.
  2. 퓨전 (Fusion): 인코딩된 벡터를 결합하여 표현합니다.
  3. 두 가지 브랜치:
     • 유틸리티 브랜치 (Utility Branch): 주어진 컨텍스트 - 행동 쌍에 대한 예상 보상 ($\mu(x, a)$) 을 예측합니다. (Huber Loss 로 학습)
     • 게이팅 브랜치 (Gating Branch): 컨텍스트에 따라 UCB(Exploration) 기반 선택이 필요한지 여부를 확률 $p(x)$로 예측합니다. (Binary Cross-Entropy Loss 로 학습)

2.3 NeuralUCB 기반 라우팅 정책

UtilityNet 의 예측을 바탕으로 행동을 선택하는 전략입니다.

• UCB 점수: $s(x, a) = \mu(x, a) + \beta \sqrt{g(x, a)^\top A^{-1} g(x, a)}$
  • $\mu(x, a)$: UtilityNet 이 예측한 평균 유틸리티.
  • 두 번째 항: 불확실성 (Exploration Bonus). 공유된 역공분산 행렬 $A^{-1}$을 사용하여 고차원 특징 공간의 불확실성을 추정합니다.
• 안전한 폴백 (Safe Fallback): 게이트 확률 $p(x)$가 임계값 $\tau_g$보다 낮으면, UCB 보너스를 적용하지 않고 단순히 예측된 평균 유틸리티가 가장 높은 안전한 행동 ($a_{safe}$) 을 선택합니다.
  • 이는 불확실성이 낮거나 탐색이 유익하지 않을 때 과잉 탐색을 방지합니다.
• 학습 프로토콜: 데이터를 20 개의 슬라이스로 나누어 순차적으로 처리하는 시뮬레이션된 온라인 환경에서 학습합니다. 각 슬라이스 후 리플레이 버퍼를 기반으로 모델을 재학습하고 $A^{-1}$을 갱신합니다.

3. 실험 및 결과 (Experimental Results)

• 데이터셋: RouterBench (36,497 개 샘플, 86 개 도메인, 11 개 후보 LLM).
• 비교 대상:
  • Random (무작위 선택)
  • Min-cost (최저 비용 모델 선택)
  • RouteLLM-BERT (지도 학습 기반 라우터)
  • Max-quality Reference (비용 무시하고 최고 품질 모델 선택 - 상한선)
• 주요 결과:
  1. 보상 (Reward) 성능: 제안된 NeuralUCB 방법은 Random 및 Min-cost 베이스라인보다 일관되게 높은 평균 보상 (약 0.590.61) 을 기록했습니다. 기존 지도 학습 방법인 RouteLLM-BERT(약 0.350.36) 보다도 훨씬 우수했습니다.
  2. 누적 보상: 시간이 지남에 따라 NeuralUCB 의 누적 보상 격차가 베이스라인 대비 지속적으로 확대되었습니다.
  3. 비용 - 품질 트레이드오프: 제안된 정책은 최고 품질 (Max-quality) 기준 대비 약 33% 수준의 추론 비용만 사용하면서도 경쟁력 있는 품질을 유지했습니다.
  4. 임베딩 분석: 다양한 텍스트 인코더 (Multilingual-E5, MPNet, Qwen 등) 를 비교한 결과, 모델 크기가 크다고 해서 항상 성능이 향상되는 것은 아니며, all-MiniLM-L6-v2나 all-mpnet-base-v2와 같은 경량 모델이 안정적이고 강력한 성능을 보였습니다.

4. 주요 기여 및 의의 (Key Contributions & Significance)

1. NeuralUCB 를 LLM 라우팅에 적용: 기존 선형 가정 (LinUCB) 이나 커널 기반 방법의 한계를 넘어, 비선형적인 쿼리 - 모델 매핑을 효율적으로 학습하고 부분 피드백 환경에서 비용 - 품질 균형을 최적화하는 새로운 프레임워크를 제시했습니다.
2. 비용 인식형 (Cost-Aware) 전략: 단순히 품질만 높이는 것이 아니라, 비용 패널티를 명시적으로 보상 함수에 통합하여 실제 배포 환경에서 더 실용적인 라우팅을 가능하게 했습니다.
3. 효율적인 탐색 전략: 게이팅 메커니즘을 통해 불필요한 탐색을 줄이고, 신뢰할 수 있는 예측이 가능할 때는 안전한 선택을 하도록 하여 학습 안정성을 높였습니다.
4. 실용적 검증: RouterBench 를 통한 오프라인 재현 (Replay) 시나리오를 구축하여, 실제 API 호출 비용 없이도 온라인 환경과 유사한 조건에서 알고리즘의 유효성을 검증했습니다.

5. 결론 및 한계

• 결론: NeuralUCB 기반 라우팅은 비용과 품질 사이의 균형을 효과적으로 맞추며, 기존 방법론보다 우수한 성능을 입증했습니다.
• 한계 및 향후 과제:
  • 많은 샘플에서 이미 '최고 품질' 모델이 '최적 보상' 모델과 일치하는 경우가 있어, 라우팅을 통해 얻을 수 있는 추가 이득 (Gain) 이 제한적일 수 있습니다.
  • 비용 제약 하에서 행동을 더 잘 구분 (Action Discrimination) 하고 탐색을 개선하기 위해 더 많은 학습 데이터와 향상된 유틸리티 모델이 필요합니다.

이 연구는 LLM 서비스의 운영 비용을 절감하면서도 사용자 경험을 유지하기 위한 지능형 라우팅 시스템의 중요한 방향성을 제시합니다.