Neural Relational Inference for Interacting Systems

1. Introduction

인과관계추론(causal inference)는 예전부터 해결하기 어려운 문제이고, 다양한 분야로 활용이 가능하다. 해당 논문은 2018년에 발표된 논문으로 graph neural network와 variational autoencoder 를 함께 사용해서 다이나믹 시스템의 인과관계를 추론 및 각 노드들의 미래를 예측하는 모델이다. 대학원생때 잠깐 공부했던 분야라 현재 트랜드는 잘 모르지만 발표 다시부터 지금까지 굉장히 많은 주목을 받은 논문인 것은 확실하다. 지금은 해당 모델도 한계점이 많이 발견되서 다양한 변형들이 많이 나오고 있다고 한다. 확률공부하다가 생각나서 다시 개념환기도 시킬 겸 리뷰를 작성해보았다.

2. Background

2.1 Causal inference

인과관계(causality)는 사실상 우리 주변에서 발생하는 모든 것들에 작용하고 있다. 우리가 팔을 움직을때, 걷기 위해 발을 움직일때 모두 관절들 사이의 인과 관계라고 볼 수도 있다. 학생들의 시험 성적을 공부시간, 공부량 등으로 모델링한다면 어느 정도의 인과관계가 존재할 것이다. 물론 공부시간이 많다고 무조건 성적이 좋은 것은 아니지만...
이렇게 인과관계를 명확히 파악할 수 있다면 다양한 활용(주식 예측 등)이 가능하다. 하지만 이것을 힘들게 하는 요인이 있는데 바로 association 이라는 것이다. 간단하게 설명하면 직접적인 영향이 causality 라면 간접적인 영향이 association 이다. 우리가 어떤 시스템을 관찰할때 X, Y라는 변수가 서로 직접적인 영향을 주고 받는다고 생각했지만, 알고보니 Z 라는 숨은 변수가 실제 원인이었던 것이다. 이럴때 X, Y는 서로 causality가 아니라 association 의 관계가 형성된다. 하지만 우리는 그것을 causality 라고 잘못해석하고 프로젝트를 진행하는 것이다. 그래서 이러한 요소를 잘 잡아내는 것이 매우 중요하다.

2.2 Graph neural network

그래프 형태의 데이터를 다루기 위한 뉴럴 네트워크이다. 서로 상호작용하는 여러 엔티티가 존재하는 시스템은 그래프로의 해석이 매우 적절한데 해당 논문도 이를 활용했다. 노드들간의 작용하는 어떠한 힘을 메시지로 모델링함으로써 효율적으로 인과관계를 예측할 수 있다.

3. Main

3.1 Encoder

모델은 전형적인 GNN의 형태를 가져왔다. Encoder에서는 각 노드들의 time step 입력을 받아서 최종적으로는 엣지 타입에 대한 hidden state($z_{ij}$)를 찾아낸다. 각 노드들의 hidden state $h_{i}^{l}$ 을 사용해서 source, destination 의 순서에 맞게 MLP($f_{e}^{l}$) 에 입력하여 엣지의 hidden state를 계산한다. 그리고 하나의 노드를 기준으로 해당 노드에 들어오는 엣지들의 hidden state를 aggregation(여기서는 sum)하여 노드의 hidden state를 업데이트한다. 이후는 미리 정해둔 업데이트 수만큼 반복된다. 최종적으로 업데이트된 엣지 hidden state에 softmax를 취해주면서 edge type를 모델링하게 된다. 

3.2 Sampling

이전 단계에서 각 엣지들의 hidden state 에 softmax를 취해서 엣지 타입을 얻어낸다고 했다. 이것은 variational autoencoder의 자연스러운 흐름 중에 하나인데 이때 gradient 계산을 위해 reparamerization trick을 사용하게 된다. 간단하게 말하면 softmax 라는 행동을 통해 우리는 일종의 sampling을 하게 된 것인데 이 행동 자체가 미분이 불가능한 행동이다. 그런데 우리는 뉴럴 네트워크를 grdient를 통해 학습하기 때문에 미분을 해야만한다. 그래서 sampling 을 통해 얻어낸 변수(여기서는 z)를 일종의 함수로 변형하여 미분 가능하게 변형하는 것을 reparamerization trick 이라고 한다. 그런데 해당 모델은 여기서 더 나아가 z 라는 것이 엣지의 타입이기 때문에 discrete 한 변수이다. 예를 들어 softmax를 통해 [0.7, 0.3] 이라는 분포를 얻었고, z를 sampling 했더니 [1, 0] 이라는 엣지 타입을 얻었다. 하지만 이것이 미분이 불가능하다는 것이다. 따라서 한가지 장치가 더 적용되는데 gumbel softmax라는 함수이다. gumbel softmax는 softmax의 변형인데, 강도를 조절할 수 있는 softmax라고 생각하면 된다.

3.3 Decoder

Decoder에서는 이전 단계에서 얻은 $z_{ij}$, 각 노드들의 입력($x_{i}^{t}$) 들을 사용해서 각 노드들의 다음 스텝을 예측한다. 이때 정확한 좌표값을 예측하는 것이 아니라 대략적인 위치 예측을 위한 평균, 분산을 계산한다. 모델 자체는 그렇게 특별한 것은 없다. 모델 최적화를 위해 해당 모델은 한가지 추가적인 장치를 적용했다. 해당 모델은 기본적으로 one-step prediction을 수행한다. 즉, 한번에 한번의 미래만 예측한다는 것이다. 이러한 방법을 선택했을때 발생할 수 있는 문제점은 너무 짧은 인과관계만 반영이 된다는 것이다. 이를 예방하기 위해 NRI는 2가지 전략을 선택했다. 첫번째는 multiple-step을 예측한다는 것이다. 두번째는 엣지 타입별로 다른 MLP를 적용한다는 것이다. 위의 수식에서 첫 번째 수식의 MLP가 엣지 타입 ($k$) 에 따라 다르게 적용된 것을 알 수 있다. 아래의 수식은 multiple-step 을 어떻게 구현했는지 설명한다. M step 마다 정답을 주고 그 사이에는 예측된 평균을 사용해서 다시 평균을 예측하는 방식이다.