Graph Attention Networks

Reading time: 52 min | Interview relevance: Very High - GAT is the canonical extension of GCN; attention mechanism questions are standard in GNN interviews | Target roles: ML Engineer, AI Engineer, Research Engineer, Graph ML Engineer

---

The Real Interview Moment

"In GCN, every neighbor gets a weight based on degree normalization," the interviewer says. "But consider a citation network. A paper by Alan Turing should probably be weighted more than a random student thesis when computing a node's representation. How would you handle that?"

This is exactly the motivation for Graph Attention Networks (GAT, Veličković et al., 2018). GCN assigns fixed weights based on graph topology alone - $1/\sqrt{d_u d_v}$ depends only on degrees, not on what the nodes are about. GAT learns to assign different importance to different neighbors based on their feature content. The result is a GNN that can focus on the most relevant parts of a node's neighborhood, and whose learned attention weights are directly interpretable.

The interviewer is testing: do you understand why GCN's fixed weights are a limitation, can you derive the GAT attention mechanism, and do you know the GATv2 improvement over the original? This lesson covers all three in depth.

---

Why GCN's Fixed Weights Are a Problem

In GCN, node $v$'s update aggregates neighbors with fixed, topology-derived weights:

$\mathbf{h}_v^{(l+1)} = \sigma\!\left(W \sum_{u \in \mathcal{N}(v) \cup \{v\}} \frac{1}{\sqrt{\tilde{d}_v \tilde{d}_u}} \mathbf{h}_u^{(l)}\right)$

The weight $\frac{1}{\sqrt{\tilde{d}_v \tilde{d}_u}}$ has three problems:

Problem 1 - Ignores feature content: a neighbor with features that are completely irrelevant to the current node's representation gets the same weight as a highly relevant neighbor, as long as their degrees are the same.

Problem 2 - Hub node suppression: high-degree hub nodes have their contribution suppressed (divided by $\sqrt{d_u}$) regardless of how relevant they are. A Turing Award winner with 500 papers ($d_u = 500$) contributes less than a graduate student with 2 papers ($d_u = 2$) if both are neighbors.

Problem 3 - Cannot adapt to task: the same weights are used regardless of what the downstream task requires. For a topic classification task, semantic similarity should dominate. For a link prediction task, structural similarity should dominate. GCN cannot differentiate.

GAT solves all three by replacing the fixed structural weight with a learned, content-dependent attention coefficient.

---

The GAT Attention Mechanism

GAT (Veličković et al., ICLR 2018) computes a scalar attention coefficient for each directed edge $(i, j)$.

Step 1: Shared Linear Transformation

Project all node features into a common space:

$\mathbf{z}_i = W \mathbf{h}_i, \quad \mathbf{z}_j = W \mathbf{h}_j$

