기술 요약: EHR 파운데이션 모델에서의 ICD 코드 계층적 모델링

문제 정의

BEHRT 및 Med-BERT와 같은 전자 건강 기록(EHR) 파운데이션 모델은 일반적으로 국제 질병 분류(ICD) 진단 코드를 평면적이고 독립적인 토큰으로 취급한다. 이러한 접근 방식은 질병을 패밀리, 하위 범주, 그리고 점진적으로 구체화되는 진단명으로 조직하는 ICD 코딩 시스템 고유의 임상적으로 유의미한 계층 구조를 간과한다. 결과적으로, 기존의 표현 학습 방법은 이러한 온톨로지를 명시적으로 활용하지 못하며, 서로 관련된 진단을 고립된 기호로 취급한다. 이는 진단명이 희소하고, 롱테일 분포를 보이며, 기관별로 특화되어 있는 임상 데이터 환경에서 특히 문제가 된다. 본 논문은 이 계층 구조를 유도 편향(inductive bias)으로 명시적으로 모델링하는 것이, 희귀 코드가 임상적으로 유사한 코드와 관련된 데이터로부터 이득을 얻게 하고, 더 넓은 케어 경로를 포착하며, 도메인 변화(domain shift) 하에서의 강건성을 개선함으로써 표현 학습을 향상시킬 수 있다는 가설을 세운다.

방법론

저자들은 현대적인 EHR 표현 학습에 ICD-10-CM 계층 구조를 주입하기 위한 두 가지 상호 보완적인 메커니즘을 제안하고 평가하며, 이를 통칭하여 계층적 ICD 모델(HICD)이라 부른다. 이들은 가장 거친 입도에서 가장 미세한 입도에 이르는 세 가지 중첩된 계층 수준($G_0, G_1, G_2$)을 정의한다.

1. HICD-BERT (토큰 레벨 주입)

이 접근 방식은 BEHRT 트랜스포머 아키텍처를 확장한다. ICD 코드를 원자적 식별자로 취급하는 대신, 문자열 접두사 절단(string-prefix truncation)을 통해 도출된 계층 특화 임베딩으로 표준 토큰 임베딩을 증강한다.

• 메커니즘: 진단 코드에 대해, 모델은 코드 문자열의 다양한 수준(예: 1글자 접두사, 2글자 접두사, 3글자 카테고리)에 해당하는 가산 임베딩(additive embeddings)을 더한다.
• 데이터 소스: 이 방법은 데이터 구동형(data-driven)으로, 훈련 시 외부 온톨로지에 접근할 필요가 없다. 계층 구조는 코드 문자열로부터 직접 도출된다.
• 훈련: 진단 시퀀스에 대한 마스크드 언어 모델링(MLM)을 사용하여 사전 훈련된 후, 다운스트림 태스크를 위해 미세 조정된다.

2. HICD-Graph (그래프 레벨 인코딩)

이 접근 방식은 트랜스포머 인코더를 초기화하기 전에 코드-그래프 레벨에서 계층 구조를 주입한다.

• 메커니즘: 하이브리드 진단 공생(co-occurrence) 그래프가 구축된다. 이는 다음을 결합한다:
  • 경험적 엣지(Empirical edges): 환자 방문 내 진단 공생에 기반한 점별 상호 정보량(PMI)에 의해 가중치가 부여된다.
    고정 가중치 엣지로 리프(leaf) 진단 노드를 ICD 계층의 조상 노드($G_0, G_1, G_2$)에 연결한다.
• 훈련: 링크 예측 목적 함수를 사용하여 그래프 합성곱 신경망(GCN)을 훈련시켜 계층 인식 코드 임베딩을 학습한다. 이 임베딩들은 환자 레벨 트랜스포머 인코더의 임베딩 테이블을 초기화한다.

실험 설정

• 데이터셋: MIMIC-IV (사전 훈련 및 인도메인 평가용) 및 eICU (교차 데이터셋 전이 평가용).
• 태스크:
  1. 30일 재입원 (MIMIC-IV): 이진 분류.
  2. 응급 입원 예측 (MIMIC-IV): 이진 분류.
  3. ICU 재입원 (eICU): 전이 평가를 위한 이진 분류 (인코더는 동결, 분류기는 미세 조정).
• 절제 연구(Ablation): 본 연구는 계층 구조가 없는 베이스라인에 대해 두 아키텍처 모두에서 세 가지 계층 수준($G_0, G_1, G_2$)의 모든 8가지 조합을 체계적으로 평가한다.

주요 결과

인도메인 성능 (MIMIC-IV)

