기술 요약: HyMaTE – EHR 표현 학습을 위한 하이브리드 Mamba 및 Transformer 모델

문제 정의

전자 건강 기록(EHR)은 환자의 건강 진행 과정을 분석하고 진단, 처치, 검사 결과 등을 포함하는 종단적 이력을 기록하는 데 있어 핵심적인 자원 역할을 한다. 그러나 EHR 데이터를 모델링하는 것은 본질적인 복잡성으로 인해 상당한 어려움을 수반한다. 즉, EHR은 길고 다변량인 시퀀스로 구성되며, 흔히 희소(sparse)하고 불규칙하게 샘플링되며, 많은 결측치를 포함하고 있다.

기존의 딥러닝 접근 방식은 이 영역에서 다음과 같은 구체적인 한계에 직면해 있다:

• Transformer 기반 모델: 장기적인 시간적 의존성을 포착하는 데는 성공적이지만, 입력 시퀀스 길이에 따른 이차 복잡도($O(N^2)$)와 메모리 요구 사항 문제를 겪는다. 이는 모델이 긴 환자 이력(수천 개의 토큰에 달하는)을 처리할 때 정확도를 저하시킬 수 있는 절단(truncation)이나 양자화(quantization) 없이 처리하는 능력을 제한한다. 또한, 표준 Transformer는 잠재적으로 변수 간의 상호작작용(채널 믹싱)을 간과할 가능성이 있는 시퀀스 믹싱(sequence mixing)을 우선시하는 경향이 있다.
• 상태 공간 모델(SSM) / Mamba: 이러한 모델은 선형 시간 시퀀스 모델링과 긴 시퀀스에 대한 탁월한 효율성을 제공한다. 그러나 전통적인 Mamba 아키텍처는 주로 시퀀스 믹싱에 집중하며, 다양한 이벤트 유형(변수) 간의 관계가 중요한 다변량 시계열 데이터에 필수적인 강력한 채널 믹싱 메커니 사망 메커니즘이 부족하다.

방법론

저자들은 SSM의 선형적 효율성과 Transformer의 맥락적 및 채널 믹싱 강점을 결합하도록 설계된 새로운 아키텍처인 HyMaTE(EHR 표현 학습을 위한 하이브리드 Mamba 및 Transformer 모델)를 제안한다.

1. 입력 표현 (Input Representation)

EHR 이벤트는 관찰 시간, 특징/변수, 관찰된 값을 구성하는 삼중항 $(t, f, v)$로 표현된다. 이 삼중항은 특징(룩업 테이블), 연속값, 그리고 시간(FFN을 통해 유도됨)에 대한 임베딩을 합산하여 $\tau$ 차원의 벡터로 임베딩된다.

2. 모델 아키텍처

아키텍처는 다음과 같은 순차적 처리 파이프라인을 따른다:

• 맥락적 어텐션 (Mamba 블록): 임베딩된 입력은 여러 Mamba 블록을 통과한다. 이 블록들은 관련 정보를 효율적으로 필터링하기 위해 이산화된 상태 공간 모델(SSM)을 활용한다. SSM 유닛은 암시적 은닉 상태를 통해 입력 시퀀스를 출력 시퀀스로 매핑하며, 병렬 처리를 위해 글로벌 컨볼루션 커널을 사용한다. 이 단계는 긴 시퀀스를 효율적으로 인코딩하는 역할을 수행한다.
• 로컬 맥락 추출 (Self-Attention): Mamba 인코딩 이후, 시퀀스 내의 로컬 맥락과 즉각적인 의존성을 추출하기 위해 셀프 어텐션 레이어가 적용된다.
• 퓨전 셀프-어텐션 (Fusion Self-Attention): Mamba 블록의 출력은 퓨전 셀프-어텐션 모듈로 전달된다. 이 레이어는 각 맥락적 임베딩에 피드 포워드 네트워크(FFN)를 적용하고 소프트맥스 함수를 사용하여 어텐션 가중치($\alpha_i$)를 계산한다. 최종 시계열 임베딩($z_T$)은 맥락적 임베딩들의 가중치 합으로 계산된다.
• 정적 변수 통합: 정적 인구통계학적 변수는 두 개의 은닉층을 가진 FFN을 통해 별도로 임베딩되어 정적 임베딩($z_d$)을 생성한다.
• 최종 표현: 시간적 임베딩($z_T$)과 정적 임베딩($z_d$)은 결합되어 최종 환자 표현($e_E$)을 형성한다.

3. 학습 전략

HyMaTE는 준지도 방식의 2단계 학습을 채택한다:

