Diffusion MRI Transformer with a Diffusion Space Rotary Positional Embedding (D-RoPE)

1. Problema y Motivación

La Imagen por Resonancia Magnética de Difusión (dMRI) es fundamental para estudiar los cambios microestructurales en el cerebro y se utiliza ampliamente en la práctica clínica. Sin embargo, el avance en el aprendizaje de representaciones de propósito general a partir de dMRI ha sido limitado debido a dos desafíos principales:

• Complejidad de los datos: A diferencia de las imágenes médicas estándar (como T1 o T2), la dMRI consiste en múltiples volúmenes cerebrales con características de atenuación que dependen de la dirección y la fuerza de los gradientes de difusión (vectores b y valores b). Esto requiere modelar conjuntamente las dependencias espaciales, de ponderación de difusión y direccionales.
• Variabilidad de protocolos: Los protocolos de adquisición varían significativamente entre estudios (diferentes números de direcciones de difusión y valores b), lo que limita la aplicabilidad de los modelos tradicionales que requieren un esquema de adquisición fijo o una interpolación previa.

Los enfoques existentes (como el muestreo a un protocolo común, modelado gráfico o convoluciones modificadas) a menudo no capturan adecuadamente las propiedades únicas de las señales de difusión ni generalizan bien a diferentes configuraciones.

2. Metodología

Los autores proponen una arquitectura basada en Transformers diseñada específicamente para dMRI, entrenada mediante un objetivo de Autoencoder enmascarado (MAE).

Arquitectura del Modelo

• Entrada: Los volúmenes de dMRI se dividen en parches 3D.
• Codificación Posicional Absoluta: Se modifica para incluir tanto la posición espacial $(x, y, z)$ (codificada sinusoidalmente) como la información del espacio de difusión (coordenadas esféricas $\rho, \theta, \phi$ del vector b), codificada mediante una capa lineal aprendible.
• Bloques de Atención: Inspirados en los transformadores de video, los bloques de atención alternan entre las dimensiones espaciales y las dimensiones de difusión para capturar las interdependencias.

Innovación Clave: D-RoPE (Diffusion Space Rotary Positional Embedding)

El núcleo de la propuesta es una generalización de la Codificación Posicional Rotatoria (RoPE) para el espacio de difusión:

• El Problema de la Distancia: La RoPE original asume una distancia lineal y ordinal ($m-n$) entre tokens. En dMRI, los tokens corresponden a puntos en un espacio esférico (dirección del gradiente) y un espacio lineal (valor b).
• La Solución: Se define una nueva métrica de distancia $D(m, n)$ que combina:
  1. La distancia lineal entre los valores b ($b_m - b_n$).
  2. El ángulo entre los vectores de difusión ($\arccos(|v_m \cdot v_n|)$), garantizando simetría de 180° (ya que direcciones opuestas producen codificaciones equivalentes).
• Fórmula: $D(m, n) = \sqrt{\gamma(b_m - b_n)^2 + \arccos^2(|v_m \cdot v_n|)}$.
• Esta distancia se utiliza para rotar los vectores de consulta y clave en el mecanismo de atención, permitiendo al modelo entender la relación geométrica entre diferentes direcciones de difusión y valores b, independientemente del número total de direcciones.

Entrenamiento (Pre-entrenamiento)

• Se utiliza un esquema de Autoencoder enmascarado (MAE) con una estrategia de enmascaramiento alternada (espacial y direccional) para forzar al modelo a aprender tanto la estructura anatómica como las relaciones entre las direcciones de difusión.
• El modelo se entrena en el conjunto de datos HCP-YA (Human Connectome Project - Young Adult) y luego se evalúa en tareas de ajuste fino (fine-tuning) y extracción de características congeladas.

3. Contribuciones Clave

1. Arquitectura Novel: Propuesta de un Transformer capaz de manejar adquisiciones de dMRI arbitrarias (número variable de direcciones y valores b) sin necesidad de re-muestreo.
2. D-RoPE: Introducción de una codificación posicional rotatoria adaptada a la geometría esférica y lineal del espacio de difusión, capturando la interdependencia entre direcciones y valores b.
3. Evaluación Exhaustiva: Validación de las representaciones aprendidas en múltiples tareas aguas abajo (regresión de edad, clasificación de sexo, clasificación de deterioro cognitivo leve y predicción de escalas cognitivas), demostrando superioridad sobre baselines.

4. Resultados

El modelo se evaluó en conjuntos de datos de HCP (Desarrollo, Envejecimiento) y ADNI (Alzheimer).

• Reconstrucción (MAE): La combinación de D-RoPE con enmascaramiento alternado obtuvo los mejores resultados en métricas de reconstrucción (PSNR, SSIM, FID), superando significativamente a los modelos sin D-RoPE. Esto indica que el modelo aprende tanto la estructura espacial como las relaciones direccionales.
• Tareas de Regresión (Edad): El modelo pre-entrenado con ajuste fino (fine-tuning de la última capa) alcanzó un coeficiente de correlación ($\rho$) de 0.94, superando a modelos entrenados desde cero (3D ResNet, ViT estándar) y a características manuales.
• Tareas de Clasificación (MCI vs. Control):
  • En la clasificación de Deterioro Cognitivo Leve (MCI) frente a Cognitivamente Normal (CN), el uso de características latentes congeladas con una cabeza MLP logró la mayor precisión equilibrada (64.7%), superando a las características manuales (61.9%) y a otros modelos de deep learning.
  • Esto es notable porque se logró con datos de entrenamiento limitados, sugiriendo que el pre-entrenamiento capturó patrones microestructurales predictivos.
• Predicción de ADAS-Cog: El modelo logró un coeficiente de correlación de 0.38 para predecir la escala de evaluación de la enfermedad de Alzheimer, superando a todas las baselines (incluyendo ViT con RoPE estándar y características manuales).

5. Significado e Impacto

• Fundación para Modelos Base: Este trabajo es un paso crucial hacia la creación de un "modelo base" (foundation model) para dMRI, capaz de armonizar análisis entre cohortes con protocolos de adquisición heterogéneos.
• Eficiencia de Datos: Demuestra que las representaciones pre-entrenadas son altamente efectivas incluso con conjuntos de datos pequeños para tareas clínicas, reduciendo la necesidad de grandes cantidades de datos etiquetados.
• Generalización: La capacidad de manejar un número arbitrario de direcciones de difusión elimina la necesidad de estandarizar protocolos de adquisición, un problema común en la investigación clínica multicéntrica.
• Biomarcadores Potenciales: Las representaciones latentes y los mapas de atención podrían revelar nuevos biomarcadores microestructurales vinculados al deterioro cognitivo y la integridad de la sustancia blanca, sin depender de modelos biofísicos explícitos.

En resumen, el artículo presenta una solución técnica robusta que integra la geometría intrínseca de los datos de difusión en la arquitectura de Transformers, logrando representaciones superiores y transferibles para aplicaciones clínicas y de investigación en neurociencia.