Attention assigns learned importance weights to inputs, enabling a graph neural network to prioritize relevant neighbors during message aggregation instead of treating every connection equally. section 11.6 section 11.6.2
A neighborhood-attention pipeline typically:
- Transforms node and neighbor features into representations used to compare relationships. section 11.6.3
- Computes attention scores, normalizes them with softmax, and forms weighted neighbor values. section 11.6.3
- Combines parallel heads so different components can learn different neighborhood patterns. section 11.6.3
Graph relationships
In a molecular graph, one attention head may emphasize chemical bond types while another captures spatial arrangements. Structural encodings can add degree-based topology information. section 11.6.3 section 11.6.3
Listing 11.3 Attention mechanism in transformer architecture
class TransformerStyleGraphAttention(nn.Module):
def __init__(self, feature_dim):
super().__init__()
self.query_transform = nn.Linear(feature_dim, feature_dim)
self.key_transform = nn.Linear(feature_dim, feature_dim)
self.value_transform = nn.Linear(feature_dim, feature_dim) #1
def forward(self, node_features, neighbor_features):
Q = self.query_transform(node_features)
K = self.key_transform(neighbor_features)
V = self.value_transform(neighbor_features) #2
attention_scores = torch.matmul(Q, K.transpose(-2, -1)) #3
attention_scores = attention_scores / math.sqrt(K.size(-1)) #3
attention_weights = F.softmax(attention_scores, dim=-1) #4
attended_values = torch.matmul(attention_weights, V) #4
return attended_values #5
Inspect a PyTorch-style implementation that computes query, key, and value projections before scaled attention and weighted aggregation. section 11.6.3
