Resumo Técnico: Patch-MoE Mamba

Declaração do Problema

A segmentação de imagens médicas depende fortemente de Redes Neurais Convolucionais (CNNs) e Transformers. No entanto, as CNNs são inerentemente limitadas por campos receptivos locais, lutando para modelar dependências de longo alcance, enquanto os Transformers frequentemente incorrem em complexidade computacional e de memória quadrática. Recentemente, Modelos de Espaço de Estados (SSMs), particularmente arquiteturas baseadas em Mamba, emergiram como alternativas eficientes com complexidade de sequência linear.

Apesar de sua promessa, os modelos de segmentação existentes baseados em Mamba enfrentam duas limitações críticas:

1. Disrupção da Estrutura Espacial Local: A maioria dos modelos converte mapas de características 2D em sequências 1D usando varredura raster pixel a pixel. Esse processo pode separar pixels espacialmente adjacentes na sequência, enfraquecendo a coerência espacial 2D local. Isso é prejudicial para tarefas de previsão densa que requerem localização precisa de limites e preservação de estrutura local.
2. Fusão Direcional Infexível: Os métodos atuais tipicamente fundem saídas de múltiplas direções de varredura via soma simples. Essa agregação fixa assume que todas as direções são igualmente importantes em cada localização espacial, falhando em adaptar-se aos diversos tamanhos, formas, texturas e complexidades de limites dos alvos médicos (por exemplo, lesões pequenas e de baixo contraste versus estruturas grandes e irregulares).

Metodologia

Os autores propõem o Patch-MoE Mamba, uma arquitetura de espaço de estados do tipo Mistura de Especialistas (MoE) ordenada por patches, projetada para abordar as limitações supracitadas. A arquitetura segue um design estilo U-Net, substituindo os blocos padrão de Espaço de Estados Visual (VSS) no VM-UNetV2 por um novo bloco Patch-MoE VSS.

1. Varredura Hierárquica Ordenada por Patches

Em vez de varrer pixels individuais raster, o método particiona mapas de características em patches locais e realiza a varredura direcional de maneira ordenada por patches.

• Mecanismo: Dado um mapa de características, a grade espacial é dividida em patches não sobrepostos de tamanho $p \times p$. Dentro de cada patch, os pixels são enumerados em ordem de linha principal antes de avançar para o próximo patch.
• Benefício: Isso garante que todos os pixels dentro de um patch permaneçam consecutivos na sequência, preservando vizinhanças espaciais locais enquanto ainda permite a modelagem de dependências de longo alcance entre patches.
• Estratégia Hierárquica: Para capturar contexto multi-escala, o modelo emprega um conjunto de tamanhos de patches $\{p^{(1)}, p^{(2)}, \dots\}$. Patches menores capturam detalhes finos de limites, enquanto patches maiores modelam estruturas mais grosseiras. Diferentes direções de varredura (frente, reverso, frente largura-altura, reverso largura-altura) podem utilizar tamanhos de patches distintos para lidar com estruturas anisotrópicas.

2. Fusão Direcional Baseada em MoE

Para substituir a soma fixa, os autores introduzem um módulo de fusão adaptativo que trata as saídas direcionais como "especialistas".

• Construção de Especialistas: O módulo utiliza cinco especialistas:
  • Quatro especialistas específicos por direção ($E_1$ a $E_4$): As saídas normalizadas dos quatro varredores Mamba direcionais.
  • Um especialista de concatenação ($E_{concat}$): Formado pela concatenação dos quatro mapas normalizados ao longo da dimensão do canal e projetando-os de volta via uma convolução $1\times1$. Este especialista captura correlações cruzadas entre direções e fornece uma visão holística e agnóstica à direção.
