📊 statistics

DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms

이 논문은 이질적이고 고차원적인 표 형식 데이터에서 최첨단 인과 발견 성능을 달성하기 위해 Mixture-of-Leaf-Experts 메커니즘을 갖춘 2단계 자기회귀 트랜스포머 아키텍처를 활용하는 새로운 파운데이션 모델인 DAG-FM을 소개한다.

원저자: Yikang Chen, Zhengkang Guan, Haoyuan Qian, Peng Cui, Yi Yang, Kun Kuang

게시일 2026-07-14

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yikang Chen, Zhengkang Guan, Haoyuan Qian, Peng Cui, Yi Yang, Kun Kuang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 미스터리를 해결하려는 탐정이라고 상상해 보세요. 누가 무엇을 일으키고 있는지를 밝혀내는 것이 임무입니다. 당신에게는 단서(데이터)가 가득 담긴 거대한 스프레드시트가 있지만, 게임의 규칙은 까다롭습니다. 때로는 단서들이 엄격하고 예측 가능한 법칙을 따르기도 하지만, 때로는 단서마다 그 법칙이 바뀌기도 합니다. 이것이 바로 **인과 발견(causal discovery)**의 세계입니다. 오랫동안 탐정들이 사용해 온 도구들은 눈을 가린 채 장갑을 끼고 루빅스 큐브를 맞추려는 것과 같았습니다. 그 도구들은 단순한 퍼즐에서 막히거나, 규칙이 복잡해지면 망가져 버리곤 했습니다.

여기에 새로운 종류의 "슈퍼 탐정" 파운데이션 모델인 DAG-FM이 등장했습니다. 이를 단 하나의 경직된 규칙책이 아니라, 거대하고 마법 같은 주방을 가진 숙련된 요리사라고 생각해보세요.

과거의 방식 vs. 새로운 셰프

이전의 탐정들은 연결 관계의 전체 지도(누가 누구를 유발하는지)를 한꺼번에 추측하려고 노력했습니다. 그것은 도시 전체의 지도를 한 번에 그리려는 것과 같았습니다. 만약 도시의 교통 규칙이 기묘하게 뒤섞여 있다면(이질적 메커니즘), 완성된 지도는 엉망이 되기 일쑤였습니다. 기존의 방법 중 일부는 직선이고 평탄한 고속도로만 달릴 줄 아는 전문가(선형 비가우시안 모델)와 같았습니다. 만약 길이 울퉁불퉁한 흙길이나 굽이굽이 휘어진 산길로 변한다면, 그 전문가들은 길을 잃고 말았습니다.

DAG-FM은 게임의 판도를 바꿉니다. 전체 지도를 한 번에 그리려 하는 대신, "역방향 보물찾기" 게임을 수행합니다. 이 모델은 두 단계로 작동합니다:

잎(Leaves) 찾기: 데이터를 보고 질문합니다. "체인의 맨 끝에 있는 사람은 누구인가? 다른 누구도 일으키지 않는 존재는 누구인가?" 그리고 체인의 마지막 사람을 하나씩 벗겨냅니다.
부모(Parents) 찾기: "잎"이 누구인지 알게 되면, 다시 질문합니다. "누가 이 잎의 줄을 당기고 있었는가?" 이를 통해 그 잎의 부모를 찾아냅니다.

이 과정은 양파 껍질을 벗기듯 전체 구조가 드러날 때까지 반복됩니다. 이는 전체 그림을 한꺼번에 추측하는 것보다 훨씬 똑똑한 방식입니다.

마법의 재료: "전문가 혼합(Mixture of Experts)"

여기서 DAG-FM이 정말 멋진 점이 드러납니다. 현실 세계에서 인과관계의 규칙은 모든 곳에서 동일하지 않습니다. 때로는 원인이 단순한 수학 방정식처럼 작용하기도 하고, 때로는 복잡하고 구불구불한 함수로 작용하기도 합니다.

DAG-FM은 뇌 속에 **혼합-리프-전문가(Mixture-of-Leaf-Experts, MoLE)**라고 불리는 전문가 팀을 보유하고 있습니다. 여러 명의 탐정이 모여 있는 방을 상상해 보세요. 한 명은 선형 퍼즐을 푸는 데 능숙하고, 다른 한 명은 노이즈 패턴을 포착하는 데 뛰어나며, 또 다른 한 명은 기묘한 곡선을 다루는 데 특화되어 있습니다. DAG-FM은 데이터를 볼 때 단순히 추측하는 것이 아니라, "이 특정 단서에 가장 적합한 탐정은 누구인가?"라고 묻습니다. 그리고 문제를 적절한 전문가에게 동적으로 배정합니다. 이를 통해 DAG-FM은 단일한 사고방식을 가진 기존의 탐정들을 혼란에 빠뜨렸던 혼란스러운 규칙의 혼합을 처리할 수 있습니다.

