From Convolution to Transformer: A Comparative Study of U-Net Variants for Brain Tumor and Retinal Vessel Segmentation

Resumo Técnico: Da Convolução ao Transformer

Definição do Problema
A segmentação de imagens médicas é fundamental para o diagnóstico auxiliado por computador, planejamento de tratamento e monitoramento de doenças. Embora a arquitetura U-Net tenha se tornado o padrão para segmentação biomédica devido à sua estrutura encoder-decoder e conexões de salto (skip connections), as variantes convencionais baseadas em convolução enfrentam limitações. Especificamente, as U-Nets padrão dependem de operações convolucionais locais, o que restringe sua capacidade de capturar dependências de longo alcance e informações contextuais globais. Essa limitação é particularmente prejudicial em tarefas complexas, como a segmentação de tumores cerebrais, onde as regiões exibem formas irregulares, aparências heterogêneas e bordas ambíguas. Além disso, o desempenho de várias modificações da U-Net (residuais, baseadas em atenção ou baseadas em transformer) pode variar significamente dependendo da modalidade de imagem (RM volumétrica vs. fundo de olho 2D) e da estrutura alvo (subregiões tumorais vs. vasos retinianos finos).

Metodologia
Este estudo conduz uma avaliação comparativa de cinco arquiteturas representativas baseadas em U-Net através de dois conjuntos de dados de referência distintos:

1. BraTS 2023: Um conjunto de dados de RM multimodal e volumétrica para segmentação de tumor cerebral (envolvendo o tumor total, tumor realçante e subregiões do núcleo do tumor).
2. DRIVE: Um conjunto de dados de imagens de fundo de olho 2D para segmentação de vasos retinianos, caracterizado por estruturas finas, ramificadas e de baixo contraste.

Os modelos avaliados representam um espectro de filosofias de design:

• 3D U-Net: Estende a arquitetura 2D original para dados volumétricos usando convoluções 3D para capturar a continuidade anatômica entre os cortes.
• Residual U-Net: Incorpora blocos residuais para melhorar a propagação do gradiente e a estabilidade do aprendizado de características, auxiliando na extração de estruturas complexas e heterogêneas.
• Attention U-Net: Introduz portões de atenção (attention gates) nas conexões de salto para suprimir informações irrelevantes de fundo e enfatizar regiões relevantes para a tarefa.
• UNETR: Substitui o encoder convolucional por um Vision Transformer (ViT) para modelar dependências globais via mecanismos de autoatenção em patches de imagem.
• Swin UNETR: Integra blocos Swin Transformer na estrutura U-Net, utilizando aprendizado de representação hierárquica e autoatenção de janela deslocada (shifted window self-attention) para capturar eficientemente tanto detalhes locais quanto o contexto global.

Todos os modelos foram treinados com consistência no pré-processamento (normalização, preenchimento/recorte) e estratégias de aumento de dados. O desempenho foi avaliado utilizando métricas quantitativas, incluindo escores Dice, Precisão, Recall, escores F1 e convergência de perda de treinamento.

Principais Contribuições
O artigo apresenta três contribuições primárias:

1. Avaliação Comparativa: Uma comparação sistemática de cinco variantes de U-Net (convolucionais, residuais, baseadas em atenção e baseadas em transformer) através de tarefas de RM cerebral volumétrica e de imagem retinal 2D.
2. Análise de Mecanismo: Uma análise de como mecanismos arquiteturais específicos — aprendizado residual, portões de atenção e autoatenção de transformer — impactam o desempenho da segmentação em diferentes cenários de imagem médica.
3. Orientação Prática: Insights práticos sobre a seleção de modelos, distinguindo entre arquiteturas mais adequadas para segmentação volumétrica complexa versus aquelas eficazes para a extração de estruturas finas.

Resultados Experimentais

• BraTS 2023 (Tumor Cerebral): O Swin UNETR alcançou o melhor desempenho geral com um escore Dice médio de 0,8965 (WT: 0,8966, ET: 0,9088, TC: 0,8841). O UNETR seguiu com um Dice médio de 0,8540. A Residual U-Net (0,7341), 3D U-Net (0,6890) e Attention U-Net (0,6522) apresentaram desempenho inferior. Os resultados indicam que a modelagem de contexto global via transformer é altamente benéfica para a segmentação volumétrica de tumores.
• DRIVE (Vasos Retinianos): O Swin UNETR novamente alcançou o maior escore Dice (0,8078), seguido de perto pela Residual U-Net (0,8069) e pelo UNETR (0,8015). A Attention U-Net obteve 0,7855, enquanto a 3D U-Net alcançou o menor Dice (0,5967), apesar de possuir a menor perda de treinamento (0,1030). Essa discrepância destaca que uma menor perda de treinamento não garante melhor qualidade de segmentação, particularmente quando a arquitetura é incompatível com a estrutura dos dados (ex: modelos 3D em estruturas 2D finas).

Significância e Alegações
O artigo afirma que o Swin UNETR é a arquitetura mais eficaz no geral, equilibrando com sucesso a preservação de detalhes espaciais locais com a modelagem de contexto de longo alcance através de seu encoder hierárquico e atenção de janela deslocada. O estudo sugere que:

• Variantes baseadas em Transformer (especificamente o Swin UNETR) são superiores para tarefas de segmentação complexas que exigem compreensão de contexto global, como a delineação de tumores cerebrais.
• O aprendizado residual permanece altamente eficaz para tarefas que exigem a preservação de detalhes estruturais finos, como a segmentação de vasos retinianos, onde a Residual U-Net apresentou um desempenho competitivo com os modelos de transformer.
• A seleção de modelos deve ser dependente da tarefa: abordagens baseadas em transformer são recomendadas para tarefas volumétricas e complexas, enquanto abordagens residuais ou híbridas podem ser suficientes para a extração de estruturas finas.

Os autores concluem que trabalhos futuros devem explorar arquiteturas híbridas que combinem a modelagem de contexto baseada em transformer com módulos convolucionais ou residuais para aumentar ainda mais a robustez e a generalização em diversos conjuntos de dados clínicos.