Graph Attention Layer (GAT) as described in the paper `"Graph Attention Networks" `. This operation can be mathematically described as: e_ij = a(W h_i, W h_j) α_ij = softmax_j(e_ij) = exp(e_ij) / Σ_k(exp(e_ik))
| 16 | ################################ |
| 17 | |
| 18 | class GraphAttentionLayer(nn.Module): |
| 19 | """ |
| 20 | Graph Attention Layer (GAT) as described in the paper `"Graph Attention Networks" <https://arxiv.org/pdf/1710.10903.pdf>`. |
| 21 | |
| 22 | This operation can be mathematically described as: |
| 23 | |
| 24 | e_ij = a(W h_i, W h_j) |
| 25 | α_ij = softmax_j(e_ij) = exp(e_ij) / Σ_k(exp(e_ik)) |
| 26 | h_i' = σ(Σ_j(α_ij W h_j)) |
| 27 | |
| 28 | where h_i and h_j are the feature vectors of nodes i and j respectively, W is a learnable weight matrix, |
| 29 | a is an attention mechanism that computes the attention coefficients e_ij, and σ is an activation function. |
| 30 | |
| 31 | """ |
| 32 | def __init__(self, in_features: int, out_features: int, n_heads: int, concat: bool = False, dropout: float = 0.4, leaky_relu_slope: float = 0.2): |
| 33 | super(GraphAttentionLayer, self).__init__() |
| 34 | |
| 35 | self.n_heads = n_heads # Number of attention heads |
| 36 | self.concat = concat # wether to concatenate the final attention heads |
| 37 | self.dropout = dropout # Dropout rate |
| 38 | |
| 39 | if concat: # concatenating the attention heads |
| 40 | self.out_features = out_features # Number of output features per node |
| 41 | assert out_features % n_heads == 0 # Ensure that out_features is a multiple of n_heads |
| 42 | self.n_hidden = out_features // n_heads |
| 43 | else: # averaging output over the attention heads (Used in the main paper) |
| 44 | self.n_hidden = out_features |
| 45 | |
| 46 | # A shared linear transformation, parametrized by a weight matrix W is applied to every node |
| 47 | # Initialize the weight matrix W |
| 48 | self.W = nn.Parameter(torch.empty(size=(in_features, self.n_hidden * n_heads))) |
| 49 | |
| 50 | # Initialize the attention weights a |
| 51 | self.a = nn.Parameter(torch.empty(size=(n_heads, 2 * self.n_hidden, 1))) |
| 52 | |
| 53 | self.leakyrelu = nn.LeakyReLU(leaky_relu_slope) # LeakyReLU activation function |
| 54 | self.softmax = nn.Softmax(dim=1) # softmax activation function to the attention coefficients |
| 55 | |
| 56 | self.reset_parameters() # Reset the parameters |
| 57 | |
| 58 | |
| 59 | def reset_parameters(self): |
| 60 | """ |
| 61 | Reinitialize learnable parameters. |
| 62 | """ |
| 63 | nn.init.xavier_normal_(self.W) |
| 64 | nn.init.xavier_normal_(self.a) |
| 65 | |
| 66 | |
| 67 | def _get_attention_scores(self, h_transformed: torch.Tensor): |
| 68 | """calculates the attention scores e_ij for all pairs of nodes (i, j) in the graph |
| 69 | in vectorized parallel form. for each pair of source and target nodes (i, j), |
| 70 | the attention score e_ij is computed as follows: |
| 71 | |
| 72 | e_ij = LeakyReLU(a^T [Wh_i || Wh_j]) |
| 73 | |
| 74 | where || denotes the concatenation operation, and a and W are the learnable parameters. |
| 75 |