훈련장: 합성 놀이터

DAG-FM이 실제 세계의 미스터리를 해결하기 전에, 저자들은 모델을 가르쳐야 했습니다. 그들은 단순히 실제 데이터를 입력하는 것에 그치지 않고, 거대한 가상 놀이터를 구축했습니다. 무작위 그래프와 무작위 노이즈를 사용하여 다양한 유형의 인과 규칙(선형, 가법 노이즈, 포스트 비선형 등)을 서로 조합하며 수백만 개의 가짜 시나리오를 생성했습니다.

결정적으로, 저자들은 이 놀이터에 특정 안전망을 설계했습니다. 만약 모델이 이 특정 규칙의 혼합물로 훈련된다면, 데이터가 아무리 지저�더라도 반드시 단 하나의 진실된 답을 찾아낼 수 있다는 것을 수학적으로 증명했습니다. 이는 마치 조종사가 모든 종류의 폭풍과 엔진 고장이 포함된 시뮬레이터에서 훈련받아, 실제 비행기를 조종할 때 어떤 상황에서도 대처할 수 있게 만드는 것과 같습니다.

결과: 속도와 영리함

저자들이 DAG-FM을 테스트했을 때, 결과는 인상적이었습니다.

가짜 데이터 상에서: 혼합된 규칙이 100가지 유형인 시뮬레이션에서, DAG-FM은 거의 모든 다른 방법들을 압도했습니다. 기존의 알고리즘이나 심지어 다른 새로운 AI 모델들보다 더 자주 연결 관계를 정확히 찾아냈고(높은 정밀도와 재현율), 오류율(낮은 에러율)도 더 적었습니다.
실제 데이터 상에서: 저자들은 7,466개의 샘플과 11개의 변수를 가진 단백질 신호 네트워크, 그리고 10,000개의 샘플과 38개의 변수를 가진 물리 시스템과 같은 실제 데이터셋으로 테스트했습니다. 다른 많은 방법이 메모리 부족으로 멈추거나 포기한 반면, DAG-FM은 계속 나아갔습니다. 데이터의 크기나 복잡성 때문에 고전하던 전통적인 도구들을 능가하며 최적의 지도를 찾아냈습니다.
효율성: 또한 매우 효율적입니다. 표준적인 24GB 비디오 메모리를 가진 컴퓨터에서, 별도의 속도 저하 기술 없이도 최대 50,000개의 샘플이나 500개의 변수를 가진 데이터셋을 처리할 수 있습니다.

아직 할 수 없는 것 (현재 기준)

이 슈퍼 탐정이 할 수 없는 일도 있다는 점을 아는 것이 중요합니다. 논문은 명시적으로 DAG-FM이 모든 단서를 가지고 있다고 가정한다(숨겨진 교란 요인이 없음)고 밝히고 있습니다. 만약 모델이 볼 수 없는 비밀 변수가 데이터에 영향을 미치고 있다면, 이 방법은 혼란에 빠질 수 있습니다. 또한 저자들은 모델이 훈련받은 규칙들에 대해서는 훌륭하게 작동하지만, 완전히 새롭고 본 적 없는 혼돈의 유형에 대해서는 여전히 테스트 중이라는 점을 언급했습니다.

결론

DAG-FM은 모든 미스터리를 즉각 해결하는 마법 지팡이는 아니지만, 거대한 도약입니다. 문제를 관리 가능한 작은 단계로 나누고, 다양한 유형의 규칙을 처리할 수 있는 전문화된 전문가 팀을 활용함으로써, 이전의 도구들이 실패했던 상황에서도 진정한 인과관계 지도를 찾아냅니다. 이는 적절한 훈련과 구조가 뒷받침된다면, AI가 마침내 세상이 실제로 작동하는 방식인 복잡하고 뒤섞인 현실을 항해할 수 있음을 시사합니다.

기술 요약: DAG-FM

문제 정의

