Efficient KernelSHAP Explanations for Patch-based 3D Medical Image Segmentation

Resumen Técnico: Explicaciones Eficientes de KernelSHAP para Segmentación 3D Médica Basada en Parches

1. Planteamiento del Problema

Los modelos de aprendizaje profundo, como nnU-Net, se han convertido en el estándar para la segmentación de imágenes médicas 3D (por ejemplo, en radioterapia). Sin embargo, su naturaleza de "caja negra" dificulta la auditoría de errores, que pueden ser sutiles y dependientes del contexto anatómico.

• Desafío Principal: Los métodos de explicabilidad basados en perturbación, como KernelSHAP, son ideales porque son agnósticos al modelo y ofrecen atribuciones robustas, pero son computacionalmente prohibitivos para la segmentación 3D basada en parches (ventana deslizante).
• Causa de la ineficiencia: Evaluar las coaliciones de características requiere miles de pasadas forward (inferencias) sobre parches superpuestos. En volúmenes 3D, esto genera una carga computacional masiva debido a la redundancia en el cálculo de parches que no se ven afectados por la perturbación.
• Falta de Consenso: No existe una práctica consolidada sobre cómo definir unidades interpretables (voxels vs. órganos vs. supervoxels) ni cómo agregar las evidencias a nivel de voxel para obtener puntuaciones significativas clínicamente.

2. Metodología Propuesta

Los autores proponen un marco de trabajo eficiente que adapta KernelSHAP a datos de CT volumétricos bajo dos principios clave: localización de la computación y reutilización de inferencias.

A. Definición de la Región de Interés (ROI) y Soporte del Campo Receptivo

• En lugar de procesar todo el volumen, el sistema se centra en una ROI definida por el usuario (ej. una componente conectada de un órgano).
• Se calcula un cuadro delimitador cúbico expandido ($B_{RF}$) que cubre no solo la ROI, sino también el soporte de entrada necesario (radio del parche) para que las predicciones dentro de la ROI sean válidas.
• Todas las perturbaciones y pasadas forward se restringen a este subvolumen, manteniendo el resto de la imagen fija.

B. Unidades Interpretables y Perturbación
El volumen dentro de $B_{RF}$ se particiona en $M$ unidades interpretables. Se comparan tres estrategias automáticas:

1. Órganos Completos (Full Organs): Cada estructura anatómica segmentada (ej. por TotalSegmentator) es una unidad. Maximiza la interpretabilidad semántica.
2. Supervoxels Regulares (FCC): Una teselación geométrica basada en una red cúbica centrada en las caras (FCC) en coordenadas físicas. Es agnóstico a la anatomía.
3. Supervoxels Híbridos (Organ-Aware): Combina la regularidad geométrica de FCC con las restricciones anatómicas. Si una celda FCC cruza límites de órganos, se subdivide. Busca un equilibrio entre granularidad y alineación anatómica.

• Operador de Perturbación: Se utiliza un enmascaramiento duro ("hard masking") donde los voxels de una unidad eliminada se reemplazan por un valor de aire ($-1024$ HU), simulando una intervención física realista.

C. Funciones de Valor y Agregación
Para evaluar cada coalición, se proponen cuatro funciones de puntuación escalares basadas en la predicción del modelo:

• Verdadero Positivo (TP): Recompensa el soporte de logits para voxels correctamente predichos.
• Falso Positivo (FP): Penaliza el soporte de logits para voxels predichos incorrectamente como positivos.
• Dice Score: Mide la similitud de Dice entre la predicción perturbada y la línea base.
• Soft Dice Score: Una versión suave que recompensa/agrega logits en voxels consistentes y penaliza los nuevos positivos.

D. Optimización: Caché de Logits de Parches (Patch Logit Caching)
Esta es la contribución técnica central para la eficiencia:

• Se realiza una inferencia inicial en la ROI no perturbada y se almacenan en caché (CPU) los logits de todos los parches de la ventana deslizante.
• Durante la evaluación de las coaliciones de KernelSHAP, si una perturbación no afecta a un parche específico, se reutiliza el logit en caché en lugar de ejecutar una nueva inferencia en la GPU.
• Esto reduce drásticamente el número de pasadas forward necesarias, acelerando el proceso sin alterar el esquema de fusión de nnU-Net.