• Comutação Consciente Espacial: Uma rede roteadora computa pesos dependentes da localização para fundir adaptativamente esses especialistas. O roteador é impulsionado por uma mistura de contexto local (via convolução por profundidade) e contexto global (via média de agrupamento adaptativa), controlado por um escalar aprendível $\alpha$.
• Estabilização Residual: Um desvio residual adiciona as saídas direcionais brutas de volta ao resultado fundido ($Z_l = \hat{Y}_l + \sum Y_l^{(i)}$). Isso garante respostas de base robustas e preserva sinais direcionais fortes, estabilizando o treinamento.

3. Infusão de Semântica e Detalhe (SDI)

As saídas do codificador são processadas por um módulo SDI (adaptado do U-Net v2) para aprimorar os mapas de características, infundindo informações semânticas de níveis superiores e integrando detalhes mais finos de níveis inferiores via produto de Hadamard.

Principais Contribuições

1. Varredura Ordenada por Patches: Um mecanismo de varredura inovador que substitui a varredura raster pixel a pixel por uma travessia ordenada por patches, preservando efetivamente a coerência espacial 2D local enquanto mantém a eficiência dos modelos de espaço de estados.
2. Fusão Direcional Baseada em MoE: Um módulo de fusão adaptativo que substitui a soma simples por uma rede de comutação aprendível. Ele pondera dinamicamente múltiplos especialistas direcionais e um especialista de concatenação com base no contexto espacial, permitindo que o modelo se adapte a escalas e complexidades variáveis de lesões.
3. Modelagem Multi-Escala Hierárquica: A integração de tamanhos de patches hierárquicos permite que a rede capture simultaneamente detalhes finos de limites e estruturas semânticas grosseiras.
4. Estabilização Residual: A inclusão de um caminho residual para as saídas direcionais brutas garante estabilidade no treinamento e previne a perda de sinais direcionais críticos.

Resultados Experimentais

Os autores avaliaram o Patch-MoE Mamba em cinco conjuntos de dados públicos de segmentação de pólipos (Kvasir-SEG, ClinicDB, ColonDB, ETIS, CVC-300) e dois conjuntos de dados de segmentação de lesões cutâneas (ISIC 2017, ISIC 2018).

• Desempenho: O Patch-MoE Mamba alcançou desempenho de última geração em todos os conjuntos de dados, superando consistentemente as bases CNN (U-Net, U-Net v2) e as bases baseadas em Mamba (VM-UNet, VM-UNetV2).
  • No desafiador conjunto de dados ETIS (conhecido por baixo contraste e pólipos pequenos), o modelo melhorou a pontuação Dice de 72,56% (VM-UNetV2) para 74,04%.
  • No ISIC 2018, alcançou uma pontuação Dice de 89,34%, superando o melhor anterior de 88,36%.
• Estudos de Ablação:
  • Substituir a varredura raster por varredura ordenada por patches apenas melhorou a Dice média de 83,46% para 84,02%.
  • Adicionar o módulo de fusão MoE aumentou ainda mais a média para 84,30%.
  • Tamanhos de patches hierárquicos (grosseiros nos estágios iniciais, finos nos estágios mais profundos) produziram os melhores resultados.
• Complexidade Computacional: Embora a própria varredura ordenada por patches não aumente o custo computacional (a contagem de tokens permanece inalterada), a inclusão do especialista de concatenação aumenta os parâmetros para ~70M e os FLOPs para ~28G. Os autores observam que essa sobrecarga é esperada devido às interações de canal mais densas exigidas pelo mecanismo MoE.

Significado e Alegações

O artigo afirma que o Patch-MoE Mamba oferece uma solução geral e eficaz para segmentação de imagens médicas. Ao abordar as limitações específicas dos modelos Mamba existentes — nomeadamente a perda de estrutura espacial local e a rigidez da fusão direcional —, a arquitetura proposta demonstra adaptabilidade superior a diversos alvos médicos com tamanhos, formas e complexidades de limites variáveis. Os resultados sugerem que a preservação de vizinhanças espaciais locais através da varredura ordenada por patches, combinada com fusão adaptativa de especialistas, aprimora significativamente a localização de limites e a precisão de segmentação, particularmente em cenários desafiadores envolvendo lesões de baixo contraste ou pequenas.