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

기술 요약: 언롤드 MRI 재구성을 위한 SO-Mamba

문제 제기
가속화된 자기공명영상 (MRI) 재구성은 대공간 영역에 걸친 해부학적 일관성을 유지하면서 undersampled k-공간 측정값으로부터 고품질 영상을 복원하는 것을 목표로 합니다. 데이터 기반 사전 지식을 통해 딥러닝이 이 분야를 발전시켰지만, 현재 아키텍처는 언롤드된 데이터 일관성 (DC) 결합 솔버에서 특정 한계에 직면해 있습니다. 합성곱 신경망 (CNN) 은 제한된 수용 영역으로 인해 장거리 해부학적 의존성을 처리하는 데 어려움을 겪는 반면, 비전 트랜스포머는 고해상도 임상 입력으로의 확장성을 저해하는 2 차 계산 복잡도에 시달립니다.

Mamba 와 같은 선택적 상태 공간 모델 (SSM) 은 선형 복잡도의 장거리 모델링을 제공하여 언롤드 재구성을 위한 매력적인 정규화기로 부각됩니다. 그러나 DC 결합 언롤드 솔버의 상태 공간 계산 내에서 재구성 증거가 어떻게 할당되어야 하는지에 관한 중요한 설계 격차가 존재합니다. 이러한 솔버에서 각 단계는 서로 다른 재구성 반복(iterate) 을 처리합니다. 모든 증거를 균일하게 처리하여 영구적인 "거주" 콘텐츠와 단계 종속적인 "비거주" 업데이트를 동일한 재귀 콘텐츠 경로에 배치하는 것은 구별된 처리가 필요한 역할들을 혼란스럽게 만들 수 있습니다. 구체적으로, 어떤 증거가 재귀 콘텐츠 (상태 궤적) 가 되어야 하는지, 그리고 어떤 증거는 단지 해당 콘텐츠의 읽기/쓰기 방식을 조절하거나 재귀 루프에 진입하지 않고 출력을 보정해야 하는지는 명확하지 않습니다.

방법론: SO-Mamba
저자들은 SO-Mamba(State-Ownership Mamba) 를 제안합니다. 이는 할당 문제를 상태 소유권 (state ownership) 으로 공식화하는 정규화기입니다. 핵심 원리는 각 언롤드 단계 내의 재구성 증거의 기능적 역할을 세 가지 명확한 경로로 분리합니다:

1. 재귀 거주성 (Recurrent Residency): 선택적 스캔 상태 공간 계산의 직접적인 콘텐츠가 되는 증거.
2. 상태 인터페이스 접근 (State-Interface Access): 상태의 읽기 및 쓰기를 제어하는 (B 및 C 인터페이스를 조절하는) 증거.
3. 비상태 출력 보정 (Non-State Output Correction): 재귀 콘텐츠 경로에 진입하지 않고 최종 출력에 기여하는 증거.

주요 구성 요소:

• 소유권 라우터 (State-Ownership Router, SOR): 이 모듈은 현재 재구성 반복의 특징 표현을 받아 두 개의 스트림으로 분할합니다:
  • 거주 운반체 ($L_k$): 고정된 비적응형 이항 운반체 프로젝션을 통해 생성됩니다. 이 스트림은 재귀 거주성을 할당받아 Mamba 콘텐츠 경로를 위한 토큰을 공급합니다. 이는 안정적이고 공간적으로 일관된 해부학적 레이아웃을 나타냅니다.
  • 비거주 증거 ($G_k$): 운반체에서 남은 증거 (운반체로부터의 국소 편차) 와 특정 비거주 특징으로 구성됩니다. 이 스트림은 재귀 콘텐츠 소스에서 제외됩니다 ($G_k \nrightarrow u_k$). 대신 다음과 같은 경로로 라우팅됩니다:
    • 상태 인터페이스 조절: $G_k$ 는 B 및 C 상태 인터페이스를 조절하는 아핀 매개변수를 예측하여 거주 콘텐츠가 처리되는 방식을 적응시킵니다.
    • 비상태 출력 구멍 (Outlet): $G_k$ 는 비상태 정제 (NSR) 블록을 통과하여 Mamba 판독값과 연결되어 최종 출력을 보정합니다.
