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

1. O Problema

A Imagem de Ressonância Magnética de Difusão (dMRI) é crucial para estudar mudanças microestruturais no cérebro e é amplamente utilizada na prática clínica. No entanto, o progresso na aprendizagem de representações de propósito geral a partir de dados de dMRI tem sido limitado. Os desafios principais incluem:

• Natureza Complexa dos Dados: Os sinais de difusão possuem propriedades únicas que dependem da direção e da força dos gradientes de sensibilidade à difusão (vetores b e valores b). Modelos tradicionais não capturam bem as dependências espaciais e direcionais conjuntas.
• Variabilidade de Protocolos: Diferentes estudos utilizam protocolos de aquisição variados (números diferentes de direções e valores b), o que dificulta o uso de modelos padronizados que exigem inputs fixos.
• Falta de Generalização: Abordagens existentes focam em tarefas específicas ou em reconstrução de alta resolução, mas carecem de modelos capazes de extrair representações transferíveis e robustas entre diferentes configurações de aquisição.

2. Metodologia

Os autores propõem uma arquitetura baseada em Transformers treinada com o objetivo de Autoencoder Máscara (MAE - Masked Autoencoder). A inovação central reside na adaptação da arquitetura para lidar com a geometria do espaço de difusão.

• Arquitetura do Modelo:

  • O input consiste em vários volumes cerebrais 3D (um para cada direção de difusão e valor b).
  • Os volumes são divididos em "patches" 3D e linearizados em tokens.
  • O modelo utiliza um encoder modificado de Vision Transformer (ViT) que alterna a atenção entre as dimensões espaciais e as dimensões de difusão.
  • Posicionamento Absoluto: Combina codificação posicional senoidal 3D para as coordenadas espaciais $(x, y, z)$ e uma camada aprendível para as coordenadas esféricas do espaço de difusão $(\rho, \theta, \phi)$.

• D-RoPE (Diffusion Space Rotary Positional Embedding):

  • Esta é a contribuição metodológica principal. O RoPE (Rotary Positional Embedding) original foi projetado para sequências 1D lineares. Como os dados de dMRI residem em um espaço esférico (direções de gradiente) e linear (valores b), o RoPE padrão é inadequado.
  • Os autores generalizam o RoPE substituindo a distância linear $(m-n)$ por uma distância geométrica no espaço de difusão $D(m, n)$.
  • A fórmula da distância considera uma combinação ponderada da diferença nos valores $b$ e o ângulo entre os vetores $b$ (direções):
    $$D(m, n) = \sqrt{\gamma(b_m - b_n)^2 + \arccos^2(|v_m \cdot v_n|)}$$
  • Isso permite que o mecanismo de atenção capture a interdependência relativa entre diferentes direções de difusão e valores b, independentemente da ordem de aquisição.

• Pré-treinamento (MAE):

  • O modelo é pré-treinado com uma estratégia de mascaramento que inclui: mascaramento espacial (75%), mascaramento de direção de difusão (50%) ou uma alternância entre ambos.
  • O objetivo é reconstruir os patches mascarados a partir dos patches visíveis, forçando o modelo a aprender tanto a estrutura espacial quanto as relações no espaço de difusão.

3. Contribuições Principais

1. Nova Arquitetura para dMRI: Um modelo Transformer capaz de processar aquisições de dMRI arbitrárias (número variável de direções e valores b) sem necessidade de reamostragem para um protocolo fixo.
2. D-RoPE: Uma generalização do Rotary Positional Embedding que modela a geometria esférica e linear do espaço de difusão, permitindo a captura de dependências direcionais complexas.
3. Avaliação Abrangente: Demonstração de que as representações aprendidas são superiores ou competitivas em diversas tarefas downstream, superando baselines totalmente treinados e métodos baseados em características manuais (handcrafted), mesmo com dados de pré-treinamento não supervisionados.

4. Resultados

O modelo foi pré-treinado no conjunto de dados HCP-YA e avaliado em tarefas downstream usando HCP-D, HCP-A e ADNI.

• Reconstrução (MAE):

  • A combinação de D-RoPE com mascaramento alternado (espacial e de difusão) obteve os melhores resultados quantitativos (maior PSNR, SSIM e menor FID).
  • As reconstruções preservaram com precisão métricas de microestrutura derivadas, como Anisotropia Fracional (FA) e Difusividade Média (MD), indicando que o espaço latente codifica informações biologicamente plausíveis.

• Tarefas Downstream:

  • Previsão de Idade e Classificação de Sexo: O modelo pré-treinado (com fine-tuning parcial ou linear probing) alcançou desempenho competitivo, superando ou igualando modelos totalmente supervisionados (como ResNet3D e ViT padrão).
  • Classificação de Comprometimento Cognitivo Leve (MCI) vs. Normal (CN): O uso de características latentes congeladas com um cabeçalho MLP resultou na maior precisão balanceada (64,7%), superando significativamente características manuais e modelos treinados do zero.
  • Predição de Escores Cognitivos (ADAS-Cog): O modelo alcançou um coeficiente de correlação de 0,38, superando todas as outras abordagens (incluindo ViT com RoPE padrão em dados brutos).

• Robustez: O modelo demonstrou capacidade de lidar com protocolos de aquisição variados, sendo capaz de processar subconjuntos aleatórios de direções de difusão sem perda significativa de desempenho.

5. Significado e Conclusão

Este trabalho representa um avanço significativo na representação de dados de dMRI.

• Modelo de Fundação (Foundation Model): A abordagem sugere o caminho para a criação de um modelo de base para dMRI, capaz de harmonizar análises entre coortes com protocolos de aquisição heterogêneos.
• Eficiência de Dados: A capacidade de obter representações robustas com pré-treinamento auto-supervisionado e fine-tuning limitado é crucial para aplicações clínicas onde os dados rotulados são escassos.
• Biomarcadores Potenciais: As representações aprendidas podem revelar novos biomarcadores microestruturais que ligam a integridade da substância branca a trajectórias cognitivas e doenças neurodegenerativas, sem a necessidade de ajustar modelos biofísicos complexos.
• Limitação Computacional: Os autores notam que o D-RoPE tem um custo computacional ligeiramente superior ($O(lh^2)$) comparado ao RoPE padrão ($O(lh)$) devido à não fatorabilidade da matriz de distância, mas o ganho em desempenho justifica o custo.

Em resumo, o D-RoPE permite que os Transformers capturem a geometria intrínseca dos dados de difusão, resultando em representações mais ricas e transferíveis para diagnóstico e pesquisa neurológica.