관측 가능한 표 형식 데이터(observational tabular data)로부터의 인과 발견(causal discovery)은 두 가지 근본적인 과제에 직면해 있다: **기저 인과 메커니즘의 이질성(heterogeneity)**과 방향성 비순환 그래프(DAG)의 고차원 조합 탐색 공간이다. 제약 기반(constraint-based), 점수 기반(score-based), 최적화 기반(optimization-based)을 포함한 전통적인 방법들은 강력한 가정이 없으면 종종 한계에 부딪힌다. 구체적으로, 함수적 인과 모델(Functional Causal Model, FCM) 기반 방법들은 모든 관계가 선형이거나 모든 노이즈가 가법적(additive)인 것과 같이 균질한 메커니즘을 가정하는 경우가 많아, 변수마다 메커니즘이 다양한 실제 시나리오로의 일반화가 제한된다. 반대로, 최근 등장하는 아모티제이션(amortized, 파운데이션) 모델들은 출력되는 DAG가 엄격하게 식별 가능한(identifiable) DAG임을 보장하는 이론적 보증이 부족하며, 사전 학습에 사용된 사전 공간(prior space) 설계에 매우 민감하여 이질적인 메커니즘을 처리하는 데 자주 실패한다.

방법론

본 논문은 관측 데이터로부터 고유하고 식별 가능한 DAG를 출력하도록 내재적으로 보장하면서, 인과 발견을 아모티제이션(amortize)하기 위해 설계된 파운데이션 모델 아키텍처인 DAG-FM을 제안한다.

1. 이론적 토대: 이질적 식별 가능성 (Heterogeneous Identifiability)

저자들은 이질적 인과 메커니즘을 허용할 수 있는 사전 공간의 설계 조건을 확립하였다. 단일 메커니즘 유형을 강제하는 기존 FCM 접근 방식과 달리, DAG-FM의 사전 공간은 구조적 인과 모델(SCM)이 다음 네 가지 서로 다른 메커니즘 패밀리의 복합 혼합물(composite mixture)이 될 수 있도록 허용한다:

LiNGAM: 비가우시안 노이즈를 가진 선형 관계.
ANM: 가법 노이즈를 가진 비선형 관계.
HNM: 이분산(heteroscedastic) 노이즈 모델.
PNL: 포스트-논리(post-nonlinear) 모델.

이론적 결과(정리 4.4 및 4.5)는 이러한 조건(인과적 충분성과 충실성 가정 하에) 하에서, 진정한 DAG가 관측 분포로부터 거의 확실하게(almost surely) 식별 가능함을 입증한다. 이는 아모티제이션 모델이 샘플 크기가 증가함에 따라 고유한 정답 DAG로 수렴함을 보장한다.

2. 아키텍처: 분해된 예측 (Decomposed Prediction)

인접 행렬(adjacency matrices)을 직접 예측하는 것의 확장성 문제(제곱 규모로 증가하며 DAG 제약이 없음)를 해결하기 위해, DAG-FM은 특화된 트랜스포머 기반 서브 모듈을 사용하여 발견 과정을 두 단계의 순차적 자기회귀(auto-regressive) 단계로 분해한다:

리프 노드 예측기 ( $Q_{\theta_1}$ ): 현재 서브 그래프 내에서 리프 노드(자식이 없는 노드)를 재귀적으로 예측하여 인과적 순서를 식별한다.
부모 노드 예측기 ( $Q_{\theta_2}$ ): 확인된 리프 노드가 주어지면, 이 모듈은 해당 노드의 부모를 식별하여 확립된 순서에 따라 상삼각(upper-triangular) 구조로 그래프를 프루닝(pruning)한다.

3. Mixture-of-Leaf-Experts (MoLE)

사전 공간 내의 다양한 메커니즘 패밀리를 처리하기 위해, 리프 예측 헤드는 Mixture-of-Leaf-Experts (MoLE) 프레임워크를 통합한다.

모델은 각 특정 메커니즘 패밀리에 대응하는 $m+1$ 개의 쌍선형(bilinear) 헤드(일반 메커니즘 전문가 및 더미 전문가 포함)를 사용한다.
Top-k 라우터가 입력 데이터 특징에 따라 적절한 전문가를 동적으로 선택한다.
이 설계는 모델이 식별 가능한 메커니즘 패밀리에 동적으로 적응할 수 있게 하여, 이론적 FCM 가정과 실제 데이터의 이질성 사이의 간극을 메운다.

4. 사전 공간 구축 (Prior Space Construction)

모델은 다음을 포함하는 사전 공간 $\Phi$ 로부터 생성된 대규모 합성 데이터셋을 통해 사전 학습된다:

