A quick question regarding the message passing logic in `GNNSparseLayer`

[cny123222]

Hey Jiale,

Hope you're doing well!

I was taking a deep dive into the ml4co/gnn4co/model/encoder/gnn_layer.py file to better understand the model architecture, and I had a quick question about the message passing implementation inside GNNSparseLayer. The design is really interesting, and I just wanted to double-check my understanding of the information flow.

In the forward method, the message Vx that gets aggregated to update the node features is created like this:

# class GNNSparseLayer
# ...
# This message is generated from the target nodes of the edges
Vx = self.V(x[edge_index[1]]) # (E, H)
# ...
# The aggregation function then aggregates messages for each node
x = Ux + self.aggregate(Vx, gates, edge_index, nodes_num) # (V, H)

The aggregate function uses torch_sparse to sum up values based on the target node index (since dim=1 corresponds to the columns, edge_index[1]).

If I'm tracing this correctly, it seems that for an edge i -> j, the message V(x_j) (generated from the target node j's features) is aggregated to update the source node i. However, since the aggregation is done based on the target node, node j would end up receiving a message that was originally generated from its own features, x_j.

For example, for an edge i -> j, Vx would contain V(x_j). The aggregate function, when updating node j, would gather all messages from incoming edges. If there's an edge k -> j, the message associated with this edge in Vx is also V(x_j).

I might be misinterpreting the flow here, as my typical understanding of MPNNs involves aggregating messages from source/neighboring nodes. I was wondering if this implementation (Vx using edge_index[1]) is an intentional design choice for a specific type of message passing, or if perhaps Vx was intended to be created from the source nodes, like self.V(x[edge_index[0]])?

No pressure to respond right away, just wanted to share my thoughts and learn more about the rationale. The codebase is incredibly well-structured, by the way!

Cheers,

Nuoyan Chen

[cny123222]

Just a quick follow-up to my previous message. I took a closer look at the GNNDenseLayer as well and noticed that a similar pattern exists in its forward method.

The message tensor Vx is constructed like this:

#  class GNNDenseLayer
Vx: Tensor = self.V(x) # (B, V, H)
Vx = Vx.unsqueeze(1).expand(-1, nodes_num, -1, -1) # (B, V, V, H)

This makes Vx[b, i, j, :] dependent on the target node j's features (self.V(x_bj)), and it's constant across the source node dimension i.

The subsequent aggregation is then performed along dim=2 (the target node dimension j):

# in GNNDenseLayer.aggregate
return torch.sum(Vx, dim=2)

When combined, this seems to result in every node i receiving the same aggregated message, which is effectively a global sum over all node features (sum_j V(x_bj)). This might prevent the layer from learning local neighborhood structures.

Just wanted to add this observation in case it's helpful. It seems the core logic is mirrored between the sparse and dense implementations.

Thanks again!