BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

Resumo Técnico de "BrainAnytime: Pré-treinamento Cross-Modal Consciente da Anatomia para Análise de Imagens Cerebrais com Disponibilidade Arbitrária de Modalidades"

Declaração do Problema

Os fluxos de trabalho de diagnóstico clínico para doenças neurodegenerativas, particularmente a Doença de Alzheimer (DA), seguem um "caminho de escalonamento de modalidades". Os clínicos geralmente iniciam com imagens estruturais de rotina (por exemplo, ressonância magnética ponderada em T1), adicionam seletivamente sequências (por exemplo, FLAIR, T2) para refinar diagnósticos diferenciais e reservam imagens moleculares (por exemplo, PET de amiloide) para casos incertos. Consequentemente, os dados dos pacientes em coortes do mundo real são caracterizados por subconjuntos de modalidades heterogêneos e frequentemente incompletos.

Os modelos de IA atuais geralmente assumem modalidades de entrada fixas, exigindo alternância arquitetônica ou re-treinamento quando a disponibilidade de dados muda. Além disso, os modelos existentes de pré-treinamento em larga escala para imagens cerebrais frequentemente sofrem de três limitações neste contexto:

1. Eles são tipicamente apenas de ressonância magnética (MRI), falhando em aprender pontes explícitas para o PET de amiloide, um biomarcador crítico.
2. Eles assumem combinações fixas de modalidades, não correspondendo à ausência dependente do estágio dos dados do mundo real.
3. O pré-treinamento padrão com Autoencoder Mascarado (MAE) baseia-se em mascaramento espacialmente uniforme, ignorando o fato de que a DA exibe vulnerabilidade regional seletiva (por exemplo, o hipocampo) em vez de patologia uniforme.

Metodologia

Os autores propõem o BrainAnytime, um framework unificado de pré-treinamento projetado para suportar a análise de imagens cerebrais sob disponibilidade arbitrária de modalidades. O framework é pré-treinado em 34.899 varreduras cerebrais 3D de cinco conjuntos de dados (A4, DOD-ADNI, AIBL, BraTS23, NACC) abrangendo MRI multi-sequência e PET de amiloide. Ele integra três componentes principais dentro de uma arquitetura compartilhada de autoencoder mascarado 3D:

1. Multi-MAE3D (Autoencoder Mascarado 3D Multi-Modal)

• Arquitetura: Um codificador Transformer compartilhado processa qualquer subconjunto de modalidades disponíveis (T1, T2, FLAIR, PET).
• Tokenização: Cada modalidade utiliza uma incorporação de patches convolucionais 3D específica da modalidade.
• Tratamento de Dados Faltantes: Modalidades ausentes são totalmente mascaradas. Para manter uma contagem constante de tokens independentemente da disponibilidade de entrada, o orçamento de tokens visíveis é redistribuído entre as modalidades observadas usando uma distribuição Dirichlet (Dir($\alpha$)).
• Decodificação: Cada modalidade possui um decodificador leve que reconstrói todos os patches. A perda de treinamento ($L_{MAE}$) é computada apenas nos patches mascarados usando erro quadrático médio com normalização por patch.

2. Destilação Cross-Modal Recíproca (RCMD)

• Objetivo: Aprender correspondências entre MRI estrutural e PET molecular.
• Mecanismo: Uma rede professor com Média Móvel Exponencial (EMA) compartilha a arquitetura do codificador. Para amostras pareadas MRI-PET, representações de nível de grupo ($z_{MRI}$ e $z_{PET}$) são computadas via pooling médio.
• Perda: Dois preditores MLP simétricos ($g_{M\to P}$ e $g_{P\to M}$) são treinados para prever características do professor cross-modais. A perda ($L_{RCMD}$) maximiza a similaridade de cosseno entre as características previstas pelo estudante e as características do professor no espaço da modalidade oposta.

3. Mascaramento de Currículo Consciente da Patologia (PACM)

• Objetivo: Priorizar a anatomia vulnerável à doença durante a reconstrução em vez de mascarar uniformemente.
• Pontuação de Importância: Uma pontuação estática é derivada de priores neuropatológicos (estadiamento de Braak) mapeados via atlas AAL116. Uma pontuação dinâmica é obtida agregando mapas de atenção de um professor de anatomia EMA. Estes são combinados para gerar pontuações de importância de patches.
• Cronograma de Currículo: Uma estratégia de três fases é empregada:
  1. Fase inicial: Mascaramento aleatório uniforme.
  2. Fase intermediária: O parâmetro de temperatura ($\tau$) é recozido via decaimento cosseno, afiando progressivamente a distribuição da máscara para favorecer regiões de alta importância para mascaramento (ou seja, forçando o modelo a reconstruir regiões críticas).
  3. Fase tardia: Temperatura fixa.
