Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

이 논문은 합성 시계열 인과 모델로 사전 학습되어 도메인별 훈련 없이도 반사실적 종단적 결과를 예측하는 동결된 사전 적합 네트워크인 CausalLongPFN을 소개하며, 암, HIV, MIMIC-III와 같은 벤치마크에서 경쟁력 있는 성능을 입증하는 동시에 전통적인 코호트별 모델링에 대한 비용 효율적인 대안을 제공한다.

핵심

당신이 의사로서 환자가 향후 몇 주 동안 특정 치료 계획에 어떻게 반응할지 예측하려고 한다고 상상해 보십시오. 문제는 환자들이 매우 복잡하다는 점입니다: 환자의 건강 상태는 매일 변하고, 과거의 치료 이력이 현재 상태에 영향을 미치며, 식단이나 스트레스와 같은 다른 요인들이 결과에 간섭합니다. 보통 이러한 예측을 수행하려면, 새로운 환자 집단을 마주할 때마다 그들에게 특화된 완전히 새로운 컴퓨터 모델을 구축해야 합니다. 이는 마치 새로운 동네로 이사할 때마다 집을 설계하기 위해 새로운 건축가를 고용하는 것과 같습니다. 이는 느리고, 비용이 많이 들며, 많은 데이터가 필요합니다.

이 논문은 CAUSALLONGPFN이라는 새로운 도구를 소개합니다. 이것을 **"보편적 의료 직관 엔진(Universal Medical Intuition Engine)"**이라고 생각하십시오. 이 엔진은 실제 환자를 만나기 전에 이미 게임의 규칙을 학습한 상태입니다.

작동 원리는 다음과 같습니다:

1. "훈련 캠프" (합성 사전 학습)

실제 환자들이 도착하기를 기다리는 대신, 제작자들은 거대한 가상 "훈련 캠프"를 구축했습니다. 이 캠프에서는 수백만 명의 가짜 환자, 즉 수백만 가지의 서로 다른 체형, 질병, 치료 반응을 가진 환자들을 시뮬레이션했습니다. 그들은 가짜 환자들에게 다음과 같은 복잡한 행동 양식을 프로그래밍했습니다:

• 어떤 환자는 천천히 회복됩니다; 다른 환자는 급격히 악화됩니다.
• 어떤 치료는 즉각적인 효과를 보이지만, 다른 치료는 지연된 효과를 보입니다.
• 어떤 환자는 숨겨진 유전적 요인에 따라 다르게 반응합니다.

AI 모델은 이 캠프에서 모든 시간을 보내며, 가짜 환자들의 결과를 예측하는 법을 배웠습니다. 모델은 단순히 정답을 암기한 것이 아니라, 치료, 시간, 그리고 생물학적 요인이 어떻게 상호작용하는지에 대한 근본적인 논리를 학습했습니다.

2. "얼어붙은 전문가" (새로운 학습 불필요)

여기 마법 같은 기술이 있습니다: 훈련 캠프를 마친 후, 연구진은 모델을 **동결(frozen)**했습니다. 모델의 뇌를 잠가버린 것입니다. 모델은 더 이상 새로운 것을 배우거나 내부 설정을 변경할 수 없습니다.

새로운 실제 환자 그룹(예: 병원의 암 환자들)이 도착했을 때, 모델은 처음부터 시작하지 않습니다. 모델은 새로 시작하는 대신, 차트를 읽는 초보 인턴처럼 행동합니다.

• 지원 궤적 (The Support Trajectories): 당신은 현재 병원의 환자 사례 몇 가지(지원 데이터)를 모델에게 보여줍니다. 이는 마치 인턴에게 해당 병원의 특수한 스타일을 이해할 수 있도록 몇 가지 케이스 파일을 읽게 하는 것과 같습니다.
• 쿼리 (The Query): 당신은 질문합니다. "만약 이 특정 환자에게 '치료 A'를 5일 동안 시행한다면 어떤 일이 벌어질까?"
• 답변: 모델은 훈련 캠프에서 배운 내용과 당신이 보여준 케이스 파일의 맥락을 결합하여 즉시 결과를 예측합니다. 이 과정에서 모델은 일반적인 재학습 과정인 "경사 하강법(gradient descent)"을 단 한 단계도 거치지 않습니다.

3. "시간 여행 시뮬레이터"

이 모델은 종단적(longitudinal) 데이터를 처리하도록 설계되었습니다. 즉, 모델은 시간을 이해합니다. 모델은 단순히 내일의 결과를 추측하는 것이 아니라, 미래를 단계별로 시뮬레이션합니다.

• 모델은 1일 차를 예측합니다.
• 그 예측값을 2일 차의 시작점으로 사용합니다.
• 이 과정을 반복하여 5일, 6일 또는 7일 후에 어떤 일이 일어날지 확인합니다.

