[ICML 2024] SleepFM: Multi-modal Representation for Sleep Across Brain Activity, ECG and Respiratory Signals

🖋️ 이 논문은 Standford 대학교에서 직접 수집한 14,000명의 환자 데이터를 가지고 Sleep Foundation Model을 구축한 내용을 담고 있다. ICML에 2024년 채택되었으며, 코드도 공개되어있다: https://github.com/rthapa84/sleepfm-codebase. Biomedical AI 분야에서 의학 저널이 아닌 AI 학회에 개제되었다는 점이 인상깊었고, 이를 위해 어떠한 방법론과 novelty를 가지고 연구를 진행했는지 살펴보면 추후에 나도 AI 학회에 논문을 제출할 때 도움이 많이 될 것 같아 리뷰해 보았다.

1. Introduction

1. Sleep 연구에서 gold standard로 쓰이는 PSG는 크게 세 가지 modality로 나눌 수 있다:
   • Brain Activity Signals (BAS) : electroencephalogram (EEG - 뇌파도), electrooculograms (EOG - 안전도), electroencephalogram (EMG - 근전도). 보통 수면 단계를 판단할 때 가장 많이 사용된다. 총 10개 채널.
   • Electrocardiogram (ECG, 심전도) : 심박동의 주기중에 일어나는 심장의 전기적 활동 상태를 측정. sleep disordered breathing events (SDB) 검출에 사용될 수 있음. 총 2개 채널.
   • Respiratory sensors : 가슴, 복부 움직임, 심박, 비강 흐름(nasal flow), 구강 흐름(oral flow)를 포함. SDB 검출에 직접적으로 활용됨. 총 7개 채널.

2. 기존 수면 연구들의 한계:
   • 수면 판독 자동화는 labeled data에 한정적으로 연구가 이루어짐.
   • 한 가지 태스크에만 적용 가능한 모델

3. Contrastive Learning
   • 수면 분야에서 CL을 활용한 연구: single channel ECG, ECG + electronic health records (EHR)
   • multi-modal CL 접근 방법으로 PSG 데이터를 분석한 것은 본 연구가 최초이다.

4. Contribution
   • Stanford Sleep Clinic에서 14,000명의 환자에게서 수집한 100,000 시간의 데이터를 활용한 최초의 contrastive learning 기반의 foundation model이다.
   • 인구통계학적 정보 (나이, 성별), 수면 단계 분류, 수면 호흡 장애 이벤트 검출의 태스크에서 SleepFM이 baseline (end-to-end CNN model)보다 좋은 성능을 보였다.
   • Leave-one-out CL 기법을 새롭게 도입해서 Pairwise CL보다 downstream task에서 성능이 좋음을 증명하였다.

2. Related Work

1. Machine Learning for Analyzing Sleep Data
   • Autoencoders, CNNs, RNNs, DNNs 등이 수면 단계 분류에 쓰임.
   • 호흡 이벤트 검출 태스크에서는 ECG, EEG와 respiratory channels 등이 주로 쓰임. multi-modal을 사용하여 (EEG, EOG, EMG) multi-task (e.g. sleep stages, arousal, leg movements, and sleep-disordered breathing) learning 모델도 있었지만, 모두 supervised learning 이었다.
     

2. Contrastive Learning
   • 컴퓨터 비전에서 시작되어 발전한 self-supervised learning 태크닉. InfoNCE, SimCLR, MoCo, SupCon 등이 있음. 이들은 대부분 image-based의 uni-modal contrastive appraoch 라고 볼 수 있음.
   • 그에 반해 Contrastive Language-Image Pretraining (CLIP) 은 image와 text embedding을 활용한 multi-modal model.
   • ConVIRT: chest radiographs 와 report 를 활용 (multi-modal)
   • 컴퓨터 비전 외 수면 의학계에서 쓰인 multi-modal contrastive learning: ECG + structured records / ECG + EHR + clinical notes
   • SleepFM은 PSG 데이터의 19개 channel을 활용한 최초의 multi-modal contrastive model이다.

3. Method

3.1. Dataset and Preprocessing

