Functional-level Uncertainty Quantification for Calibrated Fine-tuning on LLMs

Resumen Técnico: Cuantificación de la Incertidumbre a Nivel Funcional para el Ajuste Fino Calibrado en Modelos de Lenguaje Grandes

Enunciado del Problema
La cuantificación precisa de la incertidumbre (UQ) es crítica para el despliegue confiable de Modelos de Lenguaje Grandes (LLMs), particularmente en dominios críticos para la seguridad como la atención médica y el razonamiento científico. Sin embargo, los LLMs ajustados finamente exhiben frecuentemente sobreconfianza al adaptarse a tareas con datos limitados, fallando en reflejar la verdadera credibilidad de sus salidas. Los métodos existentes de UQ para el Ajuste Fino Eficiente en Parámetros (PEFT), como los basados en Adaptación de Bajo Rango (LoRA), son en gran medida a posteriori. Estiman la incertidumbre después de que el modelo ha sido entrenado mediante el análisis del espacio de parámetros (por ejemplo, mediante conjuntos o aproximaciones bayesianas) o perturbando los prompts. Estos enfoques a menudo fallan en abordar la causa raíz de la sobreconfianza: la incapacidad del modelo para capturar la incertidumbre inducida por el ajuste fino y las brechas de generalización durante el propio proceso de adaptación. Además, los métodos que dependen de múltiples pasadas hacia adelante o muestreo extensivo introducen una sobrecarga computacional significativa, limitando la escalabilidad.

Metodología: UQ4CT
Los autores proponen Cuantificación de la Incertidumbre a Nivel Funcional para el Ajuste Fino Calibrado (UQ4CT), un marco que desplaza el enfoque de la incertidumbre en el espacio de parámetros hacia la incertidumbre en el espacio funcional. En lugar de analizar los parámetros aprendidos después del entrenamiento, UQ4CT calibra la incertidumbre sobre el espacio funcional inducido por mezclas dependientes del prompt de expertos LoRA durante el proceso de ajuste fino.

La metodología consta de tres componentes centrales:

1. Descomposición Funcional Jerárquica mediante MixLoRA:
   Para evitar la explosión combinatoria de modelar todas las combinaciones posibles de expertos a través de las capas, UQ4CT emplea una descomposición jerárquica. Utiliza una arquitectura de Mezcla de Expertos (MoE) donde, en cada capa, un enrutador disperso selecciona dinámicamente un subconjunto de expertos LoRA basado en el estado oculto de entrada. Esto construye un espacio funcional flexible donde la salida del modelo es una suma ponderada de diversas funciones base (expertos LoRA) condicionadas a la entrada. El espacio funcional se representa como una composición recursiva de mezclas a nivel de capa, permitiendo que el modelo componga adaptativamente rutas de computación.

2. Estimación de la Confianza a Nivel Funcional (FLC):
   El modelo cuantifica la incertidumbre analizando la concentración de la selección de expertos. Específicamente, la Confianza a Nivel Funcional (FLC) se deriva de los pesos de enrutamiento de los top-$k$ expertos seleccionados por el enrutador en cada capa. Una alta concentración en los pesos de enrutamiento (es decir, el enrutador tiene confianza en seleccionar expertos específicos) corresponde a una alta confianza funcional, mientras que pesos dispersos indican mayor incertidumbre. Esto proporciona una señal interna y fundamentada de la certeza del modelo sin requerir múltiples pasadas de muestreo.

3. Pérdida de Calibración ($L_{cal}$):
   La innovación central es una nueva pérdida de calibración que alinea la confianza a nivel funcional con la corrección predictiva durante el entrenamiento. La pérdida se define como:
   $$L_{cal} = (1\{\text{MixLoRA}(x) = y^*\} - \text{FLC}(x))^2$$
   donde $1\{\cdot\}$ es una función indicadora para predicciones correctas. Este objetivo fuerza a que la FLC aproxime la probabilidad de que la predicción del modelo sea correcta. Al minimizar esta pérdida junto con las pérdidas estándar de entropía cruzada y de equilibrio de carga, el modelo aprende a asignar alta confianza solo cuando es probable que sea correcto y mayor incertidumbre cuando no lo es. Esto alienta al enrutador a explorar mezclas alternativas de expertos cuando la incertidumbre es alta, mejorando así la generalización.