이는 비행 시뮬레이터가 이착륙 장면만 보여주는 것이 아니라, 기상 상황이 중간에 변하더라도 조종사의 결정에 따라 전체 비행 경로를 시뮬레이션하는 것과 같습니다.

4. 왜 이것이 중요한가 (결과)

저자들은 이 "얼어붙은 전문가"를 기존 방식(각 병원마다 새로운 모델을 구축하는 방식)과 비교 테스트했습니다.

• 테스트: 암, HIV, 와파린(혈액 응고 저해제), 그리고 실제 중환자실(ICU) 기록 데이터를 사용했습니다.
• 결과: 동결된 모델은 각 특정 질병에 맞춰 특별히 훈련된 모델들과 대등하거나, 때로는 더 나은 성능을 보였습니다.
• 큰 성과: 이 모델은 "만약 ~한다면(what-if)" 시나리오를 실험적으로 테스트하기 어려운 실제 중환자실 데이터에서 특히 뛰어난 성능을 보였습니다(중환자실에서는 윤리적인 이유로 다양한 치료법을 시뮬레이션하며 실험할 수 없기 때문입니다). 모델은 관찰된 데이터만을 바탕으로 다음에 어떤 일이 일어날지를 예측할 수 있었습니다.

핵심 요약

이 논문은 새로운 의료 데이터셋마다 항상 맞춤형 모델을 만들 필요는 없다고 주장합니다. 대신, 광범위한 "만약 ~한다면" 시나리오(합성 데이터)를 통해 하나의 거대한 모델을 훈련시킨 뒤, 이를 동결된 상태의 즉시 사용 가능한 도구로 사용할 수 있습니다.

이는 모든 종류의 요리를 가상 주방에서 연습해 본 마스터 셰프를 보유하는 것과 같습니다. 새로운 식재료(새로운 환자 그룹)를 가져왔을 때, 셰프는 요리법을 다시 배울 필요가 없습니다. 그저 가져온 식재료를 보고, 자신이 이미 학습한 방대한 직관을 사용하여 즉시 요리를 준비하면 됩니다.

주의 사항: 저자들은 이 도구가 예측 및 연구를 위한 것이며, 임상 현장에서 최종적인 생사 결정을 내리기 위한 것이 아님을 분명히 하고 있습니다. 이 도구는 의사가 잠재적 결과를 이해하는 데 도움을 주지만, 여면에 다른 모든 인과 모델과 마찬가지로 표준적인 의학적 가설에 의존합니다. 이는 강력한 연구 도구이지, 의사의 판단을 대체하는 것이 아닙니다.

---

구현 및 모델 다운로드:
코드와 사전 학습된 모델 가중치를 직접 확인하고 사용하려면 아래 링크를 참조하십시오.

• Code on GitHub: https://github.com/Amirhossein-Zare/causal-long-pfn
• Pretrained model weights on Hugging Face: https://huggingface.co/Amirhossein-Zare/causal-long-pfn

기술 요약

기술 요약: 인과적 종단적 사전 적합 네트워크 (Causal Longitudinal Prior-Fitted Networks)를 이용한 반사실적 결과 예측

문제 정의

본 논문은 종단적 관측 데이터에서 미래의 치료 시퀀스에 따른 잠재적 결과(potential outcomes)를 예측하는 과제를 다룹니다. 핵심 과제는 $H_t$(시간 $t$까지 관찰된 이력)가 주어졌을 때, 계획된 미래 치료 시퀀스 $\bar{a}$에 대한 이력 조건부 잠재적 결과 $E[Y(\bar{a})_{t+\tau} | H_t]$를 추정하는 것입니다.

이 문제는 세 가지 주요 요인으로 인해 복잡해집니다:

1. 시변 혼란 (Time-varying confounding): 각 단계에서의 치료 할당은 이전 치료의 결과물인 공변량들에 의존합니다.
2. 이질적인 환자 역학 (Heterogeneous patient dynamics): 개별 궤적은 복잡하고 비선형적인 상태 진화와 잠재적 이질성을 보입니다.
3. 데이터의 한계: 관측 코호트는 새로운 도메인이나 시뮬레이터를 위해 매번 신뢰할 수 있는 딥 시퀀스 모델을 처음부터 학습시키기에는 너무 작을 수 있습니다.

기존의 종단적 인과 추정기(예: RMSN, CRN, G-Net, Causal Transformer)는 일반적으로 새로운 코호트마다 하이퍼파라미터 선택 및 성향 점수(propensity) 모델링을 포함한 별도의 지도 학습 훈련 과정을 필요로 합니다. 이러한 파이프라인은 반복적인 도메인별 훈련이 요구되는 상황에서 비용이 많이 들고 비실용적입니다.

