Resumo Técnico: BrainG3N

Declaração do Problema

A ressonância magnética (RM) cerebral tridimensional (3D) é crítica para a neurologia e neuro-oncologia clínica, mas modelos generativos capazes de operar em resolução 3D total são computacionalmente proibitivos no espaço de voxels. Consequentemente, o campo convergiu para a difusão latente, que depende de um tokenizador codificador-decodificador para comprimir volumes em um espaço latente de baixa dimensão.

Existe uma tensão fundamental nas pipelines existentes de difusão latente radiográfica 3D:

1. Utilidade Clínica: Os embeddings do codificador devem reter informações clínicas específicas necessárias para tarefas downstream (ex: classificação de doenças, predição de biomarcadores).
2. Fidelidade Anatômica: O decodificador deve reconstruir volumes de voxels anatomicamente precisos.

As abordagens atuais tipicamente treinam um par codificador-decodificador único contra um objetivo de reconstrução. Isso enviesa o codificador para a fidelidade de voxel local, muitas vezes às custas do conteúdo clínico global necessário para tarefas downstream. Além disso, os espaços latentes existentes raramente são avaliados em sua capacidade de suportar tanto a geração controlável quanto a predição clínica simultaneamente.

Metodologia

Os autores propõem o BrainG3N, um tokenizador de propósito duplo e totalmente volumétrico que desacopla o codificador e o decodificador para satisfazer ambas as demandas conflitantes.

1. Corpus de Pré-treinamento

O modelo é pré-treinado em 35.309 volumes de RM cerebral de 17.391 sujeitos únicos de 18 coortes públicas e mais de 200 locais de aquisição. Os dados abrangem quatro modalidades (T1, T2, FLAIR, T1c), dez categorias de doenças (incluindo condições neurodegenerativas e tumorais) e idades variando de 5 a 98 anos. Todos os volumes são harmonizados via registro afim ANTs para o atlas SRI24, com remoção de crânio (skull-stripping) e correção de viés, resultando em volumes isotrópicos de 160×192×160 de 1mm.

2. Arquitetura de Tokenizador de Duas Fases

A inovação central é um codificador de Autoencoder Mascarado (MAE) congelado pareado com um decodificador CNN dedicado via uma projeção linear.

• Fase 1 (Pré-treinamento): Um codificador MAE 3D (Vision Transformer de 12 camadas) é treinado para reconstruir 70% dos patches mascarados aleatoriamente a partir dos 30% de patches visíveis restantes. Isso força o codificador a modelar o contexto anatômico de longo alcance em vez de estatísticas de voxel locais.
• Fase 2 (Tokenização): O codificador MAE é congelado. Uma projeção linear ($P \in \mathbb{R}^{1152 \times 32}$) comprime os 1200 tokens por volume de 1152 dimensões para um gargalo de 32 canais. Um decodificador CNN 3D então reconstrói os voxels a partir desses tokens projetados sob uma perda $\ell_1$.
  • Racional: Treinar conjuntamente o codificador e o decodificador (como em VAEs) desvia o codificador para a fidelidade de intensidade, enquanto o congelamento do MAE preserva seus embeddings clinicamente informativos, e o gargalo reduz a dimensionalidade de entrada do modelo de difusão.

3. Difusão Latente Condicional

