Semantic Self-Distillation for Language Model Uncertainty

Resumen Técnico: Semantic Self-Distillation (SSD)

1. El Problema: Cuantificación de Incertidumbre en LLMs

Los Grandes Modelos de Lenguaje (LLMs) presentan desafíos significativos para la cuantificación de incertidumbre (UQ) de manera rigurosa. A medida que se despliegan en dominios de alto riesgo (salud, leyes, agentes autónomos), es crucial detectar errores como las alucinaciones (respuestas que no se alinean con hechos o intenciones).

• Limitaciones de los métodos actuales:
  • Métodos basados en muestreo (ej. Entropía Semántica): Son efectivos midiendo la dispersión semántica de múltiples respuestas generadas estocásticamente. Sin embargo, requieren generar múltiples completaciones, lo que introduce una latencia prohibitiva para aplicaciones en tiempo real o sistemas de agentes.
  • Métodos de sondeo (Probing): Entrenan modelos ligeros para predecir una métrica escalar de incertidumbre (ej. probabilidad de corrección) a partir de representaciones internas. Aunque son rápidos, colapsan la compleja topología de la distribución de salida en un solo número, perdiendo información estructural necesaria para tareas como la verificación posterior a la generación o la selección de respuestas.

2. Metodología: Semantic Self-Distillation (SSD)

Los autores proponen SSD, un marco que combina la eficiencia de los métodos de sondeo con la riqueza de los métodos basados en muestreo. En lugar de predecir un escalar, SSD entrena un modelo estudiante ligero para predecir una densidad condicional al prompt sobre las incrustaciones semánticas de las respuestas.

Componentes Clave del Proceso:

1. Representaciones:

   • Entrada ($h$): Se extrae el estado oculto del token final del prompt de una capa específica del modelo base ($\ell^*$).
   • Salida ($z$): Las respuestas generadas se mapean a incrustaciones semánticas ($z = \Phi(y)$) utilizando un modelo de incrustación especializado, seguido de reducción de dimensionalidad (PCA).

2. Modelo Estudiante (MDN):

   • Se utiliza una Red de Densidad de Mezcla (Mixture Density Network - MDN).
   • Dado el vector del prompt $h$, el estudiante predice una mezcla de $K$ distribuciones gaussianas sobre el espacio de incrustaciones semánticas $z$:
     $$q_\phi(z | h) = \sum_{k=1}^K \pi_k(h) \mathcal{N}(z; \mu_k(h), \text{diag}(\sigma_k^2(h)))$$
   • Entrenamiento: Se utiliza la verosimilitud condicional máxima. Se generan $S$ muestras estocásticas del modelo "maestro" (LLM base) para cada prompt, se calculan sus incrustaciones $z^{(s)}$, y se entrena el estudiante para maximizar la probabilidad de observar estas muestras dadas las representaciones del prompt.

3. Inferencia y Señales de Incertidumbre:
   Una vez entrenado, el modelo SSD realiza una sola pasada (forward pass) y ofrece múltiples señales:

   • Evaluación del Prompt (Pre-generación): Se calcula la entropía de Rényi-2 de la densidad predicha. Una alta entropía indica alta dispersión semántica y, por tanto, alto riesgo de alucinación.
   • Evaluación de la Respuesta (Post-generación): Dada una respuesta candidata $y'$, se calcula su log-verosimilitud bajo la densidad predicha. Una baja verosimilitud sugiere que la respuesta es incoherente con el contexto (fuera de dominio o alucinada).
   • Muestreo Latente: Permite generar muestras sintéticas de la distribución semántica para obtener un "consenso semántico" (media de la mezcla).

3. Contribuciones Clave

1. Nuevo Marco de Auto-distilación: Introducen SSD, que destila la distribución de respuestas semánticas muestreadas de un LLM en un estimador de densidad ligero y condicional al prompt.
2. Eficiencia sin Pérdida de Estructura: Logran un rendimiento competitivo en la predicción de alucinaciones con la misma latencia que los métodos de sondeo de un solo paso, pero preservando la estructura de la distribución completa.
3. Nuevas Primitivas de Confiabilidad: Al modelar la distribución completa, habilitan capacidades que los métodos escalares no tienen:
   • Verificación de respuestas post-generación basada en verosimilitud.
   • Detección de respuestas fuera de dominio (OOD).
   • Selección de respuestas en preguntas de opción múltiple basada en la densidad.

4. Resultados Experimentales

Los autores evaluaron SSD en los conjuntos de datos TriviaQA (conocimiento) y MMLU (razonamiento) utilizando múltiples familias de modelos (Llama, Mistral, Qwen, Gemma, SmolLM).

• Predicción de Alucinaciones (Pre-generación):

  • SSD compite favorablemente con la Dispersión del Maestro (Teacher Dispersion - TD), que requiere muestreo en tiempo de inferencia.
  • En TriviaQA, SSD superó a la dispersión del maestro en 4 de 7 modelos en AUROC y en todos en AUPRC, gracias a su capacidad de suavizar estimaciones en casos degenerados donde el maestro produce siempre la misma respuesta incorrecta.
  • Es comparable a los métodos de sondeo (PCP) pero con la ventaja de la densidad explícita.

• Evaluación Post-Generación:

  • Detección OOD: El log-verosimilitud de SSD logró un AUROC > 0.95 para distinguir respuestas coincidentes de respuestas aleatorias, actuando como un "guardián" de contenido.
  • Selección de Opción Múltiple: En MMLU, el modelo estudiante (entrenado solo con respuestas abiertas) pudo seleccionar la opción correcta basándose en la verosimilitud, acercándose al rendimiento del modelo base en generación abierta.

• Análisis de Fidelidad: Se encontró una fuerte correlación entre la fidelidad de la destilación (qué tan bien el estudiante aprende la distribución del maestro) y el rendimiento en la detección de alucinaciones.

5. Significado e Impacto

El trabajo demuestra que la incertidumbre en modelos generativos complejos puede ser destilada como una densidad explícita sobre representaciones latentes, en lugar de un simple puntaje de riesgo.

• Cambio de Paradigma: Mueve el costo computacional del tiempo de inferencia (donde la latencia es crítica) al tiempo de entrenamiento.
• Aplicabilidad General: Aunque se centra en LLMs autoregresivos, el marco es agnóstico a la modalidad y se puede aplicar a otros modelos de secuencia (ej. modelos de difusión, modelos de historiales clínicos electrónicos) donde las salidas admiten representaciones latentes significativas.
• Viabilidad para Agentes: Proporciona un mecanismo práctico para integrar la incertidumbre de nivel de distribución en sistemas de agentes que requieren baja latencia y toma de decisiones robusta (abstención, recuperación de información).

En resumen, Semantic Self-Distillation ofrece una solución elegante al dilema entre la precisión de la incertidumbre (requiriendo muestreo) y la eficiencia operativa, permitiendo que los sistemas de IA sean más confiables y seguros sin sacrificar la velocidad de respuesta.