• 30-second 단위 - 본 논문에서는 clip이라고 칭함.
• resampled to 256 Hz
• Sleep Stage: Wake, Stage1, Stage2, Stage3, REM 분류 문제
• SDB: binary label
• 전체 데이터셋을 다음과 같이 나눔: pretrain/train/validation/test as (11,261/1,265/141/1,401) - 환자 기준
• pretrained dataset은 Foundation model을 pretrain 할 때에만 쓰임.

3.2. Embedding Model

• 3개의 EfficientNet architecture를 활용한 multi-channel 1D convolution이 각각 modality의 embedding model (Encoder)로 사용됨.
• Depthwise separable convolution, Dropout layer, Residual Connection이 사용됨.

3.3. Multi-modal Contrastive Learning

• postive pair 기준: 같은 시간의 different modality - temporally aligned 30-second clips across modality
• negative pair: 그 외 training batch안의 non-matching instances는 모두 negative pair로 학습.
• 2가지 Contrastive Learning Types (+ 여기에서 새롭게 제시한 Leave-one-out CL)
  • Pairwise CL: 3개의 modality 중 2개를 고르는 조합으로 각각 contrastive loss를 계산, 최종적으로 모든 조합에서 나온 loss를 합쳐서 사용.
  • Leave-one-out CL: 3개의 modality 중 2개의 embedding을 평균 낸 것과 다른 하나의 modality를 pair로 사용.

3.4. Model Training

• Baseline model: 1D EfficientNet 구조를 사용한 CNN 기반 모델
• Downstream task를 위해 modality encoder에서 각각 training, validation, test set의 embedding을 뽑는다. 이후 Logistic Regresion classifier를 사용하여 sleep stage와 SDB 태스크를 학습시킨다.
• 코드를 보니 Logistic Regression으로 multi-class training 은 OVR(one vs. rest) 방식으로 하였음. XGBClassifier로도 테스트 해본 것 같은데 Logistic Regression이 더 잘 나와서 후자를 선택한 것 같음.

💡 나의 궁금증: 왜 Foundation Model에서는 Ridge/Logistic Regression과 같은 단순한 머신러닝 기반의 linear probing을 사용하는 것일까? Deep Neural Network (FC layer 혹은 CNN layer) 등을 사용하면 안되나? 아래는 ChatGPT의 답변이다.

• 우리는 Foundation model의 embedding이 얼마나 informative 한지, linearly separable 한 것이 궁금함.
• 만약 DNN을 사용한다면:
  • Overfit to small downstream data
  • Hide the weakness of poor embeddings by adding capacity 등의 문제가 있을 수 있음.
• 따라서 logistic/ridge regression은 model capacity가 변수가 될 수 있는 가능성을 없애므로, foundation model의 embedding quality에만 의존하여 분류 문제를 풀 수 있도록 한다.

4. Experiments and Results

🖋️ 논문의 하이라이트인 실험 파트이다. 다양한 다운스트림 태스크에서 SleepFM 월등한 성능을 보여주고, multi-modality, few-shot 성능, public dataset을 활용한 external validation을 하였다.

4.1. Demographic Attributes Classification

1. 나이를 총 4개의 그룹으로 나누어서 맞추는 분류 문제:

• 
• Leave-One-Out 방식이 성능이 가장 좋았다.

2. 성별 남 vs. 여 분류 문제:

• 

• 마찬가지로 Leave-One-Out 방식이 성능이 가장 좋았다.

4.2. Retrieval Analysis

• 하나의 modality embedding을 가지고 cosine similarity를 활용하여 가장 가까운 embedding을 찾기
• Top10 중에서 정답 pair가 있으면 count
• 총 90,000개의 랜덤하게 선택된 clip을 가지고 시행.
• 여기에서 random하게 10개를 골랐을 때 Recall@10는 10/90000 = 0.0001 (이 값보다 커야 함)

  

• 결과: Pairwise CL이 월등히 성능이 좋았음. 문제 자체가 pairwise cosine similiarty를 계산하는 것이기 때문에 training 시에 pairwise로 훈련시키는 것이 더 효과적이었을 것.
• 다만 Respiratory로 다른 modality를 retrieval 하는 태스크에서는 비교적 성능이 좋지 않았음. 그 이유로는 BAS나 ECG는 직접적인 전기 자극으로 측정하는 것인데 Repiratory channel은 간접적인 움직임을 기록하는 것이기 때문에 몸의 자세나 non-breathing related motion 등에 영향을 받을 수 있기 때문이라고 설명.

