Repository: snap-stanford/ogb
Issue: Virtual node implementation wrong?
Root cause: Algorithmic Error

Hello team.  I was reviewing some baselines for graph property prediction and had a question about the implementation for virtual nodes, which appears in a number of places like [here](https://github.com/snap-stanford/ogb/blob/master/examples/graphproppred/mol/conv.py#L189-L235).

My understanding (which may be wrong) is that we treat the virtual node as just another node in the graph that's connected to all other nodes.  The update equation for node _i_ would look something like:

<img src="https://latex.codecogs.com/svg.image?\mathbf{h}_i^{(l&plus;1)}&space;=&space;\text{MLP}[&space;(1&space;&plus;&space;\epsilon)&space;\cdot&space;\mathbf{h}_i^{(l)}&space;&plus;&space;\mathbf{h}_{\text{virt}}^{(l)}&space;&plus;&space;\sum_{j&space;\in&space;\mathcal{N}_i}&space;\text{ReLU}(\mathbf{h}_j^{(l)}&space;&plus;&space;\mathbf{e}_{ij})&space;]" title="\mathbf{h}_i^{(l+1)} = \text{MLP}[ (1 + \epsilon) \cdot \mathbf{h}_i^{(l)} + \mathbf{h}_{\text{virt}}^{(l)} + \sum_{j \in \mathcal{N}_i} \text{ReLU}(\mathbf{h}_j^{(l)} + \mathbf{e}_{ij}) ]" />

where the virtual node embedding is treated as "just another message", but is given slightly different treatment since it doesn't have a BondEncoding.

But since the public implementation adds the virtual node embedding to the other node embeddings *before* convolution ([code](https://github.com/snap-stanford/ogb/blob/master/examples/graphproppred/mol/conv.py#L199-L202)), it seems you implement an equation like:

<img src="https://latex.codecogs.com/svg.image?\mathbf{h}_i^{(l&plus;1)}&space;=&space;\text{MLP}[&space;(1&space;&plus;&space;\epsilon)&space;\cdot&space;(\mathbf{h}_i^{(l)}&space;&plus;&space;\mathbf{h}_{\text{virt}}^{(l)})&space;&plus;&space;\sum_{j&space;\in&space;\mathcal{N}_i}&space;\text{ReLU}(\mathbf{h}_j^{(l)}&space;&plus;&space;\mathbf{h}_{\text{virt}}^{(l)}&plus;&space;\mathbf{e}_{ij})&space;]" title="\mathbf{h}_i^{(l+1)} = \text{MLP}[ (1 + \epsilon) \cdot (\mathbf{h}_i^{(l)} + \mathbf{h}_{\text{virt}}^{(l)}) + \sum_{j \in \mathcal{N}_i} \text{ReLU}(\mathbf{h}_j^{(l)} + \mathbf{h}_{\text{virt}}^{(l)}+ \mathbf{e}_{ij}) ]" />

Is this intended?  Assuming not, I think a correct implementation could include changing the convolution layer itself and making the forward pass something like:

```python
def forward(self, x, edge_index, edge_attr, h_virt):
    edge_embedding = self.bond_encoder(edge_attr)
    out = self.mlp((1 + self.eps) *x + h_virt + self.propagate(edge_index, x=x, edge_attr=edge_embedding))

    return out
```

Cheers

I have a second question about the implementation.  It seems the `BondEncoder` is instantiated in the GIN/GCN modules, but `AtomEncoder` is instantiated in the GNN_Node module ([code](https://github.com/snap-stanford/ogb/blob/master/examples/graphproppred/mol/conv.py#L90)).  This means that there will be K different BondEncoders but only 1 AtomEncoder, where K is the number of GraphConv layers.  Is this intentional?  Thanks again.

Hi!

Regarding the virtual node, we just add the virtual node embedding (that captures global information) into the node embeddings in each round of message passing. https://github.com/snap-stanford/ogb/blob/master/examples/graphproppred/mol/conv.py#L199
You should think of the virtual node algorithm as h_i^(l) \leftarrow h_i^(l) + h_virtual^(l), which is then used as the input to the ordinary GIN convolution (with edge features).

Regarding the encoder, it is intended that `BondEncoder` is instantiated in each layer of message passing, while `AtomEncoder` is instantiated only in the input layer. This is because the node/atom embeddings are refined in each round of message passing (hence, no need to instantiate in every layer), while edge/bond embeddings are treated as input features.


Thanks for the response.  I think the issue is that you're adding the message before convolution rather than after.  As such, the virtual node embedding is added N times for each node update, where N is the number of neighbors.  In other words, it occurs under the sum over neighbors instead of outside of it.  Ideally what would happen is you would pass the virtual node embedding as another argument to the GIN layer, which would then add it [here](https://github.com/snap-stanford/ogb/blob/master/examples/graphproppred/mol/conv.py#L26).  

I do not think that’s an issue. That’s how we implemented the virtual node algorithm (i do not think there is established definition of the virtual node algorithm prior to our work, and we just implemented a variant of it)

Thank you for the clarification!