[논문 리뷰] Meta Transfer of Self-Supervised Knowledge: Foundation Model in Action for Post-Traumatic Epilepsy Prediction

Wenhui Cui, Haleh Akrami외 3명|PubMed|2023. 12. 21.

Neonatal and fetal brain pathology참고 문헌 27인용 수 4

한 줄 요약

이 논문은 건강한 fMRI 데이터에서의 메타학습과 자기지도 사전 훈련을 조합한 기초 모델인 메타 전이 자기지도 지식(MeTSK)을 제안한다. 이는 부족하고 이질적인 임상 데이터에 대해 일반화 능력을 향상시켜 외상성 뇌손상(TBI) 이후 간질(PTE) 예측에 기여한다. 정상에서 임상 특징으로 지식을 전이하기 위해 이중 최적화를 활용함으로써 MeTSK는 새로운 TBI 데이터셋에서 최신 기술 수준의 제로샷 성능을 달성하였으며, AUC는 0.6981 ± 0.0409를 기록하였다.

ABSTRACT

Despite the impressive advancements achieved using deep-learning for functional brain activity analysis, the heterogeneity of functional patterns and scarcity of imaging data still pose challenges in tasks such as prediction of future onset of Post-Traumatic Epilepsy (PTE) from data acquired shortly after traumatic brain injury (TBI). Foundation models pre-trained on separate large-scale datasets can improve the performance from scarce and heterogeneous datasets. For functional Magnetic Resonance Imaging (fMRI), while data may be abundantly available from healthy controls, clinical data is often scarce, limiting the ability of foundation models to identify clinically-relevant features. We overcome this limitation by introducing a novel training strategy for our foundation model by integrating meta-learning with self-supervised learning to improve the generalization from normal to clinical features. In this way we enable generalization to other downstream clinical tasks, in our case prediction of PTE. To achieve this, we perform self-supervised training on the control dataset to focus on inherent features that are not limited to a particular supervised task while applying meta-learning, which strongly improves the model's generalizability using bi-level optimization. Through experiments on neurological disorder classification tasks, we demonstrate that the proposed strategy significantly improves task performance on small-scale clinical datasets. To explore the generalizability of the foundation model in downstream applications, we then apply the model to an unseen TBI dataset for prediction of PTE using zero-shot learning. Results further demonstrated the enhanced generalizability of our foundation model.

연구 동기 및 목표

• 외상성 뇌손상(TBI) 직후에 수집된 제한적이고 매우 다양성이 큰 fMRI 데이터로부터 후천성 간질(PTE)을 예측하는 데 도전한다.
• 임상 데이터가 건강한 대조군에 비해 부족한 상황에서, 불균형 데이터로 훈련된 표준 기초 모델의 열악한 일반화 능력을 개선한다.
• 업스트림 사전 훈련 기간 동안 메타학습을 통해 정상 fMRI 패tern에서 임상적으로 관련 있는 특징으로 일반화 능력을 향상시킨다.
• 재편조 없이도 효과적인 제로샷 및 소수의 샘플에 기반한 적응을 가능하게 하여, 특히 재편조 없이도 PTE 예측과 같은 후행 임상 과제에 적용할 수 있도록 한다.

제안 방법

• 대규모 건강한 대조군 fMRI 데이터셋에서 대비 자기지도 사전 훈련을 통해 임무에 종속되지 않은 일반화 가능한 특징을 학습하기 위해 기초 모델을 사전 훈련한다.
• MAML에 영감을 받은 이중 최적화 기반의 메타학습을 통합하여, 건강한 특징에서 임상 특징으로 일반화하는 방법을 학습한다.
• 단일 업데이트 스텝 이후 소규모 임상 데이터셋에서 내부 루프 손실을 최소화하는 메타목표를 통해 모델이 임상 데이터에 신속하게 적응하도록 훈련한다.
• 결과로 도출된 MeTSK 모델을 사용하여 재편조 없이도 새로운 TBI/PTE 데이터셋에서 제로샷 추론을 수행한다. 이때 직접 특징을 추출한다.
• PTE 분류에 대한 후행 성능 평가를 위해 선형 프로빙과 제로샷 특징 평가를 적용한다.
• 원천(대조군)과 대상(임상) 도메인 간의 일반화 정도를 정량화하기 위해 도메인 유사도 분석을 수행하며, 훈련/테스트 분할 내부의 도메인 간 유사성도 분석한다.

실험 결과

연구 질문

• RQ1 메타학습은 건강한 fMRI 패턴에서 부족한 데이터셋의 임상적으로 관련 있는 특징으로 자기지도 기초 모델의 일반화 능력을 향상시킬 수 있는가?
• RQ2 새로운 TBI 데이터셋에서 MeTSK는 표준 재편조 및 기준 자기지도 모델 대비 제로샷 PTE 예측에서 어떻게 성능을 냈는가?
• RQ3 메타학습과 자기지도 학습의 통합이 이질적인 fMRI 데이터에서 도메인 이동과 내부 클래스 분산을 어느 정도 감소시키는가?
• RQ4 MeTSK에서 학습된 특징은 PTE와 관련된 뇌 네트워크 변화에 대한 해석 가능한 통찰을 제공할 수 있는가?

주요 결과

• MeTSK는 새로운 TBI 데이터셋을 사용한 제로샷 PTE 예측에서 0.6981 ± 0.0409의 최고 AUC를 기록하여 모든 기준 모델을 압도하였다.
• 모델은 도메인 간 일반화 능력 향상을 보였으며, HCP 대조군 데이터셋에서 유도된 자기지도 특징가 임상 특징과 더 높은 유사도를 보였다.
• 메타학습은 ADHD 및 ABIDE 데이터셋에서 훈련 및 테스트 세트 간 내부 클래스 도메인 유사도를 크게 향상시켜 데이터 이질성의 영향을 줄였다.
• MeTSK 특징에 대해 선형 프로빙을 수행한 결과, 재편조보다 뛰어난 성능을 보였으며, 이는 제로샷 특징가 분포 이동 상황에서도 더 견고함을 시사한다.
• 도메인 유사도 분석 결과 MeTSK 특징는 도메인 간 및 클래스 내부에서 더 나은 일반화 능력을 보였으며, 이는 후행 성능 향상의 이유를 설명한다.
• 재편조 없이도 일반화 능력을 유지를 통해, 라벨이 제한된 희귀 신경학적 질환에 대한 데이터 효율적인 배포 잠재력을 보여주었다.