Graph Attention Networks (GAT), in depth: edge scores, segment softmax, multi-head layers and GATv2

The attention score

Take node features h_i of size F. A GAT layer has a weight matrix W of shape F' by F, shared by every node, and an attention vector a of length 2F'. For an edge from neighbour j into target i, the raw score is

e_ij  = LeakyReLU( a^T [ W h_i || W h_j ] )        # || is concatenation, slope 0.2
alpha_ij = exp(e_ij) / sum_{k in N(i)} exp(e_ik)   # softmax over i's neighbourhood
h_i'  = sigma( sum_{j in N(i)} alpha_ij W h_j )    # sigma = ELU in the paper

Three details matter. First, N(i) includes i itself: the layer adds a self-loop so a node can keep its own information, and so a node with no in-edges still has a non-empty softmax. Second, the attention is masked: the softmax runs only over actual neighbours, never over all nodes. That is what makes it a graph layer rather than a transformer over a bag of nodes, and it keeps the cost proportional to the number of edges. Third, splitting a into two halves shows the score is cheaper than it looks. Write a = [att_dst ; att_src]; then

a^T [W h_i || W h_j] = att_dst . (W h_i)  +  att_src . (W h_j)
                     =        d_i         +        s_j

so you compute one scalar d per node and one scalar s per node, and the per-edge work is a single addition followed by the LeakyReLU. You never materialise the 2F'-wide concatenation. PyTorch Geometric names these halves att_dst and att_src; this article uses the same names so you can read its source alongside.

Segment softmax on an edge list

Graphs arrive as an edge list, two integer arrays src and dst, not as an adjacency matrix. A dense implementation that builds an N by N score matrix and masks it works on Cora's 2,708 nodes and runs out of memory long before a million. The scalable form is a segment softmax: group edges by their target and normalise within each group. It needs three scatter operations, a max for numerical stability, a sum for the denominator, and a sum for the output. Here is a complete multi-head layer in NumPy:

import numpy as np

def leaky_relu(x, slope=0.2):
    return np.where(x > 0, x, slope * x)

def gat_layer(X, src, dst, W, att_src, att_dst, heads):
    # X: (N, F). src[k] -> dst[k] is edge k, self-loops already added.
    # W: (F, heads*Fp). att_src, att_dst: (heads, Fp).
    N = X.shape[0]
    Fp = W.shape[1] // heads
    Wh = (X @ W).reshape(N, heads, Fp)
    s = (Wh * att_src).sum(-1)              # (N, heads): neighbour half
    d = (Wh * att_dst).sum(-1)              # (N, heads): target half
    e = leaky_relu(d[dst] + s[src])         # (E, heads): one score per edge
    m = np.full((N, heads), -np.inf)
    np.maximum.at(m, dst, e)                # per-target max, for stability
    ex = np.exp(e - m[dst])
    den = np.zeros((N, heads))
    np.add.at(den, dst, ex)                 # per-target denominator
    alpha = ex / den[dst]                   # (E, heads), sums to 1 per target
    out = np.zeros((N, heads, Fp))
    np.add.at(out, dst, alpha[..., None] * Wh[src])
    return out, alpha                       # caller concats or averages heads

The cost is O(N F F') for the projection and O(E F') for scores and aggregation, per head. Memory is dominated by the E by heads attention tensor and, in frameworks that materialise messages, the E by heads by F' message tensor. On graphs with billions of edges that message tensor, not the parameters, is what limits batch size; fused kernels in PyG and DGL avoid materialising it where they can.

Worked example: one target, four neighbours

Run the layer on a star: target node 0 with neighbours 1, 2 and 3, plus its self-loop. Use two features, W equal to the identity so Wh = h, and one head with att_dst = [0.5, -0.5] and att_src = [1.0, 0.5]. The features are h0 = (1, 0), h1 = (0, 1), h2 = (1, 1) and h3 = (2, 0).

The target half is d_0 = 0.5 times 1 minus 0.5 times 0 = 0.5. The exponentials of the four scores are 4.48, 2.72, 7.39 and 12.18, summing to 26.77, which gives the alpha column. The new state is 0.167 (1, 0) + 0.102 (0, 1) + 0.276 (1, 1) + 0.455 (2, 0) = (1.354, 0.378), before the nonlinearity. Running gat_layer on this graph with heads=1 reproduces those numbers exactly, which is the first test to write for any implementation: a star small enough to check with a calculator.

Multi-head attention and the original configuration

A single head is noisy to train, so GAT runs K independent heads, each with its own W and a, and combines them. In hidden layers the K outputs are concatenated, giving K times F' features. In the output layer they are averaged, because concatenating class logits would make no sense. The original Cora model used a first layer of 8 heads with 8 features each (64 hidden features) and ELU, and a second layer with a single head producing the class scores, with dropout of 0.6 on both the layer inputs and the attention coefficients and L2 regularisation of 0.0005. The paper reports 83.0 percent accuracy on Cora's standard split, against 81.5 for GCN. Treat that gap as modest: later re-evaluations on random splits found the ranking between GCN and GAT on small citation graphs depends on the split and tuning.

