기술 요약: CauScale – 대규모 신경망 인과 발견 (Neural Causal Discovery at Scale)

1. 문제 정의

인과 발견(Causal discovery)은 관측 데이터로부터 인과 관계를 나타내는 유향 비순환 그래프(DAG)를 추론하는 작업으로, 생물정보학에서 경제학에 이르기까지 다양한 분야에서 매우 중요합니다. 그러나 기존 방법론들은 수백 또는 수천 개의 노드를 가진 대규모 그래프에 적용할 때 심각한 확장성 병목 현상에 직면합니다.

• 제약 기반 방법론(예: PC, FCI)은 조건부 독립 검정에 의존하며, 변수의 수가 증가함에 따라 지수적인 시간 복잡도를 가집니다.
• 점수 기반 방법론(예: NOTEARS, SDCD)은 문제를 연속 최적화 문제로 재정의하지만, 새로운 데이터셋마다 매번 최적화 문제를 풀어야 하므로 대규모 환경에서 계산 비용이 매우 높습니다.
• 아모티즈드(Amortized, 제로샷) 방법론(예: AVICI)은 그래프를 직접 예측하도록 모델을 사전 학습시켜 속도 이점을 제공합니다. 그러나 변수의 수에 따라 제곱으로 증가하는 어텐션 메커즘과 축별 어텐션 맵(axis-specific attention maps) 유지를 위한 높은 메모리 요구량으로 인해 대규모 그래프에서 공간 및 시간 효율성이 떨어집니다.

핵적인 과제는 높은 정확도를 유지하면서도 과도한 메모리나 시간 비용 없이 최대 1,000개의 노드를 가진 그래프까지 효율적으로 확장 가능한 신경망 아키텍처를 개발하는 것입니다.

2. 방법론: CauScale 아키텍처

CauScale은 효율적인 아모티즈드 인과 발견을 위해 설계된 신경망 아키텍처입니다. 이는 특정 효율성 메커니즘이 통합된 이중 스트림 설계(데이터 스트림 및 그래프 스트림)를 채택하고 있습니다.

2.1. 입력 및 개요

모델은 두 가지 입력을 받습니다:

1. 데이터 ($\mathcal{D}$): 관측 변수와 이진 중재 지표(binary intervention indicators)를 결합한 텐서.
2. 그래프 사전 정보 ($\rho$): 데이터의 역공분산 행렬로 계산된 통계적 사전 정보.

아키텍처는 데이터-그래프 블록(Data-Graph Blocks)과 축소 유닛(Reduction Units)의 교차 스택을 통해 입력을 처리하며, 최종적으로 그래프 스트림 출력을 예측 헤드에 전달하여 확률적 인접 행렬을 생성합니다.

2.2. 핵심 구성 요소

A. 이중 스트림 설계 (Two-Stream Design)

• 데이터 스트림 (Data Stream): 고차원 관측치로부터 관계적 증거를 추출합니다. 초기에는 전체 샘플 차원을 처리합니다.
• 그래프 스트림 (Graph Stream): 통계적 그래프 사전 정보를 통합하고 핵심적인 구조적 신호를 보존합니다. 노드 간 관계 공간에서 작동합니다.
• 데이터-그래프 블록 (Data-Graph Block): 두 스트림을 연결하는 모듈입니다.
  • 데이터 레이어 (Data Layer): 데이터 스트림 임베딩을 업데이트합니다.
  • Data2Graph 레이어: 축 방향 어텐션(axial attention)과 풀링을 적용하여 데이터 스트림을 압축된 관측 관계 행렬($\omega$)로 압축합니다. 이는 관계적 증거를 그래프 메시지로 정제합니다.
  • 그래프 레이어 (Graph Layer): 메시지 $\omega$를 그래프 스트림에 주입하여 그래프 임베딩을 업데이트합니다. 이를 통해 그래프 스트림이 구조적 무결성을 유지하면서도 데이터로부터 유도된 신호의 이점을 얻도록 합니다.

B. 효율성 메커니즘

• 축소 유닛 (Reduction Unit): 시간 효율성을 해결하기 위해 모델은 데이터 스트림 임베딩을 관측(샘플) 차원을 따라 압축합니다. 매 $k$ 블록마다 유닛은 샘플을 크기 $r$의 청크로 그룹화하고 평균 풀링을 적용합니다. 이는 계산 부하를 $O(m)$에서 $O(m/r)$로 줄여주며, 여기서 $m$은 샘플 수입니다. 이는 인과적 신호가 주로 샘플 내 노드 간의 의존성에 기반하며, i.i.d. 샘플 간의 집계가 변수 차원을 붕괴시키는 것보다 덜 파괴적이라는 점에 착안되었습니다.
• 결합된 어텐션 가중치 (Tied Attention Weights): 공간 효율성을 위해 CauScale은 두 스트림 모두에서 결합된 어텐션 가중치(Rao et al., 2021 방식)를 채택합니다. 축별 어텐션 맵(행 어텐션의 경우 $O(R \cdot C^2)$ 메모리 필요)을 별도로 유지하는 대신, 모델은 대상 축에 대해 어텐션 가중치를 공유합니다. 이는 어텐션 맵의 메모리 점유율을 $O(C^2)$ (또는 열 어텐션의 경우 $O(R^2)$)로 줄여 GPU 메모리 사용량을 크게 낮춥니다.