무작위 그래프: 다양한 밀도를 가진 Erdős–Rényi, Barabási–Albert, Watts–Strogatz 모델.
무작위 메커니즘: 네 가지 정의된 패밀리에서 샘플링된 다양한 함수 라이브러리 (MLP, GP, Random Forest 등).
무작위 노이즈: 제어된 신호 대 잡음비(SNR)를 가진 포괄적인 분포 라이브러리 (Normal, Laplace, Student's t 등).
데이터 증강: 실제 데이터의 특성을 시뮬레이션하기 위한 비선형 워핑(warping) 및 이산화(discretization)를 수행하되, 허위 통계적 아티팩트를 방지하기 위해 엄격한 분포 정규화를 강제한다.

5. 추론 (Inference)

추론은 반복적으로 수행된다 (알고리즘 1). 모델은 리프 노드를 반복적으로 식별하고, 부모를 결정하며, 해당 노드를 집합에서 제거하고, 모든 변수가 처리될 때까지 인과적 순서를 업데이트한다. 이 과정은 사후적인 사이클 교정(cycle correction) 없이 유효한 DAG 구조를 보장한다.

주요 기여

이론적 보증: 본 논문은 이질적 인과 메커니즘을 지원하면서도 아모티제이션된 모델이 고유하고 식별 가능한 DAG로 수렴함을 이론적으로 보장하는 사전 공간의 설계 조건을 제공한다.
DAG-FM 아키텍처: 리프 노드 및 부모 노드 예측 작업으로 DAG 식별을 분리하고, 메커니즘 이질성을 처리하기 위해 트랜스포머 기반 서브 모듈과 MoLE 프레임워크를 활용하는 파운데이션 모델을 개발하였다.
실험적 성능: 이질적 메커니즘을 가진 합성 벤치마크와 복잡한 실제 데이터셋 모두에서 DAG-FM이 최첨단(state-of-the-art) 성능을 달encing함을 입증하였으며, 이는 전통적인 알고리즘 및 최신 파운데이션 모델들을 유의미하게 능가하는 수치이다.

실험 결과

합성 벤치마크: 이질적 메커니즘을 가진 $n=1000$ 샘플 및 $d=20$ 변수 데이터셋에서, DAG-FM은 가장 높은 F1-score (0.64)와 가장 낮은 정규화된 Structural Hamming Distance (nSHD) (0.19)를 달성하여 Direct-LiNGAM, AVICI, TabCausal와 같은 베이스라인을 능가하였다.
균질 부분 공간 (Homogeneous Subspaces): 엄격히 균질한 설정(예: 순수 LiNGAM)에서도 DAG-FM은 전문화된 알고리즘과 대등한 0.12의 Order Divergence Rate (ODR)를 기록하며 경쟁력을 유지하였다.
실제 세계 벤치마크:
- Sachs 데이터셋 (단백질 신호 전달): DAG-FM은 F1 0.53 및 nSHD 0.25를 달성하여 다른 모든 방법(예: FoundCause: F1 0.33; GES: F1 0.31)을 크게 앞질렀다.
- Causal Chamber: DAG-FM은 F1 0.53 및 nSHD 0.14를 달성하여 모든 베이스라인을 추월하였다.
확장성: DAG-FM은 24GB VRAM 제약 내에서 최대 50,000개 샘플 또는 500개 변수를 처리하며 강력한 확장성을 입증하였으며, 이는 많은 베이스라인(예: AVICI, CauScale)이 대규모 데이터셋에서 메모리 부족 오류를 겪는 것과 대조적이다.
강건성: 모델은 다양한 그래프 밀도, 데이터 워핑 및 이산화에 대한 회복력을 보여주었으며, 보지 못한 그래프 구조, 메커니즘 및 노이즈 분포를 포함하는 OOD(Out-of-Distribution) 시나리오에서도 강력한 성능을 유지하였다.

의의

본 논문은 DAG-FM이 FCM 기반의 식별 가능성에 대한 이론적 보증과 아모티제이션 파운데이션 모델의 확장성을 결합함으로써 중요한 진전을 이루었다고 주장한다. 직접적인 인접 행렬 예측에서 벗어나, DAG 식별을 이론적으로 근거가 있는 순차적 하위 작업으로 분해함으로써, DAG-FM은 두 가지 한계를 극복한다: 즉, 전통적 방법의 한계(이질적 메커니즘에 대한 일반화 부족)와 기존 아모티제이션 접근 방식의 한계(구조적 보증 부족 및 사전 설계에 대한 민감도)이다. 이 연구는 이론적으로 엄격하면서도 고차원의 실제 데이터에 실용적으로 적용 가능한 인과 발견의 새로운 패러다임을 제시한다.