3. Resultados Clave

A. Evaluación Cuantitativa (Curvas de Perturbación)
Se utilizaron métricas como AOPC (Área sobre la curva de perturbación) y ABPC (Área entre curvas) para medir la fidelidad de las explicaciones.

• Supervoxels Regulares (FCC): Lograron las métricas de fidelidad más altas (mayor AOPC/ABPC). Sin embargo, esto se debe a que su gran tamaño y superposición con el objetivo de segmentación provocan caídas drásticas en la puntuación al ser perturbados, lo que no necesariamente implica una explicación clínicamente útil o anatómicamente alineada.
• Órganos Completos: Proporcionaron explicaciones estables y de alto nivel, pero con poca granularidad para diagnosticar errores locales.
• Supervoxels Híbridos: Ofrecieron el mejor equilibrio. Mantuvieron la semántica anatómica y lograron una granularidad fina. Fueron particularmente efectivos para identificar los drivers de falsos positivos bajo métricas normalizadas, mostrando heterogeneidad dentro de los órganos que los métodos puramente geométricos o de órgano completo no capturan.

B. Rendimiento Computacional

• La técnica de caché de parches redujo significativamente el tiempo de cómputo.
• Ahorros: La tasa de aciertos en la caché (cache hit rate) varió entre el 15% y el 32%, dependiendo de la definición de la unidad.
• Tiempo de ejecución: Se observaron ahorros computacionales del 15% al 30% en comparación con una implementación ingenua sin caché. El método híbrido mantuvo una alta reutilización de caché (30.2%) a pesar de su mayor complejidad de partición.

C. Análisis Cualitativo

• Los mapas de atribución generados con unidades híbridas permitieron visualizar cómo subregiones específicas dentro de un mismo órgano podían tener contribuciones opuestas (positivas para la predicción correcta, negativas para falsos positivos), ofreciendo una herramienta valiosa para la depuración de modelos.

4. Contribuciones Principales

1. Marco Eficiente de KernelSHAP: Adaptación exitosa de métodos de Shapley a la inferencia por parches 3D mediante la restricción a la ROI y el soporte del campo receptivo.
2. Mecanismo de Caché de Logits: Una estrategia novedosa que reutiliza predicciones base para parches no afectados, haciendo viable el uso de KernelSHAP en volúmenes 3D clínicos.
3. Análisis de Unidades Interpretables: Estudio comparativo exhaustivo que demuestra que la elección de la unidad (geométrica vs. anatómica) implica una compensación (trade-off) clara entre fidelidad métrica (que favorece a las unidades grandes/geométricas) y interpretabilidad clínica (que favorece a las unidades anatómicas/híbridas).
4. Estrategias de Agregación: Demostración de que diferentes funciones de valor (TP, FP, Dice) revelan comportamientos distintos del modelo, siendo crucial elegir la función adecuada según la pregunta clínica (ej. entender por qué el modelo falla vs. por qué acierta).

5. Significado e Impacto

Este trabajo avanza significativamente en la IA Explicable (XAI) para medicina al resolver el cuello de botella computacional que impedía el uso de métodos basados en perturbación en segmentación 3D.

• Viabilidad Clínica: Hace que las explicaciones de Shapley sean prácticas para entornos clínicos reales, permitiendo a los radioterapeutas y físicos médicos auditar por qué un modelo segmentó un órgano de cierta manera.
• Mejora en la Confianza: Al identificar no solo qué regiones apoyan la predicción, sino también qué regiones causan falsos positivos (especialmente con unidades híbridas), se facilita la detección de sesgos del modelo o artefactos de adquisición.
• Dirección Futura: El estudio sugiere que las explicaciones deben diseñarse no solo para maximizar métricas de fidelidad abstractas, sino para alinearse con la anatomía humana y las necesidades de diagnóstico específicas, proponiendo el uso de supervoxels basados en datos (como SLIC 3D) en trabajos futuros.

En resumen, el artículo presenta una solución técnica robusta que transforma un método teórico costoso en una herramienta práctica para la validación de modelos de segmentación médica 3D, equilibrando eficiencia computacional, fidelidad matemática y utilidad clínica.