DAMamba-UNet3D: A Parameter-Efficient Mamba State Space U-Net with Dynamic Adaptive Scan for 3D Medical Image Segmentation

Resumo Técnico: DAMamba-UNet3D

Definição do Problema
A segmentação de imagens médicas 3D requer o equilíbrio entre a delineação precisa de contornos locais e um contexto espacial global coerente, tudo isso aderindo a orçamentos rigorosos de computação e parâmetros. As arquiteturas existentes enfrentam trade-offs distintos:

• U-Nets Convolucionais: Oferecem eficiência de parâmetros e mistura local $O(n)$, mas carecem de contexto global explícito, exigindo o escalonamento da profundidade para capturar dependências de longo alcance.
• Vision Transformers: Fornecem raciocínio global via autoatenção, mas incorrem em custos proibitivos de memória e computação $O(n^2)$ para grandes volumes 3D (por exemplo, um patch de $128^3$ contém $\approx 2,1$ milhões de tokens).
• Modelos de Espaço de Estados (SSMs/Mamba): Propagam informações globalmente em tempo e memória $O(n)$. No entanto, os segmentadores de SSM médicos existentes (ex: SegMamba) dependem de padrões de varredura fixos e grandes orçamentos de parâmetros ($\sim 70$ milhões de parâmetros), carecendo de adaptabilidade dependente dos dados.

Embora o Escaneamento Adaptativo Dinâmico (DAS) tenha sido introduzido para o reconhecimento de imagens naturais 2D para aprender o reordenamento dependente dos dados antes do escaneamento seletivo, ele não foi aplicado à imagem médica ou estendido para volumes 3D.

Metodologia
Os autores propõem o DAMamba-UNet3D, uma arquitetura híbrida de codificador-decodificador que integra blocos Tri-Plane 3D-DAS em uma estrutura U-Net.

1. Design da Arquitetura:

   • Estrutura Híbrida: O modelo mantém blocos convolucionais para o stem (E1), o gargalo (bottleneck) e todo o decodificador para preservar detalhes locais.
   • Integração do Codificador: Blocos Tri-plane DAMamba são inseridos exclusivamente nos estágios do codificador E2–E4.
   • Variante Compacta: O modelo principal utiliza uma largura base de $C=16$, resultando em aproximadamente 5,3 milhões de parâmetros ($\sim 13\times$ menos que o SegMamba).
   • Variante de Grande Escala (DAMamba-L): Uma variante mais larga ($\sim 70$ milhões de parâmetros) com blocos DAMamba empilhados no codificador e um gargalo convolucional, projetada para igualar a capacidade do SegMamba para comparação direta.

2. Escaneamento Adaptativo Dinâmico Tri-Plano (Tri-Plane DAS):

   • Mecanismo: Diferente dos escaneamentos fixos (ordem de linha, zigzag ou múltiplas orientações de achatamento), o DAS aprende o reordenamento espacial dependente da entrada. Uma Rede de Predição de Deslocamento (OPN) prevê deslocamentos 2D a partir do mapa de características, que são usados para reamostrar as características via interpolação bilinear antes do escaneamento seletivo.
   • Extensão 3D: O operador DAS 2D é estendido para volumes aplicando-o através de três planos anatômicos: Axial ($H \times W$), Coronal ($D \times W$) e Sagital ($D \times H$). As saídas são calculadas pela média para formar o bloco 3D DAS:
     $$DAS3D(X) = \frac{1}{3}(DAS_{ax} + DAS_{cor} + DAS_{sag})$$
   • Escaneamento Seletivo: O mapa de características reordenado é processado por um bloco de escaneamento seletivo Mamba padrão, permitindo a propagação de contexto global com custo $O(N)$ sem mapas de atenção quadráticos.

3. Treinamento e Avaliação:

   • Dataset: Conjunto de treinamento BraTS 2020 (369 sujeitos, 4 modalidades).
   • Perda (Loss): Modulated Combo Loss ($L = (L_{CE} + \delta) \cdot (L_{mod}^{Dice} + \delta)$) para focar em voxels difíceis e contornos.
   • Protocolo: Validação cruzada de cinco dobras ao nível do sujeito com inferência de volume total (janela deslizante, patches de $128^3$). Métricas incluem Dice médio e HD95 para Tumor Realçante (ET), Núcleo do Tumor (TC) e Tumor Total (WT).

Principais Contribuições

1. Primeira Aplicação Médica de DAS: Transfere o Escaneamento Adaptativo Dinâmico da classificação de imagens naturais 2D para a segmentação médica densa 3D.
2. Extensão 2D-para-3D: Eleva o operador DAS para volumes via escaneamento adaptativo tri-plano, oferecendo uma alternativa treinável ao achatamento de múltiplas orientações projetado à mão.
3. Eficiência de Parâmetros: Demonstra uma U-Net híbrida com blocos tri-plane DAS apenas no codificador, alcançando paridade próxima com modelos muito maiores com um custo de parâmetros $\sim 13\times$ menor.
4. Validação em Duas Escalas: Valida a abordagem tanto em uma escala compacta ($\sim 5,3$M) quanto em uma escala grande ($\sim 70$M), mostrando que o DAS aprendido pode superar padrões de escaneamento fixos quando a colocação arquitetônica é otimizada.

Resultados

• Modelo Compacto (DAMamba-UNet3D, $\sim 5,3$M): Alcançou um Dice médio de 0,815 $\pm$ 0,013 no BraTS 2020. Isso é quase par com o SegMamba retreinado (0,824 $\pm$ 0,014), apesar de possuir $\sim 13\times$ menos parâmetros.
• Modelo de Grande Escala (DAMamba-L, $\sim 70$M): Superou o SegMamba retreinado com um Dice médio de 0,829 $\pm$ 0,012 (uma melhoria de +0,5 ponto), com o maior ganho observado na classe Núcleo do Tumor (TC) (0,842 vs. 0,827).
• Estudos de Ablação:
  • Posicionamento: Colocar blocos DAS apenas no codificador (E2–E4) rendeu os melhores resultados (0,815). Adicionar DAS ao gargalo ou ao decodificador reduziu o desempenho (0,807 e 0,805, respectivamente), indicando que o posicionamento apenas no codificador é crítico.
  • Padrões de Escaneamento: O DAS tri-plano aprendido superou os padrões de escaneamento fixos (Raster, Snake, Cross, Tri-orientated) quando o posicionamento do codificador foi correspondido.

Significância e Alegações
O artigo afirma que o DAS tri-plano aprendido em uma U-Net híbrida é uma alternativa competitiva ao escaneamento de orientação fixa Tri-orientated Mamba (ToM) para segmentação de tumores cerebrais 3D.

• Eficiência: O modelo compacto demonstra que um alto desempenho pode ser alcançado com contagens de parâmetros significativamente reduzidas, aproveitando o escaneamento adaptativo em vez da capacidade bruta.
• Desempenho: Em grandes escalas, o mecanismo adaptativo aprendido (DAS) pode melhorar sobre as trajetórias de escaneamento fixas (ToM) sob protocolos idênticos de treinamento e avaliação.
• Modéstia: Os autores alertam que, embora os resultados no BraTS 2020 sejam promissores, o estudo é limitado a um único dataset. Eles afirmam que "validação em múltiplos datasets e testes estatísticos são necessários antes que alegações gerais sobre DAS versus ToM" possam ser feitas. A principal contribuição é a extensão do DAS para volumes médicos 3D e a demonstração de sua viabilidade em uma arquitetura híbrida eficiente em parâmetros.