Scaling Vision Transformers for Functional MRI with Flat Maps

Motivation

• Known: fMRI foundation model 연구에서는 주로 parcellation 기반(고효율이지만 정보 손실이 큼) 또는 volume 기반(정보 보존이 완전하지만 계산 비용이 큼) 표현을 사용해 self-supervised 학습을 수행해왔으며, BrainLM, Brain-JEPA 등 여러 선행 모델이 존재한다.
• Gap: 기존 연구들은 parcellation과 volume이라는 양 극단의 표현만을 탐색했을 뿐, 두 표현 사이의 절충안이 되는 표현(flat map)은 시도되지 않았으며, 또한 fMRI foundation model 분야에는 재현 가능한 공개 평가 벤치마크가 부재하여 모델 간 공정한 비교가 어려웠다.
• Why: 데이터-계산 자원이 제한된 현재 체제에서 정보 손실과 계산 효율 사이의 "goldilocks zone"을 찾는 것은 fMRI foundation model의 실질적 확장성과 성능을 좌우하는 핵심 문제이며, 표준화된 벤치마크는 향후 연구의 재현성과 공정한 비교를 가능케 한다는 점에서 중요하다.
• Approach: 저자들은 fMRI 볼륨을 cortical flat map으로 투영하여 표준 spacetime ViT patch embedding을 그대로 적용할 수 있는 CortexMAE(MAE-st 기반)를 제안하고, parcellation·volume·flat map 세 가지 표현을 동일 프레임워크에서 직접 비교하는 동시에 Brainmarks라는 공개 평가 스위트를 구축해 체계적인 스케일링 및 벤치마크 분석을 수행하였다.

Achievement

1. CortexMAE 모델 패밀리 구축: HCP의 2.1K시간 오픈 fMRI 데이터로 학습된 flat map, parcellation, volume 기반의 최초의 다중 표현 fMRI foundation model 패밀리를 제시하였다.
2. Brainmarks 벤치마크 공개: 7개 공개 데이터셋에 걸쳐 subject-level trait prediction과 cognitive state decoding을 포함하는 최초의 공개 fMRI foundation model 평가 스위트를 제공하였다.
3. 표현 방식 비교 및 최적 영역 발견: flat map이 state decoding에서, volume이 age prediction에서, parcellation이 효율성 측면에서 각각 우수함을 실험적으로 규명하였다.
4. 최초의 체계적 fMRI 스케일링 분석: 엄격한 power law scaling을 관측했으나 out-of-distribution 일반화가 약하고 모델 크기에 한계가 있음을 발견하였다.
5. 도전적인 null result 보고: subject-level trait prediction에서 어떤 모델도 명확한 state-of-the-art를 달성하지 못했고, 단순 functional connectivity baseline을 능가하기 어렵다는 결과를 제시하였다.
6. state decoding에서의 큰 폭 성능 향상: cognitive state decoding에서는 CortexMAE 계열, 특히 flat map 변형이 기존 모델들을 큰 격차로 능가함을 보였다.

How

Figure 2. MAE model architecture (He et al., 2022) adapted to

• 표준 surface-based 전처리 파이프라인(FreeSurfer, HCP pipeline, fMRIPrep 등)으로 fMRI를 그룹 템플릿 cortical surface mesh에 매핑
• pycortex를 이용해 surface 데이터를 flat surface mesh로 복사 후 정규 이미지 그리드로 리샘플링하여 2D flat map 시계열(video) 생성
• 배경(all-zero) 영역은 패치에서 제외하고 유효 픽셀에 대해서만 MSE loss 계산
• MAE-st(He et al., 2022; Feichtenhofer et al., 2022) 프레임워크를 적용해 encoder-decoder ViT를 spacetime patch(pt×p×p) 단위로 마스킹·복원 학습
• parcellation 버전은 Schaefer-400 parcellation을 사용해 각 parcel 시계열을 time-only patch로 독립적으로 임베딩(Caro et al., Dong et al. 방식 준용)
• volume 버전은 3D 큐브 패치를 직접 사용
• 세 가지 표현 방식으로 각각 CortexMAE를 학습하고, 모델 크기를 변화시키며 systematic scaling 분석 수행
• Brainmarks를 통해 trait prediction과 state decoding 과제에 대해 controlled benchmark 비교 수행

Originality

• 3D fMRI 볼륨을 cortical flat map으로 변환하여 표준 ViT/MAE 아키텍처를 수정 없이 그대로 적용한 최초의 시도
• parcellation·volume·flat map 세 표현을 동일한 모델 프레임워크에서 직접적, 통제된 비교를 수행한 최초의 연구
• fMRI foundation model에 대한 최초의 체계적 scaling law 분석
• 재현 가능한 공개 벤치마크(Brainmarks)를 제공하여 fMRI foundation model 평가의 표준화를 시도한 최초의 오픈소스 평가 스위트

Limitation & Further Study

• Flat map 변환 과정에서 surface-based 전처리에 의존하므로, 전처리 파이프라인의 품질과 표준 템플릿 정합 정확도에 성능이 좌우될 수 있음
• Scaling 분석에서 관측된 hard limit과 약한 OOD 일반화의 원인에 대한 심층 분석이 부족하며, 후속 연구에서 이를 해결하기 위한 아키텍처 개선이나 데이터 다양성 확대가 필요함
• Subject-level trait prediction에서의 null result는 deep learning 기반 fMRI foundation model 접근법 자체의 한계를 시사하며, 왜 단순 FC baseline을 능가하지 못하는지에 대한 원인 규명이 추가로 필요함
• 피질(cortex) 표면 기반 표현이므로 subcortical 구조를 포함하지 못하는 근본적 한계가 있어, 이를 보완하는 확장된 표현 방식 연구가 요구됨

Evaluation

총평: flat map이라는 간단하면서도 효과적인 아이디어로 fMRI에 표준 ViT/MAE를 직접 적용 가능하게 만들고, 공개 벤치마크까지 함께 제공함으로써 fMRI foundation model 연구의 재현성과 발전에 실질적으로 기여하는 견고한 연구이다.