Graph Attention Networks (GAT)

A Bit of History

Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio introduced Graph Attention Networks in their 2018 ICLR paper "Graph Attention Networks." The work arrived one year after Kipf & Welling's GCN (2017) and drew direct inspiration from Vaswani et al.'s revolutionary 2017 paper "Attention Is All You Need" — the same self-attention mechanism that launched the Transformer architecture. GAT's key contribution was showing that self-attention, originally designed for sequences, transfers naturally to arbitrary graph structures, and crucially requires no knowledge of the full graph structure upfront — only a node's immediate neighborhood, making it inherently well-suited to inductive settings.

Working Principle: Learned Attention Coefficients

Instead of GCN's fixed, structurally-determined normalization $\frac{1}{\sqrt{\tilde{d}_i \tilde{d}_j}}$, GAT computes a learned attention coefficient $\alpha_{ij}$ for every edge, reflecting how much node $j$'s features should influence node $i$'s updated representation:

$$e_{ij} = \text{LeakyReLU}\left( \mathbf{a}^T [W h_i \, \| \, W h_j] \right)$$

• $W \in \mathbb{R}^{d' \times d}$: a shared learnable linear transformation applied to every node's features.
• $\mathbf{a} \in \mathbb{R}^{2d'}$: a learnable attention weight vector.
• $\|$: concatenation of the transformed features of nodes $i$ and $j$.
• $e_{ij}$: an unnormalized "importance score" of neighbor $j$ to node $i$.

These raw scores are then normalized via softmax across all of $i$'s neighbors $\mathcal{N}(i)$, so they sum to 1 and can be interpreted as attention weights:

$$\alpha_{ij} = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}(i)} \exp(e_{ik})}$$

The updated node representation is then a weighted sum of neighbor features, using these learned attention weights instead of fixed structural ones:

$$h_i' = \sigma \left( \sum_{j \in \mathcal{N}(i)} \alpha_{ij} \, W h_j \right)$$

Multi-Head Attention

Following the Transformer playbook directly, GAT stabilizes learning and increases representational capacity by computing $K$ independent attention mechanisms ("heads") in parallel, then combining their outputs:

$$h_i' = \Big\|_{k=1}^{K} \sigma \left( \sum_{j \in \mathcal{N}(i)} \alpha_{ij}^{(k)} \, W^{(k)} h_j \right) \quad \text{(concatenation, for intermediate layers)}$$

At the final layer, the $K$ heads are typically averaged rather than concatenated, to produce a single output representation of the desired dimensionality.

Worked Example

Complexity Analysis

$$\text{Time per layer: } O(|E| \cdot d \cdot d') \qquad \text{Space: } O(|V| \cdot d' + |E|)$$

Crucially, computing attention coefficients is parallelizable across all edges simultaneously (like Transformer self-attention across all token pairs), making GAT highly efficient on modern GPU/TPU hardware despite the added attention computation compared to GCN.

Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class GraphAttentionLayer(nn.Module):
    def __init__(self, in_features, out_features, alpha=0.2):
        super(GraphAttentionLayer, self).__init__()
        self.W = nn.Linear(in_features, out_features, bias=False)
        self.a = nn.Linear(2 * out_features, 1, bias=False)
        self.leaky_relu = nn.LeakyReLU(alpha)

    def forward(self, h, adj):
        """
        h: node features (N, in_features)
        adj: adjacency matrix (N, N), 1 if edge exists (including self-loops)
        """
        N = h.size(0)
        Wh = self.W(h)  # (N, out_features)

        # Build all pairs [Wh_i || Wh_j] for attention score computation
        Wh_i = Wh.unsqueeze(1).repeat(1, N, 1)  # (N, N, out_features)
        Wh_j = Wh.unsqueeze(0).repeat(N, 1, 1)  # (N, N, out_features)
        combined = torch.cat([Wh_i, Wh_j], dim=-1)  # (N, N, 2*out_features)

        e = self.leaky_relu(self.a(combined).squeeze(-1))  # (N, N)

        # Mask out non-edges before softmax (set to -inf so softmax -> 0)
        mask = adj > 0
        e_masked = e.masked_fill(~mask, float('-inf'))
        attention = F.softmax(e_masked, dim=1)  # normalize per-row (per node i)

        h_prime = torch.matmul(attention, Wh)  # weighted sum of neighbor features
        return F.elu(h_prime), attention


class MultiHeadGAT(nn.Module):
    def __init__(self, in_features, out_features, num_heads=4):
        super(MultiHeadGAT, self).__init__()
        self.heads = nn.ModuleList([
            GraphAttentionLayer(in_features, out_features) for _ in range(num_heads)
        ])

    def forward(self, h, adj):
        outputs = [head(h, adj)[0] for head in self.heads]
        return torch.cat(outputs, dim=-1)  # concatenate heads


# Test on a small 4-node graph
h = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5]])
adj = torch.tensor([
    [1, 1, 0, 0],
    [1, 1, 1, 0],
    [0, 1, 1, 1],
    [0, 0, 1, 1]
], dtype=torch.float32)

gat_layer = GraphAttentionLayer(in_features=2, out_features=4)
output, attn_weights = gat_layer(h, adj)
print("GAT Layer Output:")
print(output)
print("\nAttention Weights:")
print(attn_weights)

multi_head = MultiHeadGAT(in_features=2, out_features=4, num_heads=2)
multi_output = multi_head(h, adj)
print("\nMulti-Head GAT Output shape:", multi_output.shape)

Real-World Applications

Exercises

1. Derive why softmax normalization over neighbors (rather than over the entire graph) preserves locality in the attention mechanism.
2. Compare GAT and GCN's parameter counts for a single layer with the same input/output dimensions.
3. Explain why GAT does not require the symmetric normalized Laplacian that GCN depends on.
4. Challenge: Implement sparse attention computation (only computing $e_{ij}$ for actual edges, not the full $N \times N$ matrix) to make GAT scale to graphs with millions of nodes.

Limitations