Implementing Multi-Head Attention in GATs

In the previous section, we examined the Graph Attention Network (GAT) architecture and its core component: the self-attention mechanism. This mechanism allows a node to assign different importance scores to its neighbors when aggregating features. While powerful, relying on a single attention mechanism can sometimes lead to unstable training dynamics or limit the model's capacity to capture diverse relational aspects within a node's neighborhood.

To address these limitations and enhance the robustness and expressivity of GATs, we can employ multi-head attention, a technique adapted from the success of Transformer models.

The Rationale for Multiple Heads

The fundamental idea behind multi-head attention is to run the self-attention process multiple times independently and in parallel. Each independent run is termed an "attention head." Imagine each head as a separate expert, focusing on potentially different aspects of the neighborhood information or learning distinct attention patterns. By combining the insights from these multiple experts, the model can form a more comprehensive and stable representation.

Mechanism Breakdown

Let's formalize how multi-head attention works within a GAT layer. If we want to use $K$ independent attention heads, the process for updating the representation of node $i$, denoted $h_i'$, proceeds as follows:

1. Independent Transformations and Attention: For each head $k$ (where $k = 1, ..., K$), we introduce a separate set of learnable parameters: a weight matrix $W^{(k)}$ and an attention mechanism parameter vector $\vec{a}^{(k)}$. Each head calculates its own attention coefficients $\alpha_{ij}^{(k)}$ based on its specific parameters:

   • Apply the head-specific linear transformation: $z_i^{(k)} = W^{(k)} h_i$.
   • Compute the un-normalized attention score for edge $(j, i)$ using head $k$'s attention mechanism: $e_{ij}^{(k)} = \text{LeakyReLU}(\vec{a}^{(k)T} [W^{(k)} h_i || W^{(k)} h_j])$ or equivalently $e_{ij}^{(k)} = \text{LeakyReLU}(\vec{a}^{(k)T} [z_i^{(k)} || z_j^{(k)}])$
   • Normalize the scores using the softmax function across the neighborhood $\mathcal{N}_i$ of node $i$ (including the self-loop): $\alpha_{ij}^{(k)} = \text{softmax}_j(e_{ij}^{(k)}) = \frac{\exp(e_{ij}^{(k)})}{\sum_{l \in \mathcal{N}_i \cup \{i\}} \exp(e_{il}^{(k)})}$

2. Head-Specific Feature Aggregation: Each head $k$ then computes an intermediate output representation for node $i$ by performing a weighted aggregation using its calculated attention coefficients: $h_i'^{(k)} = \sigma\left( \sum_{j \in \mathcal{N}_i \cup \{i\}} \alpha_{ij}^{(k)} W^{(k)} h_j \right)$ Here, $\sigma$ represents an activation function like ReLU or ELU. Note that the output dimension of each head $h_i'^{(k)}$ is typically $F'$, the desired output feature dimension per head.

3. Combining Head Outputs: Finally, the outputs from all $K$ heads ($h_i'^{(1)}, h_i'^{(2)}, ..., h_i'^{(K)}$) need to be combined to produce the final layer output $h_i'$. There are two common strategies:

   • Concatenation: This is the most frequent approach for intermediate GAT layers. The outputs of all heads are concatenated together: $h_i' = \Vert_{k=1}^K h_i'^{(k)} = \Vert_{k=1}^K \sigma\left( \sum_{j \in \mathcal{N}_i \cup \{i\}} \alpha_{ij}^{(k)} W^{(k)} h_j \right)$ The resulting feature vector $h_i'$ will have a dimension of $K \times F'$. If subsequent layers expect a different dimension, an additional linear transformation might be applied after concatenation.
   • Averaging: This is often used in the final layer of the GAT model, especially for node classification tasks where a specific output dimension is required. The outputs are averaged: $h_i' = \sigma\left( \frac{1}{K} \sum_{k=1}^K \sum_{j \in \mathcal{N}_i \cup \{i\}} \alpha_{ij}^{(k)} W^{(k)} h_j \right)$ In this case, the output dimension remains $F'$.

The following diagram illustrates the parallel computation across multiple heads and their subsequent aggregation.

Flow of computation in a multi-head GAT layer. Input features $h_i, h_j$ are independently processed by $K$ attention heads, each using its own weights ($W^{(k)}$) and attention mechanism ($Attn_k$) to compute attention scores ($\alpha_{ij}^{(k)}$) and aggregated features ($h_i'^{(k)}$). These head-specific outputs are then combined (e.g., via concatenation or averaging) to produce the final output $h_i'$.

Advantages of Multi-Head Attention in GATs

Using multiple heads offers several significant advantages:

1. Stabilized Learning Process: By averaging attention outputs (either explicitly in the final layer or implicitly through downstream processing after concatenation), multi-head attention can make the learning process more stable and less susceptible to noisy or poorly initialized individual attention mechanisms.
2. Increased Model Capacity and Richer Representations: Each head can learn to focus on different types of information or relationships within the neighborhood. For example, one head might attend more to nodes with similar features, while another might focus on structural importance. This allows the model to capture more complex and multi-faceted patterns in the graph structure and node features. The combination (especially concatenation) yields a richer representation.
3. Parallelizability: The computations for each head are independent, making this approach highly amenable to parallelization on modern hardware like GPUs, mitigating the increase in computational cost.

Implementation Notes

When implementing multi-head attention in GATs using libraries like PyTorch Geometric (PyG) or Deep Graph Library (DGL), you typically don't need to build the parallel heads and aggregation logic from scratch. Standard GAT layer implementations (e.g., GATConv in PyG) usually include a heads parameter.

# PyTorch Geometric Example
import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv

# Assume input features have dimension 'in_channels'
# We want output features per head 'out_channels'
# We use K heads
in_channels = 16
out_channels = 8 # Output features per head
K = 4 # Number of heads

# Intermediate layer using concatenation
gat_layer1 = GATConv(in_channels, out_channels, heads=K, dropout=0.6)

# Final layer using averaging
# Input dimension is K * out_channels from the previous layer
# Output dimension is num_classes
num_classes = 7
gat_layer_final = GATConv(K * out_channels, num_classes, heads=1, concat=False, dropout=0.6)

# Example forward pass for a node classification task
# x: node features [num_nodes, in_channels]
# edge_index: graph connectivity [2, num_edges]
def forward(x, edge_index):
    x = F.dropout(x, p=0.6, training=True)
    # Output of gat_layer1 will have dimension [num_nodes, K * out_channels]
    x = gat_layer1(x, edge_index)
    x = F.elu(x)
    x = F.dropout(x, p=0.6, training=True)
    # Output of gat_layer_final will have dimension [num_nodes, num_classes]
    x = gat_layer_final(x, edge_index)
    return F.log_softmax(x, dim=1)

Key considerations during implementation include:

• Parameter Count: Using $K$ heads means you have $K$ sets of weights ($W^{(k)}$) and attention parameters ($\vec{a}^{(k)}$), increasing the total number of parameters compared to a single-head GAT.
• Output Dimension: Be mindful of the output dimension change when using concatenation. The feature dimension becomes $K$ times the out_channels specified per head. Ensure subsequent layers correctly handle this dimensionality. Averaging maintains the out_channels dimension.
• Computational Cost: While parallelizable, the computation and memory requirements scale approximately linearly with the number of heads $K$.

By incorporating multi-head attention, GATs become more powerful and reliable, capable of learning nuanced representations from graph-structured data by attending to different neighborhood aspects simultaneously. This technique is a standard component in many state-of-the-art GAT implementations.