6. Graph Attention Networks (GATs) 기초 및 코드 구현

Contents
1. GAT 기본 개념
1.1 Linear transformation
1.2
Activation function
1.3 Softmax normalization
1.4 Multi-head attention
1.5 Improved graph attention layer
2. Numpy로 graph attention layer 구현하기
3. GAT 구현하기

 

이전 포스팅에서는 서로 다른 이웃 수를 가진 노드 degree를 정규화하여 이웃 정보를 학습하는 GCN 알고리즘을 알아보았습니다. 이번 포스팅에서 다룰 Graph Attention Network (GAT)는 GCN에서 한층 더 발전된 알고리즘입니다. GAT의 핵심 아이디어는 노드 간 상호작용을 모델링하기 위해, 그래프의 각 노드가 이웃 노드와의 상호작용에 중요도 정보를 반영할 수 있도록 하는 어텐션 메커니즘을 사용하는 것입니다. 어텐션 메커니즘은 NLP 분야에서 BERT나 GPT의 기반이 되는 transformer 모델에 사용됩니다. 2017년에 발표된 GAT는 그래프 내에서 중요한 정보에 가중치를 주어 노드의 임베딩을 더욱 효율적으로 계산하기 때문에, NLP 분야에서와 마찬가지로 뛰어난 성능을 보여주었습니다. 지금부터 기본적인 GAT의 개념을 소개하고, 코드로 구현해 보겠습니다.

1. GAT 기본 개념

GAT의 아이디어는 특정한 일부 노드가 다른 노드보다 더 중요하다는 것입니다. 사실, GCN 레이어에서도 이러한 개념을 보았습니다. 이는 normalization coefficient \( \frac{1} {\sqrt{deg(i)}\sqrt{deg(j)}}  \)에 의해 적은 이웃을 가진 노드가 더 중요한 노드라는 것을 표현하는 것입니다. 그러나 이러한 방식은 노드의 차수만을 고려한다는 한계가 존재합니다. 반면에 graph attention layer의 목표는 노드 feature의 중요 도을 고려하는 가중치 요소인 attention score를 생성하는 것입니다.

 

어텐션 스코어(attention scores)는, 노드 ℎᵢ와 이웃 노드인 ℎⱼ 사이의 어텐션 스코어를 αᵢⱼ로 나타냅니다. 따라서 graph attention의 작동 원리를 아래와 같은 식으로 나타낼 수 있습니다.

$$ h_i = \sum_{j\in N_i} a_{ij} W x_j $$

 

다음으로 Graph Attention 과정을 이해하기 위한 핵심 요소를 차근차근 알아보겠습니다.

  1. Linear transformation
  2. Activation function
  3. Softmax normalization
  4. Multi-head attention
  5. Improved graph attention layer

 

1.1 Linear transformation

어텐션 스코어는 중앙 노드 ℎᵢ와 이웃 노드 ℎⱼ 사이의 중요도를 나타냅니다. 이는 두 노드의 feature 정보를 기반으로 합니다. 그래프 어텐션 레이어에서는 이를 히든 벡터 ℎᵢ와 ℎⱼ를 concatenate 한 [ℎᵢ || ℎⱼ]로 표현합니다. 여기서 𝑊는 히든 벡터를 계산하기 위해 공유되는 가중치 행렬입니다. 이후 학습 가능한 가중치 행렬 𝐚를 통해 추가적인 선형 변환 (linear transformation)이 적용됩니다. 이 행렬 𝐚는 어텐션 계수를 생성하기 위해 가중치를 학습합니다. 이러한 과정은 아래 식으로 표현할 수 있습니다. 이는 기존 가중치 행렬과 노드 특성의 concat에 대해 어텐션 행렬을 곱해준 것으로 이해하면 되겠습니다.

 

$$ a_{ij} = W^T_{att} [W{}x_i \left| \right| Wx_j] $$

이후 이러한 attention score는 이후 activation 함수를 통과합니다.

 

2.2 Activation function

기존의 neural network 모델과 마찬가지로, 활성화 함수를 통해 비선형성을 더해줍니다. GAT에서는 ReLU 함수를 변형한 Leaky ReLU를 적용합니다.

활성화 함수를 통과한 어텐션 스코어는 아래와 같은 식으로 계산됩니다.

$$ e_{ij} = LeakyReLU(a_{ij}) $$

 

1.3 Softmax normalization

각 어텐션 스코어를 비교하기 위해 정규화가 필요합니다. 최종 어텐션 스코어는 아래와 같이 노드 i의 모든 이웃에 대한 소프트 맥스 함수를 적용하여 나타냅니다. 아래 식을 통해, 이제 노드 i와 j사이의 중요도를 나타내는 attention score를 최종적으로 알 수 있게 되었습니다.

