DAG-GNN

이번 포스트는 DAG-GNN: DAG Structure Learning with Graph Neural Networks의 논문 요약입니다.

Graph Learning

그래프 학습^{Graph Learning}이란, 주어진 관측 데이터 $\mathbf{X}$ 로부터 데이터가 생성된 그래프 $\mathcal{G}$ 의 구조를 추정하는 것이다. 그래프의 구조는 인접행렬 $A\in \mathbb{R}^{m\times m}$ 으로 나타나기 때문에, 이는 곧 인접행렬을 학습하는 문제가 된다. 그러나 일반적으로 어떠한 가정 없이 인접행렬을 학습하는 것은 매우 어려운 문제이다. 특히, 추정해야 할 그래프가 DAG(Bayesian Netowrk) 구조일 때, 이는 NP-hard problem이 된다. 이를 해결하기 위해 대개는 아래와 같은 선형 구조방정식 모형^{Linear SEM} 가정을 이용한다.

Linear SEM

그래프 $\mathcal{G}=(V,E)$ 가 $m$ 개의 노드를 갖는다고 하자. 이때, 각 노드는 $d$ 차원 확률변수에 대응한다. $X\in\mathbb{R}^{m\times d}$ 가 노드들의 결합분포로부터 얻어진 샘플이라고 할 때, 선형 구조방정식 모형은 아래와 같은 가정을 의미한다.

X = A^{T}X+Z

여기서 $Z$ 는 노이즈 벡터, 혹은 행렬을 의미한다. 노드의 위상정렬^{topological sorting}을 이용하면 인접행렬을 상삼각행렬로 만들 수 있으므로, 이는 아래와 같다.

X = (I-A)^{-1}Z

이러한 구조로부터 다음과 같은 generalized form을 생각할 수 있는데, 이것이 그래프 신경망^{Graph Neural Net}의 기반이 된다.

f_{2}^{-1}(X) = A^{T} f_{2}^{-1}(X)+ f_{1}(Z)

Model

DAG-GNN 모델은 선형 구조방정식 모형을 기반으로, VAE를 적용하여 이 과정에서 인접행렬을 학습하게 한다. 인접행렬을 학습가능한 파라미터로 정의하며, VAE의 잠재변수로는 노이즈 행렬 $Z$ 를 사용하는 것이 특징이다. 또한, 노이즈 행렬의 사전분포는 표준정규분포(혹은 Matrix normal)를 가정한다.

인코더로부터 얻은 확률분포를 $N(\mu_{e},\Sigma_{e}), \mu_{e}=(\mu_{1},\ldots,\mu_{m}), \Sigma_{e}=\mathrm{diag}(\sigma_{i}^{2})_{i=1}^{m}$ 으로 두면, ELBO는 다음과 같이 계산된다.

\begin{align} \mathcal{L}_\mathrm{ELBO} &= -D_{KL}(q(Z\vert X)\Vert p(Z)) + \mathrm{E}_{q}\log p(X\vert Z)\\ D_{KL}(q\Vert p)&= \frac{1}{2}\sum_{i=1}^m\bigg[\sigma_{i}^{2}+\mu_{i}^{2}-2\log \sigma_{i} -1\bigg] \end{align}

2번째 항인 reconstruction error의 경우 다음과 같이 데이터로부터의 몬테카를로 방법을 이용한다.

\mathrm{E}_{q}\log p(X|Z) \approx \frac{1}{N}\sum_{n=1}^{N}\sum _{i=1}^{m}-\frac{(X_{i}-\mu_{i}^{(n)})^{2}}{2\sigma_{i}^{2}}-\log \sigma_{i}^{(l)}-c

Optimization

ELBO를 최대화하는 계산으로는 학습된 인접행렬 $A$ 가 DAG 그래프에 대응된다는 것을 보장할 수 없다. 따라서 별도의 제약조건이 필요한데, 이때 사용되는 것이 acyclicity constraint이다. 이는 NOTEARS라는 알고리즘에서 처음 제안되었으며 (Zheng et al., n.d.), 아이디어는 인접행렬의 $k$ -거듭제곱이 노드 $i$ 에서 $j$ 로의 길이 $k$ 경로의 수에 대한 정보를 담고 있다는 것이다. 이로부터 DAG의 경우 인접행렬의 모든 거듭제곱에 대해 대각원소가 0이 되어야 함을 알 수 있다. 따라서, 다음과 같은 접근이 가능하다.

$\mathrm{tr}(I+\alpha A\circ A)^{m}=m$ if and only if $A$ is acyclic.

그렇기에, 이를 이용하면 최적화문제에 $h(A)=\mathrm{tr}(I+\alpha A\circ A)^{m}-m =0$ 의 제약조건을 추가할 수 있고, 이를 augmented lagrangian method를 이용해 풀 수 있다. 결과적으로, 풀어야 하는 최적화 문제는 다음과 같아진다.

\begin{align} \min_{A,\theta} &-\mathcal{L}_\mathrm{ELBO}(A,\theta)\\ s.t.\;\; h(A) &= \mathrm{tr}(I+\alpha A\circ A)^{m}-m =0 \end{align}

이를 다음과 같이 Lagrangian으로 쓸 수 있다.

\mathcal{L}_{c}(A,\theta,\lambda) = -\mathcal{L}_\mathrm{ELBO} +\lambda h(A) + \frac{c}{2}\vert h(A)\vert^{2}

학습 방법은, 모든 epoch에 대해 훈련이 종료된 뒤 $c$ 를 점진적으로 증가시켜 재학습을 진행하는 것이다. $c$ 를 업데이트하는 과정은 아래와 같다.

\begin{align} (A^{k},\theta^{k})&= \mathop{\arg\min}\limits_{A,\theta}L_{c^{k}}(A,\theta,\lambda^{k})\\ \lambda^{k+1} &= \lambda^{k}+c^{k}h(A^{k})\\ c^{k+1}&= \begin{cases}\eta c^{k},\quad \text{if}\; \vert h(A^{k})\vert \gt \gamma\vert h(A^{k-1})\vert \\ c^{k},\quad o.w.\end{cases} \end{align}

References

Yu, Y., Chen, J., Gao, T., & Yu, M. (n.d.). DAG-GNN: DAG Structure Learning with Graph Neural Networks.
Zheng, X., Aragam, B., Ravikumar, P. K., & Xing, E. P. (n.d.). DAGs with NO TEARS: Continuous Optimization for Structure Learning.