Probing Clinical Concepts in an EHR Foundation Model via Sparse Autoencoders

Motivation

• Known: EHR-FM은 MIMIC-IV 같은 대규모 임상 이벤트 스트림으로 학습되어 mortality, readmission, phenotyping 등 다운스트림 태스크에 잘 전이되는 것으로 알려져 있다. 한편 SAE는 LLM의 residual stream을 monosemantic feature로 분해하는 mechanistic interpretability 기법으로 텍스트/이미지 도메인(MedCLIP, MedGemma 등)에서 이미 적용된 바 있다.
• Gap: 기존 EHR-FM 해석 방법은 attention visualization, gradient attribution, SHAP 등 post-hoc 방식에 국한되어 개별 예측만 설명할 뿐 모델이 전체적으로 무엇을 학습했는지 드러내지 못하며, 임상 온톨로지에 매핑하기도 어렵다. 또한 SAE 기반 mechanistic interpretability는 구조화된 clinical event sequence를 다루는 EHR-FM에는 아직 적용된 적이 없다.
• Why: EHR-FM은 챗봇 LLM과 달리 고위험 임상 배치를 목표로 하기 때문에, 예측 정확도를 넘어서 모델이 병태생리를 학습했는지 허위 상관을 학습했는지 감사(audit)하는 것이 임상 안전성 측면에서 매우 중요하다.
• Approach: MIMIC-IV로 학습된 transformer 기반 EHR-FM의 residual stream에 TopK SAE를 학습시키고, LLM 기반 해석 파이프라인과 activation patching을 결합하여 학습된 feature들의 임상적 의미와 계산적 역할을 검증한다.

Achievement

Figure 4. Activation patching heatmaps (clustered). Each cell

1. Cross-system syndrome discovery: SAE feature들이 Cardiorenal, Hepatorenal, SIRS, UTI, AMI, DKA 등 ICD 챕터에 흩어져 있는 알려진 임상 증후군들을 단일 monosemantic feature로 인코딩함을 보였으며, 증후군 매칭이 등장하는 layer 깊이가 조건별로 다름(예: pancreatitis/AMI는 Layer 5, UTI/DKA는 Layer 1)을 확인했다.
2. Prevalence 기반 feature 조직화: 300개 해석 feature를 통해 SAE feature들이 ICD의 해부/병인 계층이 아니라 활성화 빈도(prevalence)라는 축을 따라 조직됨을 보였으며, 이 band 구조가 모든 layer에서 일관되게 나타남을 확인했다.
3. SAE feature의 계산적 매개 증거: 400개 source-target pair 중 69.7%에서 Layer 1의 단일 feature ablation이 Layer 3/5의 sparse한 feature 집합에 5% 이상의 다운스트림 변화를 유발함을 보여, raw lab signal이 layer를 거쳐 증후군으로 조합되는 information-flow circuit의 개입적(interventional) 증거를 제시했다.

How

Figure 4. Activation patching heatmaps (clustered). Each cell

• 6-layer, 8-head causal transformer(27.5M parameters, dmodel=256)를 MIMIC-IV v3.1(169M events, 364,627 patients)에 대해 MEDS-Torch 방식(triplet encoding: code, value, Δt)으로 학습
• Layer 1, 3, 5의 residual-stream activation(≈5M vectors/layer)에 대해 TopK SAE(dictionary size 4,096, k=32)를 5 epoch 학습, decoder column을 unit-norm 투영
• 각 alive feature에 대해 상위 30개 활성화 입력을 추적하여 31-event window를 추출, MIMIC-IV reference table로 코드를 사람이 읽을 수 있는 설명으로 변환한 뒤 LLM을 judge로 사용해 label, description, category, abstraction level, known syndrome match, confidence를 구조화된 JSON으로 산출(총 300개 feature 해석)
• 세 가지 축으로 정량 검증: (i) activation patching을 통한 interventional 검증, (ii) held-out ICD phenotype과의 AUROC 비교를 통한 external validity, (iii) sentence-transformer UMAP 클러스터링을 통한 internal consistency 검증

Originality

• SAE 기반 mechanistic interpretability를 텍스트/이미지 도메인을 넘어 구조화된 clinical event stream을 다루는 EHR-FM에 최초로 확장 적용
• 도메인 전문성이 필요한 의료 코드 시퀀스의 feature 해석 문제를 LLM 기반 파이프라인(코드→자연어 설명→LLM judge)으로 해결
• activation patching을 통해 SAE feature 간 cross-layer information-flow circuit을 규명하여, 단순 상관관계가 아닌 계산적 인과 역할을 실증적으로 검증한 접근

Limitation & Further Study

• 단일 기관(MIMIC-IV) 데이터로 학습된 하나의 EHR-FM에 대해서만 검증되어, 다른 EHR 시스템이나 더 큰 규모 모델에 대한 일반화 가능성이 불확실함
• 300개 feature(layer당 100개)만 수동/LLM 해석되어 전체 dictionary(4,096×3) 중 극히 일부만 다루며, 나머지 feature의 의미는 미확인 상태로 남음
• LLM 기반 해석 자체의 신뢰성(judge의 편향, confidence rating의 주관성)에 대한 별도 검증이 제한적임
• 후속 연구로 더 다양한 EHR 데이터셋과 더 큰 SAE dictionary, 그리고 자동화된 feature steering을 통한 임상적 개입 가능성 탐구가 필요함

Evaluation

총평: EHR-FM 해석가능성 연구의 공백을 겨냥해 SAE와 LLM 기반 해석, activation patching을 결합한 참신하고 방법론적으로 탄탄한 워크숍 논문으로, 향후 임상 배치를 위한 모델 감사(auditing) 도구로서의 잠재력을 잘 보여준다.