SO-Mamba: State-Ownership Mamba for Unrolled MRI Reconstruction

Resumo Técnico: SO-Mamba para Reconstrução de MRI Desenrolada

Declaração do Problema
A reconstrução acelerada de Imagem por Ressonância Magnética (MRI) visa recuperar imagens de alta qualidade a partir de medições subamostradas do espaço-k, preservando a coerência anatômica em grandes regiões espaciais. Embora o aprendizado profundo tenha avançado este campo por meio de priores baseados em dados, arquiteturas atuais enfrentam limitações específicas em solvers desenrolados acoplados à consistência de dados (DC). Redes Neurais Convolucionais (CNNs) lutam com dependências anatômicas de longo alcance devido a campos receptores limitados, enquanto Transformers de Visão sofrem com complexidade computacional quadrática que impede a escalabilidade para entradas clínicas de alta resolução.

Modelos de Espaço de Estados Seletivos (SSMs), como o Mamba, oferecem modelagem de longo alcance com complexidade linear, tornando-se regularizadores atraentes para reconstrução desenrolada. No entanto, existe uma lacuna crítica de design na forma como a evidência de reconstrução é atribuída dentro do cálculo do espaço de estados de um solver desenrolado acoplado à DC. Em tais solvers, cada estágio opera em uma iteração de reconstrução diferente. Tratar toda a evidência uniformemente — colocando tanto o conteúdo "residente" persistente quanto as atualizações "não residentes" dependentes do estágio na mesma rota de conteúdo recorrente — pode entrelaçar papéis que exigem tratamento distinto. Especificamente, não está claro qual evidência deve se tornar o conteúdo recorrente (a trajetória de estado) versus qual evidência deve apenas modular como esse conteúdo é escrito/lemido ou corrigir a saída sem entrar no loop recorrente.

Metodologia: SO-Mamba
Os autores propõem o SO-Mamba (Mamba de Propriedade do Estado), um regularizador que formaliza este problema de alocação como propriedade do estado. O princípio central separa os papéis funcionais da evidência de reconstrução dentro de cada estágio desenrolado em três rotas distintas:

1. Residência Recorrente: Evidência que se torna o conteúdo direto do cálculo do espaço de estados de varredura seletiva.
2. Acesso à Interface do Estado: Evidência que controla como o estado é escrito e lido (modulando as interfaces B e C).
3. Correção de Saída Não-Estatal: Evidência que contribui para a saída final sem entrar no caminho de conteúdo recorrente.

Componentes Chave:

• Roteador de Propriedade do Estado (SOR): Este módulo toma a representação de características da iteração de reconstrução atual e a divide em dois fluxos:
  • Portador Residente ($L_k$): Gerado via uma projeção de portador binária fixa e não adaptativa. Este fluxo é atribuído à residência recorrente, fornecendo os tokens para a rota de conteúdo do Mamba. Representa layouts anatômicos estáveis e espacialmente coerentes.
  • Evidência Não-Residente ($G_k$): Composta pela evidência restante (desvios locais do portador) e características específicas não residentes. Este fluxo é excluído da fonte de conteúdo recorrente ($G_k \nrightarrow u_k$). Em vez disso, é roteado para:
    • Modulação da Interface do Estado: $G_k$ prevê parâmetros afins para modular as interfaces de estado B e C, adaptando como o conteúdo residente é processado.
    • Saída de Correção Não-Estatal: $G_k$ passa por um bloco de Refinamento Não-Estatal (NSR) e é concatenado com a leitura do Mamba para corrigir a saída final.