Dropout on alpha is unusual and important. It means each training step a node sees a random subset of its neighbourhood, a form of neighbour sampling that regularises small graphs heavily. Remember it when you compare attention weights between training and evaluation modes: they differ.

Static attention and GATv2

Look again at the worked example. Because LeakyReLU is monotonic, e_ij is increasing in s_j for every target i, so the ranking of neighbours by attention depends only on s_j, a property of the neighbour alone. In the example node 3 has the largest s, so node 3 receives the most attention from every node it is connected to, whatever that node's own features. Brody, Alon and Yahav called this static attention and proved that the original GAT cannot express a ranking that depends on the query. Their fix, GATv2, moves the nonlinearity between the two linear steps:

GAT   : e_ij = LeakyReLU( a^T [W h_i || W h_j] )      # rank of j fixed across all i
GATv2 : e_ij = a^T LeakyReLU( W [h_i || h_j] )        # rank of j can depend on i

In GATv2 the target and neighbour features mix inside the LeakyReLU before the projection onto a, which makes the scoring function a small MLP over the pair and lets different targets prefer different neighbours. The cost is that the per-node shortcut disappears: the score now needs an F'-wide vector per edge rather than one scalar addition, so GATv2 uses more memory per edge. Their experiments show GATv2 matching or beating GAT across benchmarks, with the largest gains on tasks where which neighbour matters depends on who is asking. If you are starting a new model, try GATv2 first.

GATConv in PyTorch Geometric

In practice you will use a library layer. In PyTorch Geometric, GATConv(in_channels, out_channels, heads=1, concat=True, negative_slope=0.2, dropout=0.0, add_self_loops=True, edge_dim=None) implements the original layer and GATv2Conv the dynamic one. A two-layer node classifier in the paper's configuration:

import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv

class GAT(torch.nn.Module):
    def __init__(self, in_dim, n_classes, hidden=8, heads=8):
        super().__init__()
        self.conv1 = GATConv(in_dim, hidden, heads=heads, dropout=0.6)            # concat -> 64
        self.conv2 = GATConv(hidden * heads, n_classes, heads=1, concat=False,
                             dropout=0.6)                                         # mean over heads

    def forward(self, x, edge_index):
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        return self.conv2(x, edge_index)

# Inspecting attention: returns (out, (edge_index_with_self_loops, alpha))
model.eval()                                    # no dropout on alpha
out, (ei, alpha) = model.conv1(x, edge_index, return_attention_weights=True)

Two library behaviours catch people. edge_index row 0 is the source and row 1 the target, so messages flow from row 0 into row 1; reversing an asymmetric graph silently trains a different model. And the returned edge index includes the added self-loops, so align alphas with the returned index, never with the one you passed in. If your edges carry features such as bond types or transaction amounts, pass edge_dim and edge_attr so they enter the score.

Running it on real graphs

For graphs that do not fit in GPU memory, train on sampled subgraphs, as with GraphSAGE. Sampling changes the softmax denominator: a node with 5,000 neighbours sampled down to 25 normalises over 25, so the attention you see in training is not the attention it computes at full-neighbourhood inference. Either infer with the same sampler, or validate that the two agree on a holdout. Log three things per layer and head during training: the mean entropy of alpha, bucketed by node degree; the fraction of nodes whose largest alpha exceeds 0.9; and the gradient norm of a. Entropy near log(degree) everywhere means attention has collapsed to a GCN-style average and you are paying for heads that do nothing.

Failure modes

• Uniform attention. If the features carry little signal about which neighbours matter, a learns near zero and alpha becomes 1/degree. Check entropy before claiming the model "attends". A plain GCN with the same budget is the honest baseline.
• High-degree dilution. Softmax over thousands of neighbours spreads mass thinly, and a hub's state becomes an average of noise. Cap degree by sampling or add degree features.
• Overflow and NaN. Skipping the per-target max subtraction overflows exp in half precision. A target with no incoming edges and no self-loop divides by zero.
• Duplicate edges. A repeated edge appears twice in the softmax and gets double weight. Coalesce the edge list first.
• Reading alpha as an explanation. Attention weights show where mass went, not why the prediction changed; heads are combined and later layers mix everything again. Use ablations or perturbation tests before showing alpha to a stakeholder as a reason.
• Depth. Attention does not cure oversmoothing. Beyond two or three layers, add residual connections (residual=True in recent PyG) or jumping-knowledge readouts.
• Heterophily. When linked nodes tend to differ, any neighbour averaging hurts. Keep a separate path for the node's own features or use a heterophily-aware model.

Trade-offs

What to do next

1. Implement gat_layer and check it against the four-node table above, then against GATConv with copied weights.
2. Train a GCN and a GAT with the same hidden size on your graph, five seeds each, and compare means and spread before choosing.
3. Swap in GATv2Conv and keep it if it wins on validation.
4. Log attention entropy by degree bucket; if it sits at log(degree), drop to GCN.
5. Coalesce edges, confirm direction, and verify the returned self-loop alignment before inspecting any attention weights.
6. Read the attention mechanism article and multi-head attention maths to see how the same softmax-weighted sum works without a graph mask, and the softmax article for its numerics.