where $W \in \mathbb{R}^{d' \times d}$ is shared across all nodes. This is the same weight matrix used for all message computations - the transformation is global, not edge-specific.

Step 2: Compute Raw Attention Scores

$e_{ij} = \text{LeakyReLU}\!\left(\mathbf{a}^\top [\mathbf{z}_i \| \mathbf{z}_j]\right)$

where:

• $\mathbf{a} \in \mathbb{R}^{2d'}$ - a single learnable attention vector (shared across all edges)
• $\|$ - concatenation: $[\mathbf{z}_i \| \mathbf{z}_j] \in \mathbb{R}^{2d'}$
• LeakyReLU with negative slope 0.2 - allows gradient flow for negative raw scores

The attention vector $\mathbf{a}$ learns to detect which feature combinations indicate relevant neighbors. The first half of $\mathbf{a}$ reacts to the querying node's features; the second half reacts to the neighbor's features.

Step 3: Softmax Normalization

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

Softmax is taken over all neighbors of node $i$ (plus $i$ itself for the self-loop). This ensures $\sum_{j} \alpha_{ij} = 1$ - the attention weights form a proper probability distribution over the neighborhood.

Why only over the neighborhood, not all $n$ nodes? Computing attention over all pairs would be $O(n^2)$ and would ignore graph structure entirely. Restricting to the neighborhood $\mathcal{N}(i)$ makes attention $O(|E|)$ (linear in edges, not nodes squared) and preserves the message-passing structure.

Step 4: Weighted Aggregation

$\mathbf{h}_i' = \sigma\!\left(\sum_{j \in \mathcal{N}(i) \cup \{i\}} \alpha_{ij} W \mathbf{h}_j\right)$

The final representation is a weighted average of transformed neighbor features, where higher attention weights direct more information from more relevant neighbors.

Full Formula

Combining steps 2-4:

$\alpha_{ij} = \frac{\exp\!\left(\text{LeakyReLU}\!\left(\mathbf{a}^\top [W\mathbf{h}_i \| W\mathbf{h}_j]\right)\right)}{\displaystyle\sum_{k \in \mathcal{N}(i) \cup \{i\}} \exp\!\left(\text{LeakyReLU}\!\left(\mathbf{a}^\top [W\mathbf{h}_i \| W\mathbf{h}_k]\right)\right)}$

---

Multi-Head Attention

Analogous to Transformer multi-head attention, GAT uses $K$ independent attention heads to stabilize learning and capture multiple types of neighbor relationships:

$\mathbf{h}_i' = \|_{k=1}^{K}\, \sigma\!\left(\sum_{j \in \mathcal{N}(i) \cup \{i\}} \alpha_{ij}^{(k)} W^{(k)} \mathbf{h}_j\right)$

where $\|$ denotes concatenation. Each head $k$ has its own weight matrix $W^{(k)}$ and attention vector $\mathbf{a}^{(k)}$.

For the output layer, average instead of concatenate (to avoid growing the feature dimension):

$\mathbf{h}_i' = \sigma\!\left(\frac{1}{K} \sum_{k=1}^{K} \sum_{j \in \mathcal{N}(i) \cup \{i\}} \alpha_{ij}^{(k)} W^{(k)} \mathbf{h}_j\right)$

Standard architecture choices:

• Hidden layers: 8 heads, concatenate → hidden_dim × 8 features
• Output layer: 1 head (or a small number), average
• Dropout on input features and on attention weights during training

Multi-head attention provides two benefits: (1) reduces variance compared to single-head (averaging effect), and (2) allows different heads to capture different relationship types - one head might focus on semantic similarity, another on structural proximity.

---

GAT vs Transformer Self-Attention

Both use multi-head attention, but with important differences:

Aspect	GAT	Transformer Self-Attention
Scope	Only graph neighbors	All pairs (full sequence)
Complexity	$O(	E
Query-key	Concatenation + linear $\mathbf{a}^\top[\mathbf{z}_i \| \mathbf{z}_j]$	Dot product $\mathbf{q}_i \cdot \mathbf{k}_j / \sqrt{d_k}$
Separate Q/K/V	No - shared $W$	Yes - separate $W_Q$, $W_K$, $W_V$
Nonlinearity	LeakyReLU before softmax	None before softmax
Structure	Respects graph topology	Assumes positional encoding

GAT's restriction to neighbors is both a limitation (cannot attend to distant nodes in one step) and a strength (linear complexity, respects graph structure, can scale to large sparse graphs).

---

The Static Attention Problem in Original GAT

Brody, Alon, and Yahav (2021) published "How Attentive are Graph Attention Networks?" (ICLR 2022), identifying a fundamental limitation of the original GAT formulation.

The original GAT computes:

$e_{ij} = \mathbf{a}^\top \text{LeakyReLU}\!\left([W\mathbf{h}_i \| W\mathbf{h}_j]\right)$

Rewrite with $\mathbf{a} = [\mathbf{a}_1 \| \mathbf{a}_2]$ and $\mathbf{a}_1, \mathbf{a}_2 \in \mathbb{R}^{d'}$:

$e_{ij} = \mathbf{a}_1^\top \text{LeakyReLU}(W\mathbf{h}_i) + \mathbf{a}_2^\top \text{LeakyReLU}(W\mathbf{h}_j)$

Wait - this decomposition is not exact because LeakyReLU is applied to the concatenation, not to each part separately. But the key insight: the nonlinearity is applied component-wise, and after separation, the $i$-dependent term $\mathbf{a}_1^\top \text{LeakyReLU}(W\mathbf{h}_i)$ is a constant with respect to $j$.

For a fixed query node $i$, the ranking of neighbors by attention score is:

$e_{ij} \propto \mathbf{a}_2^\top \text{LeakyReLU}(W\mathbf{h}_j) + \text{constant}$

The ranking of neighbors depends only on $\mathbf{h}_j$ - the neighbor's features - not on the interaction between $\mathbf{h}_i$ and $\mathbf{h}_j$. For any two query nodes $i$ and $i'$, the ranking of their neighbors by attention score is the same. This is called static attention: the attention function cannot distinguish which neighbors are important for different query nodes.

Concretely: node $i_A$ (a biology paper) and node $i_B$ (a physics paper) might both have neighbor $j$ (a mathematics paper). In truly content-aware attention, $j$ should be weighted more for $i_B$ than $i_A$. But in original GAT, the ranking of $j$ relative to other neighbors of $i_A$ and $i_B$ is the same - the attention is static.

---

GATv2: Dynamic Attention

GATv2 (Brody et al. 2021) fixes static attention by changing where the linear transformation is applied relative to the nonlinearity:

Original GAT: $e_{ij} = \mathbf{a}^\top \text{LeakyReLU}\!\left(W[\mathbf{h}_i \| \mathbf{h}_j]\right)$

GATv2: $e_{ij} = \mathbf{a}^\top \text{LeakyReLU}\!\left(W_1 \mathbf{h}_i + W_2 \mathbf{h}_j\right)$

or equivalently: $e_{ij} = \mathbf{a}^\top \text{LeakyReLU}\!\left(W \cdot [\mathbf{h}_i \| \mathbf{h}_j]\right) \quad \text{(different matrix structure)}$

The critical difference: in GATv2, the linear transformation (matrix $W_1$, $W_2$ or equivalently $W$ applied after concatenation) is applied inside the nonlinearity to the combined input, allowing genuine feature interaction before the attention score is computed.

In GATv2, the attention score for neighbor $j$ given query node $i$ is:

$e_{ij} = \mathbf{a}^\top \text{LeakyReLU}\!\left(W_1 \mathbf{h}_i + W_2 \mathbf{h}_j\right)$

For a fixed $i$, changing $j$ changes $W_2 \mathbf{h}_j$, which interacts nonlinearly with $W_1 \mathbf{h}_i$ through the LeakyReLU. The ranking of neighbors can change depending on what $\mathbf{h}_i$ is - this is dynamic attention.

GATv2 is strictly more expressive than original GAT: there exist attention functions that GATv2 can represent but GAT cannot.

---

Heterophilic Graphs: Where GAT Shines

Homophily: connected nodes tend to have the same label ($> 50\%$ of edges connect same-class nodes). Most citation networks: homophily ≈ 0.8. GCN's low-pass filter is well-aligned with high homophily.

Heterophily: connected nodes tend to have different labels. Examples:

• Protein interaction networks: binding often occurs between proteins of different types
• Fraud detection graphs: fraudsters connect to legitimate accounts to obscure their activity
• Some social networks: people connect to others with different political views

In heterophilic settings, GCN's fixed aggregation averages dissimilar neighbors, actively hurting performance. GAT can learn to down-weight dissimilar neighbors by assigning them low attention scores, preserving the node's own distinct representation.

Empirically, on the heterophilic Texas dataset (homophily = 0.11):

• GCN: ~59% accuracy
• GAT: ~72% accuracy

The attention mechanism learns to assign near-zero weight to dissimilar neighbors, effectively making GAT behave like a selective aggregator rather than a full average.

---

Visualizing Learned Attention Weights

One of GAT's key advantages over GCN is interpretability: the attention weights $\alpha_{ij}$ tell you which neighbors the model is using for each prediction.

In the Cora citation network, analyzing learned attention weights reveals:

• Papers attending to papers from the same research sub-area get high attention
• Papers attending to papers from different areas (cross-field citations) get low attention
• Some heads learn positional patterns (attending to recently published neighbors)
• Other heads learn semantic patterns (attending to papers with similar word frequencies)

This head specialization is analogous to the attention head specialization observed in BERT.

---

Edge Features in Attention Computation

Standard GAT uses only node features to compute attention. For many applications, edges carry their own features:

• Citation type (cites, based on, extends, refutes)
• Interaction type in molecular graphs (single bond, double bond, aromatic)
• Relationship type in knowledge graphs (is-a, part-of, located-in)

To incorporate edge features $\mathbf{e}_{ij} \in \mathbb{R}^{d_e}$ into the attention computation:

$e_{ij} = \text{LeakyReLU}\!\left(\mathbf{a}^\top [W\mathbf{h}_i \| W\mathbf{h}_j \| W_e \mathbf{e}_{ij}]\right)$

where $W_e \in \mathbb{R}^{d'' \times d_e}$ is a learned edge feature projection. This is a straightforward extension: the attention vector $\mathbf{a}$ is extended to account for the edge feature component.

PyG's GATConv supports edge features via the edge_attr parameter, and GATv2Conv similarly.

---

Full Code: Manual GAT, PyG GATConv, GATv2, Attention Visualization

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv, GATv2Conv
from torch_geometric.datasets import Planetoid
from torch_geometric.utils import softmax as pyg_softmax
import numpy as np
from typing import Optional, Tuple

dataset = Planetoid(root='/tmp/cora', name='Cora')
data = dataset[0]

# ─── MANUAL GAT LAYER ─────────────────────────────────────────────────────────

class GATLayerManual(nn.Module):
    """
    Manual GAT layer implementing:
    α_ij = softmax_j(LeakyReLU(aᵀ[W·hᵢ ‖ W·hⱼ]))
    h'ᵢ = σ(Σⱼ α_ij · W·hⱼ)

    Supports multi-head attention with concatenation.
    """
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        heads: int = 1,
        dropout: float = 0.6,
        alpha: float = 0.2,
    ):
        super().__init__()
        self.heads = heads
        self.out_dim = out_dim

        # Shared linear transformation for all heads
        self.W = nn.Linear(in_dim, out_dim * heads, bias=False)
        nn.init.xavier_uniform_(self.W.weight)

        # Attention vector: [heads, 2*out_dim]
        self.a = nn.Parameter(torch.empty(heads, 2 * out_dim))
        nn.init.xavier_uniform_(self.a.unsqueeze(0))

        self.leaky_relu = nn.LeakyReLU(alpha)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Returns:
            out: (N, heads * out_dim) - concatenated multi-head output
            alpha: (num_edges, heads) - attention weights for visualization
        """
        N = x.size(0)
        src, dst = edge_index

        # Transform: [N, heads * out_dim] -> [N, heads, out_dim]
        x = self.dropout(x)
        Wx = self.W(x).view(N, self.heads, self.out_dim)

        # For each edge (src, dst): concatenate transformed features
        Wx_src = Wx[src]  # (E, heads, out_dim)
        Wx_dst = Wx[dst]  # (E, heads, out_dim)

        # Compute raw attention: aᵀ[Wx_src ‖ Wx_dst]
        # cat: (E, heads, 2*out_dim)
        # a:   (heads, 2*out_dim)
        cat = torch.cat([Wx_src, Wx_dst], dim=-1)   # (E, H, 2d)
        e = (cat * self.a.unsqueeze(0)).sum(dim=-1)  # (E, H)
        e = self.leaky_relu(e)

        # Softmax normalization per destination node, per head
        alpha = pyg_softmax(e, dst, num_nodes=N)  # (E, H)
        alpha = self.dropout(alpha)

        # Aggregate: weighted sum at each destination node
        out = torch.zeros(N, self.heads, self.out_dim, device=x.device)
        # For each edge, add alpha_ij * Wx_src to destination node's accumulator
        for h in range(self.heads):
            weighted = alpha[:, h:h+1] * Wx_src[:, h, :]  # (E, out_dim)
            out[:, h, :].scatter_add_(
                0,
                dst.unsqueeze(1).expand(-1, self.out_dim),
                weighted
            )

        # Concatenate heads: (N, heads * out_dim)
        return out.view(N, self.heads * self.out_dim), alpha

# ─── GAT MODEL WITH PyG (PRODUCTION) ─────────────────────────────────────────

class GAT(torch.nn.Module):
    """
    2-layer GAT following the original paper:
    - Layer 1: 8 heads, concat → hidden * heads features
    - Layer 2: 1 head, no concat (average) → num_classes
    - ELU activation (paper uses ELU, not ReLU)
    - Dropout on inputs and attention weights
    """
    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        out_channels: int,
        heads: int = 8,
        out_heads: int = 1,
        dropout: float = 0.6,
    ):
        super().__init__()
        self.dropout = dropout

        self.conv1 = GATConv(
            in_channels, hidden_channels,
            heads=heads, dropout=dropout, concat=True,
        )
        self.conv2 = GATConv(
            hidden_channels * heads, out_channels,
            heads=out_heads, dropout=dropout, concat=False,
        )

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

    def get_attention_weights(self, x, edge_index):
        """Extract attention weights from both layers for visualization."""
        x_in = F.dropout(x, p=self.dropout, training=False)
        out1, (ei1, alpha1) = self.conv1(
            x_in, edge_index, return_attention_weights=True
        )
        x1 = F.elu(out1)
        x1_drop = F.dropout(x1, p=self.dropout, training=False)
        out2, (ei2, alpha2) = self.conv2(
            x1_drop, edge_index, return_attention_weights=True
        )
        return {
            "layer1": {"edge_index": ei1, "alpha": alpha1},
            "layer2": {"edge_index": ei2, "alpha": alpha2},
        }

model = GAT(
    in_channels=dataset.num_features,
    hidden_channels=8,
    out_channels=dataset.num_classes,
    heads=8, out_heads=1, dropout=0.6,
)

optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

# ─── TRAINING LOOP ────────────────────────────────────────────────────────────

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

@torch.no_grad()
def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)
    accs = {}
    for split in ['train', 'val', 'test']:
        mask = getattr(data, f'{split}_mask')
        accs[split] = (pred[mask] == data.y[mask]).float().mean().item()
    return accs

best_val = 0
best_test = 0
for epoch in range(200):
    loss = train()
    accs = test()
    if accs['val'] > best_val:
        best_val = accs['val']
        best_test = accs['test']

print(f"Best Val: {best_val:.3f} | Best Test: {best_test:.3f}")
# Expected: ~0.830 on Cora (vs GCN's ~0.815)

# ─── GATv2: DYNAMIC ATTENTION ─────────────────────────────────────────────────

class GATv2(torch.nn.Module):
    """
    GATv2 (Brody et al. 2022) - dynamic attention.
    e_ij = aᵀ LeakyReLU(W₁hᵢ + W₂hⱼ)
    vs original GAT: e_ij = aᵀ LeakyReLU([W₁hᵢ ‖ W₂hⱼ])

    GATv2's ranking of neighbors can change based on the querying node.
    Original GAT: same neighbor ranking for all query nodes (static).
    """
    def __init__(self, in_channels, hidden_channels, out_channels,
                 heads=8, dropout=0.6):
        super().__init__()
        self.dropout = dropout
        self.conv1 = GATv2Conv(
            in_channels, hidden_channels, heads=heads,
            dropout=dropout, concat=True, share_weights=False,
        )
        self.conv2 = GATv2Conv(
            hidden_channels * heads, out_channels, heads=1,
            dropout=dropout, concat=False, share_weights=False,
        )

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

v2_model = GATv2(dataset.num_features, 8, dataset.num_classes)
# GATv2 typically achieves ~0.834-0.838 on Cora, slightly better than GAT

# ─── ATTENTION VISUALIZATION ──────────────────────────────────────────────────

def visualize_attention_for_node(
    model: GAT,
    data,
    node_id: int,
    layer: int = 1,
    top_k: int = 10,
) -> None:
    """
    Show which neighbors a specific node attends to most strongly.
    Useful for validating that the model learns meaningful neighbor relationships.
    """
    model.eval()
    with torch.no_grad():
        attn_data = model.get_attention_weights(data.x, data.edge_index)

    layer_key = f"layer{layer}"
    edge_index_l = attn_data[layer_key]["edge_index"]
    alpha_l = attn_data[layer_key]["alpha"]  # (num_edges, num_heads)

    # Average attention across heads
    alpha_mean = alpha_l.mean(dim=1)  # (num_edges,)

    # Find all edges where node_id is the destination
    dst_mask = edge_index_l[1] == node_id
    src_nodes = edge_index_l[0][dst_mask]
    weights = alpha_mean[dst_mask]

    # Sort by attention weight
    sorted_idx = weights.argsort(descending=True)
    src_nodes_sorted = src_nodes[sorted_idx]
    weights_sorted = weights[sorted_idx]

    node_class = data.y[node_id].item()
    print(f"\nAttention weights for node {node_id} (class {node_class})")
    print(f"{'Neighbor':>10} {'Class':>6} {'Attention':>12} {'Same class?':>12}")
    print("-" * 44)

    for i in range(min(top_k, len(src_nodes_sorted))):
        src = src_nodes_sorted[i].item()
        w = weights_sorted[i].item()
        src_class = data.y[src].item()
        same = "YES" if src_class == node_class else "no"
        bar = "█" * int(w * 50)
        print(f"{src:>10} {src_class:>6} {w:>12.4f} {same:>12}  {bar}")

    # Measure: do same-class neighbors get higher attention?
    same_class_mask = data.y[src_nodes] == node_class
    if same_class_mask.any() and (~same_class_mask).any():
        avg_same = weights[same_class_mask].mean().item()
        avg_diff = weights[~same_class_mask].mean().item()
        print(f"\nAvg attention to same-class neighbors: {avg_same:.4f}")
        print(f"Avg attention to diff-class neighbors: {avg_diff:.4f}")
        print(f"Ratio: {avg_same/avg_diff:.2f}x")

visualize_attention_for_node(model, data, node_id=42, layer=1)

# ─── DEMONSTRATING STATIC VS DYNAMIC ATTENTION ────────────────────────────────

def compare_attention_ranking(
    gat_model,
    gatv2_model,
    data,
    query_node_1: int,
    query_node_2: int,
) -> None:
    """
    Show that original GAT has static attention (same neighbor ranking for
    different query nodes) while GATv2 has dynamic attention.

    In static attention: neighbor rankings for query_node_1 and query_node_2
    should be nearly identical (same top-k neighbors by attention weight).
    In dynamic attention: rankings should differ.
    """
    for name, m in [("GAT (static)", gat_model), ("GATv2 (dynamic)", gatv2_model)]:
        m.eval()
        with torch.no_grad():
            if hasattr(m, 'get_attention_weights'):
                attn = m.get_attention_weights(data.x, data.edge_index)
                edge_index_l = attn["layer1"]["edge_index"]
                alpha_l = attn["layer1"]["alpha"].mean(dim=1)
            else:
                # For GATv2, use return_attention_weights directly
                _, (edge_index_l, alpha_l) = m.conv1(
                    data.x, data.edge_index, return_attention_weights=True
                )
                alpha_l = alpha_l.mean(dim=1)

        # Get neighbor rankings for both query nodes
        for qn in [query_node_1, query_node_2]:
            dst_mask = edge_index_l[1] == qn
            srcs = edge_index_l[0][dst_mask]
            weights = alpha_l[dst_mask]
            top3 = srcs[weights.argsort(descending=True)[:3]].tolist()
            print(f"  {name} | Query node {qn}: top-3 neighbors = {top3}")
        print()

# Example: if both query nodes are from different classes but have the same
# neighbor j, GAT will rank j identically for both; GATv2 will not.
# compare_attention_ranking(model, v2_model, data, 42, 100)

# ─── EDGE FEATURES EXTENSION ─────────────────────────────────────────────────

class GATWithEdgeFeatures(nn.Module):
    """
    GAT that incorporates edge features in attention computation.
    e_ij = LeakyReLU(aᵀ[W·hᵢ ‖ W·hⱼ ‖ Wₑ·eᵢⱼ])
    """
    def __init__(self, node_dim, edge_dim, out_dim, heads=4):
        super().__init__()
        # PyG's GATConv supports edge_dim parameter
        self.conv = GATConv(
            in_channels=node_dim,
            out_channels=out_dim,
            heads=heads,
            edge_dim=edge_dim,  # include edge features in attention
            concat=True,
        )

    def forward(self, x, edge_index, edge_attr):
        # edge_attr: (num_edges, edge_dim)
        return self.conv(x, edge_index, edge_attr=edge_attr)

# ─── PPI MULTI-LABEL CLASSIFICATION ─────────────────────────────────────────

def ppi_example():
    """
    Protein-Protein Interaction dataset: 24 graphs, 2373 nodes avg,
    20 train graphs + 2 val + 2 test.
    Multi-label node classification (121 binary labels per node).

    Original GAT achieves 97.3% micro-F1 on PPI (vs GCN's 59%).
    The inductive setting (new graphs at test time) shows GAT's
    generalization strength.

    Note: PPI is inductive - test graphs are unseen during training.
    """
    # from torch_geometric.datasets import PPI
    # from torch_geometric.data import DataLoader
    #
    # train_data = PPI(root='/tmp/ppi', split='train')
    # val_data = PPI(root='/tmp/ppi', split='val')
    # test_data = PPI(root='/tmp/ppi', split='test')
    #
    # # 3-layer GAT with skip connections for PPI
    # class GATForPPI(nn.Module):
    #     def __init__(self):
    #         super().__init__()
    #         # Layer 1: 4 heads, 256 hidden each
    #         self.conv1 = GATConv(50, 256, heads=4, concat=True)
    #         # Layer 2: 4 heads, 256 hidden each
    #         self.conv2 = GATConv(1024, 256, heads=4, concat=True)
    #         # Output: 6 heads, 121 outputs (BCEWithLogitsLoss)
    #         self.conv3 = GATConv(1024, 121, heads=6, concat=False)
    #
    #     def forward(self, x, edge_index):
    #         x = F.elu(self.conv1(x, edge_index))
    #         x = F.elu(self.conv2(x, edge_index))
    #         return self.conv3(x, edge_index)
    #
    # # Loss: BCEWithLogitsLoss (multi-label, not multi-class)
    # criterion = nn.BCEWithLogitsLoss()
    pass

---

Performance Comparison: GAT vs GCN

Dataset	GCN Test Acc	GAT Test Acc	Improvement
Cora (homophily 0.81)	81.5%	83.0%	+1.5%
CiteSeer (homophily 0.74)	70.3%	72.5%	+2.2%
PubMed (homophily 0.80)	79.0%	79.0%	0%
PPI (inductive, multi-label)	59.5% micro-F1	97.3% micro-F1	+37.8%
Texas (heterophily 0.11)	59.5%	72.2%	+12.7%

The pattern: GAT improves over GCN when either the task requires learning which neighbors matter (PPI) or when the graph has heterophily (Texas). For highly homophilic transductive tasks (Cora, PubMed), the improvement is modest.

---

GAT vs GCN: When to Use Each

Property	GCN	GAT
Neighbor weighting	Fixed (degree-based)	Learned (attention)
Parameters	Less ($W$ only)	More ($W + \mathbf{a}$ per head)
Interpretability	Low	High (attention weights)
Handles noisy neighbors	Poorly	Better (down-weights)
Computational cost	$O(	E
Homophilic graphs	Excellent	Good
Heterophilic graphs	Poor	Better
New nodes at inference	No (transductive)	Yes (inductive)

Choose GAT when: neighbors have heterogeneous relevance, interpretability matters (attention weights explain predictions), the graph has heterophily, or you need inductive inference.

Choose GCN when: graph is highly homophilic (GCN's inductive bias matches), scale demands efficiency ($K$× cheaper), or as a strong baseline before adding complexity.

---

YouTube Resources

Resource	Creator	Focus
Graph Attention Networks Explained	Aleksa Gordić	GAT derivation, PyG code, attention visualization
GAT Paper - Veličković et al. 2018	Yannic Kilcher	Original GAT paper review
GATv2: Static vs Dynamic Attention	Brody et al. presentation	ICLR 2022 talk on static attention problem
Heterophilic GNNs	Stanford CS224W	When homophily assumption fails
Attention Mechanism Intuition for GNNs	DeepMind	Graph attention and its variants

---

Common Mistakes

:::danger Common Mistake 1: Not using dropout on both input features AND attention weights The original GAT paper applies dropout to input features and to the learned attention coefficients α_ij. Dropping only one is insufficient. Input dropout prevents feature co-adaptation; attention dropout prevents the model from always focusing on the same few neighbors. The typical rate is 0.6 for both on small datasets like Cora. :::

:::danger Common Mistake 2: Using ReLU instead of ELU in GAT The original GAT paper uses ELU (Exponential Linear Unit) activation, not ReLU. ELU has a smooth negative region that allows gradient flow for slightly negative inputs, which matters when attention-weighted sums can be negative. Using ReLU in a direct reimplementation will slightly underperform the paper's reported results. :::

:::warning Common Mistake 3: Ignoring the static attention limitation of original GAT If you need content-aware neighbor selection where different query nodes should rank the same neighbor differently, use GATv2. Original GAT's static attention may appear to work on standard benchmarks but will fail on tasks requiring truly dynamic, query-dependent attention. Always prefer GATv2 unless you have a specific reason to use the original. :::

:::warning Common Mistake 4: Treating all attention weights as equally interpretable Like transformer attention, GAT attention weights reflect routing in the forward pass, not necessarily causal importance. A neighbor with high attention weight is being used by the model - but this does not mean changing that neighbor's features would change the prediction proportionally. For rigorous attribution, combine attention weights with gradient information (gradient × attention). :::

---

Interview Q&A

Q1: Explain the GAT attention mechanism. How does it differ from Transformer self-attention?

GAT computes $\alpha_{ij} = \text{softmax}_j(\text{LeakyReLU}(\mathbf{a}^\top[W\mathbf{h}_i \| W\mathbf{h}_j]))$ - a scalar attention weight for each graph edge $(i, j)$. Key differences from Transformer self-attention: (1) scope - GAT computes attention only over actual graph neighbors, not all pairs. This is $O(|E|)$, not $O(n^2)$; (2) scoring function - GAT uses a shared linear attention vector $\mathbf{a}$ applied to concatenated features via LeakyReLU, rather than the scaled dot product $QK^\top/\sqrt{d_k}$; (3) no separate value projection - the same $W$ is used for both the attention score and the value computation; (4) the softmax is normalized over the neighborhood, not the full sequence. The graph structure constrains which pairs receive attention, respecting the relational inductive bias.

Q2: What is the "static attention" problem in original GAT and how does GATv2 fix it?

In original GAT, $e_{ij} = \mathbf{a}^\top \text{LeakyReLU}(W[\mathbf{h}_i \| \mathbf{h}_j])$. Expanding: $e_{ij} = \mathbf{a}_1^\top \text{LReLU}(W_1\mathbf{h}_i) + \mathbf{a}_2^\top \text{LReLU}(W_2\mathbf{h}_j)$ - the $i$-dependent term is a constant for fixed $i$. This means the ranking of neighbors by attention score depends only on $\mathbf{h}_j$, not on the interaction between $\mathbf{h}_i$ and $\mathbf{h}_j$. Every query node $i$ ranks its neighbors in the same order - this is static attention. GATv2 fixes this with: $e_{ij} = \mathbf{a}^\top \text{LeakyReLU}(W_1\mathbf{h}_i + W_2\mathbf{h}_j)$. Now the nonlinearity is applied after summing both node representations, creating genuine feature interaction before scoring. The ranking of neighbor $j$ can change depending on what $\mathbf{h}_i$ is - dynamic attention. GATv2 is strictly more expressive and consistently matches or outperforms original GAT.

Q3: Why does GAT apply softmax over neighbors rather than all nodes?

Computing softmax over all $n$ nodes would give an attention distribution over the entire graph, ignoring graph structure and costing $O(n^2)$ per node. GAT restricts attention to $\mathcal{N}(i)$, defined by graph edges. This makes attention $O(|E|)$ (edge-linear), respects the graph topology (only connected nodes communicate), and preserves the message-passing structure that gives GNNs their relational inductive bias. The local softmax also means attention weights sum to 1 within the neighborhood - a natural normalization that prevents numerical instability as neighborhood size varies.

Q4: When would you choose GAT over GCN in a production system?

Choose GAT when: (1) neighbors have heterogeneous relevance - e.g., a fraud detection system where a suspicious account is connected to both fraudulent and legitimate nodes; you want the model to down-weight the legitimate connections for the fraud prediction; (2) interpretability matters - attention weights directly show which neighbors drove a prediction; (3) the graph is heterophilic (connected nodes often have different labels); (4) you are doing inductive learning and need to handle new nodes at inference time (standard GAT is already inductive - it computes embeddings from features, not from a stored embedding table). Choose GCN when: the graph is highly homophilic and GCN's low-pass filtering matches the task, scale demands efficiency ($K$× cheaper per layer), or you need a strong, well-understood baseline.

Q5: How do you extract and interpret attention weights from a trained GAT?

Use return_attention_weights=True in PyG's GATConv.forward(). This returns (output, (edge_index, alpha)) where alpha has shape (num_edges, num_heads). To analyze: for a specific node, filter edges where it is the destination, average alpha across heads (or inspect individual heads), and rank source nodes by their attention weight. Interpretation: high $\alpha_{ij}$ means the model is strongly using node $j$'s features to compute node $i$'s representation. In citation networks, this should correlate with semantic similarity. To validate head specialization, look for heads where high attention correlates with specific edge types (same class, same institution, similar publication date). Important caveat: high attention ≠ causal importance - pair with gradient-based attribution for rigorous feature importance claims.

Q6: Explain why GAT achieves 97.3% micro-F1 on PPI while GCN achieves only 59.5%.

PPI (Protein-Protein Interaction) is inductive, multi-label, and has low homophily. Three reasons GAT dominates: First, inductive setting - PPI evaluation uses graphs unseen during training. GCN in its standard transductive form cannot handle new graphs; GAT computes node embeddings from features, generalizing naturally. Second, heterogeneous neighbors - proteins interact for different functional reasons; some neighbors are highly relevant for specific labels, others are noise. GAT learns to weight relevant neighbors highly, while GCN averages all equally. Third, the 3-layer GAT architecture used for PPI (with skip connections) maintains a larger receptive field without collapsing through over-smoothing, because the attention mechanism automatically weights distant relevant neighbors more than noisy nearby ones. The 38-point gap is primarily attributable to the learned attention mechanism's ability to filter irrelevant interactions in a noisy, heterophilic graph.

---

Historical Context: Attention Enters Graph Learning

Attention mechanisms entered sequence modeling with Bahdanau et al. (2015) and became dominant in NLP through the Transformer (Vaswani et al., 2017). The extension to graphs was natural: instead of attending over all positions in a sequence, attend over all neighbors in a graph.

Veličković et al. (2018) published GAT at ICLR 2018, less than a year after the original Transformer paper. The key contribution was not the attention mechanism itself (borrowed from sequence modeling) but the demonstration that restricting attention to the graph neighborhood - making it $O(|E|)$ rather than $O(n^2)$ - produced a scalable, effective model that outperformed GCN on several benchmarks.

The static attention limitation (Brody et al., 2021) was a natural follow-up to the Expressiveness of GNNs paper (Xu et al., 2019), which showed that GCN and many other GNNs are strictly less expressive than the Weisfeiler-Lehman test. The GATv2 paper asked: is GAT's attention mechanism truly dynamic, or is there a hidden limitation? The answer - static attention - was surprising and immediately impactful. GATv2 is now the default GAT variant in most production deployments.

---

Key Takeaways

GAT is the natural extension of GCN when you need content-aware neighbor weighting. The attention mechanism $e_{ij} = \text{LeakyReLU}(\mathbf{a}^\top[W\mathbf{h}_i \| W\mathbf{h}_j])$ replaces GCN's fixed degree normalization with learned per-edge weights. Multi-head attention stabilizes training and captures multiple relationship types.

Always prefer GATv2 over original GAT - the static attention problem is a real limitation with a straightforward fix. Use PyG's GATv2Conv, which is well-tested and optimized. Extract attention weights for interpretability via return_attention_weights=True, but remember: high attention weight indicates usage, not causal importance.

GAT shines on heterophilic graphs, inductive settings, and tasks where noise-filtering matters (like PPI). GCN remains competitive on highly homophilic transductive benchmarks where the degree-normalization inductive bias aligns well with the task. Use GCN as the baseline, add GAT/GATv2 when you need more expressiveness or interpretability.

---

GAT in Practice: Common Hyperparameter Guide

GAT's performance is more sensitive to hyperparameter choices than GCN. This table covers the decisions that most affect results.

Hyperparameter	Default	Notes
Number of attention heads	8 (hidden), 1 (output)	More heads rarely helps beyond 8; multi-head at output layer adds noise
Hidden dimension per head	8	Total hidden = heads × dim_per_head; match to GCN hidden size
Dropout (input features)	0.6	High dropout on input features is regularization, not noise injection
Dropout (attention weights)	0.6	Applied to $\alpha_{ij}$ before aggregation; prevents head collapse
LeakyReLU negative slope	0.2	Standard; rarely needs tuning
L2 weight decay	0.0005	Same as GCN
Learning rate	0.005	Adam with default betas
Epochs	1000 + early stopping	Patience of 100 on validation accuracy

import torch
import torch.nn.functional as F
from torch_geometric.nn import GATv2Conv
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures

dataset = Planetoid(root='data/Cora', name='Cora', transform=NormalizeFeatures())
data = dataset[0]

class GATv2(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels,
                 heads=8, dropout=0.6):
        super().__init__()
        self.dropout = dropout
        # Hidden: 8 heads, each with hidden_channels//heads dims
        self.conv1 = GATv2Conv(
            in_channels, hidden_channels // heads,
            heads=heads, dropout=dropout, concat=True
        )
        # Output: 1 head, average over multi-head output
        self.conv2 = GATv2Conv(
            hidden_channels, out_channels,
            heads=1, dropout=dropout, concat=False
        )

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

model = GATv2(
    in_channels=dataset.num_features,
    hidden_channels=64,   # 8 heads × 8 dims
    out_channels=dataset.num_classes,
    heads=8,
    dropout=0.6
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

best_val, patience_count = 0, 0
for epoch in range(1000):
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        pred = out.argmax(dim=1)
        val_acc = (pred[data.val_mask] == data.y[data.val_mask]).float().mean()
        if val_acc > best_val:
            best_val = val_acc
            patience_count = 0
            torch.save(model.state_dict(), 'best_gatv2.pt')
        else:
            patience_count += 1
            if patience_count == 100:
                print(f"Early stopping at epoch {epoch}, best val: {best_val:.4f}")
                break

---

Attention Head Specialization Analysis

Multi-head attention is most valuable when different heads learn to encode different types of relationships. In citation networks, heads may specialize by research area, publication year, or citation reciprocity.

def analyze_head_specialization(model, data, num_heads=8):
    """
    For each attention head, measure its average attention weight to
    same-class vs different-class neighbors. High contrast = specialized head.
    """
    model.eval()
    with torch.no_grad():
        _, (edge_index, alpha) = model.conv1(
            data.x, data.edge_index, return_attention_weights=True
        )
    # alpha shape: (num_edges, num_heads)
    src, dst = edge_index

    results = []
    for head in range(num_heads):
        weights = alpha[:, head]
        same_class = (data.y[src] == data.y[dst])

        avg_same = weights[same_class].mean().item()
        avg_diff = weights[~same_class].mean().item()
        ratio = avg_same / (avg_diff + 1e-8)
        results.append({
            'head': head,
            'avg_same_class': round(avg_same, 4),
            'avg_diff_class': round(avg_diff, 4),
            'same_to_diff_ratio': round(ratio, 2)
        })
        print(f"Head {head}: same-class avg={avg_same:.4f}, "
              f"diff-class avg={avg_diff:.4f}, ratio={ratio:.2f}")
    return results

# High ratio (>2.0) indicates the head strongly prefers same-class neighbors.
# Ratio near 1.0 indicates the head is class-agnostic - likely encoding
# structural features (high-degree hubs, bridge nodes) rather than content.

This analysis reveals how much of GAT's performance gain comes from learning class-discriminative attention vs structural attention. In practice, you typically see 2–4 heads with high class-discrimination ratio and 4–6 heads with near-neutral ratios - the ensemble of all heads provides robustness that no single head achieves alone.

---

GAT for Knowledge Graph Completion

Beyond node classification, GAT is particularly effective for knowledge graph completion - predicting missing triples $(h, r, t)$ where $h$ is the head entity, $r$ is the relation, and $t$ is the tail entity.

The key insight: different relation types should use different attention patterns. The edge type $r$ can be incorporated into the attention computation as an additional feature, allowing GAT to learn relation-specific attention.

from torch_geometric.nn import RGATConv  # Relation-aware GAT

# Knowledge graph: each edge has a relation type
# edge_index: (2, num_triples), edge_type: (num_triples,)
# RGATConv computes separate attention per relation type

class KnowledgeGraphGAT(torch.nn.Module):
    def __init__(self, num_entities, num_relations, hidden, num_heads=4):
        super().__init__()
        self.entity_emb = torch.nn.Embedding(num_entities, hidden)
        self.conv1 = RGATConv(hidden, hidden // num_heads,
                               num_relations=num_relations,
                               heads=num_heads, concat=True)
        self.conv2 = RGATConv(hidden, hidden // num_heads,
                               num_relations=num_relations,
                               heads=num_heads, concat=False)

    def forward(self, edge_index, edge_type):
        x = self.entity_emb.weight  # (num_entities, hidden)
        h = F.elu(self.conv1(x, edge_index, edge_type))
        h = self.conv2(h, edge_index, edge_type)
        return h  # entity embeddings

    def score_triple(self, h, r_emb, t):
        """TransE-style scoring: DistMult or RotatE can also be used."""
        return (h * r_emb * t).sum(dim=-1)

This architecture (Relational GAT) is the core of modern knowledge graph systems at scale, used in biomedical knowledge graphs (drug-gene-disease), financial entity graphs, and enterprise knowledge management.

:::tip 🎮 Interactive Playground

Visualize this concept: Try the Graph Attention Network demo on the EngineersOfAI Playground - no code required.

:::