• 소유권 인식 Mamba-3 업데이트: Mamba-3 블록은 입력 토큰으로 거주 운반체를 처리하면서, 상태 공간 전이를 조건화하기 위해 비거주 스트림을 사용합니다. 이는 재귀 상태 궤적이 안정된 해부학적 콘텐츠에 의해 지배되도록 보장하는 동시에, 동적 업데이트가 조절된 인터페이스와 출력 보정을 통해 시스템에 영향을 미치도록 합니다.
• 이중 레벨 외대역 누출 진단: 역할 분리를 검증하기 위해 저자들은 두 곳에서 "외대역"(고주파수) 에너지를 측정하는 진단 지표를 도입합니다: 원시 숨겨진 상태 궤적 ($H_{Leak}$) 과 스캔 후 판독값 ($R_{Leak}$) 입니다. 이는 숨겨진 상태 저장과 판독 표현을 분리하여, 비거주 증거가 재귀 상태 궤적을 손상시키지 않으면서 출력 수준에서 효과적으로 표현됨을 검증합니다.

주요 기여

1. 문제 공식화: 본 논문은 Mamba 기반 언롤드 MRI 재구성을 위한 고유한 설계 문제로서 "단계 국소 상태 할당"을 식별하고 이를 상태 소유권으로 공식화합니다.
2. SO-Mamba 아키텍처: Mamba-3 를 기반으로 구축된 새로운 정규화기로, 각 단계 내에서 증거를 재귀 거주성, 상태 인터페이스 접근, 비상태 출력 보정으로 명시적으로 할당합니다.
3. SOR 메커니즘: 거주 운반체를 구성하고 비거주 증거를 인터페이스 조절 및 출력 보정으로 라우팅하여 소유권 계약을 강제하는 State-Ownership Router 의 도입.
4. 종합 평가: 다양한 해부학과 샘플링 패턴을 다루는 다섯 개의 공개 벤치마크 (fastMRI, CC359, ACDC, M4Raw, Prostate158) 에 대한 평가와 상태 동작에 대한 새로운 진단 분석.

실험 결과

• 성능: SO-Mamba 는 다섯 개의 데이터셋 전반에 걸쳐 CNN, 트랜스포머, Mamba 기반 베이스라인 (UNet, TransUNet, ReconFormer, LMO, HiFi-Mamba 포함) 보다 일관되게 우수한 성능을 발휘합니다.
  • fastMRI(AF=4) 에서 34.65 dB PSNR 및 0.855 SSIM을 달성합니다.
  • CC359(AF=4) 에서 36.69 dB PSNR 및 0.954 SSIM을 달성합니다.
  • 또한 방사형 (ACDC) 및 멀티코일 (M4Raw) 설정에서도 선도적인 성능을 보입니다.
• 효율성: SO-Mamba 는 8 배 가속화 시 fastMRI 에서 HiFi-Mamba(67.87G) 및 LMO(484.98G) 와 비교하여 가장 낮은 FLOPs(45.59G) 를 사용하면서도 효율적인 모델 중 최고의 정확도를 달성합니다.
• 애블레이션 연구:
  • 상태 인터페이스 접근 또는 출력 구멍을 제거하면 성능이 저하됩니다.
  • 비거주 증거가 재귀 콘텐츠 경로에 진입하도록 허용하는 것 (소유권 계약 위반) 은 성능이 급격히 떨어지는 결과 (34.65 dB 에서 31.71 dB 로 감소) 를 초래하여 분리의 필요성을 확인시켜 줍니다.
  • 누출 진단은 출력 구멍을 제거하면 외대역 콘텐츠가 재귀 판독값으로 강제됨을 보여줍니다 ($R_{Leak}$이 약 5 배 증가). 반면, 완전한 SO-Mamba 설계는 재귀 상태를 깨끗하게 유지하면서 필요한 보정을 출력에서 표현합니다.

의의 및 주장
본 논문은 SO-Mamba 가 언롤드 MRI 재구성의 근본적인 아키텍처 제약을 해결한다고 주장합니다. 즉, 영구적인 상태 궤적을 형성해야 하는 증거와 시스템의 동역학이나 출력에만 영향을 미쳐야 하는 증거를 구분해야 한다는 요구사항입니다. 상태 소유권을 강제함으로써, 이 방법은 안정적인 해부학적 구조와 단계 종속적 업데이트의 혼란을 방지하여 더 날카로운 경계와 균일한 오차 매핑을 이끌어냅니다.

저자들은 State-Ownership Router(SOR) 가 해부학 대 인공물이라는 의미론적 분해가 아닌 아키텍처 할당 규칙임을 강조합니다. 그 의의는 선형 복잡도의 SSM 이 정보 흐름을 신중하게 관리함으로써 언롤드 재구성에 효과적으로 적응될 수 있음을 입증하는 데 있으며, 높은 계산 효율성으로 최첨단 성능을 달성하는 데 있습니다. 이 연구는 향후 개선 사항이 적응형 소유권 라우팅과 더 넓은 임상 획득 설정에 있을 수 있음을 시사하지만, 현재 기여는 상태 소유권 원칙 자체의 확립에 있습니다.