기술 요약: CaMBRAIN - 인과적 상태 공간 모델을 활용한 실시간 연속 EEG 추론

문제 정의

뇌파 (EEG) 신호는 임상 및 생체역학적 환경에서 뇌 활동을 모니터링하는 데 필수적이지만, 수분에서 수 시간에 이르는 무한한 길이로 인해 딥러닝에 고유한 도전을 제기합니다. 기존 딥러닝 접근법, 특히 트랜스포머 (Transformer) 기반 모델들은 장기간 EEG 에 적용될 때 두 가지 주요 장애물에 직면합니다:

1. 계산 비효율성: 어텐션 메커니즘은 시퀀스 길이에 따라 이차적으로 스케일링 ($O(T^2)$) 되어 장기간 신호 처리에 비효율적입니다.
2. 맥락 제한: 계산 비용을 관리하기 위해 현재 모델들은 일반적으로 고정된 입력 길이를 가진 슬라이딩 윈도우를 통해 EEG 를 처리합니다. 이 접근법은 모델이 중첩된 영역에서 특징을 다시 계산하게 하고 전체 신호 기록에 대한 지속적 표현을 유지하는 것을 방지합니다. 결과적으로 이러한 모델들은 전역적 시간적 이해가 부족하며, 과거 입력을 재방문하지 않고 새로운 관측치가 이전 맥락에 기반하여 실시간 연속 추론을 수행할 수 없습니다.

또한, 기존 EEG 의 자기지도식 사전 학습 목표는 종종 신호 재구성 (예: 마스킹 오토인코더) 에 의존합니다. 이는 국소 신호 충실도 학습에는 효과적이지만, 숨겨진 상태가 스트리밍 추론에 필수적인 중요한 장거리 맥락을 유지하도록 명시적으로 훈련하지는 않습니다. 스트리밍 추론에서는 모델이 역사를 압축된 메모리로 변환해야 하기 때문입니다.

방법론

저자들은 실시간 연속 스트리밍 추론을 위해 설계된 최초의 인과적, Mamba 기반 EEG 모델인 CaMBRAIN을 제안합니다. 이 방법론은 인과적 아키텍처와 전문화된 다단계 자기지도식 학습 파이프라인이라는 두 가지 핵심 기둥 위에 구축됩니다.

1. 아키텍처: 인과적 Mamba-3

CaMBRAIN 은 EEG 신호를 처리하기 위해 단방향 인과적 상태 공간 모델 (SSM), 구체적으로 Mamba-3 아키텍처를 활용합니다.

• 입력 처리: 원시 EEG 신호는 모든 채널에 걸쳐 62.5ms(256Hz 에서 16 샘플) 의 비중첩 패치로 분할됩니다. 이러한 패치는 트랜스포머 기반 채널 임베더를 사용하여 토큰으로 임베딩됩니다.
• 지속적 숨겨진 상태: 슬라이딩 윈도우 접근법과 달리 CaMBRAIN 은 모든 과거 맥락을 점진적으로 압축하는 지속적 숨겨진 상태를 유지합니다. 새로운 패치가 도착하면 숨겨진 상태는 이전 입력을 재방문하지 않고 재귀적으로 업데이트됩니다 ($h_t = \bar{A}h_{t-1} + \bar{B}x_t$).
• 인과적 제약: 모델은 엄격히 순방향으로만 작동하여 시간 $t$에서의 예측이 과거 및 현재 입력 ($x_{\le t}$) 에만 의존하도록 보장하며, 이는 EEG 모니터링의 물리적 현실과 부합합니다.
• 복소수 전이: Mamba-3 를 활용하여 모델은 다대역 EEG 신호에 내재된 진동 역동성과 위상 관계를 더 잘 포착하기 위해 복소수 상태 전이를 사용합니다.

2. 학습 파이프라인

스트리밍을 위한 재구성 기반 목표의 한계를 해결하기 위해 저자들은 3 단계 학습 파이프라인을 도입합니다:

• 1 단계: 인과적 예측 사전 학습:
  인코더는 일반적인 시간적 구조를 주입하기 위해 하이브리드 인과적 손실을 사용하여 사전 학습됩니다. 이는 다음을 결합합니다:

  • 자기회귀 모델링 (ARM): 현재 숨겨진 상태에서 다음 패치 임베딩을 예측합니다.
  • 마스킹 재구성: 인코더가 제공하는 인과적 맥락에서 마스킹된 원시 EEG 토큰을 재구성합니다.
    이 단계는 모델이 국소 시간적 역동성과 신호 충실도를 학습하도록 보장합니다.

• 2 단계: 재구성 없는 잠재 표현 정제 (JEPA 스타일):
  장거리 예측 정보의 유지를 장려하기 위해 모델은 공동 임베딩 예측 아키텍처 (JEPA) 에서 영감을 받은 디코더 없는 학습 단계를 거칩니다.

  • 학생 - 교사 프레임워크: 학생 네트워크는 입력의 마스킹된 뷰를 처리하는 반면, 교사 네트워크 (학생의 지수 이동 평균) 는 마스킹되지 않은 입력을 처리합니다.
  • 목표: 학생 네트워크는 다음을 예측하도록 훈련됩니다:
    1. 마스킹된 잠재 변수: 마스킹된 위치에 대한 교사의 표현.
    2. 미래 잠재 변수: 현재 숨겨진 상태로부터 $K$단계 앞의 다단계 미래 표현.
       이 단계는 단순한 국소 신호 세부 사항 재구성이 아닌, 연속 추론에 필요한 정보적이고 불변적이며 시간적으로 관련 있는 특징을 숨겨진 상태가 인코딩하도록 명시적으로 강제합니다.