방법론: CAUSALLONGPFN

저자들은 광범위한 시간적 구조 인과 모델(TSCM)의 사전 분포(prior)에 걸쳐 종단적 인과 예측을 아모티제이션(amortize)하도록 설계된 사전 적합 인컨텍스트 예측기인 CAUSALLONGPFN을 소개합니다.

1. 광범위한 사전 분포를 이용한 합성 사전 학습

모델은 광범위한 TSCM의 사전 분포에서 샘플링된 합성 에피소드 상에서 전체적으로 사전 학습됩니다. 이 사전 분포는 단일 시뮬레이터를 복제하기보다는 광범위한 클래스의 종단적 인과 역학을 포괄하도록 설계되었습니다. 합성 사전의 주요 특징은 다음과 같습니다:

• 인과적 시간 그래프 (Causal Temporal Graphs): 희소한 동시적 및 지연된 의존성을 가지며, 비순환적 즉각적 그래프를 가집적합니다.
• 비선형 메커니즘 (Nonlinear Mechanisms): 상태 좌표는 다양한 기초 비선형성(identity, tanh, sinusoidal, ReLU 등)과 다양한 노이즈 분포를 사용하는 희소 비선형 자기회귀 업데이트를 따릅니다.
• 역학적 모티프 (Dynamical Motifs): 지연 효과나 조절 피드백과 같은 질적 메커니즘을 포착하기 위해 행동-메모리(action-memory), 포화(saturating), 항상성(homeostatic), 피드백 제어(feedback-control), 평활화된 출력(smoothed-readout) 채널과 같은 구조화된 모티프가 중첩됩니다.
• 혼란된 행동 정책 (Confounded Behavior Policies): 치료는 잠재적 단위 이질성($Z_i$)의 영향을 받는 상태 의존적 확률 정책으로부터 샘플링되어, 시변 치료-혼란 변수 피드백을 생성합니다.
• 결과 모델 (Outcome Models): 스칼라 결과값은 직접적이고 누적적인 치료 효과를 갖는 자기회귀적 리드아웃(readout)을 통해 생성됩니다.

2. 아키텍처

CAUSALLONGPFN은 듀얼 인코더 아키텍처를 채택합니다:

• 인과적 이력 인코더 (Causal History Encoder): 궤적 수준의 인과적 트랜스포머(마스크드 셀프 어텐션 사용)를 사용하여 종단적 시퀀스를 이력 표현으로 매핑하며, 시간 $r$에서의 표현이 오직 해당 시점까지 가용한 정보에만 의존하도록 보장합니다.
• PFN 컨텍스트 인코더 (PFN Context Encoder): 인컨텍스트 적응을 수행하는 트랜스포머 기반 컨텍스트 인코더입니다. 서포트 궤적(순서가 없는 집합으로 취급됨)과 쿼리 토큰을 공동으로 처리하기 위해 셀프 어텐션을 사용합니다. 서포트 궤적의 순서에 대한 위치 인코딩은 할당되지 않습니다.
• 가우시안 혼합 예측 헤드 (Gaussian-Mixture Prediction Head): 최종 쿼리 표현은 정규화된 다음 결과값에 대한 5개 성분의 가우시안 혼합 분포를 매개변수화하여, 점 예측과 불확실성 추정을 모두 제공합니다.

3. 인컨텍스트 예측 및 롤아웃

테스트 시 모델은 동결(frozen) 상태입니다. 모델은 다음을 입력받습니다:

• 서포트 궤적 (Support trajectories): 새로운 도메인/태스크로부터의 예시들.
• 쿼리 이력 (Query history): 시간 $t_{obs}$까지 관찰된 내용.
• 제안된 미래 치료 시퀀스: 계획된 개입.

모델은 그래디언트 업데이트, 성향 모델 적합, 또는 적대적 밸런싱 없이 예측 분포를 반환합니다. 다단계 예측($\tau > 1$)의 경우, 모델은 **자기회귀적 플러그인 롤아웃(autoregressive plug-in rollout)**을 수행합니다: 1단계 결과 분포를 예측하고, 그 혼합 평균(mixture mean)을 쿼리 이력에 삽로한 뒤, 지정된 치료 시퀀스에 따라 이 과정을 반복합니다.

주요 기여

