08. Graph Attention Networks

• Primary References:

1. 1
   Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., & Bengio, Y. (2018). Graph Attention Networks. In International Conference on Learning Representations (ICLR 2018).
2. 2
   Brody, S., Alon, U., & Yahav, E. (2022). How Attentive are Graph Attention Networks? (GATv2). In ICLR 2022.
3. 3
   Kipf, T. N., & Welling, M. (2017). Semi-Supervised Classification with Graph Convolutional Networks. In ICLR 2017.
4. 4
   Rusch, T. K., Bronstein, M. M., & Mishra, S. (2023). A Survey on Oversmoothing in Graph Neural Networks. arXiv:2303.10993.

• Why I Wrote This Note: In our published paper on Deep Multi-Modal Approach for Protein Function Prediction and Classification (IEEE ICEEICT 2024) and our work on Smart Grid Stability Prediction (IEEE ICCIT 2024), relational structure is central. Amino acid residues separated by \(80\) positions in primary sequence often fold within \(< 6\text{ \AA}\) in 3D space; attention-based message passing allows the model to dynamically weight catalytic contacts over background spatial neighbors.

• Isotropic (GCN) vs. Anisotropic (GAT) Aggregation: GCN assigns fixed structural weights \(1/\sqrt{d_i d_j}\) regardless of node features; GAT computes feature-dependent coefficients \(\alpha_{ij} = \text{softmax}_{j \in \mathcal{N}_i}(e_{ij})\), enabling inductive generalization to completely unseen graph topologies at test time.
• Why Standard GAT (2018) Is Only "Static Attention": In Veličković et al. (2018), the scoring function \(e_{ij} = \mathbf{a}^\top \text{LeakyReLU}(W h_i \,\|\, W h_j)\) applies the linear layer \(W\) and vector \(\mathbf{a} = [\mathbf{a}_1 \,\|\, \mathbf{a}_2]\) sequentially without an intermediate non-linearity between query \(i\) and key \(j\) after \(\mathbf{a}\). Consequently, \(\arg\max_j \alpha_{ij}\) is identical for every query node \(i\)!
• The GATv2 Fix (Dynamic Attention): Simply swapping the order of operations to \(e_{ij}^{\text{GATv2}} = \mathbf{a}^\top \text{LeakyReLU}\big(W_l h_i + W_r h_j\big)\) prevents \(\mathbf{a}\) and \(W\) from collapsing into a single linear projection, restoring universal approximation of dynamic query-dependent attention at zero extra parameter cost.
• Computational Complexity: Unlike dense Transformer attention (\(\mathcal{O}(|V|^2)\)), GAT masks attention strictly to edges \(E\), scaling linearly in graph sparsity: \(\mathcal{O}(|V| F F' + |E| F')\).

4.1 From Isotropic GCN to Anisotropic Graph Attention

Let \(\mathcal{G} = (\mathcal{V}, \mathcal{E})\) be a graph with \(N = |\mathcal{V}|\) nodes and node feature matrix \(\mathbf{H} = \{\mathbf{h}_1, \dots, \mathbf{h}_N\}\) where \(\mathbf{h}_i \in \mathbb{R}^F\). Including self-loops (\(i \in \mathcal{N}_i\)), a general Message Passing Neural Network (MPNN) updates node \(i\) via:

• In GCN (Kipf & Welling, 2017): \(\alpha_{ij} = \frac{1}{\sqrt{|\mathcal{N}_i|\,|\mathcal{N}_j|}}\) (purely topological, isotropic).
• In GAT (Veličković et al., 2018): \(\alpha_{ij}\) is learned end-to-end via masked self-attention over \(j \in \mathcal{N}_i\):

where \(\mathbf{W} \in \mathbb{R}^{F' \times F}\) is a shared feature transformation, \(\mathbf{a} \in \mathbb{R}^{2F'}\) is a learnable attention vector, and \(\|\) denotes concatenation.

4.2 Mathematical Proof of GAT's "Static Attention" Flaw & The GATv2 Remedy

