기술 요약: STAMP: 멀티헤드 풀링을 갖춘 시공간 어댑터

문제 제기

다양한 도메인에서 사전 훈련된 시계열 기초 모델 (TSFMs) 은 다양한 모델링 작업에서 강력한 성능을 입증해 왔습니다. 동시에 뇌의 전기적 활동의 고유한 시공간적 특성을 포착하기 위해 뇌전도 (EEG) 데이터를 위한 특정 기초 모델 (EEGFMs) 이 개발되었습니다. 그러나 EEG 작업에 있어 일반적인 TSFMs 과 EEG 전용 기초 모델 간의 비교 분석은 수행되지 않았습니다. furthermore, EEGFMs 은 종종 공간 및 시간적 의존성을 모델링하기 위해 복잡한 아키텍처를 포함하는 반면, 일반적인 TSFMs 은 일반적으로 입력을 단변량 시계열로 취급하여 채널 간의 공간적 관계와 시간적 진화가 중요한 EEG 작업에는 본질적으로 비효율적입니다. EEG 데이터에 TSFMs 을 단순하게 적용하는 것 (예: 단순 평균 풀링 사용) 은 거의 무작위 수준의 성능을 초래합니다.

방법론

저자들은 일반적인 목적의 TSFMs 위에 탑재되어 효과적인 EEG 모델링을 가능하게 하는 경량이며 유연한 어댑터인 STAMP(멀티헤드 풀링을 갖춘 시공간 어댑터) 를 소개합니다. STAMP 는 고정된 (frozen) TSFM 이 생성한 단변량 임베딩을 활용하며, 세 가지 핵심 구성 요소를 통해 EEG 데이터의 시공간적 특성을 암묵적으로 모델링합니다:

1. 위치 인코딩 (PEs): 표준 TSFM 임베딩의 공간적 및 시간적 맥락 부족을 해결하기 위해 STAMP 는 학습 가능한 위치 임베딩을 추가합니다. 여기에는 토큰별 위치 임베딩 ($p_{ij}$), 공간별 임베딩 ($s_i$), 그리고 시간별 임베딩 ($t_j$) 이 포함됩니다. 수정된 임베딩은 $\tilde{e}_{ij} = e'_{ij} + p_{ij} + s_i + t_j$로 계산되며, 여기서 $e'_{ij}$는 차원 축소된 TSFM 임베딩입니다. 제거 실험 (ablation studies) 은 세 가지 유형의 인코딩을 모두 사용하는 것이 최상의 성능을 낸다는 것을 보여줍니다.

2. 크리스크로스 게이트드 MLP(CC-GMLP): 시공간 그리드 전반에 걸친 관계를 포착하기 위해 STAMP 는 크리스크로스 트랜스포머 (CC-TF) 와 게이트드 MLP(GMLP) 에서 영감을 받은 새로운 CC-GMLP 아키텍처를 사용합니다. 이 모듈은 공간 및 시간적 관계를 모델링하는 것을 분리합니다. 입력 그리드 $\tilde{E}$를 일련의 블록을 통해 처리하는데, 여기서 시간 게이트 ($g_T$) 와 공간 게이트 ($g_S$) 유닛이 각각의 축을 따라 독립적으로 작동한 후 연결됩니다. 이를 통해 특정 신호의 시간적 진화와 무관하게 채널 간 관계를 파라미터 효율적으로 학습할 수 있습니다.

3. 멀티헤드 어텐션 풀링 (MHAP): 최종 예측을 위해 혼합된 토큰 임베딩을 집계하기 위해 STAMP 는 MHAP 를 활용합니다. 이 메커니즘은 각 토큰의 최종 작업에 대한 중요도를 학습합니다. 토큰을 저차원 공간으로 투영하고 쿼리 벡터에 대해 어텐션 점수를 계산한 후 결과를 집계합니다. 쿼리 전반에 걸친 최종 집계 단계는 각 헤드의 전반적인 중요도를 결정하여 요약 벡터를 생성합니다. 이는 최종 선형 투영 전에 혼합된 토큰 임베딩의 평균과 잔차 연결 (residual connection) 과 결합됩니다.

이 어댑터는 고정된 TSFM(주로 MOMENT Large 이지만 MOMENT Small/Base, Chronos, TSPulse 로도 테스트됨) 위에 훈련됩니다. STAMP 의 학습 가능한 파라미터 총 수는 약 75 만 개로, 일반적인 EEGFMs 에 있는 수백만 개의 파라미터의 일부에 불과합니다.

