[논문]의료 영상 분류를 위한 혁신적인 비전 Mamba 모델

의료 영상 분류 하는 새로운 기술

논문 제목: MedMamba: Vision Mamba for Medical Image Classification

초록 (Abstract): 의료 영상 분류의 새로운 지평을 여는 MedMamba

본 글에서는 의료 영상 분류 분야에 혁신을 가져올 최신 딥러닝 모델인 MedMamba에 대해 심층적으로 분석합니다. MedMamba는 State Space Model(SSM)을 기반으로 한 Mamba 아키텍처를 의료 영상 도메인에 최초로 적용한 모델입니다. 16개의 데이터셋, 10가지 이미징 모달리티, 411,007개의 이미지에 대한 광범위한 실험을 통해 MedMamba는 기존의 CNN과 ViT 모델의 한계를 극복하고 우수한 성능을 보여주었습니다. 본 글에서는 MedMamba의 핵심 기술인 SS-Conv-SSM 블록과 2D-Selective-Scan 메커니즘을 상세히 설명하고, 다양한 의료 영상 분류 작업에서의 성능 평가 결과를 분석합니다.

소개: 의료 영상 분류의 도전과 MedMamba의 등장

의료 영상 분류는 컴퓨터 보조 진단(CAD) 시스템의 핵심 과제로, 임상 진단, 질병 치료, 모니터링 등 다양한 의료 분야에서 중요한 역할을 합니다. 최근 디지털 의료 영상 기술의 발전으로 CT, 초음파, X-ray, 현미경, 내시경, MRI 등 다양한 모달리티의 의료 영상이 폭발적으로 증가하고 있습니다. 이에 따라 효율적이고 정확한 자동 분류 시스템의 필요성이 더욱 커지고 있습니다.

딥러닝 기술의 발전으로 CNN(Convolutional Neural Network)과 ViT(Vision Transformer) 등이 의료 영상 분류에 널리 사용되어 왔습니다. 그러나 이들 모델은 각각 다음과 같은 한계를 가지고 있었습니다:

1. CNN: 지역적 특징 추출에는 강점이 있으나, 장거리 의존성 모델링에 한계가 있어 복잡한 의료 영상의 전체적인 문맥을 파악하는 데 어려움이 있었습니다.
2. ViT: 자기 주의 메커니즘을 통해 장거리 의존성을 효과적으로 모델링할 수 있지만, 이 과정에서 발생하는 이차 계산 복잡도로 인해 실제 임상 환경에서의 적용이 제한적이었습니다.

이러한 문제를 해결하기 위해 연구진은 최근 자연어 처리 분야에서 주목받고 있는 State Space Model(SSM)을 기반으로 한 Mamba 아키텍처를 의료 영상 분류에 적용한 MedMamba를 제안했습니다. MedMamba는 SSM의 효율적인 장거리 의존성 모델링 능력과 선형 계산 복잡도를 활용하여 다양한 의료 영상 모달리티에 대해 우수한 성능을 보여주었습니다.

기술 1: SS-Conv-SSM 블록

MedMamba의 핵심 구성 요소인 SS-Conv-SSM(State Space-Convolution-State Space Model) 블록은 다음과 같은 특징입니다.

1. 채널 분할(Channel Split): 입력 특징 맵 x ∈ R^(H×W×C)를 두 그룹으로 나눕니다. 즉, x_i ∈ R^(H×W×(C/2)), i=1,2로 분할
2. 병렬 처리:
   • Conv-Branch: 지역적 특징을 추출하기 위해 일련의 컨볼루션 연산을 수행합니다. x_1' = BatchNorm_1(x_1) x_1'' = ReLU(BatchNorm_2(Conv_3×3(x_1'))) x_1''' = ReLU(BatchNorm_3(Conv_3×3(x_1''))) cx_1 = ReLU(PWConv(x_1'''))
   • SSM-Branch: 장거리 의존성을 모델링하기 위해 SSM 연산을 수행합니다. x_2 = LayerNorm_1(x_2) x_2' = SiLU(DWConv(Linear(x_2))) x_2'' = LayerNorm_2(SS2D(x_2')) x_2''' = SiLU(Linear(x_2)) fx_2 = Linear(x_2'' ⊗ x_2''')
3. 채널 결합(Channel Concatenation): 두 브랜치의 결과를 다시 결합하여 전체 특징 맵을 구성합니다.
4. 채널 셔플(Channel Shuffle): 그룹 간 정보 교환을 촉진하여 특징 표현력을 향상합니다.

