ACTIVA: Amortized Causal Effect Estimation via Transformer-based Variational Autoencoder

기술 요약: ACTIVA - 변분 오토인코더를 통한 진폭화된 인과 효과 추정 (Amortized Causal Effect Estimation via Variational Autoencoders)

문제 정의

관측 데이터만을 사용하여 개입 후의 분포를 추정하는 것은 의료, 경제, 금융 등의 분야에서 근본적인 과제입니다. 주요 장애물은 **인과적 모호성(causal ambiguity)**입니다. 즉, 서로 다른 인과 기제(causal mechanisms)가 동일한 관측 데이터를 생성할 수 있지만, 개입 시에는 서로 다른 결과를 초래할 수 있다는 점입니다. 또한, 많은 실제 시나리오에서는 부분적으로 명시된 개입(예: 약물이 투여되었다는 사실은 알지만 정확한 용량은 모르는 경우)이 발생하며, 이로 인해 완전히 정의된 개입 함수를 필요로 하는 방법론들은 적용이 불가능해집니다. 마지막으로, 전통적인 인과 추론 방법들은 새로운 문제 인스턴스마다 광범위한 계산을 요구하며, 다양한 태스크 전반에 걸쳐 학습된 지식을 재사용하여 계산 비용이 저렴한 추론을 가능하게 하는 진폭화(amortization) 능력이 부족합니다.

방법론: ACTIVA

저자들은 전체 개입 분포의 진폭화된 추정을 위해 설계된 트랜스포머 기반 조건부 변분 오토인코더(Conditional Variational Autoencoder, CVAE)인 ACTIVA(Amortized Causal Effect Estimation via Variational Autoencoders)를 제안합니다.

핵심 아키텍처

ACTIVA는 결합 분포 $p_\theta(V, z | D_M, do(V)) = p_\gamma(V | z)p_\eta(z | D_M, do(V))$를 모델링합니다:

• 입력: 관측 데이터 $D_M$과 개입 쿼리 $do(V)$ (대상 변수와 개입 인덱스로 지정되며, 반드시 완전한 기능적 명세일 필요는 없음).
• 잠재 공간 ($z$): 기저의 인과 모델에 대한 불확실성을 나타냅니다. 모델은 관측 데이터가 관측적으로 동등한 일련의 인과 모델들과 호환된다고 가정하며, 각 모델은 서로 다른 개입 분포를 함의합니다.
• 인코더 (사후 확률): 훈련 중에 인코더 $q_\phi(z | D_M, D_{Mdo(V)}, do(V))$는 관측 샘플과 개입 샘플 모두에 조건화되어 잠재 표현을 학습합니다.
• 사전 확률 (Prior): 조건부 사전 확률 $p_\eta(z | D_M, do(V))$는 인코더의 집계된 사후 확률을 근사하도록 학습됩니다. 결정적으로, 추론 시 사전 확률은 오직 관측 데이터와 개입 쿼리에만 조건화되며, 이를 통해 훈련 과정에서 학습된 개입 결과에 대한 지식을 효과적으로 "진폭화"합니다.
• 디코더: 트랜스포머 기반 디코더 $p_\gamma(V | z)$는 잠재 샘플을 가우시안 혼합 모델(Gaussian Mixture Model, GMM)의 파라미터로 매핑하여, 전체 개입 후 분포를 예측합니다.

주요 설계 특징

1. 치환 불변성 (Permutation Equivariance): 아키텍처는 변수(feature)와 샘플에 대해 번갈아 주의(attention)를 기울이는 다중 헤드 자기 주의(multi-head self-attention, MHSA) 블록을 활용합니다. 이는 모델이 변수 순서를 일관되게 처리하도록 보장하여, 레이블이 변경된 그래프 간에도 통계적 강도를 공유할 수 있게 합니다.
2. 약한 개입 명세 (Weak Intervention Specification): 모델은 완전한 기능적 값을 요구하는 대신 인덱스와 대상 선택기를 통해 개입을 인코딩하므로, 부분적으로 정의된 개입 환경에서도 적용이 가능합니다.
3. 혼합 모델링 (Mixture Modeling): 잠재 사전 확률과 디코더 출력 모두 혼합 모델(GMM)로 모델링되어, 모델이 다봉 분포(multi-modal distributions)를 포착하고 인과적 모호성을 명시적으로 표현할 수 있도록 합니다.

이론적 분석

