CS224W - (6) GNN Augmentation and Training

GNN 학습 파이프라인

GNN 학습 파이프라인은 위 그림과 같이 생겼습니다. 지금까지 다룬 내용은 입력 그래프부터 노드 임베딩까지의 내용이고, 본 포스트에서는 예측과 관련된 내용을 다루고자 합니다.

Different task levels require diffennt prediction heads

우선 다양한 예측 헤드가 있습니다. 예측 태스크라고도 할 수 있는데 노드 레벨, 엣지 레벨, 그래프 레벨에서 다양한 태스크가 있습니다. 각 레벨에 따라서 다른 예측 헤드를 갖게 됩니다.

예측 헤드: 노드 레벨

노드 레벨의 예측은 노드 임베딩을 이용해서 바로 예측을 생성할 수 있습니다. GNN 계산 이후 우리는 $d$ 차원의 노드 임베딩을 갖게 됩니다.

\[\left\{ h_v^{(L)} \in \mathbb{R}^d, \; \forall v \in G \right\}\]

이제 우리가 $k$-way prediction을 한다고 가정하겠습니다. 분류 문제라면 $k$ 개의 카테고리 중 하나를 예측해야 하고, 회귀 문제라면 $k$ 개의 타겟을 예측해야 하죠. 이런 태스크를 노드 임베딩을 포함하여 수식으로 나타내면 다음과 같습니다.

\[\hat{y}_v - \text{Head}_\text{node}(h_v^{(L)}) = W^{(H)} h_v^{(L)}\]

• $W^{(H)} \in \mathbb{R}^{k \times d}$: $d$ 차원의 노드 임베딩을 $k$ 차원의 타겟에 매핑하는 행렬

예측 헤드: 엣지 레벨

엣지 레벨의 예측은 두 노드의 임베딩 쌍을 활용해 생성합니다. 만약 아까와 같이 $k$-way prediction을 한다고 가정하면 예측값은 다음과 같습니다.

Prediction heads: Edge-level

\[\hat{y}_{uv} = \text{Head}_\text{edge} \left(h_u^{(L)}, h_v^{(L)} \right)\]

Concatenate + Linear

이때 $\text{Head}_\text{edge}$ 로 사용할만한 방법은 두 가지입니다. 첫 번째는 두 노드 임베딩을 붙인 후 선형 레이어를 태우는 방법입니다.

\[\hat{y}_{uv} = \text{Linear}\left( \text{Concat} \left( h_v^{(L)}, h_v^{(L)} \right) \right)\]

여기서 선형 레이어는 두 노드 임베딩을 붙인 $2d$ 차원의 임베딩을 $k$ 차원의 임베딩으로 매핑합니다.

두 번째 방법은 내적입니다.

\[\hat{y}_{uv} = \left( h_u^{(L)} \right)^T h_v^{(L)}\]

이 방법으론 두 노드간 엣지의 존재를 예측하는 등의 $1$-way prediction만 가능합니다. 만약 $k$-way prediction에 내적을 적용하려면 멀티 헤드 어텐션과 비슷한 방법을 사용해야 합니다.

\[\begin{aligned} \hat{y}_{uv}^{(1)} &= \left( h_u^{(L)} \right)^T W^{(1)} h_v^{(L)} \\ &\vdots \\ \hat{y}_{uv}^{(k)} &= \left( h_u^{(L)} \right)^T W^{(k)} h_v^{(L)} \\ \hat{y}_{uv} &= \text{CONCAT} \left( \hat{y}_{uv}^{(1)}, \cdots, \hat{y}_{uv}^{(k)} \right) \in \mathbb{R}^{k} \end{aligned}\]

예측 헤드: 그래프 레벨

그래프 레벨의 예측은 그래프 내에 있는 모든 노드의 임베딩을 사용합니다. 동일하게 $k$-way prediction 시나리오에서 예측값은 다음과 같습니다.

\[\hat{y}_G = \text{Head}_\text{Graph} \left( \left\{ h_v^{(L)} \in \mathbb{R}^d, \; \forall v \in G \right\} \right)\]

여기서 $\text{Head}_\text{Graph}(\cdot)$은 GNN 레이어에서 집계 함수와 비슷한 역할을 합니다. 그래서 보통 다음의 세 함수를 많이 사용합니다.

• Global mean pooling
  • \[\hat{y}_G = \text{Mean} \left( \left\{ h_v^{(L)} \in \mathbb{R}^d, \; \forall v \in G \right\} \right)\]
• Global max pooling
  • \[\hat{y}_G = \text{Max} \left( \left\{ h_v^{(L)} \in \mathbb{R}^d, \; \forall v \in G \right\} \right)\]
• Global sum pooling
  • \[\hat{y}_G = \text{Sum} \left( \left\{ h_v^{(L)} \in \mathbb{R}^d, \; \forall v \in G \right\} \right)\]

Hierarchical pooling

하지만 이런 함수들은 작은 그래프에서만 잘 작동합니다. 큰 그래프에서는 풀링 함수들이 정보 손실을 야기하기 때문입니다. 그래서 큰 그래프에 대해서는 계층적 풀링(hierarchical pooling) 을 사용합니다.