$$a_{ij} = softmax(e_{ij}) = \frac{exp(e_{ij})} {\sum_{k \in N_i} exp(e_{ik})}$$

 

1.4 Multi-head attention

Transformer 모델에서도 알 수 있듯이 multi-head attention은 이러한 attention score를 여러 번 구하고, 다양한 표현을 학습하는 것입니다. 즉, 이전 세 단계를 여러 번 반복하기만 하면 됩니다. 'k'는 각 연산을 통해 계산된 임베딩을 의미합니다. 이러한 여러 개의 어텐션 헤드를 결합하는 방법에는 평균과 concat 두 가지 방법이 있습니다. 보통 각 임베딩의 고유 정보를 유지하기 위해, concat 방식을 활용하고, 이는 transformer 모델에서도 마찬가지입니다.

 

1.5 Improved graph attention layer

이러한 기본 개념을 바탕으로 Brody et al. (2021)에 의해 제안된, 향상된 graph attention layer를 알아보겠습니다. 단순히 연산의 순서에 관한 변화가 있습니다. 가중치 행렬 W는 두 노드의 피처를 concatenation 한 후에 적용되고, 어텐션 가중치 행렬은 활성화 함수 후에 적용됩니다. 두 식을 비교하면 단순히 연산 순서에만 차이가 있는 것을 알 수 있습니다.

 

원래 GAT 연산

 

향상된 GAT 연산

 

2. Numpy로 graph attention layer 구현하기

어텐션 가중치 행렬을 추가한 graph attention layer는 다음과 같은 행렬 식으로 나타냅니다. 여기서 \(W_a\)는 모든 어텐션 스코어 \(a_{ij}\) 값을 행렬로 나타낸 것입니다.

$$ H = \tilde{A}^T W_a X W^T $$

 

먼저 인접 행렬 A와 4개의 노드들의 feature를 4차원으로 표현한 X를 만듭니다.

import numpy as np
np.random.seed(0)

A = np.array([
    [1, 1, 1, 1],
    [1, 1, 0, 0],
    [1, 0, 1, 1],
    [1, 0, 1, 1]
])

X = np.random.uniform(-1, 1, (4, 4))
X

### result ###

array([[ 0.09762701,  0.43037873,  0.20552675,  0.08976637],
       [-0.1526904 ,  0.29178823, -0.12482558,  0.783546  ],
       [ 0.92732552, -0.23311696,  0.58345008,  0.05778984],
       [ 0.13608912,  0.85119328, -0.85792788, -0.8257414 ]])

 

이후 일반적인 가중치 행렬 W와 attention 가중치 행렬 W_att를 랜덤 하게 만듭니다.

이때 가중치 행렬 W의 차원은 (hidden 차원의 수, 노드의 수)입니다. hidden 차원은 2로 지정하였고, 노드의 수는 4로 고정되어 있습니다. 어텐션 가중치 행렬은 히든 벡터의 concat에 적용되기 때문에 (1, hidden 차원수 * 2)로 만들어 줍니다.

W = np.random.uniform(-1, 1, (2, 4))

W_att = np.random.uniform(-1, 1, (1, 4))

 

이후 소스 노드와 이웃 노드의 hidden 벡터를 concat 합니다. 두 이웃 노드의 쌍을 얻는 간단한 방법은 COO 형식의 인접 행렬을 살펴보는 것입니다. COO 형식에서 각 행은 소스 노드와 이웃 노드를 의미합니다. NumPy는 np.where()를 사용하여 빠르고 효율적으로 수행할 수 있는 방법을 제공합니다.

connections = np.where(A > 0)
connections

### result ###

(array([0, 0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 3]),
 array([0, 1, 2, 3, 0, 1, 0, 2, 3, 0, 2, 3]))

이후 np.concatenate를 사용하여 소스 노드와 이웃 노드들의 hidden vector들을 concat 합니다.

즉 connection[0]에 해당하는 소스 노드들에 가중치 행렬 W를 곱한 것과, 이웃 노드들을 뜻하는 connections[1]에 가중치 행렬 W를 곱한 값을 concat 하는 것입니다.

np.concatenate([(X @ W.T)[connections[0]], (X @ W.T)[connections[1]]], axis=1)