본 논문은 이상적인 조건(무한한 샘플 및 최적의 표현력) 하에서의 일관성(consistency) 결과를 제공합니다. 저자들은 다음을 증명합니다:

• 학습 목적 함수는 입력 데이터와 시뮬레이터 분포에 의해 지지되는, 관측적으로 호환 가능한 모든 인과 모델의 **개입 분포에 대한 혼합(mixture)**을 목표로 합니다.
• 만약 관측적 동등 클래스가 단일 모델로 수렴하는 경우(식별 가능한 경우), 목표는 유일한 개입 후 분포로 축소됩니다.
• 학습된 사전 확률은 훈련 설정(개입 데이터가 가용한 경우)과 추론 설정(개입 데이터가 가용하지 않은 경우) 사이의 간극을 효과적으로 메우며, 훈련 중 추출된 개입 정보를 흡수합니다.

실험 결과

ACTIVA는 합성 데이터셋(Gaussian 및 Beta 노이즈, 2개 및 8개 변수)과 준합성 유전자 발현 데이터(SERGIO 시뮬레이터)를 통해 평가되었습니다. 비교 대상은 다음과 같습니다:

• Gaussian Conditional: 상관관계 기반 베이스라인.
• Linear Causal: 실제 인과 그래프에 접근할 수 있는 오라클(oracle) 베이스라인.
• MACE-TNP: 결합 분포를 예측하도록 확장된 강력한 진폭화 베이스라인.

주요 결과:

• 상관관계 베이스라인 대비 우월성: ACTIVA는 특히 **가짜 비후손 효과(spurious non-descendant effects)**를 줄이는 데 있어 Gaussian Conditional 베이스라인보다 크게 앞섭 만큼 우수한 성능을 보였습니다. 상관관계 베이스라인은 잘못된 변화를 예측하는 반면, ACTIVA는 비후손에 대한 인과 효과가 0임을 정확히 예측합니다.
• 진폭화 베이스라인과의 성능 비교: 작은 규모의 그래프와 준합성 생물학적 데이터에서, ACTIVA는 음의 로그 가능도(Negative Log-Likelihood) 및 에너지 거리(Energy Distance) 측면에서 MACE-TNP와 대등하거나 더 우수한 성능을 보였습니다. 8개 변수를 가진 더 큰 합성 그래프에서는 MACE-TNP가 때때로 더 낮은 에너지 거리를 달성하였으며, 이는 두 모델의 상호 보완적인 강점을 시사합니다.
• 다운스트림 유용성: 분포 예측은 정확한 평균 처치 효과(Average Treatment Effect, ATE) 추정으로 이어집니다. 인과 구조를 포착하는 ACTIVA의 능력은 작은 문제와 준합성 데이터에서, 실제 그래프에 대한 접근 권한이 없음에도 불구하고 오라클인 Linear Causal 베이스라인과 경쟁할 만한 성능을 보여주었습니다.
• 강건성: ACTIVA는 분포 외(OOD) 설정에서 부분적인 강건성을 보여주었으나, 특히 후손(descendant) 에너지 거리 측면에서 인 분포(in-distribution) 결과에 비해 성능이 저하되었습니다.

의의 및 주장

저자들은 ACTIVA를 인과 추론을 위한 파운데이션 모델을 향한 유망한 단계로 규정합니다. 본 논문의 의의는 다음 세 가지 영역에서 주장됩니다:

1. 진폭화된 분포 추정: 인과 그래프를 복구하거나 점 추정치를 추정하는 방식과 달리, ACTIVA는 전체 개입 후 분포의 예측을 직접 진폭화하여 새로운 데이터셋에 대한 제로샷(zero-shot) 추론을 가능하게 합니다.
2. 명시적 모호성 표현: 출력을 가능한 인과 모델들의 혼합으로 모델링함으로써, ACTIVA는 데이터를 하나의 점 추정치로 붕괴시키지 않고 관측 데이터에 내재된 불확실성을 명시적으로 표현합니다.
3. 부분적 개입 처리: 본 프레임워크는 정밀한 용량 없이 개입 레이블만 있는 경우와 같이 개입이 부분적으로만 명시된 환경에 독보적으로 적합합니다.

결론적으로, 유한 샘플 구현은 이론적 이상향에 대한 근사치이지만, 이론에 기반한 아키텍처는 단순한 관측 상관관계가 아닌 개입 후의 인과 구조를 반영하는 예측기를 성공적으로 학습함을 보여줍니다.