vgsatorras / egnn

E(n) Equivariant Graph Neural Networks

Official implementation (Pytorch 1.7.1) of:

E(n) Equivariant Graph Neural Networks
Victor Garcia Satorras, Emiel Hogeboom, Max Welling
https://arxiv.org/abs/2102.09844

Abstract: This paper introduces a new model to learn graph neural networks equivariant to rotations, translations, reflections and permutations called E(n)-Equivariant Graph Neural Networks (EGNNs). In contrast with existing methods, our work does not require computationally expensive higher-order representations in intermediate layers while it still achieves competitive or better performance. In addition, whereas existing methods are limited to equivariance on 3 dimensional spaces, our model is easily scaled to higher-dimensional spaces. We demonstrate the effectiveness of our method on dynamical systems modelling, representation learning in graph autoencoders and predicting molecular properties.

Example code

For a simple example of a EGNN implementation click here. Or copy the file models/egnn_clean/egnn_clean.py into your working directory and run:

import egnn_clean as eg
import torch

# Dummy parameters
batch_size = 8
n_nodes = 4
n_feat = 1
x_dim = 3

# Dummy variables h, x and fully connected edges
h = torch.ones(batch_size * n_nodes, n_feat)
x = torch.ones(batch_size * n_nodes, x_dim)
edges, edge_attr = eg.get_edges_batch(n_nodes, batch_size)

# Initialize EGNN
egnn = eg.EGNN(in_node_nf=n_feat, hidden_nf=32, out_node_nf=1, in_edge_nf=1)

# Run EGNN
h, x = egnn(h, x, edges, edge_attr)

If you are using the EGNN in a new application we recommend checking the EGNN attributes description that contains some upgrades not included in the paper.

N-body system experiment

Create N-body dataset

cd n_body_system/dataset
python -u generate_dataset.py --num-train 10000 --seed 43 --sufix small

Run experiments

EGNN model

python -u main_nbody.py --exp_name exp_1_egnn_vel --model egnn_vel --max_training_samples 3000 --lr 5e-4

GNN model

python -u main_nbody.py --exp_name exp_1_gnn --model gnn --max_training_samples 3000 --lr 1e-3

Radial Field

python -u main_nbody.py --exp_name exp_1_gnn --model rf_vel --n_layers 4 --max_training_samples 3000 --lr 2e-4 

Tensor Field Networks

python -u main_nbody.py --exp_name exp_1_tfn --model tfn --max_training_samples 3000 --lr 1e-3 --degree 2 --nf 32

SE3 Transformer

python -u main_nbody.py --exp_name exp_1_se3 --model se3_transformer --max_training_samples 3000 --div 1 --degree 3 --nf 64 --lr 5e-3

N-body system sweep experiment

For the experiment where we sweep over different amounts of training samples you should create a larger training dataset

cd n_body_system/dataset
python -u generate_dataset.py  --num-train 50000 --sample-freq 500 

Then you can train on in this new partition by adding --dataset nbody to the above training commands. You can choose the number of training samples by modifying the argument --max_training_samples <number of training samples>

E.g. for the EGNN for 10.000 samples

python -u main_nbody.py --exp_name exp_debug --model egnn_vel --max_training_samples 10000 --lr 5e-4 --dataset nbody

Graph Autoencoder experiment

GNN Erdos & Renyi

python -u main_ae.py --exp_name exp1_gnn_erdosrenyi --model ae --dataset erdosrenyinodes_0.25_none --K 8 --emb_nf 8 --noise_dim 0

GNN Community

python -u main_ae.py --exp_name exp1_gnn_community --model ae --dataset community_ours --K 8 --emb_nf 8 --noise_dim 0

Noise-GNN Erdos&Renyi

python -u main_ae.py --exp_name exp1_gnn_noise_erdosrenyi --model ae --dataset erdosrenyinodes_0.25_none --K 8 --emb_nf 8 --noise_dim 1

Noise GNN Community

python -u main_ae.py --exp_name exp1_noise_gnn_community --model ae --dataset community_ours --K 8 --emb_nf 8 --noise_dim 1

Radial Field Erdos&Renyi

python -u main_ae.py --exp_name exp1_rf_erdosrenyi --model ae_rf --dataset erdosrenyinodes_0.25_none --K 8 --emb_nf 8

Radial Field Community

python -u main_ae.py --exp_name exp1_rf_community --model ae_rf --dataset community_ours --K 8 --emb_nf 8

EGNN Erdos&Renyi

python -u main_ae.py --exp_name exp1_egnn_erdosrenyi --model ae_egnn --dataset erdosrenyinodes_0.25_none --K 8 --emb_nf 8

EGNN Community

python -u main_ae.py --exp_name exp1_egnn_community --model ae_egnn --dataset community_ours --K 8 --emb_nf 8

---

