기술 요약: 네트워크 데이터로부터의 인과적 표현 학습 (GraCE-VAE)

문제 정의

현재의 인과적 표현 학습 방법들은 관측된 변수들이 구조화되지 않은 특징 벡터(feature vectors)라고 가정하는 경우가 많으며, 이는 많은 과학적 영역(예: 단백질-단백질 상호작용, 경로-유전자 멤버십)에 내재된 풍부한 관계적 맥락을 간과합니다. 기존의 그래프 기반 추론 방법들은 네트워크 구조를 활용할 수 있지만, 일반적으로 인과 그래프가 알려져 있다고 가정하거나 알려진 노드에 대한 관측 가능한 개입만을 고려합니다. 반대로, 관측 및 개입 데이터를 통해 잠재적 구조 인과 모델(SCM)을 학습하는 인과적 디텐글링(disentanglement) 방법들은 측정된 엔티티들 사이의 알려진 상호작용 네트워크를 활용하지 못하는 경우가 많습니다. 이는 관측된 맥락 네트워크를 활용하면서 잠재적 인과 구조와 개입 효과를 학습하는 과정에서 발생하는 격차를 만듭니다.

본 논문은 변수들이 알려진 상호작용 네트워크(예: 생물학적 경로)와 연관되어 있는 설정에서, 개입 정보로 강화된 네트워크 데이터로부터 잠재적 인과 요인과 그 상호 의존성을 발견하는 과제를 다룹니다.

방법론: GraCE-VAE

저자들은 GraCE-VAE(Graph-Aware Causal Effects Variational AutoEncoder)를 제안합니다. 이는 그래프 신경망(GNN)과 딥 생성 모델을 결합하여 이종 네트워크 데이터로부터 잠재적 인과 그래프를 발견하는 프레임워크입니다.

아키텍처

GraCE-VAE는 구조적 맥락의 사용과 생성적 인과 모델의 사용을 분리하도록 설계된 인코더-디코더 분할 방식을 채택합니다.

1. 그래프 인식 인코더 (Graph-Aware Encoder):

   • 주요 관측 엔티티($X$, 예: 유전자)와 보조 엔티티($H$, 예: 경로)를 포함하는 이종 맥락 그래프를 처리하기 위해 GNN(구체적으로 GraphSAGE로 구현됨)을 사용합니다.
   • 인코더는 $X$의 관측된 속성, 보조 뷰 $H$, 그리고 인접 구조($A_{XX}, A_{XH}, A_{HH}$)를 조건으로 합니다.
   • 이는 외생적 노이즈 변수($Z$)에 대한 사후 분포를 매개변수화하는 구조 인식 임베딩을 생성하며, 이를 통해 암묵적 추론(amortized inference) 과정에 네트워크 맥락을 효과적으로 주입합니다.
   • 결정적으로, 맥락 그래프는 오직 추론을 위해서만 사용되며 디코더에는 입력되지 않습니다.

2. 인과적 디코더 (Causal Decoder, Latent SCM):

   • 디코더는 잠재 변수 $U$에 대한 구조적 인과 모델(SCM)을 구현합니다.
   • 잠재 요인들 사이의 인과 구조를 나타내는 방향성 비순환 그래프(DAG)를 학습합니다.
   • 개입 모델링 (Intervention Modeling): 개입은 2단계 과정을 통해 모델링됩니다. 개입 인코더는 개입 지표를 잠재 변수에 대한 소프트 선택(타겟 식별)과 제안된 메커니즘 수정(예: 가산적 이동)으로 매핑합니다. 그 후 디코더는 선택된 잠재 노드의 인과 메커니즘에 이 수정을 적용하되, 다른 노드들은 변경하지 않고 유지합니다.
   • 혼합 함수 $f_\theta$는 잠재 인과 변수 $U$를 다시 관측 공간 $X$로 매핑합니다.

학습 목적 함수 (Learning Objective)

모델은 다음 세 가지 구성 요소의 가중 합을 사용하여 학습됩니다:

1. ELBO 항: 관측 및 개입 데이터 분포를 재구성합니다.
2. 분포 정렬 항 (Distribution-Alignment Term): 최대 평균 불일치(MMD)를 사용하여 모델이 생성한 반사실적 분포를 경험적 개입 분포와 정렬합니다.
3. 정규화 (Regularization): 학습된 DAG 인접 행렬에 대한 희소성 페널티와 그래프 기반 정규화 항을 포함합니다.

식별 가능성 (Identifiability)

본 논문은 GraCE-VAE가 이전의 인과적 불일치 VAE(특히 CMVAE [34])의 식별 가능성 보증을 계승함을 입증합니다. 다음의 가정 하에:

• 선형적 개입 충실도(linear interventional faithfulness).
• 관측 데이터와 개입 데이터의 모두 가용성.
• 각 개입 체제 내에서 샘플이 i.i.d.임.
• 각 잠재 변수가 적어도 하나의 개입 대상이 됨.
• 경로 정보 $H$가 추론을 돕지만, $U$가 알려진 후에는 $X$에 대한 추가적인 생성 정보를 제공하지 않는다는 보조 뷰(auxiliary-view) 가정.

잠재적 인과 DAG와 개입 타겟은 표준 동등 클래스(동등한 요인들의 순열)까지 식별 가능합니다.

주요 기여

