Unified Medical Image Tokenizer for Autoregressive Synthesis and Understanding

1. El Problema

El campo de la inteligencia artificial multimodal ha avanzado significativamente gracias a los modelos autoregresivos (AR), que procesan datos como secuencias de tokens. Sin embargo, su aplicación en el dominio médico enfrenta dos barreras críticas:

• Falta de un tokenizador unificado: Los enfoques existentes suelen especializarse en un solo aspecto: o bien preservan la estructura visual de bajo nivel (necesaria para la reconstrucción y síntesis de imágenes, como VQGAN), o bien capturan semántica de alto nivel (necesaria para la comprensión clínica, como CLIP). Ninguno logra ambas cosas simultáneamente de manera efectiva.
• Escasez de datos emparejados: Entrenar tokenizadores unificados que combinen reconstrucción y semántica textual requiere grandes cantidades de pares imagen-texto. En medicina, estos datos son escasos y costosos de obtener, mientras que las imágenes médicas no etiquetadas son abundantes. Los métodos actuales que optimizan ambos objetivos simultáneamente sufren de interferencia de gradientes y rendimiento subóptimo debido a la falta de datos emparejados.

2. Metodología: MedITok

Los autores proponen MedITok, el primer tokenizador unificado de imágenes médicas, diseñado mediante un marco de entrenamiento de dos etapas que utiliza representaciones visuales preentrenadas como puente para resolver los conflictos entre la fidelidad estructural y la riqueza semántica.

Arquitectura

MedITok se basa en una arquitectura de Vector Quantization (VQ) con un codificador de imágenes ($E$), un cuantizador ($Q$) y un decodificador ($D$). Utiliza un código de 32,768 vectores de 8 dimensiones.

Etapas de Entrenamiento

1. Alineación de Representación Visual (Etapa 1):
   • Objetivo: Establecer una base robusta de fidelidad de reconstrucción y conciencia semántica básica.
   • Datos: Se utiliza un corpus masivo de imágenes médicas no emparejadas (33.4 millones de imágenes).
   • Mecanismo: Se optimiza la pérdida de reconstrucción (fidelidad de imagen) junto con una restricción semántica ligera guiada por un codificador de visión preentrenado (BiomedCLIP-V). Esto permite aprovechar la abundancia de datos sin etiquetas sin interferir con el aprendizaje de la estructura.
2. Alineación Semántica Textual (Etapa 2):
   • Objetivo: Inyectar semántica clínica de grano fino.
   • Datos: Se utilizan pares imagen-texto de alta calidad (2.4 millones de pares).
   • Mecanismo: Se ajusta el tokenizador para alinear los tokens de imagen con las incrustaciones textuales detalladas de los informes clínicos (usando un codificador de texto preentrenado, BiomedCLIP-T). Esto enriquece el espacio latente con conceptos clínicos específicos.

Este enfoque escalonado evita la interferencia mutua entre los objetivos de reconstrucción y comprensión, permitiendo un espacio latente unificado que soporta tanto la síntesis como la interpretación.

3. Contribuciones Clave

1. Nuevo Marco de Entrenamiento: Se introduce una estrategia de dos etapas que desacopla la alineación visual de la textual, permitiendo el uso eficiente de datos no emparejados masivos y resolviendo el problema de la interferencia de gradientes.
2. MedITok: El primer tokenizador unificado para imágenes médicas que codifica simultáneamente detalles estructurales de bajo nivel (para síntesis realista) y semántica clínica de alto nivel (para comprensión).
3. Dataset a Gran Escala: Se curó un conjunto de datos masivo que incluye más de 33 millones de imágenes médicas de 9 modalidades (CT, MRI, Rayos X, dermatoscopia, etc.) y 2.4 millones de pares imagen-texto, cubriendo una amplia gama de anatomías y patologías.
4. Rendimiento SOTA: Demostración de que MedITok supera a los tokenizadores generales y médicos existentes en múltiples tareas, estableciendo un nuevo estándar para modelos multimodales médicos.

4. Resultados Experimentales

El modelo fue evaluado en más de 30 conjuntos de datos y 9 modalidades de imagen, abarcando 4 familias de tareas:

• Codificación de Bajo Nivel (Reconstrucción): MedITok logró el mejor rendimiento en fidelidad de imagen (rFID, PSNR, SSIM) en 8 de las 9 modalidades probadas, superando incluso a modelos especializados en reconstrucción como PUMIT y MedVAE, a pesar de usar un factor de submuestreo de 16x. Preservó detalles diagnósticos críticos mejor que sus competidores.
• Codificación de Alto Nivel (Comprensión): En tareas de clasificación (probing lineal) y preguntas visuales (VQA), MedITok obtuvo los puntajes más altos en mAP y AUC, demostrando que sus tokens latentes contienen información semántica clínica rica y discriminativa.
• Síntesis de Imágenes Médicas: Integrado en un modelo generativo autoregresivo (LlamaGen), MedITok generó imágenes sintéticas de alta fidelidad y diversidad en múltiples modalidades (piel, fondo de ojo, patología, etc.), superando significativamente a modelos basados en tokenizadores generales (UniTok) o médicos puros (VQGAN).
• Comprensión Visual (VQA): En benchmarks como VQA-RAD y SLAKE, el modelo LLaVA-Med potenciado por MedITok superó a todas las variantes basadas en otros tokenizadores, logrando una mayor precisión en la identificación de hallazgos clínicos y reduciendo alucinaciones.

5. Significado e Impacto

• Fundamento para Modelos Multimodales: MedITok actúa como un componente fundamental escalable para la próxima generación de modelos de IA médica, permitiendo la construcción de sistemas que pueden tanto sintetizar imágenes médicas realistas (para entrenamiento o análisis contrafactual) como comprender y diagnosticar a partir de ellas.
• Democratización de la IA Médica: Al proporcionar un tokenizador unificado de código abierto y un marco de entrenamiento que aprovecha datos no etiquetados, reduce la barrera de entrada para desarrollar modelos de IA médica de alto rendimiento sin depender exclusivamente de costosos datos emparejados.
• Eficiencia y Generalización: La capacidad de manejar 9 modalidades diferentes con un solo tokenizador demuestra una generalización superior, eliminando la necesidad de arquitecturas específicas para cada tipo de imagen médica.

En resumen, este trabajo resuelve el dilema histórico entre la fidelidad visual y la comprensión semántica en la IA médica, ofreciendo una solución práctica y escalable que impulsa tanto la generación como el análisis de imágenes clínicas.