an interpretable vision transformer framework for automated brain tumor classification

Título: Un marco de Vision Transformer interpretable para la clasificación automatizada de tumores cerebrales

1. Planteamiento del Problema

Los tumores cerebrales representan una de las condiciones neurológicas más críticas, donde el diagnóstico temprano y preciso es vital para la supervivencia del paciente. El estándar de oro actual, la interpretación manual de imágenes por Resonancia Magnética (MRI) por parte de radiólogos, presenta varios desafíos:

• Limitaciones humanas: Es intensiva en tiempo, propensa a la variabilidad interobservador y requiere una experiencia especializada que a menudo es escasa en entornos con recursos limitados.
• Limitaciones de modelos previos: Aunque las Redes Neuronales Convolucionales (CNN) han dominado el campo, su dependencia de campos receptivos localmente espaciales limita su capacidad para modelar dependencias estructurales de largo alcance, las cuales son diagnósticamente relevantes en las imágenes de MRI cerebral.
• Necesidad de interpretabilidad: La falta de explicabilidad en los modelos de "caja negra" dificulta su despliegue clínico, ya que los médicos necesitan validar en qué regiones del cerebro se basa la decisión del modelo.

2. Metodología

El artículo propone un marco de aprendizaje profundo basado en Vision Transformer (ViT) para la clasificación de cuatro clases: glioma, meningioma, tumor pituitario y tejido cerebral sano.

• Base de Datos: Se utilizó un conjunto de datos público de 7,023 imágenes de MRI, distribuidas de manera relativamente equilibrada entre las cuatro clases. Los datos se dividieron en entrenamiento (80%), validación (10%) y prueba (10%) mediante muestreo estratificado.
• Preprocesamiento (CLAHE): Se aplicó Equalización de Histograma Adaptativa Limitada por Contraste (CLAHE) antes del entrenamiento. Operando en el espacio de color CIE LAB (solo canal de luminancia) con un límite de contraste de 2.0, esta técnica mejora el contraste local y resalta los bordes de los tumores que a menudo son invisibles con la normalización estándar, sin amplificar el ruido en regiones homogéneas.
• Arquitectura: Se empleó un ViT-B/16 (Base con parches de 16x16) preentrenado en ImageNet-21k como columna vertebral. La imagen de entrada (224x224) se divide en 196 parches, procesados por 12 bloques de codificadores Transformer con atención multi-cabeza.
• Estrategia de Entrenamiento:
  • Aumento de Datos: Se combinaron transformaciones a nivel de píxel (rotación, traslación, zoom) con técnicas de aumento a nivel de muestra: MixUp (interpolación de imágenes y etiquetas) y CutMix (parches de una imagen en otra).
  • Ajuste Fino en Dos Etapas:
    1. Calentamiento (5 épocas): El backbone se congela y solo se entrena la cabeza de clasificación.
    2. Ajuste completo (hasta 15 épocas): Se desbloquean todos los parámetros con tasas de aprendizaje discriminativas (más bajas para el backbone, más altas para la cabeza).
  • Regularización: Se utilizó Exponential Moving Average (EMA) de los pesos para estabilizar el modelo y Test-Time Augmentation (TTA) durante la inferencia, promediando las probabilidades de cinco vistas diferentes de la imagen.
• Interpretabilidad: Se implementó Attention Rollout para generar mapas de calor clínicamente interpretables. Esta técnica rastrea el flujo de información desde los parches de entrada hasta el token de clasificación [CLS] a través de las capas de atención, permitiendo visualizar qué regiones del cerebro impulsaron la predicción.

3. Contribuciones Clave

1. Alto Rendimiento con ViT: Lograr una precisión de prueba del 99.29% en un problema de clasificación de cuatro clases, superando a todas las líneas base basadas en CNN.
2. Preprocesamiento Específico para MRI: La integración de CLAHE motivada por la física de adquisición de MRI para mejorar la discriminabilidad de los bordes tumorales.
3. Estrategia de Aumento Avanzada: La combinación efectiva de MixUp y CutMix en arquitecturas Transformer para mejorar la generalización en datos médicos.
4. Pipeline de Entrenamiento Robusto: Uso de ajuste fino en dos etapas, EMA y TTA para maximizar el rendimiento y la estabilidad.
5. Interpretabilidad Nativa: Aplicación sistemática de Attention Rollout para generar mapas de calor que validan la coherencia clínica del modelo (localizando masas tumorales o simetrías anatómicas), algo crucial para la adopción clínica.

4. Resultados

• Métricas Generales: El modelo alcanzó una precisión de 99.29% y un F1-score macro de 99.25% en el conjunto de prueba de 703 imágenes.
• Rendimiento por Clase:
  • Sano: Precisión y Recall perfectos (100%).
  • Meningioma: Recall perfecto (100%), con un 97.06% de precisión.
  • Glioma: Precisión perfecta (100%) y Recall del 98.15%.
  • Tumor Pituitario: Precisión perfecta (100%) y Recall del 98.86%.
• Errores: Solo se produjeron 5 mal clasificaciones en total (3 gliomas y 2 pituitarios clasificados erróneamente como meningiomas), lo que demuestra una alta seguridad clínica, especialmente en la ausencia de confusiones entre tejido sano y tumoral.
• Comparativa: Superó a la línea base CNN más fuerte (ResNet-50/GoogLeNet) y a otros modelos ViT previos (que alcanzaban ~98.70%), reduciendo aproximadamente 120 errores por cada 7,000 pacientes en un escenario de cribado.

5. Significado e Impacto

Este trabajo demuestra que los Vision Transformers, cuando se combinan con un preprocesamiento adecuado y estrategias de entrenamiento rigurosas, superan a las CNN tradicionales en la clasificación de tumores cerebrales. Su mayor valor reside en la interpretabilidad: los mapas de calor generados por Attention Rollout confirman que el modelo se enfoca en regiones anatómicas relevantes (masas infiltrativas para gliomas, masas periféricas para meningiomas y simetría para casos sanos), lo que valida su coherencia clínica.

El marco propuesto no solo ofrece un estado del arte en precisión, sino que aborda la barrera de la "caja negra" en la IA médica, sentando las bases para futuras validaciones clínicas prospectivas y la integración de estos sistemas como herramientas de apoyo a la decisión (segundo lector) en neurorradiología.