Graph Attention Coefficients

Compute the attention coefficients of a graph attention network (GAT) layer.

A GCN weights every neighbour by a fixed number that depends only on degrees. A GAT learns how much each neighbour matters. For each edge from node i to a node j in its neighbourhood (the neighbours of i and i itself):

Z      = H @ W                                   # project the features
e[i,j] = leaky_relu(a[:d] . Z[i] + a[d:] . Z[j])   # with slope 0.2 for x < 0
alpha[i, :] = softmax of e[i, j] over j in the neighbourhood of i

Here d is the width of Z, and a is a learned vector of length 2 * d. Splitting a in two is the same as taking a . [Z[i] || Z[j]].

Return the (n, n) matrix alpha. Entries outside each neighbourhood are 0, and each row sums to 1.

Function signature

def gat_attention(A, H, W, a):

Example

Input

A = np.array([[0, 1, 0],
              [1, 0, 1],
              [0, 1, 0]])
H = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [1.0, 1.0]])
W = np.eye(2)
a = np.array([1.0, 0.0, 0.0, 1.0])
gat_attention(A, H, W, a)

Output

array([[0.26894142, 0.73105858, 0.        ],
       [0.1553624 , 0.4223188 , 0.4223188 ],
       [0.        , 0.5       , 0.5       ]])

Reasoning

With this a, the score is Z[i][0] + Z[j][1]. Node 0 attends to itself with 1 + 0 = 1 and to node 1 with 1 + 1 = 2. The softmax of (1, 2) is (0.269, 0.731). Node 0 has no edge to node 2, so that weight is 0.

Constraints and notes

• A is a 0/1 adjacency matrix with zeros on the diagonal. Always include the self loop.
• Subtract the row maximum before the exponent for stability.