[논문 리뷰] Self-Supervised Transformers for fMRI representation

Itzik Malkiel, Gony Rosenman외 2명|arXiv (Cornell University)|2021. 12. 10.

Functional Brain Connectivity Studies인용 수 10

한 줄 요약

이 논문은 4D fMRI 데이터를 라벨이 없는 상태에서 사전 훈련하여 3D 볼륨을 재구성하는 자기지도 학습 기반 Transformer 프레임워크인 TFF를 소개한다. 이후 유전자 기반 임상 과제에 대해 미세조정을 수행한다. TFF는 뇌 분할이 필요 없이 연령 및 성별 예측, 정신분열증 분류 과제에서 최신 기술 수준(SOTA) 성능을 달성하며, 다중 손실 재구성 전략을 통해 효과적인 이중 단계 사전 훈련 및 미세조정 전략 덕분에 기존의 기준 모델들을 크게 능가한다.

ABSTRACT

We present TFF, which is a Transformer framework for the analysis of functional Magnetic Resonance Imaging (fMRI) data. TFF employs a two-phase training approach. First, self-supervised training is applied to a collection of fMRI scans, where the model is trained to reconstruct 3D volume data. Second, the pre-trained model is fine-tuned on specific tasks, utilizing ground truth labels. Our results show state-of-the-art performance on a variety of fMRI tasks, including age and gender prediction, as well as schizophrenia recognition. Our code for the training, network architecture, and results is attached as supplementary material.

연구 동기 및 목표

• 자기지도 학습 사전 훈련을 통해 fMRI 과제에 대한 전이 학습을 가능하게 하는 딥 러닝 프레임워크를 개발하는 것.
• 제한된 라벨이 있는 데이터와 노이즈가 많은 라벨을 가진 고차원적이고 노이즈가 많은 fMRI 데이터에 대한 과제를 해결하는 것.
• 전체 4D fMRI 볼륨을 종단 간(end-to-end)으로 처리하여 뇌 분할이 필요 없도록 하는 것.
• 연령 예측, 성별 분류, 정신분열증 진단과 같은 임상 fMRI 과제의 성능을 향상시키는 것.
• 사전 훈련의 중요성을 검증하고 모델 아키텍처의 핵심 구성 요소를 분석하는 것.

제안 방법

• TFF는 이중 단계 훈련 전략을 사용한다: 먼저 원시 fMRI 볼륨에서 자기지도 학습 사전 훈련을 통해 3D 데이터를 재구성하고, 이후 라벨이 있는 후행 과제에 대해 미세조정을 수행한다.
• 모델은 3D fMRI 볼륨을 병렬 처리하는 데서 비롯된 볼륨의 시퀀스로 간주하고, 볼륨 내 각 볼륨을 토큰 표현으로 변환하기 위해 학습 가능한 임베딩 레이어를 사용한다.
• 다층 Transformer 인코더가 임베딩된 볼륨의 시퀀스를 처리하여 4D fMRI 스캔 전역의 장거리 공간적 및 시간적 상관관계를 포착한다.
• 사전 훈련은 L1, L2 및 감각적 손실을 조합한 다중 손실 목적 함수를 사용하여 복원된 출력에서 원래 fMRI 볼륨을 재구성한다.
• 훈련 안정성 향상과 수렴성 향상을 위해 인코더는 볼륨별 정규화와 전역 정규화로 초기화된다.
• 미세조정 단계에서는 사전 훈련된 인코더가 각 fMRI 프레임을 별도로 처리하고, 출력 결과에 [CLS] 토큰을 연결하여 시퀀스를 분류를 위해 Transformer에 통과시킨다.

실험 결과

연구 질문

• RQ1 자기지도 학습 사전 훈련 절차가 후행 과제에서 fMRI 데이터의 표현 학습에 기여하는가?
• RQ2 전체 4D fMRI 볼륨에 대한 종단 간 훈련 방식이 분할 기반 방법과 비교해 성능 및 일반화 능력에서 어떻게 다른가?
• RQ3 다중 손실 재구성 목적 함수의 각 구성 요소가 모델 성능과 수렴성에 기여하는 정도는 어떠한가?
• RQ4 사전 훈련이 소규모이고 노이즈가 많은 fMRI 데이터셋, 특히 병리학적 분류 과제에서 성능 향상에 크게 기여하는가?
• RQ5 전체 모델의 단일 단계 훈련과 비교해 이중 단계 사전 훈련 전략은 어떻게 성능에 영향을 미치는가?

주요 결과

• TFF는 연령 및 성별 예측 과제에서 최신 기술 수준의 성능을 달성하며, ST-GCN 및 DeepfMRI와 같은 기존 방법들을 능가한다.
• COBRE 정신분열증 데이터셋에서 TFF는 AUC 87.9를 기록하여 ST-GCN 대비 4점 향상되고, DeepfMRI 대비 8.3점 향상되었다.
• TFF는 TFF_vanilla 기준으로 NMSE를 약 26% 감소시켰으며, L1 손실은 절대 개선 0.4점, L2 손실은 3.2점 개선되었다.
• 사전 훈련을 생략한 TFF_vanilla 모델은 정신분열증 분류 과제에서 수렴에 실패하여 사전 훈련의 필수성을 입증한다.
• 절단 분석 결과, L1, L2 및 감각적 손실 모두와 볼륨별 정규화가 모델 수렴성과 성능 향상에 필수적임을 확인하였다.
• 학습 데이터를 감소시켜도 TFF는 모든 데이터 크기에서 뛰어난 성능 유지를 보이며, 강건성과 일반화 능력을 입증하였다.