LongMoE: Longitudinal Multimodal Learning via Trajectory-Aware Mixture-of-Experts

기술 요약: LONGMOE

문제 정의

멀티모달 임상 학습은 기존 방법론들이 결합하여 해결하지 못한 두 가지 근본적이고 상호 교차하는 과제에 직면해 있습니다:

1. 모달리티 결손(Modality Missingness): 실제 임상 코호트에서는 이질적인 데이터 수집 프로토콜과 자원 제약으로 인해, 특정 환자 방문 시점에 임의의 모달리티 하위 집합(예: 영상, 텍스트, 유전 데이터)이 가용하지 않은 경우가 빈번합니다.
2. 종단적 역동성(Longitudinal Dynamics): 관측치의 진단적 중요성은 환자의 진화하는 질병 궤적과 분리될 수 없습니다. 바이오마커의 가치는 단일 정적 스냅샷에서의 절대적인 값이 아니라, 개별적인 쇠퇴 궤적상의 위치에 따라 결정됩니다.

기존 접근 방식들은 이러한 과제들을 개별적으로 처리합니다. 결손 모달리티 프레임워크(예: Flex-MoE, FuseMoE)는 각 방문을 독립적인 정적 스냅샷으로 취급하여 시간적 맥락을 버립니다. 반대로, 종단적 모델(예: GRU 기반 접근법)은 모든 방문에서 모달리티가 완전히 가용하다고 가정하며, 체계적인 결손 상황에서 성능이 저하됩니다. 어떤 기존 모델도 불규칙한 방문 간격, 임의의 모달리티 조합, 그리고 인구 통계적 이질성의 결합된 과제를 효과적으로 처리하지 못합니다.

방법론: LONGMOE 아키텍처

저자들은 모달리티 결손과 종단적 역동성을 공동으로 해결하기 위해 설계된 통합 프레임워크인 LONGMOE(Longitudinal Mixture-of-Experts)를 제안합니다. 이 아키텍처는 환자의 종단적 방문 시퀀스를 네 가지 순차적 모듈을 통해 처리합니다:

1. 모달리티별 임베딩 (Modality-Specific Embeddings)

각 모달리티 $m_l$은 그 구조적 특성에 맞춤화된 전용 인코더 $\phi_l$에 의해 처리됩니다:

• MRI: 3D 컨볼루션 인코더와 그 뒤를 잇는 MLP 투영.
• 임상 점수, 바이오마커 및 유전체: 모달리티별 MLP.
• 시간적 임베딩: 방문 타임스탬프는 절대적 시간과 방문 간 간격을 포착하는 학습된 임베딩으로 인코딩됩니다.
  관찰되지 않은 모달리티는 초기에는 제로 벡터로 설정됩니다.

2. 모달리티 임퓨테이션 및 토큰화 (Modality Imputation and Tokenization)

이 모듈은 결측 데이터를 처리하고 통합 방문 토큰을 생성합니다:

• 맥락 인식 임퓨테이션 (Context-Aware Imputation):
  • Carry-Forward (이월): 이전 방문에서 모달리티가 관찰되었다면, 가장 최근의 임베딩을 복제합니다 (바이오마커의 강한 단기 자기상관성에 근거함).
  • 맥락적 임퓨테이션 (Contextual Imputation): 이전 관찰 기록이 없는 경우, 현재 방문의 가용 모달리티를 조건으로 하여 임베딩을 생성하는 학습된 네트워크를 사용하며, 결손 신호로부터의 오염을 방지하기 위해 제로 마스킹을 사용합니다.
• 어텐션 기반 토큰화 (Attentional Tokenization):
  • 모달리티 융합: 임퓨테이션된 모달리티 임베딩에 대해 멀티 헤드 셀프 어텐션을 적용하여 모달리티 간 의존성을 학습합니다.
  • Viscode 인코딩: 연속 시간 다중 주파수 사인형 위치 인코딩(sinusoidal positional encoding)은 방문 타임스탬프를 고정된 차원의 벡터로 매핑합니다. 이 인코딩은 비균일하게 샘플링된 시퀀스에 대해 **증명 가능한 단사성(injectivity)**을 가지며, 균일한 방문 간격을 가정하지 않고도 느린 종단적 추세와 급격한 변화율을 모두 포착합니다.
  • 융합: 모달리티 임베딩과 Viscode 인코딩은 요소별 덧셈(element-wise addition)을 통해 융합되어 방문 토큰 $z_t$를 생성합니다.

3. 궤적 인식 종단적 모델링 (Trajectory-Aware Longitudinal Modeling)

스택 형태의 트랜스포머 인코더 블록이 방문 토큰 $\{z_t\}$의 시퀀스를 처리합니다.

• 인과적 마스킹 (Causal Masking): 인과적 어텐션 마스크는 방문 $t$에서의 표현이 과거 기록 $\{v_1, \dots, v_t\}$에만 의존하도록 보장하여 미래 정보의 누출을 방지합니다.
• 출력: 이는 환자의 전체 질병 진행 궤적을 통합하는 문맥화된 임베딩 $h_t$를 생성하며, 이는 라우팅 모듈을 위한 풍부한 입력 역할을 합니다.

4. 맥락 인식 희소 MoE 라우팅 (Context-Aware Sparse MoE Routing)