• 전반적 개선: 계층 구조가 증강된 구성들이 두 모델과 태스크 모두에서 계층 구조가 없는 베이스라인보다 우수한 성능을 보였다.
• HICD-BERT: 성능 향상은 주로 가장 미세한 입도 수준($G_2$)에 의해 주도되었다. 중간 수준($G_1$) 단독으로는 성능을 유의미하게 개선하지 못했으나, 다른 수준과 결합되었을 때는 기여하였다. 최적의 구성은 태스크에 따라 달랐다 (예: 재입원의 경우 모든 수준 사용, 응급 예측의 경우 단일 미세 수준 사용).
• HICD-Graph: 계층 수준을 추가할수록 성능이 점진적으로 증가하는 단조 증가 추세를 보였다. 그래프 기반 아키텍처는 여러 입도 수준의 정보를 동시에 효과적으로 통합했다.

교차 데이터셋 전이 (MIMIC-IV에서 eICU로)

• HICD-Graph: 견고한 전이성을 입증했다. 7개 중 6개의 계층 변형이 베이스라인보다 유의미하게 개선되었으며, 모든 수준 구성이 가장 강력한 AUROC(0.587)를 달성했다. 그래프 기반의 관계적 구조가 데이터셋 변화에 더 불변적인 것으로 나타났다.
• HICD-BERT: 분포 변화 하에서 계층의 이점이 대부분 사라졌다. 가장 거친 단일 수준 구성($G_0$)만이 베이스라인 대비 유의미한 개선을 보였다. 저자들은 HICD-BERT가 학습한 순차적 토큰 패턴이 MIMIC-IV의 특정 방문 구조에 더 밀접하게 결합되어 있는 반면, 온톨로지에 기반한 그래프 구조는 더 일반화 가능하다는 점을 시사한다.

임베딩 분석

• 의미론적 클러스터링: 계층 인식 모델은 더 일관된 코드 클러스터를 생성했다. 구체적으로, HICD-Graph의 $G_1+G_2$ 구성은 임베딩 공간에서 가장 조밀한 로컬 코드 클러스터를 형성했다.
• 성능과의 상관관계: 가장 조밀하고 임상적으로 의미 있는 클러스터(예: 동일한 2글자 접두사를 공유하는 코드)를 생성한 구성이 다운스트림 예측 성능에서도 가장 강력한 성능을 달성했다. 이는 계층 구조가 임상적 관계를 존중하도록 임베딩 기하학을 형성함으로써 성능을 개선한다는 것을 시사한다.

의의 및 주장

본 논문은 ICD 코딩 시스템에 존재하는 계층 구조가 EHR 파운데이션 모델에서 과소 활용되고 있는 신호라고 주장한다. 주요 기여 및 주장은 다음과 같다:

1. 강건한 유도 편향으로서의 계층 구조: 계층 구조를 명시적으로 인코딩하는 것은 두 가지 아키텍처(트랜스포머 기반 및 그래프 기반)와 태스크 모두에서 다운스트림 예측 성능을 유의미하게 향상시킨다.
2. 입도 의존성: 계층 수준의 효용은 균일하지 않다. 최적의 수준은 모델링 접근 방식(토큰 레벨 vs 그래프 레벨)과 다운스트림 태스크 모두에 따라 달라진다.
3. 구조의 전이 가능성: 그래프 기반 계층 인코딩은 토큰 레벨 주입보다 데이터셋 간 전이가 더 견고하게 이루어진다. 이는 온톨로지 유래의 관계적 구조가 순차적 토큰 패턴보다 더 근본적인 임상 관계를 포착함을 시사한다.
4. 저비용 통합: 핵심적인 사전 훈련 목적이나 파이프라인을 변경하지 않고도 최소한의 아키텍처 수정(가산 임베딩 또는 추가적인 그래프 엣지)만으로 이점을 얻을 수 있다.
5. 일반화 가능성: 구조화된 임상 온톨로지의 가치에 관한 통찰은 ICD-10-CM을 넘어 다른 구조화된 어휘집(예: 약물에 대한 ATC 코드) 및 표현 학습 방법론에도 적용될 수 있다.

저자들은 절대적인 AUROC 개선 폭은 완만하지만, 이는 통계적으로 견고하고 일관적이라는 점을 언급하며 겸허한 태도를 유지한다. 또한 다른 임상 결과, 기관 및 EHR 시스템으로의 일반화에 관한 한계점을 인정하며, 계층 수준이 자동으로 학습된 것이 아니라 수동으로 정의되었다는 점을 명시한다.