Weight Space Correlation Analysis: Quantifying Feature Utilization in Deep Learning Models

Resumen Técnico: Análisis de Correlación en el Espacio de Pesos para Cuantificar la Utilización de Características

Planteamiento del Problema
Los modelos de aprendizaje profundo en imágenes médicas son susceptibles al "aprendizaje de atajos" (shortcut learning), donde los modelos dependen de metadatos de confusión (p. ej., modelo de escáner, protocolo de adquisición) en lugar de señales clínicas genuinas. Si bien investigaciones previas han demostrado que dicha información de confusión suele estar codificada dentro de las representaciones (embeddings) de las imágenes (p. ej., Gichoya et al., 2022; Glocker et al., 2023), persiste una brecha crítica: la mera presencia de una característica en el espacio de embeddings no demuestra que el modelo la utilice activamente para las predicciones finales. Los métodos existentes, como el análisis de distribuciones de embeddings o la generación de conjuntos de prueba contrafactuales, no abordan directamente si la cabeza de clasificación aprovecha estos factores de confusión. Este artículo aborda la pregunta fundamental: ¿Se alinea activamente el límite de decisión del modelo con factores clínicamente irrelevantes?

Metodología: Análisis de Correlación en el Espacio de Pesos (WSC)
Los autores proponen el análisis de Correlación en el Espacio de Pesos (WSC, por sus siglas en inglés), una metodología interpretable para cuantificar la utilización de características mediante la medición de la alineación entre las cabezas de clasificación de una tarea clínica primaria y las tareas de metadatos auxiliares. El método procede en tres pasos:

1. Representación Lineal: El método trata los pesos de la cabeza de clasificación lineal ($W$) como "vectores de atención" que actúan sobre el espacio de embeddings latente. Para una tarea con $C$ clases, la matriz de pesos $W \in \mathbb{R}^{C \times d}$ define las direcciones de decisión para cada clase.
2. Proyección sobre la Variedad de Datos: Para asegurar que las comparaciones se realicen en dimensiones significativas, los pesos del clasificador se proyectan sobre la variedad intrínseca de los datos mediante el Análisis de Componentes Principales (PCA). Los autores extraen los $k$ componentes principales superiores (imponiendo un mínimo de 50 y capturando el 99% de la varianza) para formar una matriz de proyección $P$. Los pesos se transforman como $W' = W P^\top$. Este paso filtra las dimensiones ortogonales y no informativas.
3. Cuantificación de la Alineación: El método calcula la similitud de coseno entre los vectores de pesos proyectados de la tarea primaria ($W'_A$) y las tareas de metadatos auxiliares ($W'_m$). Una alta correlación entre los límites de decisión de una tarea clínica y una tarea de metadatos (p. ej., tipo de escáner) sirve como un indicador cuantitativo de shortcut learning, lo que indica que el modelo depende de las mismas características latentes para ambas predicciones.

Contribuciones Clave

• Métrica Novedosa: Introducción del análisis WSC para distinguir entre la codificación de metadatos y su utilización activa en la inferencia.
• Validación mediante Sesgo Inducido: Los autores validaron el método creando un "Conjunto de Datos de Sesgo Inducido" donde se introdujo artificialmente una fuerte correlación entre las clases de planos fetales y los modelos de escáner. El análisis WSC detectó con éxito la alta alineación resultante en los vectores de pesos, mientras que el "Conjunto de Datos Balanceado" (sin correlación) mostró una baja alineación, confirmando la sensibilidad del método al shortcut learning.
• Aplicación a Modelos Clínicos: La metodología se aplicó al modelo SA-SonoNet para la predicción de Parto Pretérmino Espontáneo (sPTB), un modelo de vanguardia en ecografía obstétrica.

Resultados Experimentales
El estudio utilizó dos conjuntos de datos de ecografía clínica privados: un conjunto de datos de clasificación de anatomía fetal y un conjunto de datos de cuello uterino para la predicción de sPTB.

• Codificabilidad vs. Utilización: En el conjunto de datos fetal, los modelos base confirmaron que los metadatos (escáner, ID del hospital) eran altamente codificables en los embeddings de la imagen (alto AUROC al realizar sondeos). Sin embargo, el análisis WSC reveló que los vectores de pesos para la tarea de clasificación de planos primarios eran mayoritariamente ortogonales (baja correlación) a los de la tarea de clasificación de escáner. Esto se mantuvo incluso para modelos multitarea entrenados explícitamente para codificar metadatos, sugiriendo que la tarea primaria no dependía de atajos del escáner.
• Detección de Atajos Inducidos: Cuando el conjunto de datos fue sesgado para correlacionar planos específicos con escáneres específicos, los valores de WSC entre los pesos de los planos y los del escáner aumentaron significativamente. Esto confirmó que el WSC puede detectar cuándo un modelo es forzado a adoptar una estrategia de atajo.
• Análisis de SA-SonoNet: Aplicado al modelo de predicción de sPTB, el análisis mostró:
  • Relevancia Clínica: Los pesos del clasificador de sPTB estaban moderadamente correlacionados con factores clínicamente relevantes como la longitud cervical y el espaciado de píxeles.
  • Desacoplamiento de Confundidores: Los pesos estaban desacoplados de factores de adquisición clínicamente irrelevantes, específicamente el tipo de escáner.
  • Advertencia: Aunque el peso al nacer mostró un alto WSC con sPTB, los autores señalan que esto era funcionalmente insignificante porque la cabeza de sondeo falló al predecir el peso al nacer de manera efectiva, indicando que la información no estaba codificada de forma robusta.

Significancia y Reivindicaciones
El artículo sostiene que el análisis WSC proporciona una herramienta necesaria para verificar la confiabilidad de los modelos en imágenes médicas. Al inspeccionar si un clasificador utiliza características ajenas a la señal clínica genuina, los investigadores pueden ir más allá de las métricas de rendimiento simples para comprender la lógica de decisión interna. Los autores concluyen que, aunque los factores de confusión suelen estar presentes en los embeddings de las imágenes médicas, modelos como SA-SonoNet pueden diseñarse o entrenarse para que sus límites de decisión permanezcan alineados con señales clínicamente significativas y desacoplados de correlaciones espurias como los modelos de escáner. Esto ofrece un método para confirmar que un modelo está aprendiendo la patología pretendida en lugar de explotar artefactos del conjunto de datos.

Limitaciones
Los autores reconocen que el método asume una relación lineal entre el espacio de embeddings y las predicciones de clase, lo cual se alinea con las convenciones estándar del aprendizaje profundo pero podría no capturar interacciones no lineales complejas. Además, el método actualmente analiza atributos de metadatos individuales; los atajos en el mundo real pueden involucrar combinaciones no lineales de múltiples confundidores, algo que investigaciones futuras pretenden investigar. El artículo también señala que un bajo WSC no garantiza la robustez ante cambios de dominio causados por otros factores, como cambios en la calidad de la imagen o desplazamientos de características, pero sí descarta el shortcut learning como la causa de tales fallos.