Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Calibration

Resumen Técnico: Reconfiguración inspirada en difusión de Transformadores para la calibración de incertidumbre

Enunciado del Problema
Los transformadores preentrenados se han convertido en la columna vertebral de los modelos fundacionales a gran escala en los dominios de visión, lenguaje y voz. Sin embargo, estos modelos carecen típicamente de un mecanismo principiado para la propagación de la incertidumbre a través de sus pilas de transformación de características. En consecuencia, a menudo asignan alta confianza a predicciones incorrectas, lo que plantea riesgos de seguridad en aplicaciones de alto riesgo. Si bien existen técnicas de calibración de incertidumbre (UC) para redes más pequeñas, estas suelen depender de aproximar los posteriors del modelo mediante inferencia variacional o muestreo posterior, lo cual se vuelve computacionalmente prohibitivo e inexacto para modelos a gran escala.

Los enfoques existentes que intentan reparametrizar los bloques de atención como procesos gaussianos (GP) enfrentan una limitación crítica: tratan los bloques de forma independiente. Esto no tiene en cuenta las correlaciones entre las transformaciones de características establecidas durante el preentrenamiento. Como resultado, estos métodos a menudo fuerzan una compensación donde una mejor calibración de la incertidumbre conlleva el costo de un rendimiento predictivo significativamente degradado. Además, la intercalación de bloques de atención probabilísticos con componentes estimados puntualmente (como MLPs y Normalización de Capa) interrumpe el flujo de propagación de la incertidumbre, impidiendo el establecimiento de una ruta estocástica bien definida.

Metodología
Los autores proponen DIRECTOR (Reconfiguración de TransfORmers Inspirada en Difusión), un marco que reconfigura los transformadores preentrenados en un proceso de difusión unificado para habilitar una propagación de incertidumbre principada.

1. Reconfiguración Estructural como Ruta de Probabilidad:
   El método comienza re-particionando la arquitectura del transformador de modo que cada bloque de transformación termine con un módulo de Atención Multicabeza (MHSA). Basándose en la conexión establecida entre MHSA y Procesos Gaussianos, la salida de cada bloque MHSA se interpreta como una distribución gaussiana sobre las características intermedias.
   Crucialmente, los segmentos deterministas (MLP y Normalización de Capa) que típicamente interrumpen el flujo de incertidumbre se reinterpretan no como funciones observadas de priores desconocidos, sino como parametrizaciones adicionales de la transición gaussianizada neural inducida por el MHSA reparametrizado como GP. Esta reestructuración asegura que toda la secuencia de transformaciones de características forme un proceso estocástico coherente $\{X_t\}_{t=0}^T$ con transiciones gaussianas, alineando efectivamente la arquitectura del transformador con un proceso de difusión en tiempo inverso.

2. Destilación en un Modelo de Difusión Unificado:
   En lugar de tratar la transición de cada bloque de forma independiente, DIRECTOR destila la secuencia observada de transformaciones de características en un núcleo de transición espacio-temporal unificado. Esto se logra entrenando un modelo basado en difusión (implementado como un DiT de un solo bloque) para aprender el proceso en tiempo inverso que mapea la distribución de características final de vuelta a la distribución de datos de entrada.
   El objetivo de entrenamiento combina dos componentes:

   • Pérdida Consciente de la Incertidumbre ($L_1$): Un término de divergencia de Kullback-Leibler (KL) que alinea el núcleo de transición aprendido con las transiciones observadas desde el backbone preentrenado congelado. Esto asegura que el modelo capture las correlaciones inherentes entre los bloques de atención.
   • Pérdida Consciente del Rendimiento ($L_2$): Una pérdida de predicción estándar (por ejemplo, entropía cruzada) aplicada a muestras extraídas del proceso de difusión aprendido. Esto asegura que el modelo reconfigurado mantenga o mejore la precisión predictiva.

Contribuciones Clave

• Reconfiguración Probabilística: El artículo introduce una reconfiguración estructural novedosa que transforma un transformador preentrenado en una secuencia de transiciones gaussianas neuralizadas, generalizando el modelo en un proceso de difusión parametrizado por un núcleo de transición unificado.
• Propagación Unificada de la Incertidumbre: Al destilar la ruta de características del transformador en un único modelo de difusión, el método captura las correlaciones entre bloques establecidas durante el preentrenamiento, habilitando una propagación de incertidumbre principada sin la intratabilidad de calcular posteriors completos del modelo.
• Calibración Post-hoc Eficiente: El enfoque funciona como una incrustación post-hoc del razonamiento probabilístico. El modelo de transición unificado es significativamente más pequeño que el backbone original (por ejemplo, 2.7M de parámetros frente a 6.24M para ViT), ofreciendo una mayor eficiencia de memoria mientras preserva las capacidades predictivas del modelo original.

Resultados Experimentales
Los autores evaluaron DIRECTOR en benchmarks de visión (CIFAR-10, CIFAR-100, CIFAR-10-C) y benchmarks de lenguaje (IMDB, CoLA, CoLA-OOD).

• Calibración y Precisión: DIRECTOR supera consistentemente a los transformadores conscientes de la incertidumbre existentes (incluidas las variantes KEP-SVGP) y a las líneas base estándar (Escalado de Temperatura, Dropout de Monte Carlo, Ensembles Profundos). En CIFAR-10, DIRECTOR logró una mayor precisión (84.74% frente a 82.68% para KEP-7/7) y un Error de Calibración Esperado menor (5.18% frente a 6.95%). Se observaron mejoras similares en todos los conjuntos de datos, demostrando que el método resuelve la compensación típica entre calibración y rendimiento.
• Robustez: Bajo desplazamientos de distribución (CIFAR-10-C y CoLA-OOD), DIRECTOR mantuvo una precisión predictiva competitiva mientras mejoraba significativamente las métricas de calibración en comparación con los modelos preentrenados.
• Detección OOD: El método mostró un rendimiento superior en tareas de detección de distribución fuera de distribución (OOD), logrando puntuaciones AUROC y AUPR más altas.
• Escalabilidad: Los experimentos en modelos más grandes (ViT-B-16 y Qwen-2.5) confirmaron que DIRECTOR escala de manera efectiva, ofreciendo una calibración fiable con menos parámetros que los modelos originales.

Significado y Afirmaciones
El artículo afirma que DIRECTOR ofrece un paradigma escalable y principiado para la calibración de incertidumbre en grandes modelos preentrenados. Al reinterpretar la pila de transformación de características del transformador como un proceso de difusión, el método incrusta estructuras expresivas conscientes de la incertidumbre directamente en la arquitectura interna del modelo. Los autores afirman que este enfoque permite un despliegue fiable en aplicaciones sensibles al riesgo, asegurando que las puntuaciones de confianza reflejen fielmente la precisión empírica, todo ello manteniendo o mejorando el rendimiento predictivo y mejorando la eficiencia de memoria. El trabajo abre una nueva dirección para mejorar la fiabilidad de los modelos fundacionales sin sacrificar sus capacidades centrales.