UNITES-Lab/Flex-MoE

NeurIPS 2024 (Spotlight)에 채택된 "Flex-MoE: Modeling Arbitrary Modality Combination via the Flexible Mixture-of-Experts"의 공식 구현체입니다.

• 저자: Sukwon Yun, Inyoung Choi, Jie Peng, Yangfan Wu, Jingxuan Bao, Qiyiwen Zhang, Jiayi Xin, Qi Long, Tianlong Chen

업데이트 (2025년 6월 9일):

• 🚀 이미지 양상 (modality)에 대한 더 쉽고 사용자 친화적인 지원을 요청하는 여러 건의 요청을 받았습니다. 데이터 로더 (dataloader)가 전처리된 이미지 버전을 지원하도록 업데이트되었습니다. 자세한 내용은 여기를 참조하세요: data/adni/README.md

멀티모달 학습 (Multimodal learning)은 다양한 분야에서 점점 더 중요해지고 있으며, 의료 분야에서 자주 관찰되는 이미지, 텍스트, 개인화된 기록과 같은 다양한 소스의 데이터를 통합할 수 있는 능력을 제공합니다. 그러나 일부 양상 (modalities)이 누락된 시나리오에서, 기존의 많은 프레임워크는 임의의 양상 조합 (arbitrary modality combinations)을 수용하는 데 어려움을 겪으며, 종종 단일 양상이나 완전한 데이터에 크게 의존합니다. 이러한 잠재적 양상 조합에 대한 간과로 인해 실제 상황에서의 적용 가능성이 제한됩니다. 이 과제를 해결하기 위해, 우리는 데이터 누락에 대한 강건성 (robustness)을 유지하면서 임의의 양상 조합을 유연하게 통합하도록 설계된 새로운 프레임워크인 Flex-MoE (Flexible Mixture-of-Experts)를 제안합니다. Flex-MoE의 핵심 아이디어는 먼저 관찰된 양상 조합과 그에 상응하는 누락된 양상을 통합하는 새로운 누락 양상 뱅크 (missing modality bank)를 사용하여 누락된 양상을 처리하는 것입니다. 그 다음은 독특하게 설계된 희소 MoE (Sparse MoE) 프레임워크가 이어집니다. 구체적으로, Flex-MoE는 먼저 모든 양상을 가진 샘플을 사용하여 전문가 (experts)를 학습시켜 일반화된 라우터 (generalized router)를 통해 일반화된 지식을 주입합니다 (

conda create -n flex-moe python=3.10 -y
conda activate flex-moe
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia -y
...

*자세한 내용은 data/adni/README.md를 참조하십시오

ADNI 데이터셋에 접근하려면 먼저 ADNI 웹사이트(https://adni.loni.usc.edu/)를 방문하여 다음 링크에서 데이터 접근 권한을 신청하십시오: https://ida.loni.usc.edu/collaboration/access/appApply.jsp?project=ADNI.

접근 권한을 얻으면, IDA에 로그인하여 각 양상 (Modality)별로 필요한 파일을 다운로드하십시오.

• 이미지 (Image):

• UCSF - Cross-Sectional FreeSurfer (7.x) (전처리 버전 - 업데이트됨!)

• {Subject ID}_{Scan date} (원본 버전)

• 유전체 (Genomic):

• ADNI 1 SNP 유전자형 데이터 - PLINK

• ADNI GO/2 SNP 유전자형 데이터 - 세트 1 - 9를 위한 전체 PLINK

• ADNI GO/2 SNP 유전자형 데이터 - 세트 10 - 15를 위한 전체 PLINK

• ADNI3 SNP 유전자형 데이터 세트 1 - PLINK

• ADNI3 SNP 유전자형 데이터 세트 2 - PLINK

• 임상 (Clinical):

• MEDHIST

• NEUROEXM

• PTDEMOG

• RECCMEDS

• VITALS

• 생체 시료 (Biospecimen):

• APOERES

• UPENNBIOMK_ROCHE_ELECSYS

• 이미지 (Image):

파일이 준비되면, 각 양상 (Modality)별 전처리 코드를 실행하십시오.

기존 데이터 로더 (Dataloader) 형식을 기반으로 다른 데이터셋에서 Flex-MoE를 실행하려면, 아래 두 부분을 커스텀하십시오.

main.py

• main.py의 114행과 같이, 데이터셋 전용 modality_dict를 정의하십시오.

• 양상 (Modality) 이름을 설정할 때는 (1) 소문자를 사용하고, (2) 각 양상의 첫 글자가 고유하도록 하십시오.

• 예시:
  if args.data == 'your_new': modality_dict = {'a_modality':0, 'b_modality': 1, 'c_modality': 2}

• data.py의 114행과 같이

• (1) 데이터셋에 대한 적절한 경로를 설정하고, (2) 각 모달리티 (modality)에 대해 고유한 전처리 (preprocessing) (예: 초기 채우기 (initial filling), MinMaxScaler)를 맞춤 설정해 주세요.

• 예시:
  def load_and_preprocess_new_data(args, modality_dict): # Paths a_modality_path = './data/your_dataset/a_modality' b_modality_path = './data/your_dataset/b_modality' c_modality_path = './data/your_dataset/c_modality' label_df = pd.read_csv('./data/your_dataset/your_label_csv_file', index_col='your_index') labels = label_df['your_label'].values.astype(np.int64) with open('./data/your_dataset/your_index_train_valid_test_split.json') as json_file: data_split = json.load(json_file) id_to_idx = {id: idx for idx, id in enumerate(label_df.index)} common_idx_list = [] observed_idx_arr = np.zeros((labels.shape[0], args.n_full_modalities), dtype=bool) # Load modalities (a_modality as an example) if 'A' in args.modality or 'a' in args.modality: if args.initial_filling == 'mean': path = a_modality_path + '_mean.csv' else: path = a_modality_path + '.csv' ...

실행 예시 (Run Example)

python main.py \
--data adni \
--modality IG \
...

IGCB 결과 (IGCB Result)

sh run_adni_igcb.sh

@misc{yun2024flexmoemodelingarbitrarymodality,
title={Flex-MoE: Modeling Arbitrary Modality Combination via the Flexible Mixture-of-Experts},
author={Sukwon Yun and Inyoung Choi and Jie Peng and Yangfan Wu and Jingxuan Bao and Qiyiwen Zhang and Jiayi Xin and Qi Long and Tianlong Chen},
...