사전학습 네트워크로 가속하는 인과 구조 탐색

초록

본 논문은 사전학습된 Prior‑Fitted Networks(PFN)를 이용해 데이터 의존적 로그우도 추정을 amortize함으로써 인과 구조 탐색의 정확성을 크게 향상시킨다. PFN 기반 우도 추정기가 기존 MLP 기반 추정기보다 변동성이 적고 편향이 낮으며, 이를 정책 그래디언트(PPO)와 결합해 DAG 샘플링을 강화학습 형태로 최적화한다. 실험 결과, 합성·반합성·실제 데이터셋 모두에서 구조 회복 성능이 기존 PC, DCDI 등 베이스라인을 능가한다.

상세 분석

이 연구는 인과 구조 학습을 두 단계(우도 추정, 구조 최적화)로 분리하고, 첫 번째 단계의 근본적인 한계인 우도 추정 오차를 PFN을 통해 완화한다는 점에서 혁신적이다. PFN은 대규모 변환기 모델을 사전학습하여 다양한 베이지안 네트워크 사후분포를 근사하도록 설계되었으며, 테스트 시에는 학습 데이터를 컨텍스트에 포함시켜 단일 포워드 패스로 조건부 예측분포 qθ(y|x,Dtrain)를 출력한다. 이 과정은 전통적인 MLP 기반 추정기가 필요로 하는 개별 모델 학습을 없애고, 데이터 규모가 작아도 안정적인 로그우도 값을 제공한다.

구조 최적화는 Gumbel‑Softmax와 SoftSort를 이용해 연속적인 확률 행렬 P, W 로부터 이산적인 DAG A를 샘플링하고, 샘플링된 그래프에 대해 각 변수 i 의 부모 집합 paA(i) 에 대해 PFN을 “인‑컨텍스트” 방식으로 피팅한다. 이렇게 얻은 log p_i(D_est|D_train,paA(i)) 들을 합산해 전체 스코어 ŝ(A,D) 를 계산하고, 이를 보상으로 PPO가 그래프 posterior πθ(A) 를 업데이트한다. 정책 그래디언트 접근법을 채택함으로써 우도 추정이 비미분 가능하더라도 전체 파이프라인을 효율적으로 학습할 수 있다.

실험에서는 5노드·5엣지 Erdős‑Rényi 그래프를 기반으로 30개의 부트스트랩 데이터셋을 생성해 PFN과 MLP의 우도 추정 정확도를 비교하였다. PFN은 평균·중위수 변동성이 크게 낮고, NLL도 감소했으며, 특히 125샘플 이하를 제외한 모든 규모에서 정답 그래프를 정확히 복원했다. 구조 회복 측면에서는 SHD 기준으로 합성 데이터에서 기존 PC·DCDI보다 현저히 낮은 값을 기록했고, SERGIO와 Causal Chambers 같은 반합성·실제 데이터에서도 경쟁력 있는 성능을 보였다. 다만 PFN의 사전 분포가 실제 도메인과 불일치할 경우 성능 저하가 관찰돼, 도메인‑특화 사전 학습이나 파인튜닝이 향후 과제로 제시된다.

이 논문은 “우도 추정만을 amortize하고 구조 최적화는 기존 이론적 보장을 유지한다”는 설계 철학을 통해, 최근 제안된 완전 amortized 방식이나 대형 언어 모델 기반 접근법이 갖는 이론적 불확실성을 회피한다. 또한, PFN을 활용한 베이지안 추정이 작은 데이터 환경에서도 강건함을 보여, 인과 탐색을 실험실을 넘어 실제 과학·공학 문제에 적용할 가능성을 크게 확장한다.