최종적으로 SS-Conv-SSM 블록의 출력은 다음과 같이 표현됩니다: y = x ⊕ g(f^(-1)(fx_1, fx_2))

여기서 g는 채널 셔플 연산, f^(-1)은 채널 결합 연산을 나타냅니다.

이러한 구조를 통해 MedMamba는 의료 영상의 지역적 특징과 전역적 특징을 동시에 효율적으로 모델링할 수 있습니다.

기술 2: 2D-Selective-Scan (SS2D)

MedMamba는 VMamba에서 제안된 2D-Selective-Scan(SS2D) 메커니즘을 채택했습니다. SS2D의 주요 특징은 다음과 같습니다:

1. Cross-Scan Module (CSM):
   • 4방향 스캐닝 전략을 사용하여 2D 특징 맵을 효과적으로 처리합니다.
   • 좌상단에서 우하단, 우하단에서 좌상단, 우상단에서 좌하단, 좌하단에서 우상단 방향으로 스캔합니다.
   • 이를 통해 각 픽셀이 모든 방향의 정보를 통합할 수 있어 전역적인 수용 영역을 확보할 수 있습니다.
2. S6 블록:
   • Mamba의 선택적 메커니즘을 적용하여 입력에 따라 SSM의 파라미터를 조정합니다.
   • 관련 정보를 선별적으로 유지하고 불필요한 정보를 필터링합니다.
3. 선형 복잡도:
   • 자기 주의 메커니즘과 달리 선형 계산 복잡도를 가져 대규모 의료 영상 데이터셋에 효율적으로 적용할 수 있습니다.

SS2D의 동작은 다음과 같은 단계로 이루어집니다:

1. Scan Expanding: 입력 이미지를 4가지 방향으로 unfolding 하여 시퀀스로 변환합니다.
2. S6 Processing: 변환된 시퀀스를 S6 블록으로 처리하여 특징을 추출합니다.
3. Scan Merging: 4개 방향의 출력 특징을 다시 병합하여 최종 2D 특징 맵을 구성 합니다.

이러한 과정을 통해 SS2D는 2D 의료 영상 데이터에 대해 효과적으로 장거리 의존성을 모델링하면서도 계산 효율성을 유지할 수 있습니다.

수학적 설명: MedMamba의 핵심인 State Space Model(SSM)은 다음과 같은 연속 시스템 수식으로 표현됩니다:

h'(t) = Ah(t) + Bx(t)   |  y(t) = Ch(t) 여기서 A ∈ R^(N×N)는 상태 행렬, B ∈ R^(N×1)와 C ∈ R^(N×1)는 투영 매개변수입니다.

이 연속 시스템은 다음과 같이 이산화 됩니다: A = exp(∆A)   |  B = (∆A)^(-1)(exp(∆A) - I) · ∆B

여기서 ∆는 타임스케일 파라미터입니다. 이산화 된 시스템은 다음과 같이 표현됩니다: h'(t) = Ah(t) + Bx(t) y(t) = Ch(t)

최종적으로 SSM은 전역 컨볼루션을 사용하여 출력을 계산합니다: K = (CB, CAB, ..., CA^(L-1)B) y = x * K

여기서 K ∈ R^L은 구조화된 컨볼루션 커널이며, L은 입력 시퀀스 x의 길이입니다.

이러한 수학적 기반을 통해 MedMamba는 효율적으로 장거리 의존성을 모델링하면서도 선형 복잡도를 유지할 수 있습니다.

결론: MedMamba: 의료 AI의 미래를 선도하는 혁신적 모델

MedMamba는 의료 영상 분류 분야에 혁신적인 접근 방식을 제시했습니다. 16개의 데이터셋, 10가지 이미징 모달리티, 411,007개의 이미지에 대한 광범위한 실험을 통해 MedMamba는 최신 모델들과 비교하여 경쟁력 있는 성능을 보여주었습니다. 특히 파라미터 수와 계산 복잡도 측면에서 효율성을 유지하면서도 우수한 성능을 달성했다는 점이 주목할 만합니다.

MedMamba는 의료 영상 분석을 위한 새로운 기준을 제시했으며, 앞으로 더 강력한 SSM 기반 AI 알고리즘과 의료 분야 응용 시스템 개발에 중요한 통찰을 제공할 것으로 기대됩니다. 향후 연구에서는 MedMamba를 다양한 의료 영상 작업에 적용하고, 모델의 해석 가능성을 개선하는 데 초점을 맞출 필요가 있습니다.