1. 구조화된 맥락을 위한 프레임워크: 측정된 엔티티(유전자)와 그룹 노드(경로)를 포함하는 이종 네트워크로부터의 사전 지식에 의해 유도되는, 잠재적 인과 그래프와 개입 효과를 공동으로 학습하는 변이형 그래프 오토인코더를 제안합니다. 완전히 관측된 변수나 주어진 인과 그래프를 가정하는 방법들과 달리, GraCE-VAE는 잠재 구조를 발견합니다.
2. 이론적 명확화: 생성 모델 클래스가 기존의 인과적 불일치 VAE와 일치하고, 그래프 인식 인코더가 암묵적 추론만을 수정하기 때문에, GraCE-VAE가 새로운 디코더 패밀리에 대한 새로운 정리를 요구하지 않고도 기존의 식별 가능성 보증을 계승한다는 점을 명확히 합니다.
3. 경험적 검증: 세 가지 CRISPR 섭동 데이터셋(norman, replogle-small, replogle-large)에 대한 실험을 통해, 구조화된 생물학적 맥락을 활용하는 것이 특히 보지 못한 섭동 조합에 대한 개입 결과 예측을 개선함을 보여줍니다. 절제 연구(ablation studies)를 통해 그래프 인식 인코더와 인과적 디코더가 모두 성능에 필수적임을 확인했습니다.

실험 결과

실험은 유전적 섭동과 관련된 세 가지 데이터셋에 대해 수행되었습니다:

• 단일 개입 (Single Interventions): GraCE-VAE는 강력한 $R^2$ 성능을 달察했으며, 구조를 무시하는 CMVAE나 SENA와 같은 베이스라인을 능가했습니다. 또한 CMVAE-multihot 및 VGAE(인과적 디코더가 없는 절제 모델)와 대등하거나 더 우수한 성능을 보였습니다.
• 보지 못한 이중 개입 (Unseen Double Interventions): 모델은 단일 개입으로 학습되었고 보지 못한 이중 개입에 대해 평가되었습니다. GraCE-VAE는 CMVAE와 SENA를 크게 앞질러 복잡한 조합을 예측하는 데 있어 우수한 성능을 보였으며, CMVAE의 0.7382 대비 평균 0.7845의 $R^2$를 달성했습니다. 이는 구조화된 생물학적 맥락을 인코딩하는 GNN과 인과적 디코더 사이의 시너지를 강조합니다.
• 절제 연구 (Ablation Studies):
  • GNN 설계: 다양한 GNN 아키텍처(GCN, GAT, GraphSAGE)를 테스트하였으며, GraphSAGE가 가장 좋은 성능을 보였습니다. 레이어 깊이를 늘리는 것은 추가적인 개선을 가져오지 않았습니다.
  • 맥락 수준 (Context Levels): 더 풍부한 엣지 정보(유전자-유전자, 경로-유전자, 경로-경로)를 포함하는 것이 일반적으로 단일 개입 예측을 개선했으나, 이중 개입에서는 이득이 정체되거나 약간 감소하는 경향을 보여, 맥락의 풍부함과 과적합 사이의 트레이드오프를 시사했습니다.
  • 손상된 그래프 (Corrupted Graphs): 생물학적 엣지의 50%를 무작위 엣지로 교체했을 때 성능이 일관되게 저하되었으며, 이는 의미 있는 생물학적 토폴로지가 이득을 주도한다는 것을 확인시켜 줍니다.
  • 사전 지식 (Priors): 단일 세포 기초 모델(scFM)의 임베딩이나 대안적인 개입 인코더를 추가해도 기본 GraCE-VAE 구성보다 일관되게 우수한 성능을 내지는 못했습니다. 이는 이득이 단순히 더 많은 사전 지식을 추가하는 것이 아니라, 그래프 구조와 인과 모델링의 구체적인 통합에서 온다는 것을 시사합니다.

의의 및 주장

본 논문은 GraCE-VAE가 관계적 정보와 개입 기반 인과 추론을 결합함으로써 복잡한 네트워크 시스템에서 잠재적 인과 요인을 밝혀내는 새로운 접근법을 제공한다고 주장합니다. 주요 의의는 다음과 같습니다:

• 격차 해소: 측정된 엔티티들 사이의 관측된 맥락 네트워크를 활용하면서 잠재적 인과 구조를 학습하는 방법의 부재를 해결합니다.
• 일반화: 구조화된 생물학적 맥락을 활용하여 보지 못한 개입 조합(예: 이중 섭동)으로 외삽하는 탁월한 능력을 입증하며, 이는 과학적 발견을 위한 핵심 역량입니다.
• 이론적 건전성: 표준 가정 하에서 엄격한 식별 가능성 보증을 유지하며, 그래프 구조를 인코더에 포함하는 것이 잠재적 인과 모델의 이론적 회복 가능성을 저해하지 않음을 증명합니다.

저자들은 개입 데이터 커버리지에 대한 의존성, 단순화된 개입 기능 형태(가산적 이동)의 사용, 그리고 식별성을 위해 필요한 특정 보조 뷰 가정 등의 한계점을 언급합니다. 그들은 학습된 DAG가 생물학적으로 타당한 가설(예: 알려진 세포 주기 또는 분화 경로와 일치)을 제공하지만, 이를 확정적인 인과 검증이라기보다는 반사실적 시뮬레이션을 위한 가설 공간으로 간주해야 한다고 결론짓습니다.