import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor from typing import Union from torch_geometric.nn.conv import MessagePassing from torch_geometric.nn.inits import reset from torch_geometric.typing import OptPairTensor, Size from torch_geometric.utils import scatter from .utils import create_activation class ACM_GIN(MessagePassing): """Single ACM-GIN convolution layer with edge-aware message passing. The message from node j to node i incorporates the intermediate edge feature e_ij_prime (precomputed by the outer model) alongside the scalar spatial weight a_ij used for degree normalization: m_ij = a_ij * ReLU(H_j + e_ij_prime) """ def __init__( self, nn_lowpass: torch.nn.Module, nn_highpass: torch.nn.Module, nn_fullpass: torch.nn.Module, nn_lowpass_proj: torch.nn.Module, nn_highpass_proj: torch.nn.Module, nn_fullpass_proj: torch.nn.Module, nn_mix: torch.nn.Module, T: float = 3.0, **kwargs, ): kwargs.setdefault("aggr", "add") super().__init__(**kwargs) self.nn_lowpass = nn_lowpass self.nn_highpass = nn_highpass self.nn_fullpass = nn_fullpass self.nn_lowpass_proj = nn_lowpass_proj self.nn_highpass_proj = nn_highpass_proj self.nn_fullpass_proj = nn_fullpass_proj self.nn_mix = nn_mix self.sigmoid = torch.nn.Sigmoid() self.softmax = torch.nn.Softmax(dim=1) self.T = T self.reset_parameters() def reset_parameters(self): reset(self.nn_lowpass) reset(self.nn_highpass) reset(self.nn_fullpass) reset(self.nn_lowpass_proj) reset(self.nn_highpass_proj) reset(self.nn_fullpass_proj) reset(self.nn_mix) def forward( self, x: Union[Tensor, OptPairTensor], edge_index: Tensor, edge_weight: Tensor, edge_feat: Tensor, size: Size = None, ) -> Tensor: """Forward pass of a single ACM-GIN layer. Args: x: Node features [N, hidden_dim] or (x_src, x_dst) pair. edge_index: Edge indices [2, E]. edge_weight: Scalar spatial distance per edge [E] (first column of the original edge_attr, used for degree normalization). edge_feat: Intermediate edge features [E, hidden_dim], i.e. e_ij_prime precomputed by the edge MLP in the outer model. size: Optional bipartite graph size. """ if isinstance(x, Tensor): x: OptPairTensor = (x, x) # propagate_type: (x: OptPairTensor, edge_weight: Tensor, edge_feat: Tensor) out = self.propagate( edge_index, x=x, edge_weight=edge_weight, edge_feat=edge_feat, size=size ) # Degree here is the sum of edge weights, not the neighbour count as in # standard GIN. Nodes whose weights sum to zero are handled below. deg = scatter(edge_weight, edge_index[1], 0, out.size(0), reduce="sum") deg_inv = 1.0 / deg deg_inv.masked_fill_(deg_inv == float("inf"), 0) out = deg_inv.view(-1, 1) * out x_r = x[1] assert x_r is not None, ( "Target node features (x_r) must not be None for ACM_GIN" ) out_lowpass = (x_r + out) / 2.0 out_highpass = (x_r - out) / 2.0 # compute embeddings for each filter out_lowpass = self.nn_lowpass(out_lowpass) out_highpass = self.nn_highpass(out_highpass) out_fullpass = self.nn_fullpass(x_r) # compute importance weights per filter alpha_lowpass = self.sigmoid(self.nn_lowpass_proj(out_lowpass)) alpha_highpass = self.sigmoid(self.nn_highpass_proj(out_highpass)) alpha_fullpass = self.sigmoid(self.nn_fullpass_proj(out_fullpass)) alpha_cat = torch.concat([alpha_lowpass, alpha_highpass, alpha_fullpass], dim=1) alpha_cat = self.softmax(self.nn_mix(alpha_cat / self.T)) out = alpha_cat[:, 0].view(-1, 1) * out_lowpass out = out + alpha_cat[:, 1].view(-1, 1) * out_highpass out = out + alpha_cat[:, 2].view(-1, 1) * out_fullpass return out def message(self, x_j: Tensor, edge_weight: Tensor, edge_feat: Tensor) -> Tensor: """Edge-aware message: m_ij = a_ij * ReLU(H_j + e_ij_prime).""" return edge_weight.view(-1, 1) * F.relu(x_j + edge_feat) def __repr__(self) -> str: return ( f"{self.__class__.__name__}(" f"nn_lowpass={self.nn_lowpass}, " f"nn_highpass={self.nn_highpass}, " f"nn_fullpass={self.nn_fullpass})" ) class ACM_GIN_model(nn.Module): """Multi-layer ACM-GIN model with edge-aware message passing. Both node and edge features are projected into hidden_dim at the start (via ``node_input_proj`` and ``edge_input_proj``). This ensures uniform dimensions throughout, so every edge MLP receives 3 * hidden_dim and the edge residual connection is valid from layer 0 onward. At each layer k the model: 1. Computes intermediate edge features via an edge MLP: e_ij_prime = MLP_edge(H_i || H_j || E_ij) 2. Updates edge state with a residual connection: E_ij^(k) = E_ij^(k-1) + e_ij_prime 3. Passes messages using the scalar spatial weight and the intermediate edge features: m_ij = a_ij * ReLU(H_j + e_ij_prime) 4. Applies ACM channel mixing (low/high/full-pass) on the aggregated messages. """ def __init__( self, in_dim, out_dim, num_layers, hidden_dim, edge_in_dim, batchnorm, activation="relu", ): super(ACM_GIN_model, self).__init__() self.num_layers = num_layers self.hidden_dim = hidden_dim self.gnn_batchnorm = batchnorm self.out_dim = out_dim # Project raw node and edge features into hidden_dim so that both # live in the same space from the very start. This enables the # edge residual connection at every layer (including layer 0) and # keeps all edge MLP input dimensions uniform at 3 * hidden_dim. self.node_input_proj = nn.Linear(in_dim, hidden_dim) self.edge_input_proj = nn.Linear(edge_in_dim, hidden_dim) self.ACM_convs = nn.ModuleList() self.nns_lowpass = nn.ModuleList() self.nns_highpass = nn.ModuleList() self.nns_fullpass = nn.ModuleList() self.nns_lowpass_proj = nn.ModuleList() self.nns_highpass_proj = nn.ModuleList() self.nns_fullpass_proj = nn.ModuleList() self.nns_mix = nn.ModuleList() self.edge_mlps = nn.ModuleList() self.activation_name = activation for i in range(self.num_layers): # --- Edge MLP for this layer --- # Both nodes and edges have been projected to hidden_dim before # the loop, so the input is always (src + dst + edge_state) = # 3 * hidden_dim for every layer. edge_mlp_in = 3 * hidden_dim if self.gnn_batchnorm: self.edge_mlps.append( nn.Sequential( nn.Linear(edge_mlp_in, hidden_dim), nn.BatchNorm1d(hidden_dim), create_activation(activation), nn.Linear(hidden_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), create_activation(activation), ) ) else: self.edge_mlps.append( nn.Sequential( nn.Linear(edge_mlp_in, hidden_dim), create_activation(activation), nn.Linear(hidden_dim, hidden_dim), create_activation(activation), ) ) # --- Projection modules to compute importance weights --- for channel_proj_module in [ self.nns_lowpass_proj, self.nns_highpass_proj, self.nns_fullpass_proj, ]: if i == self.num_layers - 1: channel_proj_module.append(nn.Linear(self.out_dim, 1)) else: channel_proj_module.append(nn.Linear(self.hidden_dim, 1)) # --- Weights mixing module as attention mechanism --- self.nns_mix.append(nn.Linear(3, 3)) # --- GIN channel MLPs --- # After node_input_proj, all nodes are hidden_dim, so # local_input_dim is always hidden_dim. local_input_dim = self.hidden_dim if i == self.num_layers - 1: local_out_dim = self.out_dim else: local_out_dim = self.hidden_dim for channel_module in [ self.nns_lowpass, self.nns_highpass, self.nns_fullpass, ]: if self.gnn_batchnorm: sequential = nn.Sequential( nn.Linear(local_input_dim, self.hidden_dim), nn.BatchNorm1d(self.hidden_dim), create_activation(self.activation_name), nn.Linear(self.hidden_dim, local_out_dim), nn.BatchNorm1d(local_out_dim), create_activation(self.activation_name), ) else: sequential = nn.Sequential( nn.Linear(local_input_dim, self.hidden_dim), create_activation(self.activation_name), nn.Linear(self.hidden_dim, local_out_dim), create_activation(self.activation_name), ) channel_module.append(sequential) self.ACM_convs.append( ACM_GIN( nn_lowpass=self.nns_lowpass[i], nn_highpass=self.nns_highpass[i], nn_fullpass=self.nns_fullpass[i], nn_lowpass_proj=self.nns_lowpass_proj[i], nn_highpass_proj=self.nns_highpass_proj[i], nn_fullpass_proj=self.nns_fullpass_proj[i], nn_mix=self.nns_mix[i], ) ) def reset_parameters(self): for m in self.modules(): if isinstance(m, nn.Linear): m.reset_parameters() elif isinstance(m, nn.BatchNorm1d): m.reset_parameters() def forward(self, x, edge_index, edge_attr, batch=None, return_hidden=False): """Forward pass through all ACM-GIN layers with edge updates. `batch` is accepted for API parity with ACM_GINEConv_model (GraphNorm); it is unused here. Args: x: Node features [N, in_dim]. edge_index: Edge indices [2, E]. edge_attr: Edge features [E, edge_in_dim]. The first column (index 0) is the scalar spatial distance used for degree normalization; the full vector evolves through layers via the edge MLPs. return_hidden: If True, also return all intermediate node states. Returns: x: Final node embeddings [N, out_dim]. outs: (optional) List of node states after each layer. """ # Extract scalar spatial distance for degree normalization # (stays fixed across layers) edge_weight = edge_attr[:, 0] # Project node and edge features from raw dims to hidden_dim x = self.node_input_proj(x) edge_state = self.edge_input_proj(edge_attr) outs = [] for i in range(self.num_layers): # Step 1: Compute intermediate edge features src, dst = edge_index edge_mlp_input = torch.cat([x[src], x[dst], edge_state], dim=-1) e_ij_prime = self.edge_mlps[i](edge_mlp_input) # Step 2: Update edge state with residual (safe at every layer # because both edge_state and e_ij_prime are hidden_dim) edge_state = edge_state + e_ij_prime # Step 3-4: Edge-aware ACM message passing x = self.ACM_convs[i]( x=x, edge_index=edge_index, edge_weight=edge_weight, edge_feat=e_ij_prime, ) outs.append(x) if return_hidden: return x, outs else: return x if __name__ == "__main__": acm_gin = ACM_GIN_model(46, 46, 2, 256, 74, True) print(sum(p.numel() for p in acm_gin.parameters() if p.requires_grad))