기술 요약: 인-컨텍스트 표 예측을 위한 인과적 순서 학습

1. 문제 제기

표 기반 기초 모델 (TFMs) 과 같은 인-컨텍스트 학습 아키텍처는 대규모 사전 학습을 활용하여 표 데이터에서 강력한 예측 성능을 달성했습니다. 그러나 이러한 모델들은 주로 연관적이며, 명시적인 인과 방향성보다는 상관 구조에 의존합니다. 결과적으로, 이러한 모델들은 일반화되지 않는 비인과적 통계적 연관성을 이용할 수 있어 분포 변화나 개입 하에서 예측이 신뢰할 수 없게 될 수 있습니다.

확립된 인과 발견 방법들이 방향성 비순환 그래프 (DAG) 를 복원하기 위해 존재하지만, 이들은 종종 데이터셋 특정적이며 계산 집약적이고 현대적인 예측 아키텍처와 분리되어 있습니다. 반면, 최근의 감가상각적 인과 발견 접근법은 합성 그래프에 대한 지도 학습에 의존하며, 실제 인과 레이블이 필요합니다. 결측 데이터와 개입을 처리할 수 있는 비지도 인-컨텍스트 학습 프레임워크 내에서 인과 구조 발견을 예측 아키텍처에 체계적으로 통합하는 데에는 간극이 존재합니다.

2. 방법론: TABORDER

저자들은 인과적 변수 순서를 동시에 추론하고 강건한 예측을 수행하도록 설계된 트랜스포머 기반 아키텍처인 TABORDER를 제안합니다. 핵심 혁신은 인과 순서 제약 어텐션의 사용입니다.

아키텍처 설계

TABORDER 는 결측 항목이 마스크 $M$으로 표시된 데이터셋 $D$를 처리하는 이중 트랜스포머 시스템으로 작동합니다:

1. 임베딩: 각 셀 $x_i^{(r)}$과 그 결측 지시자가 벡터 표현으로 임베딩됩니다.
2. 순서 추론 분기: 표준 멀티헤드 셀프 어텐션 트랜스포머가 임베딩을 처리하여 각 변수 $X_i$에 대한 스칼라 점수 $s_i$를 추론합니다. 낮은 점수는 인과 순서에서 더 앞선 위치를 나타냅니다. 이러한 점수들은 행을 가로질러 평균화되어 전역 순서 $\hat{\pi}$를 생성합니다.
3. 순서 제약 예측 분기: 두 번째 트랜스포머가 각 셀의 조건부 분포를 모델링합니다. 결정적으로, 이 트랜스포머의 열별 어텐션은 학습된 점수 $s$에서 파생된 마스크에 의해 제약됩니다. 구체적으로, 변수 $X_i$는 $s_j \leq s_i$인 변수 $X_j$에만 어텐션할 수 있습니다. 이는 결합 분포의 순서화된 분해를 강제합니다:
   $$p_\theta(x^{(r)} | D, \pi) = \prod_{k=1}^d p_\theta(x_{\pi(k)}^{(r)} | x_{\pi(1:k-1)}^{(r)}, D)$$
   이는 모델이 타겟을 예측할 때 "반인과적" 특징 (미래 원인) 을 사용하는 것을 방지하여, 허위 상관관계에 대한 민감도를 줄입니다.

결측 데이터 및 불확실성 처리

모델은 인과적 부모의 결측성을 명시적으로 고려합니다. 가법 노이즈 모델 (ANM) 가정 ($X_i = f(X_{pa(i)}) + N_i$) 하에서, 모델은 조건부 평균 $\hat{\mu}$와 분산 $\hat{\sigma}^2$을 디코딩합니다.

• 분산 증가: 인과적 부모가 결측인 경우, 모델은 학습 가능한 양 $\delta_{i,j}$만큼 대상 변수의 예측 분산을 증가시킵니다. 이는 관련 입력이 관찰되지 않을 때 증가하는 불확실성을 포착합니다.
• 학습 목적 함수: 모델은 마스크된 항목의 로그 가능도를 최대화함으로써 엔드 투 엔드로 학습됩니다. 목적 함수는 순서 학습과 예측을 결합합니다: 마스크된 데이터를 더 잘 설명하는 (잔차 분산을 최소화하는) 순서들은 더 높은 가능도를 받습니다.

이론적 정당성

저자들은 가법 노이즈 모델 (ANM) 프레임워크 하에서 가능도 기반 접근법을 정당화합니다. 그들은 가법 인과 모델의 경우, 가능도 최적 순서가 실제 기본 DAG 의 위상적 순서와 일치함을 보여줍니다. 또한, 결측성 인식 가능도 비대칭성에 관한 정리를 증명합니다: 분산 증가를 통해 결측성을 고려하는 것은 올바른 인과 방향과 잘못된 방향 사이의 가능도 격차를 증폭시켜, 결측 데이터가 존재함에도 불구하고 인과 순서를 식별하는 것을 용이하게 합니다.

