Multi-head attention lets a graph attention network (GAT) learn from multiple views of a node’s neighbors at once. Rather than treating every connected node equally, each head computes attention that helps determine which neighbors matter when updating a node’s features. section 4.5.3
A GAT layer uses the heads in a short pipeline:
- For each connected node pair, compute an attention score from their transformed features. section 4.5.3
- Normalize scores across a node’s neighbors, then use them to weigh and combine neighboring features. section 4.5.3
- Run multiple heads so the layer can learn from multiple perspectives simultaneously. section 4.5.3
A two-head GAT layer
In PyTorch Geometric, set `heads=2` on the first `GATConv`. If that layer produces 64 features per head, the following batch-normalization and convolution layers receive 128 features. The output convolution uses one head to produce a single representation for each node. section 4.4.1
Listing 4.7 GAT-based architecture
class GAT(torch.nn.Module):
def __init__(self, hidden_layers=32, heads=1, dropout_p=0.0):
super().__init__()
torch.manual_seed(2022)
self.conv1 = GATConv(data.num_node_features,\
hidden_layers, heads, dropout=dropout_p) #1
self.bn1 = nn.BatchNorm1d(hidden_layers*heads) #2
self.conv2 = GATConv(hidden_layers * heads, \
1, dropout=dropout_p)
def forward(self, data, dropout_p=0.0):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
x = self.bn1(x)
x = F.relu(x)
x = F.dropout(x, training=self.training)
x = self.conv2(x, edge_index)
return torch.sigmoid(x)
The GAT architecture shows where `heads` is set and why the next layers use `hidden_layers * heads` input dimensions. section 4.4.1
More heads carry a memory cost: GAT calculates attention scores for every head and every edge, while autograd retains tensors during training. For a large graph, mini-batch training with PyG’s `NeighborLoader` can replace training on the entire graph at once. section 4.4.1
