UCCI: Calibrated Uncertainty for Cost-Optimal LLM Cascade Routing

Resumen Técnico: UCCI – Incertidumbre Calibrada para el Enrutamiento de Cascada de LLMs Optimizado por Costos

Enunciado del Problema
Los Modelos de Lenguaje Grande (LLMs) ofrecen un alto rendimiento en tareas de predicción estructurada, pero los costos de inferencia escalan abruptamente con el tamaño del modelo. La inferencia en cascada —enrutando consultas fáciles a un modelo pequeño y escalando las difíciles a un modelo grande— promete una reducción de costos. Sin embargo, los enrutadores desplegados actualmente a menudo dependen de puntuaciones de confianza no calibradas (por ejemplo, entropía de token cruda o márgenes) que son sensibles a la redacción del prompt y a los cambios en la carga de trabajo. En consecuencia, el ajuste de umbrales suele ser un ejercicio manual, específico por carga de trabajo, que no generaliza, creando una brecha entre la promesa teórica de las cascadas y el rendimiento en producción.

Metodología
El artículo propone UCCI (Inferencia de Cascada Calibrada por Incertidumbre), un marco de enrutamiento que trata el diseño de la cascada como un problema de decisión calibrado en lugar de un ejercicio de ajuste. La metodología procede en tres pasos:

1. Extracción de Incertidumbre a Nivel de Token: Para una secuencia generada, el sistema calcula una puntuación de incertidumbre escalar $u(x)$ basada en el margen entre las probabilidades del siguiente token de top-1 y top-2 ($p_{t,1} - p_{t,2}$) a través de todos los tokens. Este margen se agrega para representar la incertidumbre general de la salida del modelo pequeño.
2. Calibración Isotónica: La puntuación de incertidumbre cruda $u(x)$ no es una probabilidad. UCCI mapea esta puntuación a una probabilidad de error por consulta calibrada $\hat{p}(x)$ utilizando regresión isotónica en un conjunto de calibración retenido. Este paso transforma la señal mal calibrada en una estimación fiable de la probabilidad de que la salida del modelo pequeño sea incorrecta ($P(e(x)=1 | u(x))$).
3. Minimización de Costos con Restricciones: Se selecciona una política de umbral $\pi_\theta$ en un conjunto de validación. El umbral $\theta^*$ se elige para minimizar el costo de inferencia esperado sujeto a una restricción de precisión objetivo ($\tau$). La política enruta las consultas con $\hat{p}(x) \leq \theta$ al modelo pequeño y escala aquellas con $\hat{p}(x) > \theta$ al modelo grande.

Contribuciones Teóricas
El artículo proporciona dos resultados teóricos clave:

• Optimalidad de las Políticas de Umbral: Bajo tres suposiciones explícitas (el costo del modelo grande > el costo del modelo pequeño; la precisión del modelo grande es invariante a qué consultas se escalan; y acceso a la probabilidad de error calibrada verdadera), los autores demuestran que una política de umbral sobre la probabilidad de error calibrada es óptima en costos.
• Complejidad de Muestra: El artículo establece que la regresión isotónica logra un Error de Calibración Esperado (ECE) de $O(n^{-1/3})$ con $n$ ejemplos de calibración, proporcionando un límite teórico sobre la eficiencia de datos del paso de calibración.

Configuración Experimental y Resultados
El marco fue evaluado en una carga de trabajo de Reconocimiento de Entidades Nombradas (NER) en producción que involucra 75,000 consultas etiquetadas manualmente de un sistema empresarial de gestión de fotos.

• Modelos: Un modelo pequeño de 4 mil millones de parámetros ($f_s$) y un modelo grande de 12 mil millones de parámetros ($f_\ell$), ambos ajustados por instrucción y servidos en GPUs H100.
• Medición de Costos: Los costos se derivaron de la latencia de extremo a extremo medida en GPUs H100, arrojando una relación de costos de $c_\ell/c_s \approx 3.02$.
• Líneas Base: El estudio comparó UCCI contra la inferencia de un solo modelo, el umbral de entropía no calibrado, el enrutamiento conformal dividido y un umbral aprendido al estilo de FrugalGPT.

Hallazgos Clave:

• Reducción de Costos: En un micro-F1 objetivo de 0.91, UCCI redujo los costos de inferencia en un 31% (IC 95%: [27%, 35%]) en comparación con el uso exclusivo del modelo grande.
• Calidad de Calibración: La regresión isotónica redujo el Error de Calibración Esperado (ECE) de 0.12 (no calibrado) a 0.03.
• Rendimiento vs. Líneas Base: En el mismo punto de operación (F1 = 0.91), UCCI logró el costo más bajo (2.08 unidades normalizadas) en comparación con el umbral de entropía (2.31), la predicción conformal (2.18) y el enrutamiento al estilo FrugalGPT (2.24).
• Validación de Suposiciones: La suposición de que la precisión del modelo grande permanece invariante al enrutamiento se probó empíricamente; la precisión en las consultas escaladas (0.928) fue casi idéntica a la precisión global (0.932), lo que respalda el marco teórico.

Significado y Afirmaciones
El artículo argumenta que el principal valor de ingeniería en los sistemas de cascada radica en calibrar la puntuación de enrutamiento en lugar de ajustar manualmente el umbral final. Al convertir una señal de margen de token ruidosa en una estimación probabilística de error mediante un ajuste isotónico pequeño, UCCI logra un rendimiento de enrutamiento casi óptimo sin arquitecturas complejas de aprendizaje para diferir o datos de preferencia.

Los autores enfatizan que sus resultados se basan en enrutamiento de extremo a extremo utilizando salidas reales de modelos y latencia de hardware medida, evitando el sesgo de selección inherente a las simulaciones basadas en precisiones globales o precios nominales de API. Aunque el marco es agnóstico a la tarea, los autores señalan que los ahorros de costos específicos (31%) y las métricas de calibración son específicos de la carga de trabajo. Reconocen limitaciones respecto a la suposición de invarianza en dominios donde las consultas "difíciles" para el modelo pequeño son desproporcionadamente difíciles para el modelo grande, y señalan que la calibración estática puede requerir adaptación para implementaciones de transmisión con cambios de distribución.