### result ###
array([[ 0.37339233,  0.38548525,  0.37339233,  0.38548525],
       [ 0.37339233,  0.38548525,  0.85102612,  0.47765279],
       [ 0.37339233,  0.38548525, -0.67755906,  0.73566587],
       [ 0.37339233,  0.38548525, -0.65268413,  0.24235977],
       [ 0.85102612,  0.47765279,  0.37339233,  0.38548525],
       [ 0.85102612,  0.47765279,  0.85102612,  0.47765279],
       [-0.67755906,  0.73566587,  0.37339233,  0.38548525],
       [-0.67755906,  0.73566587, -0.67755906,  0.73566587],
       [-0.67755906,  0.73566587, -0.65268413,  0.24235977],
       [-0.65268413,  0.24235977,  0.37339233,  0.38548525],
       [-0.65268413,  0.24235977, -0.67755906,  0.73566587],
       [-0.65268413,  0.24235977, -0.65268413,  0.24235977]])

이후 어텐션 행렬 W_att를 곱해주어 선형 변환을 실시합니다.

연결이 존재하는 pair가 12개 이므로, 각 pair의 중요도를 나타내는 12개의 값을 구했습니다.

a = W_att @ np.concatenate([(X @ W.T)[connections[0]], (X @ W.T)[connections[1]]], axis=1).T
a
array([[-0.1007035 , -0.35942847,  0.96036209,  0.50390318, -0.43956122,
        -0.69828618,  0.79964181,  1.8607074 ,  1.40424849,  0.64260322,
         1.70366881,  1.2472099 ]])

이러한 a에 활성화 함수를 적용합니다.

def leaky_relu(x, alpha=0.2):
    return np.maximum(alpha*x, x)

e = leaky_relu(a)
e

### result ###
array([[-0.0201407 , -0.07188569,  0.96036209,  0.50390318, -0.08791224,
        -0.13965724,  0.79964181,  1.8607074 ,  1.40424849,  0.64260322,
         1.70366881,  1.2472099 ]])

활성화 함수를 통과한 a값을 다시 인접행렬 A와 같은 형태로 복구합니다. A와 마찬가지로, 연결이 없는 노드쌍은 0으로 나타내고, 연결이 존재하는 노드 쌍에 e 값이 채워집니다. 앞서 정의한 connections에 존재하는 위치 값을 활용합니다.

E = np.zeros(A.shape)
E[connections[0], connections[1]] = e[0]
E

### result ###
array([[-0.0201407 , -0.07188569,  0.96036209,  0.50390318],
       [-0.08791224, -0.13965724,  0.        ,  0.        ],
       [ 0.79964181,  0.        ,  1.8607074 ,  1.40424849],
       [ 0.64260322,  0.        ,  1.70366881,  1.2472099 ]])

이러한 attention score에 softmax 정규화를 실시합니다. 커스텀 softmax 함수를 정의하고 최정 W_alpha 값을 구했습니다.

def softmax2D(x, axis):
    e = np.exp(x - np.expand_dims(np.max(x, axis=axis), axis))
    sum = np.expand_dims(np.sum(e, axis=axis), axis)
    return e / sum

W_alpha = softmax2D(E, 1)
W_alpha

### result ###
array([[0.15862414, 0.15062488, 0.42285965, 0.26789133],
       [0.24193418, 0.22973368, 0.26416607, 0.26416607],
       [0.16208847, 0.07285714, 0.46834625, 0.29670814],
       [0.16010498, 0.08420266, 0.46261506, 0.2930773 ]])

이러한 어텐션 매트릭스는 그래프에 존재하는 모든 노드 pair에 가중치로 적용합니다. 최종적으로 임베딩 차원이 2인 행렬 H를 각 노드의 표현으로 구합니다.

H = A.T @ W_alpha @ X @ W.T

