X-WIN: Building Chest Radiograph World Model via Predictive Sensing

Título: X-WIN: Construindo um Modelo de Mundo para Radiografias Torácicas via Sensoriamento Preditivo

1. O Problema

A radiografia de tórax (CXR) é a técnica de imagem mais utilizada para diagnóstico de doenças torácicas. No entanto, como as CXRs são imagens 2D projetadas, elas sofrem de sobreposição estrutural, o que impede a captura da anatomia 3D completa dos órgãos. Em contraste, a Tomografia Computadorizada (CT) fornece estruturas internas 3D detalhadas, mas é significativamente mais cara, irradia mais o paciente e tem menor acessibilidade, especialmente em regiões subdesenvolvidas.

Existe uma lacuna de capacidade entre os modelos de IA atuais para CXR e os radiologistas humanos: enquanto os médicos podem cognitivamente reconstruir um modelo 3D interno a partir de vistas frontal e lateral, os modelos de aprendizado de máquina existentes (incluindo modelos de "mundo" anteriores como o CheXWorld) são limitados a representações 2D e não incorporam conhecimento volumétrico 3D.

2. Metodologia: X-WIN

Os autores propõem o X-WIN (X-ray World Intelligence Network), um novo modelo de mundo que destila conhecimento volumétrico de CTs para melhorar a representação de CXRs. A ideia central é treinar um modelo para prever novas projeções 2D de um volume 3D sob diferentes transformações no espaço 3D, utilizando o conceito de sensoriamento preditivo.

Arquitetura e Fluxo:

• Rede Dual: O sistema consiste em dois componentes principais:
  1. Um codificador ($f_\theta$) que recebe uma projeção de raio-X rotineira (contexto) e um preditor de visão ($g_v$) condicionado a uma "ação" (rotação da fonte de raio-X).
  2. Um codificador alvo ($f_{\theta'}$), atualizado via Exponential Moving Average (EMA), que recebe múltiplas projeções do mesmo volume CT para supervisionar a previsão.
• Design de Ação: A "ação" é definida como a rotação da fonte de raio-X (e.g., guinada/yaw) para gerar novas projeções a partir do volume CT. O modelo deve prever a representação latente dessa nova projeção com base na projeção original e no ângulo de rotação.

Funções de Perda (Loss Functions):
Para treinar o modelo, são utilizadas três estratégias principais de perda:

1. Alinhamento Contrastivo Guiado por Afinidade (Affinity-guided Contrastive Alignment):
   • Utiliza uma perda contrastiva (InfoNCE) para alinhar as representações previstas com as reais.
   • Introduz um regularizador baseado em afinidade que relaxa o alinhamento rígido. Reconhece que projeções do mesmo volume CT contêm informações correlacionadas ricas, permitindo que o modelo aprenda similaridades sutis entre pares negativos, em vez de apenas repeli-los.
2. Modelagem de Imagem Mascarada (MIM - Masked Image Modeling):
   • Aplicada tanto em CXRs reais quanto simuladas.
   • O modelo deve reconstruir patches mascarados, forçando-o a aprender características locais e contextuais finas de estruturas anatômicas e anomalias.
3. Adaptação de Domínio Preservadora de Estrutura:
   • Utiliza um classificador de domínio para reduzir a lacuna entre o domínio simulado (projeções de CT) e o domínio real (CXRs reais).
   • Incentiva que as representações do domínio simulado sejam estatisticamente similares às do domínio real, enquanto preserva a informação estrutural necessária para a tarefa de previsão.

3. Contribuições Principais

1. Primeira Integração 3D em Modelo de Mundo para CXR: O X-WIN é a primeira abordagem a integrar conhecimento espacial 3D em um modelo de mundo para radiografias, destilando conhecimento volumétrico de CTs.
2. Novo Objetivo de Alinhamento: Introdução de uma perda de alinhamento contrastivo guiada por afinidade, que melhora a codificação de características discriminativas e explora correspondências ricas entre diferentes projeções.
3. Desempenho SOTA e Reconstrução 3D: O modelo alcança desempenho state-of-the-art em tarefas de interpretação de CXR (usando linear probing e few-shot fine-tuning) e demonstra a capacidade de renderizar projeções 2D para reconstruir volumes 3D de CT, provando a compreensão interna da estrutura anatômica.

4. Resultados Experimentais

O modelo foi treinado com 371.951 CXRs (MIMIC-CXR) e 32.371 volumes CT (NLST).

• Avaliação por Linear Probing: O X-WIN superou modelos fundamentais existentes (como CheXFound, Ark+, CheXWorld) e modelos de visão-linguagem em seis benchmarks padrão (VinDr, CheXpert, NIH-CXR, RSNA, JSRT, COVIDx).
  • Exemplo: No VinDr, o X-WIN (ViT-Large) alcançou 0.925 de AUROC, superando o CheXWorld (0.903) e o Ark+ (0.906).
• Avaliação por Few-Shot Fine-Tuning: Em tarefas de detecção de pneumonia/COVID-19 com poucos dados (4, 8, 16 amostras por classe), o X-WIN demonstrou alta adaptabilidade, alcançando 0.993 de AUROC com 16 shots, superando todos os concorrentes.
• Reconstrução 3D: O modelo foi capaz de reconstruir volumes 3D de CT a partir das projeções renderidas, atingindo um PSNR de 27.87 dB e SSIM de 0.789, preservando a estrutura global e detalhes locais razoáveis.
• Análise de Ablação:
  • A combinação de perdas contrastivas e MIM foi crucial.
  • A rotação direta (yaw) funcionou melhor que rotações passo-a-passo ou rotações 3D completas (yaw, pitch, roll).
  • Um tamanho de passo de rotação de 3° foi otimizado para equilibrar densidade de amostragem e diversidade.

5. Significado e Conclusão

O trabalho X-WIN representa um avanço significativo na interseção entre modelagem de mundo e imagem médica. Ao ensinar o modelo a "imaginar" como uma radiografia mudaria se a fonte de raio-X se movesse no espaço 3D, o sistema internaliza uma compreensão da anatomia tridimensional, superando as limitações da projeção 2D.

Isso permite:

• Melhor diagnóstico de doenças em CXRs, mesmo com oclusão de estruturas.
• Redução da dependência de CTs caros e irradiantes para tarefas que podem ser inferidas a partir de CXRs enriquecidas com conhecimento 3D.
• Um caminho para sistemas de IA mais explicáveis e robustos na análise radiográfica, aproximando a capacidade da IA da cognição espacial dos radiologistas humanos.

O código e os dados estão disponíveis publicamente, facilitando a reprodução e o avanço futuro na área.