Resumo Técnico: Mamba Goes HoME

Declaração do Problema

A segmentação de imagens médicas tridimensionais (3D) é crítica para o diagnóstico e planejamento de tratamento em diversas modalidades, como Tomografia Computadorizada (TC), Ressonância Magnética (RM) e Ultrassom (US). No entanto, as arquiteturas existentes enfrentam trocas significativas:

• Redes Neurais Convolucionais (CNNs) oferecem complexidade computacional linear, mas sofrem com campos receptivos limitados, dificultando a captura de padrões espaciais globais essenciais para compreender a relação hierárquica entre lesões locais e estruturas anatômicas globais.
• Vision Transformers (ViTs) utilizam atenção global para modelar dependências de longo alcance, mas incorrem em complexidade computacional quadrática em relação ao número de tokens. Isso os torna computacionalmente proibitivos para dados volumétricos 3D de alta resolução.
• Modelos de Espaço de Estados Seletivos (SSMs), como o Mamba, fornecem complexidade linear para modelagem de sequências longas, oferecendo uma alternativa eficiente aos ViTs. No entanto, os SSMs padrão não são inerentemente projetados para lidar de forma adaptativa com os diversos padrões locais encontrados em dados médicos multiescala complexos.
• Estruturas de Mistura de Especialistas (Mixture-of-Experts - MoE) podem rotear dinamicamente características para subredes especializadas para lidar com padrões locais, mas a combinação da eficiência global dos SSMs com a adaptabilidade localizada do MoE em imagens médicas 3D permanece amplamente inexplorada.

Metodologia: Mamba-HoME

Os autores propõem o Mamba-HoME, uma nova arquitetura que integra o backbone de Modelo de Espaço de Estados (SSM) Mamba com uma camada de Mistura de Especialistas Suave Hierárquica (Hierarchical Soft Mixture-of-Experts - HoME). O design visa modelar eficientemente as hierarquias de padrões locais-para-globais com complexidade computacional linear.

1. Mistura de Especialistas Suave Hierárquica (HoME)

A inovação central é uma camada de roteamento de tokens de dois níveis projetada para equilibrar a especialização local com a integração de contexto global:

• Nível 1: Agrupamento e Roteamento Local. As sequências de entrada são particionadas em grupos locais. Uma camada de Mistura de Especialistas Suave (SMoE) roteia os tokens dentro desses grupos para especialistas locais especializados. Isso permite a extração de características localizadas enquanto preserva a localidade espacial e reduz o overhead de memória em comparação com o roteamento global.
• Nvez 2: Agregação Global. As saídas do primeiro nível são agregadas e passadas por uma segunda camada de SMoE global. Isso facilita a fusão de informações entre grupos e o refinamento do contexto global.
• Escalabilidade: Ao realizar o roteamento dentro de grupos antes da agregação global, o método evita os custos de memória proibitivos do roteamento global em dados 3D de alta resolução (onde a contagem de tokens pode atingir milhões), mantendo uma complexidade computacional de $O(Nd)$.

2. Bloco Mamba-HoME

A arquitetura incorpora a camada HoME em uma rede codificador-decodificador em forma de U (inspirada no SegMamba):

• Backbone: Utiliza camadas Mamba para modelagem eficiente de dependências de longo alcance com complexidade de tempo linear.
• Convolução Espacial com Portão (Gated Spatial Convolution - GSC): Integrada para capturar priors espaciais locais antes de achatar os dados em sequências 1D para o processamento Mamba.
• Tanh Dinâmica (Dynamic Tanh - DyT): Substitui a Normalização de Camada (Layer Normalization) padrão por uma não-linearidade ponto a ponto ($f(x) = w \cdot \tanh(\alpha \cdot x) + b$). Isso estabiliza os gradientes e reduz o overhead computacional ao evitar cálculos de média e variância, acelerando o treinamento e a inferência.
• Progressão Hierárquica: O número de especialistas aumenta monotonicamente com a profundidade ($E_1 < E_2 < \dots$), enquanto os tamanhos dos grupos diminuem ($K_1 > K_2 > \dots$), permitindo um processamento progressivamente mais fino à medida que a rede se aprofunda.