The following overfit eperiments are for (p=0.2). p can be modified by replacing the 0.2 value from the dataset name (e.g. erdosrenyinodes_0.2_overfit) to other values.

GNN Erdos&Renyi overfit

python -u main_ae.py --model ae --dataset erdosrenyinodes_0.2_overfit --epochs 10001 --test_interval 200 --K 16 --emb_nf 16 2>&1 | tee outputs_ae/$EXP.log &

Noise-GNN Erdos&Renyi overfit

python -u main_ae.py --model ae --dataset erdosrenyinodes_0.2_overfit --epochs 10001 --test_interval 200 --noise_dim 1 --K 16 --emb_nf 16 2>&1 | tee outputs_ae/$EXP.log &

EGNN Erdos&Renyi overfit

python -u main_ae.py --model ae_egnn --dataset erdosrenyinodes_0.2_overfit --epochs 10001 --test_interval 200 --K 16 --emb_nf 16 2>&1 | tee outputs_ae/$EXP.log &

QM9 experiment

properties --> [alpha | gap | homo | lumo | mu | Cv | G | H | r2 | U | U0 | zpve]
learning rate --> 1e-3 for [gap, homo lumo], 5r-4 for the rest

python -u main_qm9.py --num_workers 2 --lr 5e-4 --property alpha --exp_name exp_1_alpha
python -u main_qm9.py --num_workers 2 --lr 1e-3 --property gap --exp_name exp_1_gap
python -u main_qm9.py --num_workers 2 --lr 1e-3 --property homo --exp_name exp_1_homo
python -u main_qm9.py --num_workers 2 --lr 1e-3 --property lumo --exp_name exp_1_lumo
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property mu --exp_name exp_1_mu
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property Cv --exp_name exp_1_Cv
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property G --exp_name exp_1_G
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property H --exp_name exp_1_H
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property r2 --exp_name exp_1_r2
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property U --exp_name exp_1_U
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property U0 --exp_name exp_1_U0
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property zpve --exp_name exp_1_zpve

Acknowledgements

The Robert Bosch GmbH is acknowledged for financial support.

egnn's Issues

Question about velocity updates in N-body system experiment

Dear author,

I have a question about the updates of velocity in your N-body experiments.
According to equation (7), both the velocity and coordinates are updated after each EGCL layer:

However, in the corresponding E_GCL_vel class, only the coordinates are updated while the velocity is kept the same:

 def forward(self, h, edge_index, coord, vel, edge_attr=None, node_attr=None):
      row, col = edge_index
      radial, coord_diff = self.coord2radial(edge_index, coord)

      edge_feat = self.edge_model(h[row], h[col], radial, edge_attr)
      coord = self.coord_model(coord, edge_index, coord_diff, edge_feat)


      coord += self.coord_mlp_vel(h) * vel
      h, agg = self.node_model(h, edge_index, edge_feat, node_attr)
      # coord = self.node_coord_model(h, coord)
      # x = self.node_model(x, edge_index, x[col], u, batch)  # GCN
      return h, coord, edge_attr

I was wondering if there is any motivation or reasoning for this or whether this helps achieve better results.

Thank you!

extrapolation on other datasets

I am wondering whether you could provide a tutorial concerning how to extrapolate egnn on another dataset, saying alchemy.
You, your group, have done quite a good job so that the code could be run so automatically. But on the other hand, this set some obstacles when I am trying to extrapolate your model on another datasets.

Egnn flow

Hello @vgsatorras
I am not sure if the new work "E(n) Equivariant Normalizing Flows" will be released here or in separate repo?
I hope you tell me when this will be released?
Thanks

Clean code for interpretation

I was studying your paper and, when looking into the implementation, came upon the following in .\models\egnn_clean\egnn_clean.py:

    self.edge_mlp = nn.Sequential(
        nn.Linear(input_edge + edge_coords_nf + edges_in_d, hidden_nf),
        act_fn,
        nn.Linear(hidden_nf, hidden_nf),
        act_fn)

    self.node_mlp = nn.Sequential(
        nn.Linear(hidden_nf + input_nf, hidden_nf),
        act_fn,
        nn.Linear(hidden_nf, output_nf))

    layer = nn.Linear(hidden_nf, 1, bias=False)
    torch.nn.init.xavier_uniform_(layer.weight, gain=0.001)

    coord_mlp = []
    coord_mlp.append(nn.Linear(hidden_nf, hidden_nf))
    coord_mlp.append(act_fn)
    coord_mlp.append(layer)
    if self.tanh:
        coord_mlp.append(nn.Tanh())
    self.coord_mlp = nn.Sequential(*coord_mlp)

I was somewhat confused as why self.coord_mlp was different from self.edge_mlp and self.node_mlp. I had to wrestle with myself until I was convinced they are the same.

Therefore, for clarity reasons, I would recomend the following refactor:

    self.coord_mlp = nn.Sequential(
        nn.Linear(hidden_nf, hidden_nf),
        act_fn,
        layer,
        nn.Tanh() if self.tanh else nn.Identity())