• 자기 지도 사전 학습 (Self-Supervised Pre-training): 모델은 마스크된 예측(masked forecasting) 과제를 사용하여 더 큰 데이터셋에서 사전 학습된다. 모델은 관찰된 이력을 바탕으로 미래의 변수 값을 예측하는 법을 배우며, 결측치를 처리하기 위해 마스크된 평균 제곱 오차(MSE) 손실을 활용한다. 이 단계는 라벨이 제한된 데이터로부터 견고한 표현을 학습할 수 있게 한다.
• 지도 미세 조정 (Supervised Fine-tuning): 사전 학습된 모델은 특정 임무별 헤드를 최종 임베딩에 부착하여 특정 하위 임상 예측 작업(예: 사망률, 재원 기간)에 대해 미세 조정된다.

주요 기여

1. 하이브리드 아키텍처: HyMaTE는 효율적인 선형 시간 시퀀스 인코딩을 위한 Mamba 블록과, 순수 SSM의 채널 믹싱 한계 및 순수 Transformer의 계산적 병목 현상을 해결하기 위한 Transformer 기반 어텐션 메커니즘을 전략적으로 결합한다.
2. 퓨전 메커니즘: 퓨전 셀프-어텐션 레이어의 도입은 로컬 맥락(셀프 어텐션으로부터)과 글로벌 시퀀스 정보(Mamba로부터)를 하나의 통합된 표현으로 합성할 수 있게 한다.
3. 해석 가능성: 모델은 퓨전 레이어의 어텐션 가중치와 최종 예측 밀집 레이어의 계수 분석을 통해 해석 가능성을 명시적으로 포함하며, 이를 통해 시간적 및 글로벌 특징 중요성에 대한 통찰을 모두 제공한다.
4. 준지도 학습: 자기 지도 예측 사전 학습 단계의 통합은 라벨이 부족한 EHR 데이터셋에서 고품질의 표현을 학습할 수 있게 하며, 이는 의료 환경에서 중요한 이점이 된다.

실험 결과

HyMaTE는 세 가지 다양한 데이터셋인 PhysioNet Challenge 2012(ICU 사망률), MIMIC-IV(사망률, 재원 기간, 재입원), 그리고 개인 소아 체중 관리 데이터셋에서 평가되었다.

• 성능: HyMaTE는 SeFT, STraTS, EHR-Mamba, DuETT를 포함한 베이스라인 모델들을 일관되게 능가했다.
  • MIMIC-IV 사망률 예측에서 HyMaTE는 0.907의 AUROC와 0.648의 AUPRC를 달성하였으며, 두 번째로 우수한 모델(EHR-Mamba)보다 AUROC에서 0.026 앞섰다.
  • PhysioNet 2012에서 0.868의 AUROC를 달성했다.
  • 소아 체중 감소 과제에서 0.752의 AUROC를 달성하여 EHR-Mamba(0.704)를 크게 앞질렀다.
• 절제 연구 (Ablation Studies): 주요 구성 요소를 제거했을 때 성능 저하가 발생했다:
  • 입력 임베딩을 제거했을 때 주목할 만한 하락이 발생하여, 삼중항 표현의 중요성을 강조했다.
  • 셀프 어텐션 또는 어텐션 퓨전을 제거하면 동일한 성능 저하가 나타나, 효과적인 특징 혼합을 위한 이 구성 요소들의 상호 의존성을 확인했다.
  • 자기 지도 사전 학습을 제거하면 성능이 크게 저하되었으며, 특히 소아 데이터셋에서 그러했는데, 이는 준지도 접근 방식의 필요성을 입증한다.
  • Mamba 블록을 제거하면 성능이 급격히 감소하여, 선형 시간 시퀀스 처리의 가치를 확인시켜 주었다.
• 확장성: 모델은 긴 시퀀스(최대 8,654개의 관측치)를 처리하는 데 있어 견고함을 입증하였으며, 더 긴 시퀀스 길이에서도 성능 저하가 최소화되었다.

의의 및 주장

본 논문은 HyMaTE가 계산 효율성과 표현의 풍부함 사이의 균형을 효과적으로 맞춤으로써 실제 의료 애플리케이션을 위한 확장 가능하고 일반화 가능한 솔루션을 제공한다고 주장한다. Transformer의 이차 복잡도와 Mamba의 채널 믹싱 한계를 극복함으로써, HyMaTE는 장기 의존성과 미세한 변수 간 관계를 모두 포착하는 EHR 데이터의 통합된 표현을 제공한다.

저자들은 어텐션 가중치와 밀집 레이어 계수로부터 도출되는 모델의 해석 가능성—예를 들어 사망률에 대한 결정적 요인으로서 NIMAP과 심박수를 식서하는 것—이 기존의 임상적 발견과 일치함을 강조하며, 이를 통해 투명성과 잠재적인 임상적 유용성을 높인다고 강조한다. 이 연구는 HyMaTE를 종단적 의료 데이터의 복잡성을 처리하면서도 광범-환자 이력을 처리하는 데 필요한 효율성을 유지할 수 있는 EHR 표현 학습의 진일보로 자리매김한다.