• Seleção: Patches são selecionados usando o truque Gumbel-top-k, garantindo que patches de alta importância tenham maior probabilidade de serem mascarados e, portanto, reconstruídos.

Treinamento e Ajuste Fino:
A perda total é $L = L_{MAE} + \lambda \cdot \mathbb{1}_{[paired]} \cdot L_{RCMD}$. Para tarefas a jusante, apenas o codificador estudante é usado com um cabeçalho de tarefa leve. Modalidades ausentes permanecem zeradas e com bloqueio de atenção durante a inferência, garantindo que o modelo lide com subconjuntos arbitrários sem alterações arquitetônicas.

Contribuições Principais

1. Framework Unificado: Um único modelo capaz de aceitar qualquer subconjunto de modalidades (de uma única varredura T1 a uma avaliação multimodal completa) sem re-treinamento ou alternância arquitetônica.
2. Aprendizado Cross-Modal: Aprende explicitamente correspondências estruturais-moleculares entre MRI e PET via destilação recíproca, fechando a lacuna entre evidências estruturais e patológicas.
3. Pré-treinamento Consciente da Anatomia: Introduz uma estratégia de mascaramento de currículo que direciona o esforço reconstrutivo do modelo para a neuroanatomia saliente à DA, abordando as limitações do mascaramento uniforme em MAEs padrão.
4. Pré-treinamento Escalável: Pré-treinado com sucesso em um conjunto de dados heterogêneo em larga escala (34.899 varreduras) imitando padrões de ausência clínica.

Resultados Experimentais

O modelo foi avaliado em quatro tarefas a jusante (CN vs. DA, CN vs. MCI, regressão MMSE, previsão de idade) em cinco configurações de modalidades motivadas clinicamente (variando de apenas T1 a T1+T2+FLAIR+PET).

• Desempenho: O BrainAnytime superou modelos específicos de modalidade (por exemplo, 3D ResNet, M3T), métodos de fusão de dupla modalidade, modelos fundacionais cerebrais em larga escala (BrainIAC, SAM-Brain3D) e abordagens de aprendizado multimodal incompleto (Flex-MoE, FuseMoE).
• Métricas Chave:
  • CN vs. DA: Alcançou uma precisão média de 89,4% (vs. 84,2% para a próxima melhor linha de base, Flex-MoE), representando uma melhoria relativa de 6,2% sobre a linha de base mais forte de modalidade ausente.
  • CN vs. MCI: Alcançou uma precisão média de 67,2% (vs. 60,6% para MoE-retriever), uma melhoria relativa de 7,0%.
  • Regressão: Demonstrou desempenho superior na previsão de MMSE e idade em todas as configurações de modalidades.
• Robustez: Um estudo de robustez intra-sujeito confirmou que o desempenho melhorou consistentemente à medida que mais modalidades foram adicionadas (por exemplo, F1 CN vs. DA aumentou de 65,0% com apenas T1 para 80,4% com entrada completa), demonstrando a capacidade do modelo de aproveitar efetivamente os dados disponíveis.
• Ablação: A remoção individual de RCMD ou PACM resultou em quedas de desempenho, confirmando a necessidade de ambas a destilação cross-modal e o mascaramento consciente da anatomia.

Significado

O artigo afirma que o BrainAnytime aborda uma lacuna crítica entre a realidade clínica (aquisição de dados escalonada e incompleta) e os paradigmas atuais de IA (modelos de entrada fixa). Ao permitir que um único modelo aceite flexivelmente qualquer imagem disponível, ele alinha-se com o caminho de escalonamento de modalidades clínico. O framework demonstra que o pré-treinamento com destilação cross-modal e mascaramento consciente da patologia produz um modelo fundacional robusto que supera significativamente as linhas de base existentes, particularmente em cenários onde os dados estão incompletos ou as modalidades estão ausentes. Esta abordagem oferece uma solução prática para implantar IA em ambientes clínicos do mundo real onde dados multimodais completos raramente estão disponíveis para cada paciente.