About the implemenation of Eq4 of the EGNN paper.

Hi,

I wonder about the implementation of Eq4 of the EGNN paper.
According to equation 4 from the paper, the position update takes into accounts all interactions among the nodes in the graph. In that sense, the meaning of normalizer C=1/(M-1) makes more sense. However, according to the
code, the aggregations to update the coord are done only for the existing edges. May I ask to clarify which one is the expected behavior of EGNN?

def coord_model(self, coord, edge_index, coord_diff, edge_feat):
        row, col = edge_index
        trans = coord_diff * self.coord_mlp(edge_feat)
        if self.coords_agg == 'sum':
            agg = unsorted_segment_sum(trans, row, num_segments=coord.size(0))
        elif self.coords_agg == 'mean':
            agg = unsorted_segment_mean(trans, row, num_segments=coord.size(0))
        else:
            raise Exception('Wrong coords_agg parameter' % self.coords_agg)
        coord = coord + agg
        return coord

Thank you :)
Junyoung

Question about the AE experiment

Hi, I am a little confused by the AE experiment. The input data for the AE experiment is a fully-connected graph. I understand that the "attention" parameter is turned on while feeding a fully-connected graph and let the model infer edges, but it is only the case in the QM9 experiment. In the AE experiment, the attention parameter is set to False and the given graph is still fully-connected. Would you mind explaining a little bit about it?

Thanks for your time!

About training egnn on qm9 dataset

Hello! I'm really interested in your work, and tested it on qm9 dataset following your setting.
However, I found a strange situation that when training about 80 epochs, the loss suddenly burst from about 0.05 or lower to 1 and never drop again. This usually happenes when I'm training on gap/lumo/homo property.
Have you ever encountered this situation? If so do you know how to deal with it? Thank you a lot!

Edit: I think it's because CosAnealing lr scheduler, but i'm not sure
update: tried other kind of scheduler but still failed

Experiment results on qm9 dataset

Dear author,

I am trying to reproduce the results on the qm9 dataset in Table 3 but encounter something that I don't really understand. Specifically, when I ran your code, I achieved a very low MAE on the chemical property U_0 (train MAE ~ 0.001, test MAE ~ 0.011 - much lower than the test result 11 in the paper), which makes me think that the unit used in the code is eV instead of meV as reported in the paper. Can you help me confirm if what I think is true or not?

Another thing is that for the chemical property ZPVE, the test MAE was much higher than the train MAE (about 0.002 for train vs 1.61 for test - but this number is pretty close to the test result 1.55 in the paper), which confuses me about the unit used in the code for this property (I guess it is meV, not eV as in the above case). It would be great if you can provide more details with regard to this situation.

Thank you :)
Khanh Nguyen.

Some questions about egnn

Hi,

I wonder about the implementation of two parts in the egnn codes.

First, in E_GCL code,

    def forward(self, h, edge_index, coord, edge_attr=None, node_attr=None):
        row, col = edge_index
        radial, coord_diff = self.coord2radial(edge_index, coord)

        edge_feat = self.edge_model(h[row], h[col], radial, edge_attr)
        coord     = self.coord_model(coord, edge_index, coord_diff, edge_feat)
        h, agg    = self.node_model(h, edge_index, edge_feat, node_attr)
        # coord = self.node_coord_model(h, coord)                                                                                
        # x = self.node_model(x, edge_index, x[col], u, batch)  # GCN                                                            
        return h, coord, edge_attr

the edge_feat is not return of forward method and the edge_attr, which is input of forward, is just returned.
Can I use the edge_feat as the one of the returns of forward method?

Second, Can I predict the h, edge_attr with the EGNN Autoencoder?
In example main_ae.py code, the decoder simply consists of linear transformation and just predicts adjacency matrix.
What I want to do is to predict the h_x (node attribute) and edge attribute. So did I just use the output of autoencoder as same way of adjacency matrix?

Best regards,
Jinyoung

Environment file

Hi can you share your environment used to run code in this repo? It doesn't work with my local environment

In place operation in `coord_model()` of `egnn_clean.py`

Thanks for sharing the code for this paper!

I ran into some autograd issues with the following in place operation:

egnn/models/egnn_clean/egnn_clean.py

Line 81 in dcd3761

coord += agg

Is this a bug or was it done intentionally for some other reason?

Question about graph edges in QM9 experiment

Dear author,

In the QM9 experiement, you mentioned you were "not provided with an adjacency matrix" and used "the extension of our model from Section 3.3 that infers a soft estimation of the edges. "

But as far as I know, it's not hard to get the molecular graphs with edges between atoms from the QM9 dataset. And I was wondering why you chose to infer the soft edges with molecular graphs given. Was there any benefits of this operation?

Many thanks,
Lixin Yang