3. 주요 기여

1. 인과 발견과 예측의 통합: TABORDER 는 예측 트랜스포머 프레임워크 내에서 인과적 순서를 비지도 방식으로 추론함으로써, 인과 구조 학습과 인-컨텍스트 표 예측 간의 간극을 메웁니다.
2. 순서 제약 어텐션: 학습된 인과 순서에서 선행하는 변수들로만 정보 흐름을 제한하는 미분 가능한 학습 어텐션 마스크를 도입하여, 예측이 오직 잠재적 원인에만 의존하도록 보장합니다.
3. 결측성 인식 인과 학습: 인과적 부모의 결측이 예측 분산을 명시적으로 증가시키는 메커니즘을 설계하여, 이론적 및 실증적으로 인과 방향을 구별하는 데 도움이 됩니다.
4. 비지도 학습: 감가상각적 인과 발견 방법과 달리, TABORDER 는 학습 중 실제 인과 그래프를 요구하지 않으며, 마스크된 데이터에 대한 가능도 최대화를 통해서만 최적 순서를 학습합니다.

4. 실험 결과

저자들은 알려진 인과 구조를 가진 합성 데이터와 실제 세계 벤치마크에서 TABORDER 를 평가했습니다.

• 인과 순서 추론: 비선형 가우시안 프로세스 ANM 에서 생성된 합성 데이터에서 TABORDER 는 CAM, SCORE, NOTEARS 와 같은 전문 인과 발견 알고리즘과 비교 가능한 위상적 발산 ($d_{TOP}$) 을 달성하지만, 단일 순전파를 통해 거의 즉각적으로 순서를 추론합니다. 실제 세계 Sachs 단백질 신호 전달 데이터셋에서 TABORDER 는 합의 그래프와 비교하여 낮은 발산 ($d_{TOP} \approx 0.21$) 으로 인과 순서를 발견하여 다른 베이스라인보다 우수한 성능을 보였습니다.
• 예측 성능 및 대체: 결측값 대체 작업 (OpenML CC18 및 CTR23 벤치마크) 에서 TABORDER 는 특히 높은 결측률 ($\geq 40\%$) 에서 경쟁력 있는 성능을 달성하여 KNN, MICE, GAIN 과 같은 표준 방법을 능가했습니다.
• 개입 강건성: $Y \to Z$ 메커니즘에 개입이 있는 3 변수 체인 ($X \to Y \to Z$) 에서 TABORDER 는 $X$와 $Z$가 주어졌을 때 중재자 $Y$의 정확한 예측을 유지합니다. 반면, 표준 TFM(예: TABPFN) 과 XGBoost 는 $Y$와 $Z$ 사이의 깨진 비인과적 연관성에 의존하므로 성능이 크게 저하됩니다.
• 생물학적 통찰력: 다양한 개입 조건 하의 Sachs 데이터셋에서 TABORDER 는 생물학적으로 타당한 재구성을 반영하는 순서를 추론합니다. 예를 들어, 상류 개입 (예: PKC 억제) 은 추론된 순서에서 전역적 변화를 일으키는 반면, 하류 개입은 국소적 변화를 일으켜 알려진 신호 전달 경로 역학과 일치합니다.

5. 중요성과 주장

이 논문은 TABORDER 가 인과 방향성을 명시적으로 모델링함으로써 분포 변화와 개입에 더 강건한 표 데이터용 인-컨텍스트 학습을 위한 원칙적인 접근법을 제공한다고 주장합니다.

• 강건성: 학습된 인과 순서의 선행자에만 예측을 의존하도록 강제함으로써, 모델은 개입 하에서 실패하는 허위 상관관계를 이용하는 것을 피합니다.
• 효율성: 데이터셋마다 최적화 문제를 해결하는 고전적 인과 발견과 달리, TABORDER 는 단일 순전파 내에서 구조를 효율적으로 추론하여 하류 작업에 확장 가능합니다.
• 이론적 기반: 이 연구는 표준 ANM 가정 하에서 결측 데이터가 존재할 때의 가능도 최대화가 인과 발견의 대리가 될 수 있음을 확립하여, 비지도 순서 학습을 위한 이론적 기초를 제공합니다.

저자들은 현재 구현이 완전한 DAG 가 아닌 위상적 순서에 초점을 맞추고 있으며 실수형 특징을 가정한다는 한계를 인정합니다. 그들은 이 방법이 강건하지만, 추론된 구조는 도메인 검증을 필요로 하는 가설이며, 향후 연구는 더 풍부한 학습 사전 지식과 혼합형 데이터를 탐구할 수 있음을 시사합니다.