주요 기여

• 최초의 비교 분석: 이 논문은 EEG 전용 작업에 있어 EEG 전용 기초 모델과 일반적인 TSFMs 간의 최초의 비교 분석을 제공합니다.
• 새로운 어댑터 아키텍처: TSFM 을 EEG 데이터로 재훈련할 필요 없이 단변량 TSFM 임베딩을 다변량 시공간 EEG 데이터에 성공적으로 적응시키는 STAMP 의 도입.
• 효율성: STAMP 는 CBraMod(약 0.74M 대 29M) 와 LaBraM(5.8M) 과 같은 최첨단 EEGFMs 과 비교 가능한 성능을 달성하면서 학습 가능한 파라미터를 훨씬 적게 사용합니다.
• 종합적 평가: 이 방법은 모션 이미지, 정신 산술, 감정 인식, 이벤트 분류 등 다양한 임상 작업을 포함하는 8 개의 벤치마크 EEG 데이터셋에서 평가되었습니다.

결과

• 기선 대비 성능: STAMP 는 대부분의 데이터셋에서 EEGFMs(CBraMod, LaBraM) 와 유사하거나 더 나은 성능을 보이며, 비기초 EEG 모델 (EEG Conformer, ST-Transformer) 보다 엄격하게 우월한 성능을 발휘합니다.
• 제거 실험 (Ablation Studies):
  • 위치 인코딩: 세 가지 위치 인코딩 유형 (토큰, 공간, 시간) 을 모두 포함하는 것이 일관되게 성능을 향상시켰으며, PEs 를 제외할 경우 가장 나쁜 결과가 나왔습니다.
  • 토큰 믹싱: CC-GMLP 아키텍처는 기본 GMLP 와 트랜스포머 기반 믹싱 전략 (B-TF, CC-TF) 모두를 데이터셋 전반에서 능가했습니다.
  • 토큰 집계: 대부분의 데이터셋에서 MHAP 와 평균 풀링이 유사한 성능을 보였으나, MHAP 는 BCIC-IV-2a 데이터셋에서 상당한 성능 향상을 제공하여 복잡한 토큰 관계를 포착하는 데 유용함을 시사합니다.
• TSFM 호환성: STAMP 는 EEG 데이터로 사전 훈련되지 않은 모델 (Chronos, TSPulse) 을 포함한 다양한 TSFMs 의 임베딩과 함께 잘 작동하여 어댑터의 일반화 능력을 입증했습니다.
• 한계점: STAMP 는 감정 인식 데이터셋 (SEED-V 및 FACED) 에서 다른 방법들에 비해 열악한 성능을 보였습니다. 저자들은 현재 TSFMs 이 세분화된 감정 인식과 관련된 특징을 추출하지 못할 수 있다고 가설을 세웠는데, 이는 Chronos 임베딩으로 전환하더라도 성능이 크게 향상되지 않았기 때문입니다.

의미 및 주장

이 논문은 STAMP 가 일반적인 목적의 시계열 기초 모델이 최소한의 추가 학습 가능한 파라미터로 EEG 작업에 효과적으로 적응될 수 있음을 보여준다고 주장합니다. 저자들은 TSFMs 이 "아웃 오브 더 박스 (out of the box)" 상태에서 EEGFMs 에 비해 성능이 낮아 보일지라도, 상대적으로 적은 적응—특히 위치 인코딩과 토큰 믹싱 전략의 추가—만으로도 경쟁력 있는 결과를 달성할 수 있다고 논합니다.

의미 있는 점은 기존 대규모 사전 훈련된 TSFMs 을 활용하여 제로 스크래치 (from scratch) 로 거대한 도메인 전용 기초 모델을 훈련하는 대신 EEG 솔루션 개발의 비용과 시간을 줄일 수 있는 잠재력에 있습니다. 저자들은 초기 임베딩 생성에 계산 자원이 필요하지만, 이는 사전에 그리고 병렬로 수행될 수 있다고 지적합니다. 그들은 STAMP 가 EEG 데이터 모델링을 위한 유연하고 파라미터 효율적인 경로를 제공한다고 결론지으며, 동시에 현재 TSFMs 이 어려움을 겪는 세분화된 감정 인식과 같은 작업에 대해서는 추가 탐구가 필요함을 인정합니다.