논문 요약: 구조화된 EHR 기반 모델에서의 토큰화 트레이드오프 (Tokenization Tradeoffs in Structured EHR Foundation Models)

이 논문은 구조화된 전자 건강 기록 (EHR) 기반 모델 (Foundation Models) 의 성능과 효율성에 결정적인 영향을 미치는 토큰화 (Tokenization) 전략의 설계 선택을 체계적으로 분석합니다. 저자들은 소아 환자 데이터를 기반으로 다양한 토큰화 방식을 비교 실험하여, 어떤 방식이 하류 임상 예측 작업에서 더 우수한 성능을 발휘하면서도 계산 비용을 절감할 수 있는지 규명했습니다.

---

1. 문제 제기 (Problem)

EHR 기반 모델은 진단, 절차, 약물, 검사 결과 등 시간 순서가 있는 임상 이벤트 시퀀스를 학습하여 환자 표현을 생성합니다. 이러한 시퀀스를 이산적인 토큰으로 변환하는 토큰화 과정은 모델이 어떤 정보를 보존할지, 얼마나 효율적으로 인코딩할지, 그리고 어떤 관계를 데이터에서 학습해야 하는지 결정합니다.

기존 연구에서는 토큰화 설계 (이벤트 인코딩, 시간 인코딩, 워크플로우 주석 등) 가 하류 작업의 성능과 계산 효율성에 미치는 영향을 체계적으로 비교한 사례가 부족했습니다. 특히 데이터가 제한적인 의료 환경에서, 복잡한 관계를 토큰화 단계에서 미리 계산할지 (Pre-compute), 아니면 모델이 학습하도록 할지 (Learn from data) 에 대한 트레이드오프가 명확하지 않았습니다.

2. 방법론 (Methodology)

2.1 데이터 및 실험 설계

• 데이터: 토론토 소아병원 (SickKids) 의 소아 EHR 데이터 (약 200 만 명, 1.69 억 건의 이벤트) 를 사용하여 사전 학습 (Pretraining) 을 수행했습니다. 외부 검증으로는 성인을 대상으로 한 MIMIC-IV (BIDMC) 데이터를 사용했습니다.
• 실험 설계: 2 × 2 × 2 **계인 실험 (Factorial Experiment)**을 통해 8 가지 토큰화 조건을 생성하고 비교했습니다.
  1. 이벤트 인코딩 (Event Encoding):
     • Joint (결합): 임상 코드와 값 (예: 포도당 수치) 을 하나의 토큰으로 결합.
     • Factorized (분해): 코드와 값을 별도의 토큰으로 분리하여 순차적으로 입력.
  2. 시간 인코딩 (Time Encoding):
     • Time-Positions: 환자의 나이나 이벤트 간격을 로터리 포지셔널 임베딩 (RoPE) 을 통해 위치 정보로 인코딩.
     • Time-Tokens: 이벤트 사이에 시간 간격을 나타내는 별도의 토큰을 삽입.
  3. 워크플로우 주석 (Workflow Annotations):
     • With Workflow: 검사 의뢰, 채취, 결과 도출 등 다단계 워크플로우를 별도의 이벤트로 표현.
     • Without Workflow: 최종 결과 시간 등 주요 시점 하나만 이벤트로 표현.

2.2 모델 및 평가

• 모델: 28 레이어, 768 차원의 디코더 전용 트랜스포머 (115.6M 파라미터) 를 사용.
• 평가: SickKids 데이터에서 74 개의 임상 예측 작업 (수혈, 절차, 영상, 검사 결과, 약물, 임상 결과 등) 에 대해 AUROC 를 측정했습니다.
• 외부 검증: MIMIC-IV 데이터 (13 개 작업) 에 대해 SickKids 에서 학습된 모델을 고정 (Frozen) 하고 선형 프로브 (Linear Probe) 만 학습시켜 전이 학습 성능을 평가했습니다.
• 추가 실험: 토큰 효율성과 '로컬 바인딩 (Local Binding)' 문제의 영향을 분리하기 위한 Ablation Study를 수행했습니다.