• 3 단계: 하류 미세 조정:
  사전 학습된 모델은 표준 지도 학습을 사용하여 특정 작업 (예: 발작 감지, 아티팩트 분류) 에 대해 미세 조정됩니다. 주목할 점은 CaMBRAIN 이 고정된 세그먼트에 걸쳐 토큰을 풀링하거나 집계할 필요 없이 분류를 위해 최종 숨겨진 상태만 활용한다는 것입니다.

주요 기여

1. 인과적 추론 공식화: 이 논문은 슬라이딩 윈도우 및 양방향 트랜스포머 접근법과 구별되는 단방향 EEG 추론 공식을 제안하며, 신호를 독립적인 세그먼트가 아닌 연속 스트림으로 취급합니다.
2. CaMBRAIN 아키텍처: 불필요한 슬라이딩 윈도우 재계산을 지속적 숨겨진 상태로 대체하는 인과적 Mamba-3 인코더로, 임의의 길이의 신호에 대해 선형 시간 복잡도로 효율적이고 확장 가능한 추론을 가능하게 합니다.
3. 표현 수준 학습: 신호 재구성을 넘어선 새로운 다단계 자기지도식 프레임워크입니다. 자기회귀 예측과 마스킹 및 미래 잠재 변수 예측을 결합함으로써 숨겨진 상태가 중요한 장거리 맥락을 유지하도록 명시적으로 훈련합니다.

결과

CaMBRAIN 은 TUAR(아티팩트 감지), TUAB(비정상 EEG 감지), MAT(정신적 스트레스 감지), CHB-MIT(발작 감지) 의 네 가지 다양한 EEG 데이터셋에서 평가되었습니다.

• 성능: CaMBRAIN 은 모든 벤치마크에서 최첨단 (SOTA) 또는 경쟁력 있는 결과를 달성했습니다.
  • TUAR: AUROC 0.936을 달성하여 LUNA-H(0.921) 와 같은 더 큰 파운데이션 모델을 능가했습니다.
  • CHB-MIT: AUROC 0.921을 달성하여 불균형이 심한 발작 데이터에서 강력한 성능을 입증했습니다.
  • MAT: 비교된 방법 중 가장 높은 균형 정확도 (0.778) 와 AUROC(0.876) 를 달성했습니다.
• 계산 효율성: 모델은 스트리밍 시나리오에서 상당한 효율성 향상을 보여주었습니다.
  • 16Hz 스트리밍 업데이트 속도에서 CaMBRAIN 은 1.23 GFLOPs/s의 계산량을 필요로 했습니다.
  • 이는 REVE 가 194.79 GFLOPs/s 를 필요로 한 것과 같은 다른 파운데이션 모델 베이스라인에 비해 6.1 배에서 158 배의 지속적 계산 감소로, 정확도는 일치하거나 초과했습니다.
• 애블레이션 연구:
  • 지속적 상태: 5 초마다 재설정하는 것 대비 지속적 숨겨진 상태를 사용하는 것은 CHB-MIT 에서 AUROC 를 +0.015 향상시켰으며 발작 시작 국소화도 크게 개선했습니다.
  • 학습 단계: 전체 2 단계 사전 학습 (1 단계 + 2 단계) 은 표준 미세 조정이나 1 단계만 사용하는 것보다 훨씬 큰 향상 (+2.30 AUROC) 을 보여주어 스트리밍 작업에 대한 표현 수준 정제의 필요성을 확인시켰습니다.

중요성 및 주장

이 논문은 CaMBRAIN 이 EEG 신호의 고유한 인과적이고 연속적인 특성과 아키텍처 및 학습 목표를 정렬함으로써 EEG 모델링의 패러다임 전환을 나타낸다고 주장합니다.

• 실시간 능력: 이 모델은 이차적 스케일링의 계산적 패널티나 슬라이딩 윈도우의 중복 없이 가변 길이 EEG 신호에 대한 장거리 연속 추론을 수행할 수 있는 최초의 모델을 가능하게 합니다.
• 임상적 관련성: 이 접근법은 임상가가 침대 옆에서 EEG 를 모니터링하는 방식을 반영하여, 긴 간격으로 분리된 사건을 감지하기 위해 전체 신호 기록의 압축된 메모리를 유지합니다.
• 효율성: 상태 공간 모델의 선형 시간 스케일링과 표현 수준 학습 목표를 활용함으로써 CaMBRAIN 은 높은 처리량과 낮은 지연 시간을 달성하여 계산 자원이 제한된 실시간 임상 및 웨어러블 애플리케이션에 적합합니다.

저자들은 CaMBRAIN 이 여러 데이터셋에서 새로운 벤치마크를 설정하지만, TUAB 비정상 EEG 감지에서는 가장 강력한 파운데이션 모델에 비해 우월함보다는 경쟁력 있는 수준으로 남아 있음을 인정하며, 더 크고 이질적인 데이터셋 (예: REVE) 과 일치하도록 사전 학습 코퍼스를 확장하는 것이 필요한 미래 방향임을 시사합니다.