Um Transformer de Difusão de Correspondência de Fluxo (Flow-Matching Diffusion Transformer - DiT) é treinado nas sequências de tokens projetadas ($z' \in \mathbb{R}^{1200 \times 32}$).

• Condicionamento: O DiT aceita seis variáveis de condicionamento: doença (8 vias), sexo, modalidade (4 vias), local de aquisição (19 vias), idade (contínua) e status de mutação IDH1 (binário).
• Mecanismo: As condições são roteadas através de modulação adaLN-Zero. A Orientação Livre de Classificador (Classifier-Free Guidance - CFG) é implementada descartando aleatoriamente as condições (probabilidade $p=0.1$) durante o treinamento, exceto a modalidade, que é sempre especificada.
• Variante Longitudinal: Uma versão especializada substitui o interpolante ruído-para-dados padrão por uma ponte latente de linha de base-para-seguimento, permitindo a previsão específica do paciente ao longo de horizontes de tempo solicitados ($\Delta t$).

Principais Contribuições

1. Espaço de Embedding de Propósito Dual: Os autores demonstram que um único codificador MAE congelado pode servir como um tokenizador tanto para tarefas clínicas downstream quanto para geração controlável.
2. Desacoplamento Arquitetural: Ao congelar o codificador MAE e utilizar uma projeção linear + decodificador CNN, o método evita o compromisso entre conteúdo clínico e fidelidade de reconstrução inerente ao treinamento conjunto.
3. Pré-treinamento Escalável: O modelo é pré-treinado em um corpus massivo, multimodal e multissite (mais de 35 mil volumes), significativamente maior do que os conjuntos de dados generativos 3D típicos.

Resultados Experimentais

1. Representação Clínica (Linear Probing)

O codificador congelado foi avaliado em um benchmark de teste linear de 23 tarefas (sem ajuste fino do codificador).

• Desempenho: O modelo superou ou igualou três modelos de fundação de estado da arte (SOTA) (BrainIAC, BrainSegFounder, MedicalNet) em 21 de -23 tarefas.
• Métricas Chave:
  • Status de Mutação IDH1: AUC 0.937 (vs. 0.805 para BrainIAC).
  • Grau do Tumor: AUC 0.921.
  • Regressão de Idade Cerebral: Erro Absoluto Médio (MAE) de 4.43 anos.
  • Predição de Sexo: AUC 0.967.
• Comparação: A abordagem de MAE congelado superou um modelo baseado em fatias 2D (BrainDINO) em tarefas volumétricas, posicionando a abordagem de auto-supervisão totalmente volumétrica como superior para dados 3D.

2. Geração Controlável

Um DiT treinado nesses embeddings foi avaliado por sua capacidade de gerar amostras condicionadas a variáveis específicas.

• Controlabilidade: Sondas de dados reais aplicadas às amostras geradas confirmaram que o modelo recupera com sucesso os atributos solicitados (ex: estado de doença, sexo, idade) sob Orientação Livre de Classificador (CFG).
  • Exemplo: O condicionamento de idade mostrou uma correlação de Pearson de $r=0.93$ entre a idade solicitada e a gerada.
  • Contrafatuais: Manter as sementes de ruído constantes enquanto variava as condições produziu mudanças visíveis e anatomicamente consistentes (ex: aumento ventricular em casos de Alzheimer) enquanto preservava a identidade.
• Fidelidade: O 3D-FID agrupado entre as amostras geradas e os volumes reais foi de 107.3, competitivo com o trabalho generativo médico 3D anterior. A lacuna entre os gerados e os reais foi atribuída principalmente à compressão de detalhes de alta frequência do scanner pelo tokenizador, em vez de falha do gerador.

3. Previsão Longitudinal

O modelo foi testado em extrapolação temporal intra-paciente usando dados do ADNI.

• Configuração: Dado um latente de linha de base e um horizonte de tempo ($\Delta t$), o modelo previu um exame de acompanhamento.
• Resultados: O modelo recuperou aproximadamente 27% da magnitude real do envelhecimento (inclinação 0.268) com uma correlação de Pearson de $r=0.716$ entre o $\Delta t$ previsto e o solicitado.
• Localização: As mudanças anatômicas foram corretamente localizadas nos ventrículos e sulcos, consistentes com padrões conhecidos de neurodegeneração, embora a magnitude tenha sido atenuada devido ao regularizador de ponte estocástica que interpola em direção à média condicional.

Significância e Alegações

O artigo estabelece que um único espaço de embedding de RM cerebral 3D pode suportar simultaneamente:

1. Alta performance em tarefas clínicas downstream (superando modelos de fundação existentes em teste linear).
2. Modelagem generativa controlável (suportando geração condicional através de seis variáveis e previsão longitudinal específica do paciente).

Os autores argumentam que esta abordagem valida o uso de codificadores MAE congelados como tokenizadores para difusão radiográfica 3D, superando as limitações de tokenizadores impulsionados pela reconstrução que sacrificam o conteúdo clínico pela fidelidade de voxel. O trabalho permite aplicações como a geração de dados sintéticos para coortes sub-representadas, simulação de trajetórias de doenças contrafatuais e o fornecimento de leituras baseadas apenas em imagem para biomarcadores clínicos (ex: status de IDH1, idade cerebral) sem a necessidade de procedimentos invasivos.

Limitações observadas pelos autores:

• As descobertas são específicas para RM cerebral; a transferência para outras modalidades (CT, RM de corpo inteiro) não foi testada.
• A previsão longitudinal recupera a direção e a localização do envelhecimento, mas apenas ~27% da magnitude.
• A comparação com outros modelos de fundação envolve potenciais discrepâncias de distribuição devido aos diferentes corpora de pré-treinamento.
• Confundidores de aquisição (fabricante do scanner, força de campo) são detectáveis nos embeddings, necessitando de avaliação pareada por coorte para aplicações downstream.