Rhamba: Region-Aware Hybrid Attention-Mamba Framework for Self-Supervised Learning in Resting-State fMRI

기술 요약: Rhamba – 영역 인식 하이브리드 어텐션–Mamba 프레임워크

문제 제기
자기지도 학습 (SSL) 은 대규모 신경영상 분석을 위한 유망한 길을 제시하지만, 기존 접근법은 상당한 한계에 직면해 있습니다. 현재 방법론들은 종종 신경해부학적 구조의 명시적 통합이 부족하여, 미세한 공간 정보를 희생할 수 있는 무작위 볼록 마스킹 또는 일반적 그래프 구조에 의존합니다. 또한, 휴지기 상태 기능적 자기공명영상 (rs-fMRI) 데이터를 모델링하기 위한 아키텍처 선택은 트레이드오프를 수반합니다: 트랜스포머 기반 모델은 전역 상호작용을 포착하지만 긴 시퀀스에 대해 이차 복잡성을 겪는 반면, Mamba 와 같은 상태 공간 모델 (SSM) 은 선형 시간 효율성을 제공하지만 복잡한 전역 상호작용을 모델링할 유연성이 부족할 수 있습니다. 더 나아가, 사전 학습 재구성 목표와 하위 작업 성능 간의 관계, 특히 영역 인식 마스킹 전략이 표현의 질에 미치는 영향에 대해서는 아직 충분히 탐구되지 않았습니다.

방법론
저자들은 대규모 rs-fMRI 분석을 위해 설계된 영역 인식 사전 학습 프레임워크인 Rhamba를 제안합니다. 이 방법론은 세 가지 핵심 설계 원칙에 기반합니다:

1. 영역 인식 토큰화 및 마스킹:

   • 해부학적 정렬: 프레임워크는 자동 해부학적 라벨링 (AAL3) 아틀라스를 활용하여 4D fMRI 볼륨 (MNI152 공간으로 정규화됨) 을 해부학적으로 정의된 관심 영역 (ROI) 에 정렬합니다.
   • 패치 생성: 볼륨은 겹치지 않는 $6 \times 6 \times 6$ 볼록 패치로 분할됩니다.
   • 마스킹 전략: 패치 내 해부학적 동질성의 정도에 따라 세 가지 서로 다른 마스킹 전략이 구현됩니다:
     • Any (임의): 패치가 특정 거대 영역에 속하는 최소 하나의 볼록을 포함하는 경우 패치가 마스킹됩니다.
     • Majority (다수): 패치의 볼록 중 50% 이상이 목표 거대 영역에 속하는 경우에만 패치가 마스킹됩니다.
     • Pure (순수): 지배적인 해부학적 라벨이 라벨링된 볼록의 최소 70% 를 차지하는 경우에만 패치가 마스킹되어 높은 공간적 동질성을 보장합니다.

2. 하이브리드 어텐션–Mamba 아키텍처:

   • 백본은 시프트드 윈도우 Mamba(SWM) 블록을 활용하는 기반 모델인 NeuroSTORM에 기반합니다.
   • 저자들은 어텐션과 상태 공간 메커니즘의 최적 통합을 결정하기 위해 네 가지 아키텍처 변형을 실험합니다:
     • Mamba-only: 순수 상태 공간 모델링.
     • Alternate: 교차 배치된 Mamba 및 어텐션 블록.
     • Attention–Mamba (AM): 인코더의 어텐션 블록과 디코더의 Mamba.
     • Mamba–Attention (MA): 인코더의 Mamba 블록과 디코더의 어텐션.
   • 이 모델은 사전 학습 중 누락된 시공간 패치를 재구성하는 마스킹 오토인코더 (MAE) 목표를 사용합니다.

3. 학습 및 평가 파이프라인:

   • 사전 학습: 세 가지 마스킹 전략과 네 가지 아키텍처 변형을 사용하여 ABIDE 데이터셋 (1,041 명) 에서 수행되었습니다.
   • 파인튜닝: 사전 학습된 모델은 COBRE 및 ADHD-200 데이터셋을 사용하여 이진 분류 작업 (조현병 vs 대조군; ADHD vs 대조군) 에 적응되었습니다.
   • 해석 가능성: 모델 예측에 기여하는 뇌 영역을 식별하기 위해 attribution 맵을 생성하는 데 SmoothGrad Squared 가 통합된 IG-SQ(Intergrated Gradients with SmoothGrad Squared) 가 사용되었습니다.

주요 기여

