The Attention Mechanism in GATs

Some GNN architectures, such as GCN and GraphSAGE, treat all neighbors of a node with equal or structurally-determined importance. A GCN, for example, averages the features of neighboring nodes using weights derived from node degrees. This is computationally efficient, but it's a significant limitation. In many graphs, some neighbors are far more important than others. Think of a citation network where a foundational paper cited by a new article should carry more weight than a tangential reference.

Graph Attention Networks (GATs) address this by allowing each node to learn the relative importance of its neighbors. Instead of using fixed weights, GATs compute attention coefficients for each edge, effectively learning how much "attention" a node should pay to each of its neighbors during the aggregation step. This mechanism is inspired by the successful attention models used in natural language processing, particularly the Transformer.

Calculating Attention Coefficients

The core of the GAT layer is the self-attention mechanism, which computes a score for every edge in the graph. This score, an attention coefficient $e_{ij}$, quantifies the importance of node $j$'s features to node $i$.

The process begins with the input features for two connected nodes, $h_i$ and $h_j$, where $h \in \mathbb{R}^{N \times F}$ and $N$ is the number of nodes and $F$ is the number of features.

1. Linear Transformation: First, a shared learnable linear transformation, parameterized by a weight matrix $W$, is applied to every node's feature vector. This projects the features into a higher-level representation.

   $$z_i = W h_i$$

2. Scoring Mechanism: The GAT paper proposes a simple single-layer feedforward network to compute the attention coefficient. The transformed feature vectors of the two nodes, $z_i$ and $z_j$, are concatenated. This combined vector is then multiplied by a learnable weight vector $a$, followed by a LeakyReLU non-linearity.

   $$e_{ij} = \text{LeakyReLU}(a^T[z_i || z_j])$$

This calculation produces a single scalar score $e_{ij}$ for the edge from $j$ to $i$. This score is un-normalized and represents the raw importance of their connection. This process is performed for every neighbor $j$ of node $i$.

Normalizing with Masked Softmax

These raw attention scores $e_{ij}$ are not directly usable because they are difficult to compare across different nodes. A node with many neighbors might have scores with a different scale than a node with few neighbors. To make them comparable and turn them into a probability distribution, we normalize them using the softmax function.

A significant detail here is masked attention. The softmax is only applied over the set of neighbors of node $i$, denoted by $\mathcal{N}_i$. This injects the graph structure into the mechanism, ensuring that a node only computes attention scores with its immediate neighbors.

The final attention weight $\alpha_{ij}$ is calculated as:

$$\alpha_{ij} = \text{softmax}_j(e_{ij}) = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}_i} \exp(e_{ik})}$$

The resulting weights $\alpha_{ij}$ are positive and sum to 1 for all neighbors $j$ of node $i$. They represent a distribution of attention that node $i$ places on its neighbors.

The flow of the GAT attention mechanism for a single node i and its neighbors. The process computes, normalizes, and applies attention weights to create an updated node representation.

Applying the Attention Weights

With the normalized attention weights $\alpha_{ij}$ in hand, the final step is to compute the new feature vector for node $i$, denoted as $h'_i$. This is done by taking a weighted sum of the linearly transformed features of all its neighbors.

$$h'_i = \sigma\left(\sum_{j \in \mathcal{N}_i} \alpha_{ij} z_j\right) = \sigma\left(\sum_{j \in \mathcal{N}_i} \alpha_{ij} W h_j\right)$$

Here, $\sigma$ represents a final non-linearity, such as ReLU or ELU, applied to the aggregated features. This entire process constitutes a single GAT layer.

Stabilizing with Multi-Head Attention

The learning process of self-attention can sometimes be unstable. To mitigate this and to allow the model to capture different types of relationships, GATs employ multi-head attention. This involves running several independent attention mechanisms, or "heads," in parallel.

Each head $k$ has its own set of parameters ($W^{(k)}$ and $a^{(k)}$) and computes its own set of attention weights $\alpha_{ij}^{(k)}$. Each head produces a separate embedding. These embeddings are then combined to form the final output. There are two common ways to combine them:

1. Concatenation (for intermediate layers): The outputs from all $K$ heads are concatenated, creating a larger feature vector.

   $$h'_i = \Big\|_{k=1}^K \sigma\left(\sum_{j \in \mathcal{N}_i} \alpha_{ij}^{(k)} W^{(k)} h_j\right)$$

   The resulting output has a dimension of $K \times F'$, where $F'$ is the output dimension of a single head.

2. Averaging (for the final layer): For the final output layer of the network (e.g., for classification), concatenation is not sensible. Instead, the embeddings from the different heads are averaged.

   $$h'_{i, \text{final}} = \sigma\left(\frac{1}{K}\sum_{k=1}^K \sum_{j \in \mathcal{N}_i} \alpha_{ij}^{(k)} W^{(k)} h_j\right)$$

By using multiple heads, the model can jointly attend to information from different representation subspaces at different positions. For example, one head might focus on community structure while another focuses on local connectivity patterns, making the overall model more expressive and stable.