### result ### 
array([[-1.10126376,  1.99749693],
       [-0.33950544,  0.97045933],
       [-1.03570438,  1.53614075],
       [-1.03570438,  1.53614075]]

 

3. GAT 구현하기

앞선 섹션 2에서 graph attention layer를 밑바닥부터 구현해 보았습니다. 이번 섹션에서는 이러한 레이어를 통해 graph attention network를 구현하는 과정을 알아보겠습니다. Cora 데이터셋을 사용하고, PyG에서 제공하는 GAT 레이어를 활용하겠습니다.

 

먼저 Cora 데이터셋을 불러옵니다.

from torch_geometric.datasets import Planetoid

# Import dataset from PyTorch Geometric
dataset = Planetoid(root=".", name="Cora")
data = dataset[0]

GAT 클래스에서 GAT 레이어는 PyG에서 제공하는 레이어를 간편하게 사용했습니다. 또한 위에서 설명한 향상된 GAT인 GATv2COnv 레이어를 불러왔습니다.

 

첫 번째 GAT 레이어의 multi-heads는 8개로 지정하고, 두 번째 GAT의 head는 하나로 설정했습니다.

 

오버피팅을 피하기 위해 dropout 레이어를 적용했습니다. 또한 코드에서는 활성화 함수로 ELU를 적용합니다.

나머지 과정은 이전 포스팅과 동일합니다.

import torch
torch.manual_seed(1)
import torch.nn.functional as F
from torch_geometric.nn import GATv2Conv, GCNConv
from torch.nn import Linear, Dropout


def accuracy(y_pred, y_true):
    """Calculate accuracy."""
    return torch.sum(y_pred == y_true) / len(y_true)


class GAT(torch.nn.Module):
    # multi-head GAT에서 첫번째 레이어의 head는 8개가 좋음. 이후는 달라짐
    def __init__(self, dim_in, dim_h, dim_out, heads=8):
        super().__init__()
        self.gat1 = GATv2Conv(dim_in, dim_h, heads=heads)
        self.gat2 = GATv2Conv(dim_h*heads, dim_out, heads=1)

    def forward(self, x, edge_index):
        h = F.dropout(x, p=0.6, training=self.training) # Preventing overfitting
        h = self.gat1(h, edge_index)
        h = F.elu(h)
        h = F.dropout(h, p=0.6, training=self.training)
        h = self.gat2(h, edge_index)
        return F.log_softmax(h, dim=1)

    def fit(self, data, epochs):
        criterion = torch.nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(self.parameters(), lr=0.01, weight_decay=0.01)

        self.train()
        for epoch in range(epochs+1):
            optimizer.zero_grad()
            out = self(data.x, data.edge_index)
            loss = criterion(out[data.train_mask], data.y[data.train_mask])
            acc = accuracy(out[data.train_mask].argmax(dim=1), data.y[data.train_mask])
            loss.backward()
            optimizer.step()

            if(epoch % 20 == 0):
                val_loss = criterion(out[data.val_mask], data.y[data.val_mask])
                val_acc = accuracy(out[data.val_mask].argmax(dim=1), data.y[data.val_mask])
                print(f'Epoch {epoch:>3} | Train Loss: {loss:.3f} | Train Acc: {acc*100:>5.2f}% | Val Loss: {val_loss:.2f} | Val Acc: {val_acc*100:.2f}%')

    @torch.no_grad()
    def test(self, data):
        self.eval()
        out = self(data.x, data.edge_index)
        acc = accuracy(out.argmax(dim=1)[data.test_mask], data.y[data.test_mask])
        return acc

# Create the Vanilla GNN model
gat = GAT(dataset.num_features, 32, dataset.num_classes)
print(gat)

# Train
gat.fit(data, epochs=100)
### result ###

GAT(
  (gat1): GATv2Conv(1433, 32, heads=8)
  (gat2): GATv2Conv(256, 7, heads=1)
)
Epoch   0 | Train Loss: 1.969 | Train Acc: 15.00% | Val Loss: 1.96 | Val Acc: 11.80%
Epoch  20 | Train Loss: 0.259 | Train Acc: 96.43% | Val Loss: 1.10 | Val Acc: 67.60%
Epoch  40 | Train Loss: 0.163 | Train Acc: 98.57% | Val Loss: 0.90 | Val Acc: 70.80%
Epoch  60 | Train Loss: 0.205 | Train Acc: 98.57% | Val Loss: 0.96 | Val Acc: 69.00%
Epoch  80 | Train Loss: 0.130 | Train Acc: 100.00% | Val Loss: 0.91 | Val Acc: 70.80%
Epoch 100 | Train Loss: 0.148 | Train Acc: 99.29% | Val Loss: 0.90 | Val Acc: 73.00%

테스트셋에 적용 결과 다음과 같이 향상된 결과를 보여주었습니다.

# Test
acc = gat.test(data)
print(f'GAT test accuracy: {acc*100:.2f}%')

### result ###
GAT test accuracy: 82.00%

 

GNN 이론이 점점 어려워지고 있지만, attention과 기법에 대한 이해가 있다면 쉽고 직관적으로 이해할 수 있을 것이라고 생각합니다. 이번 포스팅에서는 GAT(Graph Attention Networks)에서 소개하는 선형 변환부터 다중 헤드 어텐션까지 네 가지 주요 단계로 내부 동작을 살펴보았습니다. 이후 Cora 데이터셋에 모델 아키텍처를 적용하는 것까지 구현했습니다.

 

다음 포스팅은 그래프에서 스케일 문제를 해결하기 위한 GraphSAGE와 함께 대규모 그래프 관리에 특화된 새로운 아키텍처를 소개드리겠습니다. 감사합니다.

반응형