CoRe: Joint Optimization with Contrastive Learning for Medical Image Registration

1. Problema

O registro de imagens médicas é uma tarefa fundamental para estabelecer correspondências anatômicas e semânticas densas entre imagens (obtidas em diferentes momentos, sujeitos ou modalidades). O desafio principal reside em duas fontes de complexidade:

• Variabilidade na aparência das imagens: Diferenças em protocolos de aquisição, hardware e modalidades tornam comparações diretas de pixels (baseadas em intensidade) pouco confiáveis.
• Deformações anatômicas não lineares: Comuns em exames longitudinais ou entre pacientes, essas deformações podem ser causadas por patologias ou intervenções cirúrgicas, exigindo modelos que capturem variações espaciais complexas.

Métodos existentes muitas vezes falham em equilibrar a extração de características robustas com a precisão da estimativa do campo de deformação. Abordagens supervisionadas exigem dados de treinamento alinhados com "ground truth" (caros e difíceis de obter), enquanto métodos de aprendizado auto-supervisionado muitas vezes separam a fase de pré-treinamento de características da fase de registro, perdendo a sinergia entre os dois objetivos.

2. Metodologia: CoRe

O artigo propõe o CoRe (Contrastive learning for medical image Registration), uma nova abordagem que integra aprendizado contrastivo auto-supervisionado diretamente no pipeline de registro de imagens deformáveis através de um framework de otimização conjunta.

Arquitetura e Fluxo

O framework consiste em:

1. Extrator de Características ($G$): Uma rede convolucional que extrai representações das imagens fixa ($I_f$) e móvel ($I_m$).
2. Cabeça de Projeção ($P$): Projeta as características em embeddings.
3. Módulo de Otimização Diferenciável ($H$): Infere o campo de deslocamento ($u$) minimizando a dissimilaridade de características e a suavidade do campo.

Otimização Conjunta e Perdas

Diferente de métodos que pré-treiam o extrator separadamente, o CoRe otimiza simultaneamente o extrator de características sob duas funções de perda:

1. Perda de Registro ($L_{reg}$): Baseada em um esquema de auto-treinamento com pseudo-rótulos. O campo de deslocamento gerado pelo otimizador é refinado iterativamente e usado como "ground truth" para supervisionar a rede. A perda é o Erro Quadrático Médio (MSE) entre o campo previsto e o pseudo-rótulo.
2. Perda Contrastiva com Restrição de Equivariância ($L_c$):
   • Em vez de usar apenas invariância (ignorar transformações), o CoRe impõe equivariância. Isso garante que transformações geométricas no espaço de entrada resultem em mudanças previsíveis e estruturadas no espaço de embeddings.
   • Aplica-se transformações afins aleatórias ($T_g$) às imagens. O modelo é treinado para manter a consistência entre as características da imagem original transformada e as características da imagem transformada.
   • Utiliza-se a perda InfoNCE para maximizar a similaridade entre pares positivos (mesma localização anatômica sob diferentes transformações) e minimizar a similaridade com pares negativos.

A função de perda total é definida como:
$$L = L_{reg} + \alpha \cdot L_c$$
Onde $\alpha$ equilibra a contribuição de ambos os objetivos.

Pseudocódigo (Resumido)

O treinamento envolve:

• Aplicar transformações afins aleatórias às imagens de entrada.
• Extrair características das imagens originais e aumentadas.
• Calcular o campo de deslocamento e a perda de registro ($L_{reg}$).
• Calcular a perda contrastiva ($L_c$) entre as características das imagens aumentadas e as características transformadas das imagens originais.
• Atualizar os pesos do extrator e do projetor via retropropagação conjunta.

3. Contribuições Principais

1. Framework de Otimização Conjunta: Integração direta do aprendizado contrastivo auto-supervisionado no pipeline de registro, alinhando a aprendizagem de características com a tarefa de registro downstream.
2. Sinergia de Objetivos: Demonstração de que a otimização simultânea das perdas de contraste e registro supera o treinamento com qualquer um dos objetivos isoladamente.
3. Robustez Semântica: O método gera embeddings densos e semanticamente significativos que são robustos a variações de intensidade e deformações de tecidos, sem depender de dados de treinamento rotulados ou pares multimodais alinhados.
4. Validação Abrangente: Testado em cenários intra-paciente e inter-paciente, superando métodos convencionais, baseados em aprendizado e híbridos.

4. Resultados Experimentais

O método foi avaliado em dois conjuntos de dados públicos:

• AbdomenCT: Registro inter-paciente (30 exames de TC abdominal).
• RadChestCT: Registro intra-paciente (371 pares de exames de TC torácica longitudinal).

Métricas de Desempenho:

• Coeficiente de Similaridade de Dice (DSC): Medida de sobreposição de estruturas anatômicas.
• SDLogJ: Desvio padrão do logaritmo do determinante Jacobiano (medida de suavidade/fisicalidade do campo de deformação).
• Tempo de Inferência ($T_{inf}$).

Comparação:
O CoRe superou todos os métodos comparados (DEEDs, NiftyReg, VoxelMorph, LapIRN, SAMConvex, RegCyc).

• AbdomenCT: CoRe alcançou 52.1 ± 7.5 de DSC (vs. 49.2 do SAMConvex e 51.1 do RegCyc).
• RadChestCT: CoRe alcançou 89.4 ± 6.6 de DSC (vs. 88.5 do DEEDs e 88.1 do SAMConvex).
• A suavidade dos campos (SDLogJ) permaneceu comparável aos melhores métodos, indicando que a melhoria na precisão não comprometeu a fisicalidade da deformação.

Estudo de Ablação:

• A otimização conjunta superou o uso isolado de perda de registro ou perda contrastiva.
• Curiosamente, no contexto de dados de TC (monomodal), o uso de aumentações de intensidade não lineares (comuns em aprendizado contrastivo) não trouxe benefícios adicionais e até piorou o desempenho em comparação com o uso apenas da restrição de equivariância geométrica. Isso sugere que, para dados de TC padronizados, a robustez geométrica é mais crítica que a invariância de intensidade.

5. Significado e Conclusão

O trabalho CoRe representa um avanço significativo ao demonstrar que a integração de restrições de equivariância geométrica via aprendizado contrastivo diretamente no processo de otimização de registro melhora a qualidade das características aprendidas.

• Impacto Clínico: Oferece uma ferramenta mais precisa para rastrear a progressão de doenças e avaliar intervenções terapêuticas, especialmente em cenários complexos como o registro inter-paciente abdominal.
• Inovação Técnica: Elimina a necessidade de pré-treinamento separado ou dados de "ground truth" perfeitos, criando um ciclo virtuoso onde o aprendizado de representações e a estimativa de deformação se reforçam mutuamente.
• Eficiência: Mantém tempos de inferência rápidos (cerca de 1-2 segundos), tornando-o viável para aplicações clínicas.

Em suma, o CoRe estabelece um novo estado da arte ao provar que a aprendizagem de representações e a otimização de registro não devem ser etapas desconectadas, mas sim processos interdependentes e otimizados conjuntamente.