Does Your ViT Still Need U-Net for Segmentation?

Resumo Técnico: O seu ViT ainda precisa de U-Net para segmentação?

Declaração do Problema
A segmentação de imagens médicas tem sido dominada por mais de uma década por arquiteturas do tipo encoder–decoder U-Net. Embora esses designs abordem com sucesso a perda de detalhes espaciais por meio de conexões de salto (skip connections) e decodificação simétrica, eles dependem de operações convolucionais que são limitadas por campos receptivos restritos, tornando difícil a modelagem de contexto global. Embora os Vision Transformers (ViTs) ofereçam mecanismos de autoatenção para capturar dependências de longo alcance, os métodos iniciais baseados em ViT (ex: TransUNet, SwinUNet) mantiveram decoders de U-Net porque as representações pré-treinadas de ViT eram consideradas insuficientes para uma predição densa precisa.

Avanços recentes em pré-treinamento em larga escala (ex: DINOv2, SigLIP) redefiniram significativamente as capacidades de representação dos ViTs, permitindo que aprendam características visuais poderosas a partir de massivos conjuntos de dados não rotulados. Essa evolução levanta uma questão crítica: o paradigma da U-Net ainda é necessário para a segmentação de imagens médicas quando backbones de ViT modernos e fortemente pré-treinados estão disponíveis? Se não, como um framework apenas de encoder (encoder-only) deve ser projetado para prever máscaras de segmentação de forma eficaz sem um decoder?

Metodologia: EoSeg
Os autores propõem o EoSeg (Encoder-only Segmentation), um framework que elimina o decoder da U-Net em favor de um design baseado em consultas (query-based) construído sobre backbones de ViT pré-treinados modernos. A metodologia é desenvolvida através de uma exploração sistemática de escolhas arquiteturais fundamentais:

1. Seleção de Backbone: O estudo avalia vários backbones de ViT pré-treinados (incluindo ViT-L/14, SigLIP, MTP, DINOv2 e DINOv3) em um cenário de apenas encoder. O DINOv2 é selecionado como o backbone ideal, alcançando o maior desempenho de baseline (81,86% mDice no Synapse) sem qualquer decoder.
2. Geração de Máscara:
   • Deconvolução vs. Baseado em Consulta: Inicialmente, um módulo de deconvolução foi testado para fazer o upsampling de características de baixa resolução do ViT. No entanto, os autores descobriram que uma abordagem de predição de máscara baseada em consulta era superior. Em vez de decodificação pixel a pixel, embeddings de consulta aprendíveis atuam como sondas semânticas globais.
   • Mecanismo: Vetores de consulta são projetados em um espaço de embedding de máscara e interagem com o mapa de características globais com upsampling via um produto escalar canal a canal para gerar máscaras binárias diretamente.
3. Modelagem de Consulta Multinível: Inspirado pelas conexões de salto em U-Nets, o framework aproveita as representações dos últimos três blocos Transformer em vez de apenas o bloco final. Consultas aprendíveis são injetadas antes desses blocos e refinadas através de autoatenção com os tokens da imagem, permitindo que o modelo agregue informação semântica de múltiplos estágios.
4. Fusão de Blocos Aprendível: Para combinar as predições dos três blocos intermediários, um módulo de fusão aprendível é introduzido. Um portão de fusão (fusion gate) utiliza a representação da consulta para prever pesos adaptativos para a saída de cada bloco, permitindo que o modelo determine dinamicamente a contribuição de cada nível para a predição final.
5. Configuração de Consulta: O framework adota um paradigma de classificação baseado em consulta onde cada consulta prevê uma máscara e um rótulo de classe. Experimentos determinaram que atribuir duas consultas por categoria produz o desempenho ideal, equilibrando a necessidade de capturar variações intra-classe sem superparametrização.
6. Estratégia de Supervisão: Contrariamente às práticas comuns em segmentação médica, a supervisão profunda (deep supervision — aplicando perda às saídas intermediárias) descobriu-se que não melhora o desempenho neste setup específico de apenas encoder e foi excluída do design final.

Principais Contribuições

• Reavaliação Arquitetural: O artigo fornece uma avaliação extensiva demonstrando que um decoder do tipo U-Net não é mais um pré-requisito para segmentação de imagens médicas de alta qualidade quando se utiliza backbones de ViT pré-treinados em larga escala e modernos.
• Framework EoSeg: Os autores apresentam o EoSeg, um framework de apenas encoder baseado em consultas que integra modelagem de consulta multinível e fusão de blocos aprendível para permitir a predição direta de máscaras a partir de características do encoder.
• Validação Abrangente: O design proposto é validado em sete datasets de referência abrangendo cinco modalidades distintas de imagens médicas: CT (Synapse), MRI (ACDC), Histopatologia (GlaS, MoNuSeg), Endoscopia (Kvasir-Seg) e Dermoscopia (ISIC-2016, ISIC-2017).

Resultados Experimentais
O EoSeg supera consistentemente métodos baseados em CNN, Transformer e híbridos de encoder–decoder em todos os métodos testados:

• Synapse (CT): Alcançou 85,50% de Dice médio (mDice), superando o melhor anterior (EViT-UNet com 80,87%) em 4,63 pontos. Também alcançou pontuações máximas para órgãos específicos como a Aorta (90,80%) e o Baço (94,13%).
• ACDC (MRI): Alcançou 91,73% de mDice, excedendo o melhor anterior em 0,24 pontos.
• GlaS (Histopatologia): Alcançou 93,27% de mDice e 87,79% de mIoU, superando o melhor anterior em 0,83 e 1,29 pontos, respectivamente.
• MoNuSeg (Histopatologia): Alcançou 80,51% de mDice e 67,45% de mIoU.
• Kvasir-Seg & ISIC: Demonstraram desempenho superior na segmentação de pólipos e lesões, alcançando altos scores de mIoU e mDice com bordas suaves e estruturas coerentes.

A análise qualitativa confirma que o EoSeg produz bordas mais suaves e menos artefatos em comparação com métodos concorrentes, preservando efetivamente detalhes estruturais finos sem a restauração de características baseada em decoder.

Significância e Alegações
O artigo afirma que a dominância do paradigma U-Net na segmentação de imagens médicas não é mais justificada pelas limitações dos backbones de ViT iniciais. Em vez disso, as fortes representações aprendidas através de pré-treinamento em larga escala permitem designs de apenas encoder que podem prever diretamente máscaras de segmentação precisas.

Os autores enfatem que o sucesso do EoSeg não se deve meramente à remoção do decoder, mas sim à adoção de um paradigma de predição baseado em consulta apropriado combinado com modelagem multinível e fusão adaptativa. Essas descobertas sugerem que futuros frameworks de segmentação de imagens médicas construídos sobre modelos de visão pré-treinados em larga escala podem se afastar de arquiteturas de decoder complexas, confiando, em vez disso, no poder de representação intrínseco dos ViTs modernos.