• Atualização Mamba-3 Consciente de Propriedade: O bloco Mamba-3 processa o portador residente como os tokens de entrada, enquanto usa o fluxo não residente para condicionar as transições do espaço de estados. Isso garante que a trajetória de estado recorrente permaneça dominada por conteúdo anatômico estável, enquanto atualizações dinâmicas influenciam o sistema por meio de interfaces controladas e correções de saída.
• Diagnóstico de Vazamento de Banda Externa em Dois Níveis: Para validar a separação de papéis, os autores introduzem uma métrica diagnóstica que mede a energia de "banda externa" (alta frequência) em dois locais: a trajetória bruta do estado oculto ($H_{Leak}$) e a leitura pós-varredura ($R_{Leak}$). Isso separa o armazenamento do estado oculto da expressão de leitura, verificando que a evidência não residente não corrompe a trajetória de estado recorrente, mas é efetivamente expressa no nível de saída.

Contribuições Principais

1. Formulação do Problema: O artigo identifica "alocação de estado local ao estágio" como um problema de design distinto na reconstrução de MRI desenrolada baseada em Mamba e o formula como propriedade do estado.
2. Arquitetura SO-Mamba: Um novo regularizador construído sobre o Mamba-3 que atribui explicitamente evidência à residência recorrente, acesso à interface do estado e correção de saída não estatal dentro de cada estágio.
3. Mecanismo SOR: A introdução do Roteador de Propriedade do Estado, que constrói um portador residente e roteia evidência não residente para modulação de interface e correção de saída, fazendo cumprir o contrato de propriedade.
4. Avaliação Abrangente: Avaliação em cinco benchmarks públicos (fastMRI, CC359, ACDC, M4Raw, Prostate158) cobrindo diversas anatomias e padrões de amostragem, juntamente com uma nova análise diagnóstica do comportamento do estado.

Resultados Experimentais

• Desempenho: O SO-Mamba supera consistentemente as bases baseadas em CNN, Transformer e Mamba (incluindo UNet, TransUNet, ReconFormer, LMO e HiFi-Mamba) em cinco conjuntos de dados.
  • No fastMRI (AF=4), alcança 34,65 dB PSNR e 0,855 SSIM.
  • No CC359 (AF=4), alcança 36,69 dB PSNR e 0,954 SSIM.
  • Também lidera em configurações radiais (ACDC) e de múltiplos coils (M4Raw).
• Eficiência: O SO-Mamba alcança a melhor precisão entre modelos eficientes enquanto utiliza os menores FLOPs (45,59G) em comparação com HiFi-Mamba (67,87G) e LMO (484,98G) no fastMRI com aceleração de 8×.
• Estudos de Ablação:
  • Remover o acesso à interface do estado ou a saída de correção degrada o desempenho.
  • Permitir que evidência não residente entre na rota de conteúdo recorrente (violando o contrato de propriedade) causa uma queda acentuada no desempenho (de 34,65 dB para 31,71 dB), confirmando a necessidade da separação.
  • O diagnóstico de vazamento mostra que remover a saída de correção força o conteúdo de banda externa para a leitura recorrente (aumentando $R_{Leak}$ em ~5×), enquanto o design completo do SO-Mamba mantém o estado recorrente limpo enquanto expressa correções necessárias na saída.

Significado e Alegações
O artigo alega que o SO-Mamba aborda uma restrição arquitetônica fundamental na reconstrução de MRI desenrolada: a necessidade de distinguir entre evidência que deve formar a trajetória de estado persistente e evidência que deve apenas influenciar a dinâmica do sistema ou a saída. Ao fazer cumprir a propriedade do estado, o método previne o entrelaçamento de estruturas anatômicas estáveis com atualizações dependentes do estágio, levando a limites mais nítidos e mapas de erro mais uniformes.

Os autores enfatizam que o Roteador de Propriedade do Estado (SOR) é uma regra de atribuição arquitetônica e não uma decomposição semântica de anatomia versus artefatos. O significado reside em demonstrar que SSMs de complexidade linear podem ser efetivamente adaptados para reconstrução desenrolada gerenciando cuidadosamente o fluxo de informações, alcançando desempenho de última geração com alta eficiência computacional. O trabalho sugere que melhorias futuras podem residir no roteamento adaptativo de propriedade e em configurações mais amplas de aquisição clínica, mas a contribuição atual é o estabelecimento do próprio princípio de propriedade do estado.