Causality-Encoded Diffusion Models for Interventional Sampling and Edge Inference

1. 연구 배경 및 문제 정의 (Problem)

• 기존 확산 모델 (Diffusion Models) 의 한계: 확산 모델은 복잡한 분포를 유연하게 추정하는 강력한 생성 모델이지만, 본질적으로 인과 구조 (Causal Structure) 를 인코딩하지 않습니다. 따라서 방향성 비대칭성을 학습하지 못해, 구조적 인과 모델 (SCM) 에서 요구하는 개입 (Intervention, do-연산자) 쿼리에 대한 답변이나 인과 분석을 직접적으로 지원하지 못합니다.
• 차원의 저주 (Curse of Dimensionality): 기존 확산 모델의 수렴 속도는 전체 차원 (Ambient Dimension) 에 의존하는 경향이 있어, 고차원 데이터에서 추론이 어렵습니다.
• 인과적 추론의 필요성: 방향성 비순환 그래프 (DAG) 가 주어졌을 때, 이를 활용하여 개입 분포를 샘플링하고, DAG 내의 엣지 (인과 관계) 에 대한 통계적 검정을 수행할 수 있는 프레임워크가 필요합니다.

2. 제안 방법론: 인코딩된 확산 모델 (CEDM)

저자들은 인과성 인코딩 확산 모델 (Causality-Encoded Diffusion Models, CEDM) 을 제안합니다. 이는 알려진 DAG 구조를 확산 과정의 역방향 역학 (Reverse Dynamics) 에 통합하는 프레임워크입니다.

2.1. 인과성 인코딩 확산 과정

• 구조적 인과 모델 (SCM) 통합: 각 변수 $X_j$가 부모 노드 $X_{pa(j)}$와 독립 노이즈 $U_j$에 의해 생성된다고 가정합니다 ($X_j = g_j(X_{pa(j)}, U_j)$).
• 조건부 스코어 학습: 기존 확산 모델이 전체 결합 분포의 스코어 함수 ($\nabla \log p_t(x)$) 를 학습하는 반면, CEDM 은 DAG 분해에 부합하는 조건부 스코어 함수 ($\nabla_{x_j} \log p_t(x_j | x_{pa(j)})$) 를 각 노드별로 학습합니다.
• 위상적 순서 샘플링: 알고리즘은 DAG 의 위상적 순서 (Topological Order) 를 따라 노드를 하나씩 생성합니다. 각 노드 $j$를 역확산 (Reverse Diffusion) 할 때, 이미 생성된 부모 노드 값들을 고정된 조건으로 사용하여 해당 노드의 역 SDE 를 풉니다.

2.2. 개입 샘플링 (Interventional Sampling)

• do-연산자 구현: 특정 변수 집합 $S$에 대해 $do(X_S = x^*_S)$를 수행하려면, 역확산 과정에서 $S$에 속한 노드들을 $x^*_S$로 고정 (Freeze) 하고, 나머지 노드들만 조건부 스코어를 통해 생성합니다.
• 결과: 이는 인과적 개입에 따른 분포 $p(x | do(X_S = x^*_S))$를 정확하게 재현하며, 관찰 분포와 구별되는 인과적 효과를 시뮬레이션할 수 있게 합니다.

2.3. 엣지 추론 (CEDMI)

• 재샘플링 기반 검정: 제안된 CEDM 을 이용하여 CEDM 기반 추론 (CEDMI) 절차를 개발했습니다.
  1. 특정 엣지 $(k, j)$가 존재하지 않는다는 귀무가설 ($H_0$) 하에서, 해당 엣지를 제거한 그래프를 가정합니다.
  2. 이 그래프에서 CEDM 을 훈련하고, 관측된 부모 변수들을 조건으로 하여 $X_j$를 재샘플링 (Resampling) 합니다.
  3. 재샘플링된 데이터와 관측 데이터 간의 조건부 의존성을 측정하여 $p$-값을 계산합니다.
• 통계량: Azadkia-Chatterjee 계수의 다변량 확장인 MCODEC (Multivariate Conditional Dependence Coefficient) 를 사용하여 조건부 독립성 ($X_j \perp X_k | X_{pa(j)\setminus k}$) 을 검정합니다. 이 통계량은 튜닝 파라미터가 필요 없으며 순위 기반 (Rank-based) 입니다.

3. 주요 기여 (Key Contributions)