3. 주요 기여 및 결과 (Key Contributions & Results)

3.1 토큰화 선택이 성능과 효율성에 미치는 독립적 영향

• 이벤트 인코딩: Joint Encoding이 Factorized Encoding 보다 74 개 작업 중 73 개에서 더 높은 성능을 보였습니다. 또한, Joint Encoding 은 Factorized 방식보다 **39.5% 적은 사전 학습 FLOPs(연산량)**를 요구했습니다.
• 시간 인코딩: Time-Positions가 Time-Tokens 보다 74 개 작업 중 71 개에서 우세했으며, 9.6% 적은 FLOPs를 사용했습니다. Time-Tokens 은 시퀀스 길이를 불필요하게 늘려 성능을 저하시켰습니다.
• 워크플로우: 워크플로우 주석을 포함하면 SickKids 내부에서는 성능이 향상되었으나 (64/74), 35.6% 의 추가 계산 비용이 발생했고, 외부 데이터 (MIMIC) 로 전이될 때는 유의미한 이점이 없었습니다.

3.2 메커니즘 분석: 로컬 바인딩 효율성 (Local Binding Efficiency)

Ablation 실험을 통해 Joint Encoding 의 우월성이 '로컬 바인딩 효율성'에서 기인함을 규명했습니다.

• Factorized 방식: 모델이 인접한 토큰 (코드 + 값) 간의 관계를 학습해야 합니다. 예를 들어, 같은 '높은 분위수' 토큰이라도 '포도당' 뒤에 오면 고혈당을, '크레아티닌' 뒤에 오면 신장 부전을 의미하는 맥락 의존적 관계를 학습해야 합니다. 이는 데이터가 제한적인 EHR 환경에서 학습하기 어렵습니다.
• Joint 방식: 코드와 값을 하나의 토큰으로 결합하여 사전에 관계를 고정하므로, 모델이 복잡한 바인딩을 학습할 필요가 없어 효율성이 극대화됩니다.

3.3 외부 일반화 (Generalization)

• Joint Encoding의 이점은 MIMIC 데이터 (성인, 다른 병원) 로 전이될 때도 유지되었습니다 (OOV 비율이 70% 이상임에도 불구하고).
• 반면, Time-Positions와 Workflow의 이점은 병원별 또는 인구집단별 특성 (예: 소아 vs 성인 연령 분포) 에 종속적이어서 전이되지 않았습니다.

4. 의의 및 결론 (Significance & Conclusion)

1. 구조적 인코딩의 중요성: EHR 기반 모델에서는 복잡한 임상적 구조 (코드 - 값 관계) 를 토큰화 단계에서 미리 계산하여 단일 토큰으로 만드는 것이, 모델이 데이터에서 이를 학습하는 것보다 성능과 효율성 모두에서 우월합니다. 이는 데이터가 부족한 의료 AI 환경에서 특히 중요합니다.
2. 계산 비용 절감: 더 나은 성능을 내면서도 사전 학습에 필요한 연산량을 크게 줄일 수 있어, 자원 효율적인 모델 개발이 가능합니다.
3. 실무적 시사점:
   • Joint Event Encoding과 **Positional Time Encoding (RoPE)**은 EHR 기반 모델 개발을 위한 권장 기본 설정 (Default) 입니다.
   • 워크플로우 주석은 특정 기관의 업무 흐름을 반영하므로 전이 학습 시 주의가 필요하며, 외부 데이터에 대한 일반화 성능을 기대하기 어렵습니다.
   • 토큰화 전략은 모델의 성능과 비용에 직접적인 영향을 미치므로, 향후 연구에서 토큰화 설계의 투명성과 체계적인 비교가 필수적입니다.

이 연구는 EHR 기반 모델의 성능 향상을 위한 핵심 레버 (Lever) 로서 토큰화 전략의 중요성을 입증하고, 데이터 효율적인 모델 설계를 위한 구체적인 지침을 제시했습니다.