최종 레이어는 환자별 이질성과 모달리티 패턴을 처리하기 위해 희소 혼합 전문가(Sparse Mixture-of-Experts, SMoE) 아키텍처를 채택합니다:

• 전문가 아키텍처: 레이어는 다음으로 구분되는 $K=8$개의 전문가를 포함합니다:
  • 개별 전문가 (Individual Experts): 구별되는 환자 하위 집단에 특화된 서브 네트워크.
  • 공유 전문가 (Shared Experts): 인구 수준의 공통 패턴을 포착하고 과도한 전문화를 방지하기 위해 모든 입력에 접근 가능한 전문가.
• 라우팅: 라우터 네트워크는 궤적 임베딩 $h_t$를 점수 벡터로 매핑합니다. 하드 탑-1(Hard top-1) 라우팅은 가장 관련성이 높은 단일 전문가(항상 활성화되는 공유 전문가 포함)를 선택합니다.
• 집계: 최종 예측은 라우터의 신뢰도에 따라 가중치가 부여된 모든 방문에 대한 전문가 출력의 어텐션 집계를 통해 얻어집니다.
• 정규화: 로드 밸런싱 손실(load-balancing loss)은 라우터 붕괴를 방지하여 전문가 활용이 균등하게 이루어지도록 합니다.

주요 기여

1. 통합 프레임워크: LONGMOE는 멀티모달 임상 학습에서 모달리티 결손과 종단적 시간 역동성을 공동으로 다루는 최초의 모델입니다.
2. 새로운 아키텍처 구성 요소:
   • 과거 데이터와 현재 맥락을 모두 활용하는 맥락 인식 임퓨테이션 모듈.
   • 교차 모달 어텐션과 불규칙한 시퀀스를 위한 증명 가능한 단사적 다중 주파수 시간 인코딩을 결합한 어텐션 기반 토큰화 모듈.
   • 질병 진행을 모델링하는 궤적 인식 인코더.
   • 관찰된 모달리티 조합과 환자 궤적에 따라 전문가를 동적으로 선택하는 맥락 조건부 SMoE 라우팅 레이어.
3. 강건성: 본 프레임워크는 결손되거나 약한 동시 모달리티 상황에서도 강건함을 입증하면서도, 전체 모달리티 설정에서 경쟁력 있는 성능을 유지합니다.

실험 결과

저자들은 ADNI (Alzheimer's Disease Neuroimaging Initiative), OASIS-3, MIMIC-IV의 세 가지 데이터셋에서 LONGMOE를 평가했습니다.

• 성능: LONGMOE는 ADNI와 MIMIC-IV의 모든 모달리티 조합에서 최고의 AUC-ROC를 달성했습니다. OASIS-3에서는 7개 조합 중 4개에서 선두를 차지했으며, 특히 다른 베이스라인들이 무너지는 영상 및 유전체 지배적 환경에서 뛰어난 성능을 보였습니다.
• 결손 모달리티 강건성:
  • ADNI에서 LONGMOE는 단 하나의 모달리티만 존재할 때도 높은 성능(AUC $\approx$ 0.94–0.97)을 유지한 반면, TF와 같은 베이스라인은 무작위 확률(0.50) 수준으로 무너졌고 다른 모델들도 크게 저하되었습니다.
  • 궤적 인식 인코더는 결손된 동시 모달리티 정보를 종단적 맥락으로 효과적으로 대체하여, 모달리티가 제거됨에 따라 성능 저하 곡선이 거의 평탄하게 유지되는 결과를 낳았습니다.
• 절제 연구 (Ablation Studies): 종단적 인코더(LE) 또는 Viscode 인코딩(VE)을 제거했을 때 가장 큰 성능 저하가 발생하여, 이들의 결정적인 역할을 확인했습니다. 공유 전문가와 로드 밸싱 메커니즘 또한, 특히 심각한 결손이 있거나 코호트 규모가 작은 설정에서 필수적임이 입증되었습니다.
• 복잡도: LONGMOE는 궤적 인식 트랜스포머로 인해 정적 베이스라인(예: Flex-MoE의 183K 파라미터 대비 4.4M 파라미터)보다 더 많은 파라미터 수와 순전파 지연 시간을 가집니다. 그러나 저자들은 이것이 오프라인 배치 추론에는 실용적이지만, 실시간 병상 지원에는 제한적일 수 있다고 언급했습니다.

의의 및 주장

본 논문은 LONGMOE가 종단적 인식을 갖춘 멀티모달 임상 학습을 위한 강력한 토대를 구축한다고 주장합니다. 이 모델의 주요 의의는 결측 데이터 처리와 질병 진행 모델링 사이의 간극을 메우는 데 있습니다. 환자의 이력을 독립적인 스냅샷의 연속이 아닌 연속적이고 불규칙한 궤적으로 취급함으로써, 현재 방문에 핵심 모달리티가 부족하더라도 차별적인 신호를 추출할 수 있습니다.

저자들은 LONGMOE가 복잡하고 불완전한 임상 기록을 처리하고 오프라인 배치 추론을 수행하는 데 매우 효과적이지만, 궤적 인식 인코더의 계산 비용 증가로 인해 실시간 병상 지원에 대한 적용은 제한적이라고 명시했습니다. 또한, 인코더가 의미 있는 임베딩을 생성하기 위해 최소 하나의 과거 관찰 기록을 필요로 하므로, 첫 방문 환자에 대해서는 종단적 맥락을 활용할 수 없습니다.