Lab: Selective Attention with GAT

In this lab, we will implement a Graph Attention Network (GAT) to solve the node classification task. Unlike GCN, which averages neighbors, GAT will learn to "pay attention" to specific citation links that are more relevant for determining a paper's category.

---

1. The Power of "Selective Listening"

In the Cora dataset, not all citations are equal. A "Neural Networks" paper might cite a "Mathematics" paper for a specific proof, but its citations to other "Neural Networks" papers are much more important for its own classification. GAT learns this distinction automatically.

---

2. Model Architecture: Multi-Head GAT

We will use 8 attention heads in the first layer to stabilize the learning process, followed by a single-head output layer.

gat_model.py

import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv

class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        # Layer 1: 8 Attention Heads
        self.conv1 = GATConv(in_channels, hidden_channels, heads=8, dropout=0.6)
        
        # Layer 2: Final Output (1 head)
        self.conv2 = GATConv(hidden_channels * 8, out_channels, heads=1, 
                             concat=False, dropout=0.6)

    def forward(self, x, edge_index):
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv1(x, edge_index).elu() # ELU is common for GAT
        
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        
        return F.log_softmax(x, dim=1)

---

3. Training & Performance

GAT typically requires more tuning than GCN but can achieve slightly higher accuracy due to its higher capacity.

Optimizer	Learning Rate	Weight Decay	Dropout
Adam	0.005	5e-4	0.6

Performance on Cora

• GCN Test Acc: ~81.5%
• GAT Test Acc: ~83.0%

---

4. Visualizing Attention Weights

One of the "pro" features of GAT is interpretability. We can extract the attention coefficients $\alpha_{ij}$ to see which edges the model considers most important.

# Extract attention weights from the first layer
out, (edge_index, alpha) = model.conv1(data.x, data.edge_index, return_attention_weights=True)

# alpha contains the weight for every edge across all 8 heads
print(f"Attention weights shape: {alpha.shape}") 

Interpretability

By plotting these weights as a heatmap on the graph, you can visually confirm that the model is assigning higher weights to "intra-class" edges (edges within the same category) compared to "inter-class" edges.

---

Hands-on Resources

• Google Colab: Full GAT Interpretation Lab
• Source Code: Available in the labs/ directory.