C. 예측 헤드 (Prediction Head)

최종 그래프 임베딩은 분해된 예측 헤드에 의해 처리됩니다. 순서가 없는 각 노드 쌍 $\{i, j\}$에 대해 모델은 '에지 없음', $i \to j$, 또는 $j \to i$의 세 가지 상태에 대한 로짓(logits)을 예측합니다. 기존 연구들과 달리, CauScale는 값비싼 제약 최적화를 피하기 위해 순방향 패스 중에 비순환 제약을 명시적으로 강제하지 않으며, 대신 학습된 표현과 필요 시 사후(post-hoc) 사이클 제거 전략에 의존합니다.

3. 주요 기여

1. 확장 가능한 사전 학습: 신경망을 통한 인과 발견 사전 학습이 훈련 시 최대 500개, 추론 시 최대 1,000개의 노드를 가진 그래프까지 확장될 수 있음을 보여준 최초의 연구 중 하나입니다.
2. CauScale 아키텍처: 이중 스트림 설계, 축소 유닛, 결합된 어텐션을 활용하여 정확도를 희생하지 않으면서 시간 및 메모리 효율성을 동시에 최적화하는 새로운 신경망 아키텍처를 제시합니다.
3. 포괄적 검증: 합성 및 반합성(단일 세포 유전자 발현) 데이터셋에 대한 광범한 실험을 통해 다양한 그래프 규모, 인과 메커니즘 및 분포 외(OOD) 설정에서 우수한 성능을 입증했습니다.

4. 실험 결과

저자들은 CauScale을 제약 기반(FCI), 점수 기반(NOTEARS, SDCD), 아모티즈드(AVICI, SEA) 베이스라인과 비교 평가했습니다.

• 정확도:

  • 분포 내(in-distribution) 합성 데이터($n=100$)에서 CauScale은 99.6% mAP(평균 정밀도)를 달성했습니다.
  • 복잡한 메커니즘(예: 다항 함수)을 가진 분포 외(OOD) 데이터에서 84.4% mAP를 기록하며, SEA(36.2%) 및 SDCD(41.9%)와 같은 베이스라인을 크게 앞질렀습니다.
  • SERGIO-GRN 단일 세포 데이터셋에서 모든 인과 발견 베이스라인 대비 가장 우수한 성능을 보였습니다.
  • 훈련 시 최대 규모(500 노드)를 초과하는 1,000개 노드 그래프에서도 높은 성능(96.6% mAP)을 유지했습니다.

• 효율성:

  • 시간: CauScale은 기존 방법론보다 4배에서 13,000배 빠른 추론 속도를 달성했습니다. 예를 들어, $n=1000$ 그래프에서 NOTEARS가 10,896초, AVICI가 3.34초가 걸린 반면, CauScale은 0.83초가 소요되었습니다.
  • 공간: CauScale은 500개 노드 그래프로 훈련하고 1,000개 노드 그래프로 추론하는 데 성공했습니다. 반면, AVICI는 메모리 제약으로 인해 SERGIO-GRN 데이터셋($n=100$)에서 "Out of Memory" 오류를 일으켰습니다.

• 절제 연구 (Ablation Studies):

  • **축소 유닛(Reduction Unit)**을 제거하면 더 큰 훈련 세트에서 OOM 오류가 발생하고 추론 속도 및 메모리 효율성이 크게 저하되었습니다.
  • **그래프 스트림(Graph Stream)**을 제거했을 때 정확도가 가장 크게 하락하여, 이중 스트림 설계의 중요성을 입증했습니다.
  • **결합된 어텐션(Tied Attention)**은 정확도를 유지하거나 오히려 향상시키면서 일반적인 어텐션 대비 6배의 속도 향상을 제공했습니다.

5. 의의 및 주장

본 논문은 CauScale이 대규모 데이터 집약적 과학 분야에서 인과 구조 학습을 계산적으로 접근 가능하게 만드는 실질적인 진전이라고 주장합니다.

• 확장성: 연구진은 이론적인 대규모 인과 발견의 필요성과 실제적인 메모리 및 연산의 한계 사이의 간극을 성공적으로 메웠으며, 기존 신경망 방법론들이 실패하는 1,000개 노드 규모의 그래프 추론을 가능하게 했습니다.
• 일반화: 모델은 OOD 그래프 구조에 대해 강력한 일반화 능력을 보여주었으나, OOD 노이즈 분포 및 메커니즘 함수에 대해서는 다소 민감함을 보여, 향후 다양한 노이즈 생성에 집중할 필요가 있음을 시사합니다.
• 한계점: 저자들은 시뮬레이션된 훈련 데이터(SCM)에 대한 의존도가 실제 환경에서의 일반화 성능을 제한할 수 있음을 솔직하게 언급하며, 예측된 그래프는 고위험 결정 전 도메인 전문가의 검증과 후속 실험을 거쳐야 하는 가설로 취급되어야 함을 강조합니다.

요약하자면, CauScale은 시간 및 공간의 병목 현상을 극복하여 대규모 인과 그래프에 딥러닝을 적용할 수 있게 하는, 확장 가능하고 효율적이며 정확한 신경망 인과 발견 프레임워크를 제공합니다.