Semantic Token Clustering for Efficient Uncertainty Quantification in Large Language Models

1. El Problema

Los Modelos de Lenguaje Grandes (LLMs) han demostrado capacidades excepcionales, pero carecen de garantías de veracidad y a menudo sufren de un problema de sobreconfianza: generan respuestas que suenan plausibles pero son incorrectas con alta confianza.

La cuantificación de la incertidumbre es crucial para identificar estas salidas no fiables. Sin embargo, los métodos existentes presentan limitaciones significativas:

• Métodos basados en muestreo: Requieren generar múltiples respuestas (muestreo estocástico) para medir la dispersión semántica (ej. Entropía Semántica), lo que introduce una gran sobrecarga computacional.
• Métodos con modelos auxiliares: Utilizan modelos externos (como NLI) para verificar la coherencia, lo que aumenta la latencia y la complejidad.
• Métodos basados en logits simples: A menudo ignoran la consistencia semántica, ya que un solo token puede tener baja probabilidad aunque existan otros tokens semánticamente equivalentes con alta probabilidad.

El objetivo es desarrollar un método que cuantifique la incertidumbre de manera eficiente, utilizando una única generación y sin depender de modelos externos, aprovechando las representaciones internas del LLM.

2. Metodología: Semantic Token Clustering (STC)

Los autores proponen STC, un método que agrupa tokens en clusters semánticamente consistentes para agregar sus masas de probabilidad. El proceso se divide en dos etapas:

A. Etapa de Pre-cálculo (Offline)

• Agrupación de Embeddings: Se utilizan los embeddings de entrada y salida de los tokens del vocabulario del modelo.
• Clustering No Supervisado: Se aplica un algoritmo de agrupamiento (específicamente Agglomerative Clustering) para agrupar tokens que comparten significado semántico (ej. "TV", "tv", "televisión" se agrupan en el mismo cluster).
• Filtrado: Se excluyen palabras funcionales (stopwords) y números arábigos (ya que embeddings similares no garantizan equivalencia matemática).
• Resultado: Se crea un mapa de clusters semánticos que se utiliza durante la inferencia sin costo adicional.

B. Etapa de Inferencia (Online)

Durante la generación de una sola respuesta, el método calcula la incertidumbre paso a paso:

1. Identificación de Tokens: Para cada paso de decodificación $i$, se identifica el token generado $y_i$.
2. Agregación de Probabilidad: En lugar de usar solo la probabilidad de $y_i$, se suman las probabilidades de todos los tokens en el vocabulario que pertenecen al mismo cluster semántico que $y_i$.
3. Coincidencia de Prefijos (Prefix Matching): Para abordar la fragmentación de tokens (tokenización), se verifica si un token candidato es un prefijo de la secuencia generada posterior. Esto asegura que variaciones de tokenización (ej. "tele" + "vision" vs "television") se consideren consistentes.
4. Cálculo de la Puntuación de Incertidumbre:
   • Se calcula la masa de probabilidad agrupada $\hat{p}_c$ para cada paso.
   • La puntuación final de incertidumbre $S(x, y)$ se define como:
     $$S(x, y) = 1 - \prod_{i=1}^{n} \hat{p}_c(y_i | x, y_{<i})$$
   • Una puntuación más alta indica mayor incertidumbre (menor confianza agregada).

3. Contribuciones Clave

1. Uso de Representaciones Internas: STC explota directamente la estructura semántica codificada en los embeddings del LLM, sin necesidad de entrenamiento adicional o datos etiquetados.
2. Eficiencia Computacional:
   • Requiere una sola generación (no muestreo múltiple).
   • No utiliza modelos auxiliares.
   • La parte más costosa (clustering) se realiza offline.
   • La sobrecarga en tiempo de inferencia es mínima (se puede ejecutar en CPU).
3. Implementación Autocontenida: Funciona con cualquier LLM de caja blanca (white-box) disponible, sin necesidad de fine-tuning.

4. Resultados Experimentales

Los autores evaluaron STC en tres conjuntos de datos (TriviaQA, Natural Questions, WebQuestions) y varios modelos (Llama-2/3, Mistral, Qwen2.5).

• Rendimiento (AUROC): STC logra un rendimiento comparable o superior a los métodos state-of-the-art (como Entropía Semántica, EigenScore y CCP). En muchos casos, supera a los baselines, especialmente en modelos más pequeños.
• Eficiencia:
  • Comparado con métodos basados en muestreo (que requieren múltiples generaciones), STC reduce la sobrecarga de tiempo de ejecución en un 98% en comparación con métodos como CCP.
  • Mientras que otros métodos pueden tardar cientos de segundos por pregunta, STC añade solo unos pocos segundos (o menos de un 2% de sobrecarga relativa).
• Estudio de Ablación:
  • La eliminación de la agrupación por embeddings o la coincidencia de prefijos degrada el rendimiento, demostrando que ambos componentes son complementarios y necesarios para capturar la consistencia semántica y superficial.

5. Significado y Limitaciones

Significado:
El trabajo presenta un avance crucial hacia la implementación práctica de la cuantificación de incertidumbre en entornos de producción. Al eliminar la necesidad de muestreo múltiple y modelos externos, STC hace viable la detección de alucinaciones en aplicaciones de baja latencia y recursos limitados, manteniendo una alta precisión.

Limitaciones:

1. Acceso a Logits/Embeddings: Requiere acceso a las representaciones internas del modelo, lo que impide su uso directo en modelos de caja negra (closed-source) sin acceso a la API interna.
2. Polisemia: Al basarse en embeddings estáticos (independientes del contexto), puede agrupar tokens con significados divergentes si el contexto no se considera. Sin embargo, los autores notan que los LLMs suelen asignar probabilidades muy bajas a tokens semánticamente incompatibles en un contexto dado, mitigando parcialmente este ruido.
3. Calibración: El método no calibra explícitamente las puntuaciones de incertidumbre, aunque esto podría hacerse en una etapa posterior.

En conclusión, STC ofrece un equilibrio óptimo entre precisión y eficiencia, proporcionando una solución robusta para evaluar la fiabilidad de las respuestas de los LLMs en tiempo real.