[논문 리뷰] SwiFT: Swin 4D fMRI Transformer

P. Kim, Junbeom Kwon외 7명|arXiv (Cornell University)|2023. 07. 12.

Functional Brain Connectivity Studies인용 수 10

한 줄 요약

SwiFT는 원시 4D fMRI 데이터로부터 엔드 투 엔드 시공간 표현을 학습하기 위해 4D Swin Transformer를 도입하여 대규모 데이터셋에서 성별, 연령, 그리고 인지 지능을 효율적으로 예측 가능하게 하며, 이점인 자기지도 사전학습 및 해석 가능한 통찰을 제공합니다.

ABSTRACT

Modeling spatiotemporal brain dynamics from high-dimensional data, such as functional Magnetic Resonance Imaging (fMRI), is a formidable task in neuroscience. Existing approaches for fMRI analysis utilize hand-crafted features, but the process of feature extraction risks losing essential information in fMRI scans. To address this challenge, we present SwiFT (Swin 4D fMRI Transformer), a Swin Transformer architecture that can learn brain dynamics directly from fMRI volumes in a memory and computation-efficient manner. SwiFT achieves this by implementing a 4D window multi-head self-attention mechanism and absolute positional embeddings. We evaluate SwiFT using multiple large-scale resting-state fMRI datasets, including the Human Connectome Project (HCP), Adolescent Brain Cognitive Development (ABCD), and UK Biobank (UKB) datasets, to predict sex, age, and cognitive intelligence. Our experimental outcomes reveal that SwiFT consistently outperforms recent state-of-the-art models. Furthermore, by leveraging its end-to-end learning capability, we show that contrastive loss-based self-supervised pre-training of SwiFT can enhance performance on downstream tasks. Additionally, we employ an explainable AI method to identify the brain regions associated with sex classification. To our knowledge, SwiFT is the first Swin Transformer architecture to process dimensional spatiotemporal brain functional data in an end-to-end fashion. Our work holds substantial potential in facilitating scalable learning of functional brain imaging in neuroscience research by reducing the hurdles associated with applying Transformer models to high-dimensional fMRI.

연구 동기 및 목표

• 고차원 4D fMRI에서 ROI 기반 전처리 없이 뇌 다이나믹스를 더 잘 포착하기 위해 엔드투엔드 학습을 직접 수행하도록 동기를 부여합니다.
• 메모리 및 계산 효율이 높은 로컬 윈도우 어텐션을 갖춘 4D Swin Transformer를 fMRI 용으로 개발합니다.
• 대규모 데이터(HCP, ABCD, UKB)에서 엔드투엔드 SwiFT가 성별, 연령 및 지능 예측 성능을 향상시킨다는 것을 보여줍니다.
• 대상 fMRI 작업에 대한 대조적 자기지도 사전 학습의 타당성과 이점을 입증합니다.
• 예측에 기여하는 뇌 영역을 식별하기 위한 해석 가능성 분석을 제공합니다.

제안 방법

• fMRI 부피의 시간 및 3D 공간 차원에서 작동하도록 Swin Transformer를 4D로 확장합니다.
• 효율적인 국소 상호작용을 위해 4D 윈도우 기반 자기어텐션(4D W-MSA)과 4D 시프트 윈도우 어텐션(4D SW-MSA)을 사용합니다.
• 세 공간 차원에서 패치 분할과 패치 병합을 구현하되 시간 차원은 그대로 유지합니다.
• 각 단계 뒤에 더해지는 절대 4D 포지셔널 임베딩을 채택하여 공간 및 시간 좌표를 부호화합니다.
• 전체 토큰 간의 상호작용을 위한 최종 글로벌 어텐션 스테이지로 엔드투엔드 학습을 가능하게 합니다.
• 대상 성능 향상을 위해 인스턴스 대비 손실과 로컬-로컬 시간적 대비 손실의 두 가지 대조적 자기지도 사전 학습 목표를 활용합니다.
• 고정된 4D Swin Transformer 백본과 최종 MLP 헤드를 사용한 가중치 효율적 학습을 수행합니다.

(a) Overall architecture of SwiFT

실험 결과

연구 질문

• RQ1 엔드 투 엔드 4D Swin Transformer가 원시 fMRI 데이터로부터 시공간 뇌 다이나믹스를 효과적으로 학습할 수 있는가?
• RQ2 SwiFT가 ROI 기반 및 2단계 Transformer/CNN 기반 기초모형을 대규모 데이터셋에서 성별 분류 및 연령/지능 예측에서 능가하는가?
• RQ3 대조적 자기지도 사전 학습이 SwiFT의 다운스트림 fMRI 예측 작업을 개선할 수 있는가?
• RQ4 해석 가능한 기여도에 따라 성별 분류에 가장 크게 기여하는 뇌 영역은 어디인가?
• RQ5 TFF와 같은 기존 4D fMRI 모델과 비교하여 SwiFT의 효율성(매개변수, FLOPs, 처리량)은 어떤가?

주요 결과

• SwiFT는 HCP, ABCD, UKB 데이터셋 전반에서 성별 분류 및 연령/지능 예측에서 최근 기초모형을 지속적으로 능가했다.
• 인스턴스 대비와 로컬-로컬 시간적 대비 손실을 포함한 자기지도 사전 학습은 다운스트림 성능을 향상시킬 수 있으며, 데이터셋과 작업에 따라 효과가 다르게 나타난다.
• 통합 그래디언트 기반 해석은 성별 차 differentiations 문헌과 일치하는 뇌 영역(mPFC, PCC, 전두상회(precuneus)) 등을 식별하고 연령대에 따라 다른 영역들을 나타낸다.
• SwiFT는 글로벌 어텐션 기반 트랜스포머(TFF)보다 매개변수 및 계산 효율이 더 높으면서 예측 성능도 더 좋음을 보여준다.
• 이 모델은 원시 4D fMRI 데이터에서 엔드투엔드 학습을 지원하여 ROI 기반 특징 추출과 2단계 학습 파이프라인의 필요성을 줄인다.
• 더 긴 입력 시간 시퀀스가 일부 작업(특정 코호트의 지능 등)에서 성능을 향상시킬 수 있지만 효과는 작업 및 데이터세트에 따라 다르다.

(b) Successive 4D Swin Transformer Blocks