[인턴] EEG 신호로 수면 단계 분류하기: Chronos와 RAG를 활용한 실험

들어가며

문제 정의: 수면 단계 분류가 왜 어려운가?

수면 단계란?

• Wake (깨어있음): 의식이 있는 상태

• N1 (얕은 수면 1단계): 잠이 들기 시작

• N2 (얕은 수면 2단계): 본격적인 수면

• N3 (깊은 수면): 서파수면, 회복에 중요

• REM (렘수면): 꿈을 꾸는 단계

기술적 도전 과제

1. 개인차가 엄청납니다: 두개골 두께, 전도도, 개인의 생리학적 특성 때문에 같은 수면 단계여도 신호 모양이 사람마다 다릅니다.

2. 데이터가 복잡합니다: 100Hz 샘플링 레이트로 30초면 3,000개의 데이터 포인트입니다.

3. 라벨링 비용이 높습니다: 전문가가 일일이 판독해야 하죠.

핵심 아이디어: Time-Aware RAG 시스템

전체 아키텍처

EEG 신호
  ↓
Chronos Encoder (시계열 임베딩 추출)
  ↓
512차원 벡터
  ↓
Two-Stage Retrieval (유사한 수면 패턴 검색)
  ↓
LLM In-Context Learning (검색된 예시로 분류)
  ↓
수면 단계 예측

1. Chronos로 EEG 신호를 dense vector로 변환

2. 유사도 검색으로 비슷한 과거 사례를 찾고

3. LLM이 그 사례를 참고해서 최종 판단

왜 Chronos인가?

• 사전학습: 엄청나게 많은 시계열 데이터로 학습됨

• Zero-shot Transfer: 의료 신호를 따로 학습 안 해도 됨

• 고정된 임베딩: 512차원 벡터로 시계열을 표현

실험 과정: 단계별 구현

1단계: 데이터 준비

• 형식: Parquet (CSV보다 빠르고 압축률 좋음)

• 샘플링: 100Hz (1초에 100개 샘플)

• 전처리: Z-Score 정규화 (평균 0, 표준편차 1)

# 30초 에폭 = 100Hz × 30s = 3000 샘플
context_tensor_cpu = torch.tensor(
    segment['target'].values,  # Z-score normalized EEG
    dtype=torch.float32
).unsqueeze(0)

2단계: Chronos 임베딩 추출

# ❌ 이렇게 하면 에러!
input_ids = pipeline.tokenizer.context_input_transform(tensor_gpu)

# ✅ 이렇게 해야 함
input_ids, attention_mask, scale = pipeline.tokenizer.context_input_transform(
    context_tensor_cpu  # CPU 텐서로!
)
# 그 다음에 GPU로 이동
input_ids = input_ids.to(device)
attention_mask = attention_mask.to(device)

# (batch, seq_len, 512) → (batch, 512)
pooled = torch.mean(embeddings, dim=1)

3단계: Two-Stage Retrieval (핵심!)

Stage 1: User Selection (사용자 선택)

# 각 사용자의 "프로필" = 모든 에폭의 평균 벡터
user_profile = np.mean(user_embeddings, axis=0)

# 현재 사용자와 가장 비슷한 사용자 찾기
similarity = cosine_similarity(target_profile, source_profiles)
best_user = argmax(similarity)

Stage 2: Instance Selection (에폭 선택)

# 선택된 사용자의 에폭들 중 가장 유사한 것 찾기
similar_epochs = find_top_k_similar(current_epoch, best_user_epochs, k=3)

4단계: LLM In-Context Learning

당신은 수면 단계 분류 전문가입니다.
현재 사용자(SC4001)와 매우 유사한 뇌파 패턴을 가진
사용자(SC4002, 유사도=0.95)의 데이터를 검색했습니다.

검색된 예시:
1. [유사도=0.92] → N2 (얕은 수면)
2. [유사도=0.89] → N2 (얕은 수면)
3. [유사도=0.87] → N3 (깊은 수면)

현재 입력의 수면 단계를 예측하세요.

검증 & 시각화

t-SNE로 임베딩 품질 확인

# 512차원 → 2차원 투영
tsne = TSNE(n_components=2, perplexity=30, random_state=42)
X_embedded = tsne.fit_transform(embeddings)

• 왼쪽 플롯 (수면 단계별): 같은 단계끼리 뭉쳐있음 → 생리학적 특성을 캡처

• 오른쪽 플롯 (사용자별): 사용자별로도 뭉쳐있음 → 개인차도 보존

Forecasting으로 파이프라인 검증

forecast = pipeline.predict(
    context_tensor_cpu,
    prediction_length=100,  # 다음 1초 예측
    num_samples=20          # 확률적 예측
)

배운 교훈

1. CPU/GPU 경계를 명확히 하라

• 에러 메시지를 잘 읽자: "CUDA operation not supported" 같은 힌트

• 단계별로 디바이스 체크: print(tensor.device) 습관화

2. Foundation Model은 Zero-shot으로도 강력하다

• 사전학습 데이터의 다양성: 다양한 시계열 패턴 학습

• 아키텍처의 범용성: T5의 Attention 메커니즘

3. 검색 전략이 성능을 좌우한다

• 노이즈 감소

• 검색 속도 향상 (후보군 축소)

• 해석 가능성 증가 ("이 사람과 비슷해서 참고했습니다")

4. 정규화는 선택이 아니라 필수

• Chronos는 정규화된 데이터로 사전학습됨

• 원본 스케일의 EEG는 분포가 완전히 다름

한계와 개선 방향

현재 시스템의 한계

1. 정량 평가: 테스트 셋에서 Accuracy 측정

2. 베이스라인 비교: CNN, LSTM 같은 전통적 방법과 비교

3. 클래스 불균형 처리: N2가 압도적으로 많음

개선 아이디어

검색된 예시:
1. [유사도=0.92] → N2 (주파수 파워: 델타 70%, 세타 20%)

SC4002 (유사도 0.95): N2, N2, N3 → 66% N2
SC4005 (유사도 0.91): N2, N2, N2 → 100% N2
→ 최종 예측: N2 (신뢰도 높음)

마치며

• Foundation Model의 범용 표현력

• RAG의 예시 기반 학습

• Two-Stage Retrieval의 도메인 적응

References & Resources

• Chronos 논문: Ansari et al. (2024), "Chronos: Learning the Language of Time Series"

• Sleep-EDF 데이터셋: PhysioNet (https://physionet.org/)

• AASM 매뉴얼: Berry et al. (2012), 수면 단계 분류 표준

• RAG 개념: Lewis et al. (2020), "Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks"

코드 저장소