1. 종단적 인과 예측을 위한 사전 적합 모델: CAUSALLONGPFN은 계획된 종단적 치료 시퀀스 하에서 이력 조건부 잠재적 결과 예측을 위한 최초의 PFN 스타일 모델입니다. 이 모델은 테스트 시 적응이 필요 없는 동결된 모델로 작동합니다.
2. 종단적 인과 태스크를 위한 합성 사전 분포: 저자들은 치료-혼란 변수 피드백, 잠재적 이질성, 비선형 역학, 지연/누적 효과, 그리고 체제 변화(regime changes)를 포함하는 다양한 태스크를 생성하는 TSCM 사전 분포를 설계했습니다.
3. 종단적 인컨텍스트 추론을 위한 아키텍처: 인과적 트랜스포머 이력 인코더와 PFN 컨텍스트 인코더, 그리고 가우시안 혼합 헤드를 결합한 새로운 듀얼 인코더를 제안합니다.
4. 자기회귀적 반사실적 롤아웃: 결정론적 플러그인 롤아웃을 통해 학습된 1단계 예측기를 다단계 예측으로 확장했습니다.
5. 제로샷 평가: 분기 가능한 반사실적 벤치마크와 실제 데이터(MIMIC-III)에 대해 단일 동결 모델을 도메인 학습 기반 베이스라인(MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer)과 비교하는 종합적인 평가를 수행했습니다.

결과

모델은 암 종양 성장, 와파린(Warfarin) PK/PD, HIV 치료 역학, MIMIC-III ICU 궤적의 네 가지 벤치마크에서 평가되었습니다.

• 도메인 균형 성능 (Domain-Balanced Performance): CAUSALLONGPFN은 최상의 도메인 균형 1단계 정규화 RMSE(0.222)를 달성하여 MSM과 RMSN을 근소한 차이로 앞섰습니다. 5단계 예측의 경우, RMSN과 G-Net에 이어 전체 3위를 기록했으나, MSM, CRN 및 트랜스포머 기반 베이스라인보다는 우수한 성능을 보였습니다.
• 도메인별 성능:
  • MIMIC-III (Factual): 모델은 1단계 및 5단계 예측 모두에서 1위를 차지하며, 도메인 특화 훈련 없이도 실제 임상 궤적에 대한 강력한 전이 능력을 입증했습니다.
  • 반사실적 벤치마크 (Cancer, HIV, Warfarin): 모델은 1단계 태스크에서 2위 또는 3위를 기록하며 경쟁력을 유지했습니다. 그러나 장기 호흡의 반사실적 태스크(예: 5단계 암 예측)에서는 특정 도메인에 적합한 전문화된 재귀 모델(RMSN, CRN)이 더 낮은 오차를 기록했는데, 이는 충분한 타겟 도메인 데이터가 존재하여 특정 적합(fitting)이 가능할 때의 이점을 시사합니다.
• 불확실성 교정 (Uncertainty Calibration): 가우시안 혼합 헤드는 유용한 분포 정보를 제공했습니다. 교정 성능은 도메인마다 달랐는데, 와파린은 가장 좋은 교정 성능을 보였고, MIMIC-III는 높은 이질성으로 인해 더 넓은 구간을 보였습니다.

의의 및 주장

본 논문은 다음과 같은 경우에 광범위한 합성 사전 학습이 반복적인 도메인별 훈련에 대한 유용한 동결 대안이 될 수 있음을 주장합니다:

• 전문화된 모델을 재학습시키는 것이 비용이 많이 들거나 비실용적인 경우.
• 새로운 코호트에 대한 빠른 적응이 필요한 경우.
• 반사실적 감독(counterfactual supervision)을 사용할 수 없는 경우(MIMК-III와 같은 실제 사실적 예측 태스크).

저자들은 CAUSALLONGPFN이 관측 데이터를 해석하는 데 필요한 표준 인과 가정(일관성, 양의성, 순차적 교환 가능성)을 제거하는 것이 아니라, 추정 문제를 아모티제이션(amortize)하는 것이라고 강조합니다. 결과는 충분히 광범한 합성 사전 분포가 치료-반응 태스크 전반에 걸쳐 재사용 가능한 구조를 포착할 수 있어, 모델을 강력한 범용 인컨텍스트 예측기로 만든다는 것을 시사합니다. 다만, 저자들은 특정 도메인의 충분한 타겟 데이터와 검증 신호가 존재하여 장기 호서(long-horizon) 반사실적 예측을 수행할 때는 도메인별 훈련이 여전히 가치가 있다는 점을 겸허히 언급했습니다.

이 연구는 독립적인 임상 의사 결정 시스템이라기보다 인과적 시퀀스 모델링 및 가설 생성을 위한 연구 도구로서의 위치를 점하며, 인과적 가정이나 사전 분포의 지원이 불충분할 때 예측을 과신하지 말 것을 경고합니다.

코드 및 모델 가용성

본 연구의 구현 코드는 GitHub(https://github.com/Amirhossein-Zare/causal-long-pfn)에서, 사전 학습된 모델 가중치는 Hugging Face(https://huggingface.co/Amirhossein-Zare/causal-long-pfn)에서 공개되어 있습니다.