• Rhamba 프레임워크: rs-fMRI 를 위해 해부학적으로 유도된 마스킹과 하이브리드 어텐션–Mamba 백본을 통합하는 영역 인식 하이브리드 아키텍처의 도입.
• 해부학적으로 유도된 마스킹: 무작위 볼록 마스킹을 넘어 구조화되고 영역 정렬된 교란을 위한 토큰화 및 마스킹을 위한 체계적 접근법의 제안 (AAL3 파셀레이션을 사용).
• 아키텍처 분석: 상태 공간 및 어텐션 메커니즘의 교차 배치가 전역 기능적 상호작용을 보존하면서 효율적인 긴 시퀀스 모델링을 가능하게 함을 입증.
• 벤치마킹: SwiFT(트랜스포머 기반) 및 NeuroSTORM(Mamba 기반) 을 포함한 최신 기반 모델에 대한 포괄적 평가.

결과

• 마스킹 전략 vs 재구성: 마스킹 특이성에 기반한 재구성 손실에 명확한 순서가 있었습니다: Any > Majority > Pure. 더 엄격한 마스킹 (Pure) 은 더 낮은 재구성 손실을 초래했는데, 이는 동질적인 패치가 재구성이 더 쉽음을 나타냅니다.
• 마스킹 전략 vs 하위 작업 성능: 재구성 손실 추세는 하위 분류 성능으로 직접 전환되지 않았습니다. 마스킹 전략 간의 차이는 미미하고 데이터셋에 의존적이었습니다. 예를 들어, ADHD-200 데이터셋에서는 "Any"와 "Pure" 전략이 유사한 성능을 보인 반면, COBRE 에서는 "Any"와 "Majority"가 유사했습니다.
• 아키텍처 성능: Mamba–Attention (MA) 구성이 두 데이터셋 모두에서 일관되게 가장 높은 평균 AUROC 를 달성했습니다.
  • COBRE: Rhamba-MA(후두엽 다수) 는 71.32% 정확도와 68.44% AUROC를 달성하여 SwiFT(58.38% / 61.30%) 및 NeuroSTORM 변형들을 능가했습니다.
  • ADHD-200: Rhamba-MA(변연계 다수) 는 71.44% AUROC를 달성하여 SwiFT(69.06%) 와 NeuroSTORM 을 능가했으나, 정확도는 SwiFT 와 비교 가능했습니다.
• 영역별 분석: 최고 성능은 단일 지배적 영역보다는 마스킹 전략과 아키텍처 간의 상호작용에 의존했습니다. 예를 들어, COBRE 에서는 측두엽 영역이 MA 와 함께 가장 좋은 성능을 보인 반면, Limbic 영역은 ADHD-200 에서 MA 와 함께 최고 성능을 보였습니다.
• 해석 가능성: Attribution 맵은 각 장애에 대한 알려진 신경생물학적 프로필과 일치하는 영역을 강조했습니다 (예: ADHD 의 경우 DMN 및 전두정엽 네트워크; 조현병의 경우 DMN 및 시상피질 네트워크).

의의 및 주장
이 논문은 Rhamba 가 대규모 fMRI 표현 학습에서 해석 가능성, 확장성, 성능 간의 균형을 잡을 수 있는 유연한 프레임워크를 제공한다고 주장합니다. 주요 발견 사항은 다음과 같습니다:

• 재구성 vs 표현: 재구성 난이도는 표현의 질에 대한 신뢰할 수 있는 대리 지표가 아닙니다. 더 높은 재구성 손실을 유발하는 이질적 마스킹 (Any) 은 영역 간 통합을 장려함으로써 더 이전 가능한 특징의 출현을 촉진합니다.
• 하이브리드 우월성: MA 구성은 분산된 네트워크 상호작용과 장거리 시간적 의존성을 모두 포착하기 위해 전역 컨텍스트 모델링 (어텐션을 통해) 과 효율적인 시퀀스 역학 (Mamba 를 통해) 을 결합하는 것의 중요성을 강조합니다.
• 강건성: 영역 인식 사전 학습은 해부학적 정렬이 변하는 이질적인 다중 사이트 데이터셋 (ADHD-200 등) 에서도 효과적이며, 이는 프레임워크가 엄격한 해부학적 대응을 넘어 기능적으로 관련된 패턴을 포착함을 시사합니다.
• 미래 방향: 이 연구는 사전 학습 목표와 아키텍처 설계를 함께 고려해야 함을 시사합니다. 표현 학습 행동과 근본적인 신경생물학적 메커니즘 간의 명확한 구분이 존재한다면, 해부학적으로 정보화된 자기지도가 기능적으로 관련된 뇌 조직을 탐구하기 위한 원칙적인 도구로 작용할 수 있음을 제안합니다.