1. 방법론적 혁신:

   • 인과 구조를 존중하는 생성 프레임워크인 CEDM 과 이를 활용한 개입 샘플링, 엣지 검정 절차 (CEDMI) 를 제안했습니다.
   • 기존 방법론이 가정한 선형성이나 가법적 노이즈 가정을 완화하고, 일반적인 비선형 및 비가법적 설정에서도 작동합니다.

2. 통계적 보장 (Theoretical Guarantees):

   • 수렴 속도: 관찰 및 개입 분포 추정 오차의 수렴 속도가 전체 차원이 아닌 최대 국소 차원 (Maximum Local Dimension, $q = \max_j (d_j + d_{pa(j)})$) 에 의해 결정됨을 증명했습니다. 이는 희소 그래프에서 차원 저주를 극복하는 구조적 차원 축소 효과를 의미합니다.
   • 유형 I 오류 통제: CEDMI 에 대한 점근적 유형 I 오류 (Type I error) 통제 보장을 증명했습니다.

3. 실증적 검증:

   • 다양한 시뮬레이션 (Chain, Hub, Random DAG) 과 실제 유세포 분석 (Flow Cytometry) 데이터를 통해 기존 방법 (RePaint, VACA, CCIT 등) 대비 우수한 개입 분포 회복 능력과 검정력 (Power) 을 입증했습니다.

4. 실험 결과 (Results)

• 분포 회복 (Distribution Recovery):

  • 시뮬레이션에서 CEDM 은 관찰되지 않는 개입 분포를 회복하는 데 있어 RePaint(일반 확산) 나 VACA(변분 오토인코더) 보다 최소 평균 제곱 오차 (MMD) 를 보여주었습니다.
  • 특히 희소 그래프 (Chain, Hub) 에서 국소 차원 의존성으로 인해 성능 향상이 두드러졌습니다.

• 추론 성능 (Inference Performance):

  • 크기 (Size): CEDMI 는 명목 수준 (0.05) 에서 잘 보정된 (Well-calibrated) 오류율을 보여주었습니다. 반면, 기존 방법들 중 일부 (RCoT, RCIT 등) 는 과도하게 보수적이지 않거나 (Anti-conservative) 오류 통제가 불안정했습니다.
  • 검정력 (Power): 신호 강도가 증가하거나 표본 크기가 커짐에 따라 CEDMI 는 높은 검정력을 보였으며, 특히 작은 표본에서도 경쟁 방법들보다 우월한 성능을 발휘했습니다.

• 실제 데이터 적용 (Flow Cytometry):

  • Sachs et al. 의 T 세포 신호 전달 네트워크 데이터에 적용하여, 기존 합의 네트워크 (Consensus Network) 와 데이터 기반 베이지안 네트워크 간의 이견이 있는 엣지들을 검정했습니다.
  • Erk $\to$ Akt 연결에 대해 통계적으로 유의미한 증거를 발견했는데, 이는 기존 문헌에는 명시되지 않았으나 생물학적으로 타당한 상호작용 (Cross-talk) 을 지지하는 결과로 해석되었습니다.

5. 의의 및 결론 (Significance)

이 논문은 확산 모델을 단순한 생성 도구를 넘어 인과 추론을 위한 강력한 프레임워크로 확장했다는 점에서 의의가 있습니다.

• 구조적 차원 축소: 인과 그래프 구조를 모델 아키텍처에 직접 통합함으로써, 고차원 데이터에서도 효율적인 학습과 추론이 가능함을 이론적으로 증명했습니다.
• 개입 예측 능력: 관찰 데이터만으로 학습된 모델이 개입 (Intervention) 하의 분포를 정확하게 예측할 수 있음을 보여주어, 정책 평가 및 실험 설계 등 실제 응용 분야에 큰 잠재력을 제공합니다.
• 검증 가능한 인과 추론: 기존 인과 발견 방법들이 직면한 비선형성 및 복잡한 의존성 문제를 해결하면서도, 통계적으로 엄밀한 엣지 검정을 가능하게 하여 신뢰할 수 있는 인과 네트워크 구축을 지원합니다.

결론적으로, CEDM 은 인과 구조를 활용한 차원 축소와 확산 모델의 유연성을 결합하여, 관찰 데이터로부터 인과적 메커니즘을 더 정확하게 이해하고 예측하는 새로운 패러다임을 제시합니다.