[CausalFM] Foundation Models for Causal Inference via PFNs

https://arxiv.org/pdf/2506.10914

(Sep 2025)

https://github.com/yccm/CausalFM

GitHub - yccm/CausalFM: PyTorch Implementation on Paper: Foundation Models for Causal Inference via Prior-Data Fitted Networks

---

결정판이 나왔다!!
(물론 향후 더 발전된 model이 나오겠지만)

Do-PFN, CausalPFN을 보면서 느꼈던 한계점들을 명시적으로 지적하고, 보완하여서

다양한 causal inference setting에 적용할 수 있는 PFN-based foundation model을 만든다.

1. 첫번째로, 나는 Amortized bayesian inference setting에서 prior fitting이 굉장히 중요하다고 생각되는데,
그렇다면 prior distribution을 정의하고, prior dataset을 sampling하는 과정이 합리적으로 구체적이어야 한다고 생각한다.

즉, "prior design이 induced된 posterior로 하여금 유효한 Bayesian Causal Inference가 가능하도록 하여야 한다"는 것이다. 

- Do-pfn는 SCM prior를 제시하였는데, 이를 통한 data generating process는 잘 이해가 되었고, 납득이 갔다. 
다만 기본적인 structure, simple한 structural equation을 사용하였는데, 좀더 다양하고 복잡한 SCM prior를 generate하는 procedure가 아쉬웠다.

- CausalPFN의 data generating process는 tabular dataset (table)을 가져와서, covariate, potential outcome을 assign하는 방식이었는데 (내가 이해한 게 맞다면), 이 방식은 솔직히 전혀 납득이 되지 않고, 별로 사용하고 싶지 않았다.. (내가 오해하거나 틀린 걸 수도 있다). 
왜냐하면 underlying data generating process (ground-truch SCM)을 모르기 때문이다.

그리고, 두 논문 모두에서 느낀 건, 이렇게 prior를 generate하는 게 맞나..? 아쉬운 느낌이 남았다.

=> 본 논문 (CausalFM)은 이런 점을 시원하게 긁어준다. 

Systematic하게 prior data를 generating하는 procedure를 제시하고, valid한 prior에 대한 기준을 명확히 제시한다. 

그리고 Bayesian neural network를 사용한 prior로써 다양한 상황 하에서 Bayesian causal inference를 가능하게 했다는 게 매우 인상적이다!

2. Prior-fitting을 통해서 test time inference 시에 in-context learning으로 causal inference를 수행하는 것은, 앞선 모델들과 같다. 

우수한 performance를 보여준다. 

중요한 건, back-door adjustment만 가능했던 앞선 model들과 달리, front-door, instrumental variable adjustment, unobserved confounders 까지 가능하게끔 했다.

(이는 본 논문의 주장이고, 내가 본 Do-pfn에서도 기본적인 dag structure에 front-door, mediator, unobserved confounder 등도 포함하여 실험 결과를 제시했는데, 왜 본 논문에서는 Do-pfn이 back-door adjustment만 가능하다고 주장하는지 이해를 못했다..)

3. Bayesian causal inference 답게 uncertainty quantification을 구체적으로 명시했다.

4. identifiability guarantee를 구체적으로 명시했다.

특히, identification step과 estimation step를 구분하였다는 게 인상적이다.

전반적으로 causal inference의 기본 철학을 굉장히 충실하게 반영하였다. 

---

여담으로 단점은..

figure가 너무 없다.. ㅠㅠ revised version이 나올까요..? 

 

상당히 detail하게 task setting, prior design, training, architecture 등을 서술해주는데, 

과정이 복잡하기 때문에 figure가 아쉽다.. ㅜㅜ

fi..fi...figure.... 좀..

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---