4.3. Downstream Classification Tasks

1. Sleep Stage classification 성능
   

2. SDB classification 성능 (Binary)

   

• Supervised CNN: pretraining + training dataset을 모두 활용하여 훈련
• SleepFM (Leave-On-Out / Pairwise): pretraining으로 FM 모델 훈련, training data로 linear probing.
• 결과: SleepFM 중 Leave-One-Out 이 성능이 좋았다.
• 추가적으로, 각각의 single modality embedding만을 가지고 위 태스크를 진행. BAS는 sleep stage classification에서 좋은 성능을, repiratory signal은 SDB event detection에서 좋은 성능을 보여주었다. 각각의 single modality만을 가지고도 어느 정도의 성능은 나오더라.
• 서로 다른 age와 gender group안에서도 각각의 태스크 성능을 비교해 보았을 때, 큰 차이가 없었다.

4.4. Few-Shot Evaluation

• few shot의 성능을 비교하기 위해 model이 보는 환자의 수를 k=1 부터 최대(1265명)까지 늘려가며 비교하였음.
• Supervised CNN: few-shot example만 가지고 훈련
• SleepFM: 이미 pretrained 된 모델에 few-shot만 가지고 logistic regression 모델 훈련
• 예상할 수 있듯이, 결과는 SleepFM, Leave-One-Out이 가장 좋았다.
  

4.5. Benefit of Multi-Modal Pretraining

• 3개의 modality를 전부 사용하는 것이 가장 좋은 성능을 보였음.
• 2 modality와 single modality를 가지고 pretraining 을 통해 Foundation Model을 학습
• 이후 3개, 2개, 1개의 modality를 가지고 학습한 Foundation Model에서 각각 BAS embedding을 뽑아 sleep stage classification을 위한 logistic regression 을 학습, 동일하게 SDB detection에서는 respiratory embedding만 뽑아서 학습하여 성능을 비교하였다.
  
• Sleep Stage에서는 BAS가 가장 연관성이 높으므로 BAS / BAS+RESP / BAS+ECG / BAS+ECG+RESP 를 활용한 Foundation Model 성능을 비교하였고,
• SDB에서는 Respiratory signal이 가장 연관성이 높으므로 RESP / RESP + BAS / RESP + ECG / RESP + BAS + ECG 를 활용한 Foundation Model 성능을 비교하였음.
• Few-shot evaluation으로 training dataset에 쓰인 환자 수를 점차 늘려가며 성능 비교
• 결론적으로 2개의 modality를 결합한 BAS-ECG 모델과 RESP-ECG 모델의 성능이 3개 modality를 사용한 것과 비슷한(혹은 더 좋은) 성능을 보였는데, 이를 통해 ECG signal이 좀더 pretraining 시에 representation을 더욱 풍부하게 만들어준다는 것을 알 수 있다.
• 또한 BAS나 RESP를 단독으로 pretraining 에 사용시에는 성능이 지속적으로 좋지 않았다.

4.6. External Validation

• pretraining stage에서는 Stanford 에서 수집한 데이터셋을 활용하고, downstream task에서는 PhysioNet2018 데이터를 활용하여 generalizability 를 보여주었다.
• 비교: Supervised-CNN은 external dataset으로만 supervised learning 을 진행하였음.
• 100개의 test set에 대한 성능 비교 결과: SleepFM이 Supervised CNN보다 좋았다.
• 즉, domain이 다른 병원의 데이터를 하나도 보지 않은 pretrained foundation model이 새로운 site의 데이터셋에도 adaptation할 수 있다는 것을 보였다.
  

5. Discussion and Conclusion

Future Work

• Standford에서 수집한 one institution sleep data만을 가지고 실험을 진행하였다. -> 좀 더 다양한 site를 가지고 generalizability를 측정하는 것이 유의미할 것.
• Downstream task로 sleep stage와 SDB detection만을 사용했지만, arousal detection, periodic leg movements, disease 등을 분류하는 태스크도 가능할 것이다.
• Contrastive Learning을 제외하고도 다른 SSL method를 활용하여 성능을 비교하는 것.