Decomposing \(\mathbf{a} = [\mathbf{a}_l \,\|\, \mathbf{a}_r]\) into two halves \(\mathbf{a}_l, \mathbf{a}_r \in \mathbb{R}^{F'}\), the original GAT pre-nonlinearity term is:

Notice that because \(\text{LeakyReLU}(u + v)\) is monotonically increasing in \(v\) for any fixed scalar \(u = s_l(\mathbf{h}_i)\), the neighbor \(j^\star\) that maximizes \(e_{ij}^{\text{GAT}}\) is:

which does not depend on the query node \(\mathbf{h}_i\) at all! Every node \(i\) in a complete neighborhood is forced to rank its neighbors in the exact same global order determined by \(\mathbf{a}_r^\top \mathbf{W}\mathbf{h}_j\).

GATv2 (Brody et al., ICLR 2022) resolves this by applying the \(\text{LeakyReLU}\) non-linearity before projecting with \(\mathbf{a} \in \mathbb{R}^{F'}\):

Because \(\mathbf{a}^\top \text{LeakyReLU}(\cdot)\) is a universal approximator (a 2-layer MLP over \([\mathbf{h}_i \,\|\, \mathbf{h}_j]\)), each query node \(i\) can now induce a completely distinct ranking over keys \(j \in \mathcal{N}_i\).

flowchart LR
    subgraph GAT_2018["Standard GAT (Static Ranking)"]
        H1["W h_i || W h_j"] --> A1["Linear Dot: a_l^T W h_i + a_r^T W h_j"] --> LR1["LeakyReLU(scalar)"]
    end
    subgraph GATv2_2022["GATv2 (Dynamic Query-Key Interaction)"]
        H2["W_l h_i + W_r h_j ∈ ℝ^{F'}"] --> LR2["Elementwise LeakyReLU(vector)"] --> A2["Vector Projection: a^T (·)"]
    end

4.3 Multi-Head Aggregation & Over-Smoothing Dynamics

To stabilize attention learning, \(K\) independent attention heads are concatenated in intermediate layers and averaged in the final prediction layer:

4.4 PyTorch Implementation of Multi-Head GATv2 Layer

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

class MultiHeadGATv2Layer(nn.Module):
    # Dynamic Graph Attention (GATv2, Brody et al., ICLR 2022) with masked adjacency softmax
    def __init__(
        self,
        in_features: int,
        out_features: int,
        num_heads: int = 4,
        concat: bool = True,
        negative_slope: float = 0.2,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.out_features = out_features
        self.concat = concat

        self.W_left = nn.Linear(in_features, num_heads * out_features, bias=False)
        self.W_right = nn.Linear(in_features, num_heads * out_features, bias=False)
        self.attn_vec = nn.Parameter(torch.empty(1, num_heads, out_features))
        self.leaky_relu = nn.LeakyReLU(negative_slope)
        self.dropout = nn.Dropout(dropout)

        nn.init.xavier_uniform_(self.W_left.weight)
        nn.init.xavier_uniform_(self.W_right.weight)
        nn.init.xavier_uniform_(self.attn_vec)

    def forward(self, h: torch.Tensor, adj_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # h: (N, F_in), adj_mask: (N, N) binary/boolean adjacency matrix including self-loops
        N = h.size(0)
        proj_l = self.W_left(h).view(N, self.num_heads, self.out_features)   # (N, K, F_out)
        proj_r = self.W_right(h).view(N, self.num_heads, self.out_features)  # (N, K, F_out)

        # Pairwise broadcast sum: (N, 1, K, F_out) + (1, N, K, F_out) -> (N, N, K, F_out)
        pairwise_sum = self.leaky_relu(proj_l.unsqueeze(1) + proj_r.unsqueeze(0))
        e_ij = (pairwise_sum * self.attn_vec.unsqueeze(0)).sum(dim=-1)  # (N, N, K)

        # Mask non-existent edges with -inf prior to neighborhood softmax
        mask = (adj_mask == 0).unsqueeze(-1)  # (N, N, 1)
        e_ij = e_ij.masked_fill(mask, float("-inf"))

        alpha = F.softmax(e_ij, dim=1)  # (N, N, K) normalized over neighbors j
        alpha_drop = self.dropout(alpha)

        # Weighted message aggregation over j: einsum('ijk,jkf->ikf')
        out = torch.einsum("ijk,jkf->ikf", alpha_drop, proj_r)  # (N, K, F_out)
        if self.concat:
            return F.elu(out.reshape(N, self.num_heads * self.out_features)), alpha
        return out.mean(dim=1), alpha

1. 1
   Deep Multi-Modal Protein Function Prediction (IEEE ICEEICT 2024): Proteins are naturally represented both as \(1\text{D}\) amino acid sequences and \(3\text{D}\) residue contact graphs (\(C_\alpha\text{–}C_\alpha\) distance \(< 8\text{ \AA}\)). Using GATv2 over residue contact graphs allows the model to highlight conserved catalytic triads and binding pockets via high \(\alpha_{ij}\) weights, which we fuse with sequence embeddings for Gene Ontology (GO) classification.
2. 2
   Smart Grid Stability Prediction (IEEE ICCIT 2024): Power grids are physical graphs where generator and load buses interact over transmission lines with varying admittance. Incorporating line impedance \(\mathbf{e}_{ij}\) into GATv2 edge attention models cascading instability far more faithfully than flat tabular MLPs.

• Always Use GATv2 Instead of GAT (2018): There is zero architectural reason to use the original 2018 GAT formulation today. By sharing \(\mathbf{W}_l = \mathbf{W}_r\) in GATv2, parameter count is \(50\%\) smaller than standard GAT's \(\mathbf{a} \in \mathbb{R}^{2F'}\) while strictly dominating it in expressivity and robustness to edge noise.
• Attention Entropy as a Diagnostic: Monitoring the average neighborhood attention entropy \(\mathcal{H}(\boldsymbol{\alpha}_i) = -\sum_{j \in \mathcal{N}_i} \alpha_{ij}\log\alpha_{ij}\) normalized by \(\log |\mathcal{N}_i|\) tells you immediately whether GAT is actually learning anisotropic structure: if \(\tilde{\mathcal{H}}(\boldsymbol{\alpha}_i) \approx 0.99\), your GAT has collapsed to a uniform GCN with \(4\times\) the compute overhead!

• Implement MultiHeadGATv2Layer with dynamic query-key interaction and return of attention weights \(\alpha_{ij}\).Verified
• Extend MultiHeadGATv2Layer to accept continuous edge attributes \(\mathbf{e}_{ij} \in \mathbb{R}^{D_e}\) (3D Euclidean distance + bond type for protein graphs).Pending
• Measure Dirichlet Energy \(E_{\text{Dir}}(\mathbf{H}^{(\ell)}) = \frac{1}{N}\sum_{(i,j)\in\mathcal{E}} \|\mathbf{h}_i^{(\ell)} - \mathbf{h}_j^{(\ell)}\|_2^2\) across layers \(\ell = 1 \dots 8\) with and without residual connections.Pending

"If the attention weights over a node's neighbors equal \(1/\deg(i)\), you are paying the price of a Transformer to run a simple graph average."

Edge-Aware GATv2 Scoring with Residual Connection:

• Connected Research Notes:
• Mixture-of-Experts Architecture — Viewing sparse top-\(k\) expert routing as dynamic bipartite graph attention between tokens and experts.
• Understanding Predictive Uncertainty — Using normalized neighborhood attention entropy \(\tilde{\mathcal{H}}(\boldsymbol{\alpha}_i)\) to diagnose uniform attention collapse.
• LoRA: Mathematical Foundations — Adapting large pretrained molecular/protein GNNs via low-rank updates on \(\mathbf{W}_l, \mathbf{W}_r\).