계층적 풀링은 풀링을 여러 단계에 걸쳐서 진행합니다. 우선 GNN 레이어 두 개를 각각 GNN-A, GNN-B라고 하겠습니다. GNN-A는 노드 임베딩을 계산하고, GNN-B는 각 노드들이 속할 군집을 계산합니다. 이 두 레이어는 병렬적으로 실행할 수 있습니다.

각 풀링 레이어에서는 GNN-B에서 배정한 군집마다 GNN-A에서 계산한 노드 임베딩을 집계합니다. 그 다음 각 군집에 대해 하나의 새 노드를 생성하고, 군집 간의 엣지를 유지하여 새로 풀링된 네트워크를 생성합니다.

GNN 학습

지도 학습과 비지도 학습

지도 학습에서의 라벨은 특정 상황에서 비롯됩니다. 예를 들어 노드 라벨의 경우 인용 네트워크에서 노드가 속한 주제 영역이 무엇인지로 정할 수 있을겁니다. 엣지 라벨은 트랜잭션 네트워크에서 어떤 엣지가 사기와 관련된 엣지인지로 정할 수 있을거구요. 그래프 라벨은 분자 그래프에서 그래프의 약물 유사성이 될 수 있습니다. 항상 라벨은 작업하기 쉬운 노드/엣지/그래프 라벨로 만들어야 합니다.

그런데 우리는 대부분의 상황에서 그래프 구조는 갖고 있지만 별도의 라벨이 없는 경우가 많습니다. 이런 경우에는 준지도 학습(self-supervised learning) 을 고려해야 합니다. 노드 라벨의 경우 군집 계수나 페이지랭크같은 노드 통계량을 사용할 수 있고, 엣지 라벨의 경우 두 노드의 엣지를 숨기고 실제 그 엣지의 예측 여부로 사용할 수 있습니다. 그래프 라벨은 두 그래프가 같은 구조인지와 같은 통계량을 사용할 수 있습니다.

손실 함수

이전 포스트에서 언급했듯 분류 문제에서는 보통 크로스 엔트로피(cross entropy)를 사용합니다. $k$-way prediction에서 $i$번째 데이터 포인트에 대해 크로스 엔트로피는 다음과 같습니다.

\[\text{CE}\left( y^{(i)}, \hat{y}^{(i)} \right) = -\sum^k_{j=1} y_j^{(i)} \log\left(\hat{y}_j^{(i)} \right)\]

회귀 문제의 경우 L2 loss를 많이 사용합니다.

\[\text{MSE}\left( y^{(i)}, \hat{y}^{(i)} \right) = \left( y_j^{(i)} - \hat{y}_j^{(i)} \right)^2\]

평가 메트릭

회귀 문제에 대해서는 RMSE나 MAE를 많이 사용합니다. 멀티클래스 분류 문제에서는 정확도를, 이진 분류 문제에서는 정확도, Precision, Recall, ROC/AUC 등을 많이 사용합니다.

GNN 예측 태스크 설정

고정 스플릿(Fixed split)의 경우 데이터셋을 한 번만 나눕니다.

• 학습 데이터: GNN 파라미터 최적화 용도
• 검증 데이터: 모델과 하이퍼파라미터 튜닝
• 테스트 데이터: 최종 성능 확인할 때까지 보류

하지만 테스트 데이터가 실제로 있을지 보장하기 어렵습니다. 그래서 보통 랜덤 스플릿(random split)을 통해 학습 데이터와 검증 데이터, 테스트 데이터를 모두 나눕니다. 그리고 여러 랜덤 시드 값에 대해서 평균 성능을 구하죠.

사실 그래프 형태의 데이터를 나누는 것은 일반적인 경우와는 다릅니다. 예를 들어 이미지 분류 문제에서 이미지 데이터셋은 서로에게 영향을 주지 않습니다. 하지만 그래프는 하나의 노드가 다른 노드에 영향을 줄 수 있습니다. 각각의 노드가 메시지 전달에 참여하기 때문입니다. 따라서 어떤 노드가 없다면 다른 노드의 임베딩은 변하게 되죠.

Transductive setting

그래프 데이터를 나눌 때는 두 가지 방법이 있습니다. 첫 번째는 전이적 설정(transductive setting) 입니다. 입력 그래프를 전체 데이터로 간주합니다. 이때 노드 라벨만 분리합니다. 위 그림처럼 학습 시에는 전체 그래프를 이용해 임베딩을 계산하고 노드 1과 노드 2로 라벨을 학습합니다. 검증할 때는 전체 그래프를 이용해 임베딩을 계산하고 노드 3과 노드 4를 이용해 평가합니다. 그래서 노드 예측과 엣지 예측 태스크에만 적용이 가능합니다.

Inductive setting

두 번째는 귀납적 설정(inductive setting) 입니다. 귀납적 설정에서는 엣지를 분리해 여러 개의 그래프로 나눕니다. 위 그림처럼 학습 데이터, 검증 데이터, 테스트 데이터 각각의 노드만 연결되어 있고, 다른 데이터셋끼리의 엣지는 모두 분리합니다. 학습 시에도 노드 1, 노드 2에 대한 임베딩만 계산하고 노드 1, 노드 2에 대한 라벨만 사용합니다. 이런 방법을 통하면 확인이 안된 그래프에 대한 일반화도 가능하고 전이적 설정과 다르게 그래프 예측 태스크에도 적용 가능합니다.