Principais Contribuições

1. Camada HoME: Introdução de uma camada de roteamento de tokens MoE de dois níveis que captura eficientemente as hierarquias de padrões locais-para-globais ao agrupar tokens para roteamento de especialista local seguido pelo refinamento de especialista global.
2. Arquitetura Unificada: Design do Bloco Mamba-HoME, que integra os SSMs do Mamba com o HoME, combinando o processamento de sequências longas com eficiência de memória e roteamento de especialistas hierárquico.
3. Rede de Segmentação de Múltiplos Estágios: Incorporação desses componentes em uma arquitetura em forma de U especificamente adaptada para segmentação de imagens médicas 3D, melhorando a representação contextual através das dimensões espaciais e de profundidade.

Resultados Experimentais

Os autores avaliaram o Mamba-HoME em quatro conjuntos de dados públicos (PANORAMA, AMOS, FeTA 2022, MVSeg) e um conjunto de dados de TC próprio, cobrindo as modalidades CT, RM e US.

• Desempenho: O Mamba-HoME alcançou resultados de estado da arte (SOTA) em precisão de segmentação (Coeficiente de Similaridade de Dice - DSC) e precisão de borda (Distância de Hausdorff 95% - HD95) em todas as modalidades testadas.
  • No conjunto de dados PANORAMA, alcançou um DSC de 56,7% (pré-treinado) e 54,8% (do zero), superando baselines como SegMamba e Swin UNETR.
  • No conjunto de dados de RM fetal FeTA 2022, atingiu um DSC de 87,7% (pré-treinado).
  • No conjunto de dados de US 3D MVSeg, alcançou um DSC de 85,0% (pré-treinado).
• Eficiência: Apesar de possuir uma contagem de parâmetros maior (170,1M) em comparação com alguns competidores, o modelo demonstrou baixo uso de memória GPU durante a inferência, uma vantagem crítica para dados 3D de alta resolução. Embora a inferência tenha sido aproximadamente 30% mais lenta que alguns baselines, os autores observam que este é um compromisso atraente pelos ganhos significativos em precisão.
• Generalização: O modelo demonstrou forte generalização cross-modal. Um modelo pré-treinado em dados de CT e RM adaptou-se com sucesso a tarefas de Ultrassom 3D com ajuste fino mínimo, superando métodos SOTA em ambas as métricas de DSC e HD95.
• Estudos de Ablação:
  • O desempenho ideal foi encontrado com contagens específicas de especialistas ($E_1 \in \{4, 8, 12, 16\}$, $E_2 \in \{8, 16, 24, 32\}$) e tamanhos de grupo ($K \in \{2048, 1024, 512, 256\}$).
  • A normalização Tanh Dinâmica (DyT) melhorou a velocidade de treinamento e inferência em ~6% sem sacrificar a precisão.
  • O pré-treinamento supervisionado em grandes conjuntos de dados de CT e RM aumentou significamente o desempenho e a robustez através de diversas regiões anatômicas e modalidades.

Significância e Alegações

O artigo afirma que o Mamba-HoME representa um avanço significativo na segmentação de imagens médicas 3D ao unir efetivamente a eficiência dos SSMs e a adaptabilidade das estruturas MoE. Os autores sustentam que seu design hierárquico permite que os modelos capturem as hierarquias espaciais inerentes de local-para-global dos dados médicos de forma mais eficaz do que as abordagens anteriores baseadas em CNN ou Transformers.

O trabalho destaca o potencial de representações de características agnósticas de modalidade, demonstrando que um modelo pré-treinado em CT e RM pode generalizar-se efetivamente para tarefas de Ultrassom, abordando o desafio da variabilidade de dados entre diferentes modalidades de imagem. Os autores posicionam o Mamba-HoME como uma solução escalável para dados 3D de alta resolução que mantém eficiência competitiva de memória e computação enquanto entrega precisão de segmentação superior.

Limitações Reconhecidas: Os autores observam que a escalabilidade para conjuntos de dados extremamente grandes (ex: >10.000 exames) e o comportamento sob aprendizado autossupervisionado de grande escala (>200.000 exames) permanecem inexplorados, identificando estes como direções para trabalhos futuros.