MedSAE: Dissecting MedCLIP Representations with Sparse Autoencoders

Resumen Técnico: MEDSAE – Disecando las Representaciones de MedCLIP con Autoencoders Escasos

Planteamiento del Problema
La integración de la Inteligencia Artificial (IA) en la atención sanitaria, particularmente a través de Modelos de Lenguaje Multimodal Grandes (MLLMs) como MedCLIP, ha mejorado las capacidades diagnósticas pero ha introducido desafíos significativos en cuanto a la interpretabilidad. Aunque estos modelos logran un alto rendimiento en tareas como el análisis de radiografías de tórax, sus mecanismos internos permanecen como "cajas negras" opacas. Operan en un régimen de superposición, donde el número de características aprendidas excede la dimensionalidad del espacio de representación, provocando que múltiples características se multiplexen a través de neuronas compartidas. Esto complica la atribución de conceptos clínicos específicos a activaciones individuales del modelo. Las técnicas existentes de interpretabilidad mecánica no han sido escaladas o adaptadas suficientemente para tareas complejas de visión y lenguaje médico, dejando una brecha entre la IA médica de alto rendimiento y la transparencia requerida para la fiabilidad clínica.

Metodología
Los autores proponen MedSAE, un marco que aplica Autoencoders Escasos (SAE) al espacio latente de MedCLIP para desentrañar las características superpuestas en conceptos comprensibles para humanos, aproximadamente monosémicos. El pipeline consta de tres etapas principales:

1. Extracción de Incrustaciones y Entrenamiento:

   • Los autores utilizan el codificador de visión de MedCLIP (congelado) para generar incrustaciones a partir del conjunto de datos CheXpert (radiografías de tórax).
   • Para abordar un posible desalineamiento entre modalidades, las incrustaciones se normalizan centrando (restando las medias por modalidad) y escalando para garantizar normas consistentes.
   • Se entrena un SAE basado en ReLU sobre estas incrustaciones normalizadas. El SAE mapea vectores de entrada densos ($x$) a una representación latente dispersa ($z$) utilizando una arquitectura de codificador-decodificador con una penalización $L_1$ para imponer dispersión:
     $$L := \|x - \hat{x}\|_2^2 + \lambda\|z\|_1$$
   • El objetivo es aprender un conjunto sobredeterminado de direcciones latentes que separen las características multiplexadas.

2. Evaluación de la Monosémanticidad:

   • Para validar que las neuronas del SAE corresponden a conceptos únicos, los autores calculan la correlación de Pearson entre la activación de cada neurona y las etiquetas clínicas codificadas como multi-hot (por ejemplo, Edema, Neumonía).
   • Calculan la entropía de la distribución de correlación para cada neurona a través de todas las etiquetas. Una entropía más baja indica que una neurona está fuertemente concentrada en un solo concepto (alta monosémanticidad), mientras que una entropía más alta sugiere una activación distribuida entre múltiples conceptos.

3. Nombrado Automatizado y Validación:

   • Reconociendo la dificultad de construir un vocabulario médico robusto para la coincidencia basada en decodificadores, los autores emplean MedGemma, un modelo fundacional de visión y lenguaje, para la identificación automatizada de conceptos.
   • Para cada neurona del SAE, las imágenes con mayor activación se envían a MedGemma mediante indicaciones estructuradas para generar una explicación de concepto compartida (por ejemplo, "Edema pulmonar severo").
   • Validación: Para asegurar que estos nombres generados no sean alucinaciones, se realiza una tarea de detección. Se presenta a MedGemma la descripción del concepto generado y un conjunto equilibrado de imágenes activadoras y no activadoras, y se le pide que clasifique qué imágenes coinciden con el concepto. Esto proporciona una medida cuantitativa de la alineación semántica.

Contribuciones Clave

• Marco de Evaluación Novel: El artículo introduce un método de evaluación cuantitativa que combina métricas basadas en correlación con el nombrado automatizado de conceptos y su validación utilizando MedGemma. Esto va más allá del simple análisis de activaciones para proporcionar etiquetas interpretables por humanos.
• Validación Empírica de la Linealidad: Los autores demuestran que las incrustaciones de MedCLIP satisfacen aproximadamente la hipótesis de superposición lineal. Mediante una comprobación de sanidad que involucra imágenes compuestas, mostraron que las incrustaciones compuestas pueden reconstruirse linealmente a partir de sus componentes con una alta similitud del coseno (0.97), justificando la aplicación de herramientas de interpretabilidad lineal como los SAE.
• Desentrañamiento de Características: El estudio muestra que los SAE pueden extraer con éxito características médicamente significativas e interpretables de las representaciones internas de un modelo de visión y lenguaje médico, superando a las incrustaciones crudas de MedCLIP en términos de selectividad de clase.

Resultados
Los experimentos realizados en el conjunto de datos CheXpert arrojan los siguientes hallazgos:

• Mejora en la Monosémanticidad: Las neuronas del SAE exhiben una entropía más baja (2.25) en comparación con las incrustaciones crudas de MedCLIP (2.38), lo que indica una mayor alineación con conceptos únicos.
• Mayor Precisión en la Interpretabilidad: El marco de nombrado automatizado identificó 21 neuronas SAE distintas con una precisión de detección superior al 70% (por ejemplo, "Derrame pleural derecho con posibles loculaciones", "Cardiomegalia severa con congestión vascular pulmonar significativa"). En contraste, las incrustaciones crudas de MedCLIP produjeron solo dos interpretaciones de alta precisión de este tipo.
• Robustez de los Hallazgos: Una condición de control donde los pares neurona-concepto se barajaron aleatoriamente resultó en una caída significativa en el rendimiento de detección, confirmando que las interpretaciones significativas surgen de la estructura del SAE y no de artefactos estadísticos.
• Equilibrio entre Dispersión y Reconstrucción: El modelo seleccionado logró un equilibrio favorable con una tasa de activación $L_0$ del 0.20% y una fracción de varianza explicada (FVE) de 0.98, a pesar de una tasa del 30% de neuronas "muertas".

Significado y Afirmaciones
El artículo posiciona a MedSAE como un paso escalable hacia representaciones clínicamente fiables en la IA médica. Al conectar modelos de alto rendimiento con la transparencia, el trabajo ofrece una vía para descifrar los mecanismos computacionales en conceptos médicos comprensibles para humanos. Los autores afirman que este es el primer análisis empírico de la interpretabilidad mecánica en MLLMs médicos utilizando SAE.

Sin embargo, los autores mantienen un tono modesto respecto a las limitaciones. Reconocen que su enfoque se basa en la hipótesis de representación lineal, que podría no sostenerse universalmente, y que su análisis actual se limita a una sola capa (MedCLIP-ResNet). Señalan que las unidades inactivas podrían reflejar tensiones entre capacidad y dispersión en lugar de un fallo del enfoque, y que la validación con intervención humana por parte de expertos médicos sigue siendo necesaria para evaluar la relevancia clínica y la seguridad. El trabajo se presenta como un paso fundamental para motivar más investigación hacia sistemas de IA médica transparentes, robustos y de confianza.