UCCI: Calibrated Uncertainty for Cost-Optimal LLM Cascade Routing

기술 요약: UCCI – 비용 최적화 LLM 캐스케이드 라우팅을 위한 보정된 불확실성

문제 제기
대규모 언어 모델 (LLM) 은 구조화된 예측 작업에서 높은 성능을 제공하지만, 추론 비용은 모델 크기에 따라 급격히 증가합니다. 캐스케이드 추론은 쉬운 쿼리를 작은 모델로 라우팅하고 어려운 쿼리는 큰 모델로 에스컬레이션하여 비용 절감을 약속합니다. 그러나 현재 배포된 라우터들은 종종 프롬프트 문구와 작업 부하 변화에 민감한 보정되지 않은 신뢰도 점수 (예: 원시 토큰 엔트로피 또는 마진) 에 의존합니다. 결과적으로 임계값 조정은 종종 작업별 수동 작업에 그쳐 일반화에 실패하며, 캐스케이드의 이론적 가능성과 실제 운영 성능 사이의 격차를 초래합니다.

방법론
본 논문은 캐스케이드 설계를 튜닝 작업이 아닌 보정된 의사결정 문제로 간주하는 라우팅 프레임워크인 UCCI(Uncertainty-Calibrated Cascade Inference) 를 제안합니다. 방법론은 세 단계로 진행됩니다:

1. 토큰 수준 불확실성 추출: 생성된 시퀀스에 대해 시스템은 모든 토큰에 걸쳐 상위 1 위와 상위 2 위 다음 토큰 확률 ($p_{t,1} - p_{t,2}$) 간의 마진을 기반으로 스칼라 불확실성 점수 $u(x)$ 를 계산합니다. 이 마진은 작은 모델 출력의 전반적인 불확실성을 나타내기 위해 집계됩니다.
2. 등방성 보정 (Isotonic Calibration): 원시 불확실성 점수 $u(x)$ 는 확률이 아닙니다. UCCI 는 홀드아웃 보정 세트에서 등방성 회귀 (isotonic regression) 를 사용하여 이 점수를 보정된 쿼리별 오류 확률 $\hat{p}(x)$ 로 매핑합니다. 이 단계는 오보정된 신호를 작은 모델의 출력이 잘못될 확률 ($P(e(x)=1 | u(x))$) 에 대한 신뢰할 수 있는 추정치로 변환합니다.
3. 제약 비용 최소화: 검증 세트에서 임계값 정책 $\pi_\theta$ 가 선택됩니다. 임계값 $\theta^*$ 는 목표 정확도 제약 ($\tau$) 하에서 기대 추론 비용을 최소화하도록 선택됩니다. 이 정책은 $\hat{p}(x) \leq \theta$ 인 쿼리는 작은 모델로 라우팅하고 $\hat{p}(x) > \theta$ 인 쿼리는 큰 모델로 에스컬레이션합니다.

이론적 기여
본 논문은 두 가지 주요 이론적 결과를 제공합니다:

• 임계값 정책의 최적성: 세 가지 명시적 가정 (큰 모델 비용 > 작은 모델 비용; 에스컬레이션된 쿼리에 관계없이 큰 모델의 정확도는 불변; 실제 보정된 오류 확률에 대한 접근) 하에서, 저자들은 보정된 오류 확률에 대한 임계값 정책이 비용 최적임을 증명합니다.
• 샘플 복잡도: 본 논문은 등방성 회귀가 $n$ 개의 보정 예시로 $O(n^{-1/3})$ 의 기대 보정 오차 (ECE) 를 달성함을 규명하여, 보정 단계의 데이터 효율성에 대한 이론적 상한을 제공합니다.

실험 설정 및 결과
본 프레임워크는 기업용 사진 관리 시스템에서 75,000 개의 수동 라벨링된 쿼리를 포함하는 생산 환경 명명 개체 인식 (NER) 작업 부하에서 평가되었습니다.

• 모델: 4B 파라미터 작은 모델 ($f_s$) 과 12B 파라미터 큰 모델 ($f_\ell$) 로 구성되었으며, 둘 다 지시 미세 조정 (instruction-tuned) 되었으며 H100 GPU 에서 서빙되었습니다.
• 비용 측정: 비용은 H100 GPU 에서 측정된 엔드투엔드 지연 시간에서 도출되었으며, 비용 비율 $c_\ell/c_s \approx 3.02$ 를 산출했습니다.
• 베이스라인: 연구는 UCCI 를 단일 모델 추론, 보정되지 않은 엔트로피 임계값 설정, 분할-준합동 (split-conformal) 라우팅, 그리고 FrugalGPT 스타일의 학습된 임계값과 비교했습니다.

주요 발견:

• 비용 절감: 목표 마이크로-F1 이 0.91 일 때, UCCI 는 큰 모델만 사용하는 것에 비해 추론 비용을 31% (95% 신뢰구간: [27%, 35%]) 감소시켰습니다.
• 보정 품질: 등방성 회귀는 기대 보정 오차 (ECE) 를 보정되지 않은 0.12 에서 0.03으로 감소시켰습니다.
• 베이스라인 대비 성능: 동일한 운영 지점 (F1 = 0.91) 에서 UCCI 는 엔트로피 임계값 설정 (2.31), 준합동 예측 (2.18), FrugalGPT 스타일 라우팅 (2.24) 에 비해 가장 낮은 비용 (2.08 정규화 단위) 을 달성했습니다.
• 가정 검증: 라우팅에 관계없이 큰 모델의 정확도가 불변이라는 가정이 실증적으로 테스트되었습니다. 에스컬레이션된 쿼리에 대한 정확도 (0.928) 는 전역 정확도 (0.932) 와 거의 동일하여 이론적 프레임워크를 지지했습니다.

의의 및 주장
본 논문은 캐스케이드 시스템의 주요 공학적 가치는 최종 임계값을 수동으로 튜닝하는 것이 아니라 라우팅 점수를 보정하는 데 있다고 주장합니다. 작은 등방성 피팅을 통해 노이즈가 있는 토큰 마진 신호를 확률적 오류 추정치로 변환함으로써, UCCI 는 복잡한 학습-기반-지연 (learning-to-defer) 아키텍처나 선호도 데이터 없이도 거의 최적의 라우팅 성능을 달성합니다.

저자들은 그들의 결과가 전역 정확도나 명목상 API 가격을 기반으로 한 시뮬레이션에 내재된 선택 편향을 피하기 위해 실제 모델 출력과 측정된 하드웨어 지연 시간을 사용한 엔드투엔드 라우팅에 기반하고 있음을 강조합니다. 프레임워크는 작업 중립적이지만, 저자들은 특정 비용 절감 (31%) 과 보정 지표는 작업 부하에 특정적임을 지적합니다. 그들은 "어려운" 쿼리가 작은 모델에게는 어렵고 큰 모델에게는 불균형적으로 어려운 영역에서 불변 가정과 관련된 한계를 인정하며, 정적 보정은 분포 변화가 있는 스트리밍 배포를 위해 적응이 필요할 수 있음을 언급합니다.