Contribuciones Clave

• Perspectiva a Nivel Funcional: El artículo introduce un cambio de la cuantificación de la incertidumbre en el espacio de parámetros al espacio funcional, calibrando explícitamente la confianza del modelo sobre los mapeos entrada-salida que puede realizar durante el ajuste fino.
• Ajuste Fino Consciente de la Calibración: Se propone una nueva pérdida de calibración que alinea dinámicamente los pesos de mezcla de LoRA dependientes del prompt con la corrección predictiva, mitigando la sobreconfianza sin degradar la precisión.
• Arquitectura Eficiente: El método aprovecha una arquitectura jerárquica MixLoRA que captura diversas relaciones funcionales con una sobrecarga computacional mínima, requiriendo solo una pasada hacia adelante durante la inferencia.
• Garantía Teórica: Los autores proporcionan una prueba de que optimizar el riesgo de calibración sobre la distribución de datos produce una FLC óptima que corresponde a la verdadera probabilidad de corrección predictiva.

Resultados Experimentales
Los autores evaluaron UQ4CT en cuatro benchmarks de opción múltiple (ARC-E, ARC-C, OBQA, ClimateQA) y dos tareas de QA generativa de respuesta abierta (TruthfulQA, TriviaQA) utilizando Llama-3.1-8B y Mistral-7B.

• Rendimiento In-Distribution: UQ4CT redujo el Error de Calibración Esperado (ECE) en más de 25% en todos los benchmarks en comparación con bases sólidas (incluyendo LoRA, MC Dropout, Deep Ensembles y Bayesian LoRA), manteniendo al mismo tiempo una precisión competitiva o superior. Cabe destacar que UQ4CT logró estos resultados con una sola pasada hacia adelante, mientras que métodos como BLoB (N=10) requirieron diez pasadas y aún así produjeron un ECE más alto.
• Desplazamiento de Distribución: Bajo desplazamientos de distribución (probados en subtareas de MMLU en CS, Ingeniería, Derecho y Salud), UQ4CT mantuvo una calibración superior y una precisión competitiva, demostrando una mayor robustez frente a entradas fuera de distribución en comparación con las bases.
• Tareas Generativas: En la generación de respuesta abierta, UQ4CT logró el ECE más bajo en TruthfulQA y TriviaQA, lo que sugiere que la perspectiva de calibración en el espacio funcional se extiende más allá de configuraciones de estilo de clasificación.
• Eficiencia: El análisis de memoria indica que el uso máximo de memoria GPU de UQ4CT es solo ~7.5% superior al de LoRA estándar, lo que lo hace computacionalmente viable para despliegues a gran escala.

Significado y Afirmaciones
El artículo afirma que UQ4CT aborda una limitación fundamental en las prácticas actuales de ajuste fino: la falta de conciencia de la incertidumbre durante el proceso de adaptación. Al integrar la calibración de la incertidumbre directamente en el objetivo de ajuste fino mediante señales a nivel funcional, el método asegura que la confianza del modelo refleje mejor su fiabilidad predictiva real. Los autores posicionan este trabajo como un paso fundamental hacia LLMs ajustados finamente dignos de confianza, particularmente en entornos con escasez de datos donde la sobreconfianza plantea riesgos significativos. Señalan modestamente que, aunque el enfoque se extiende a la generación de respuesta abierta utilizando proxies basados en F1, escalar hacia razonamiento de largo alcance o generación de múltiples pasos sigue siendo una dirección para trabajos futuros.