Implementation of E(n)-Equivariant Graph Neural Networks, in Pytorch

Dear authors,

thanks for your great work! I saw your example, which is easy to understand. But I notice that during training, in each iteration, it seems it supports the case where batch-size > 1, but all the graphs have the same adj_mat. do you have better solution for that? thanks

https://github.com/lucidrains/egnn-pytorch/blob/e35510e1be94ee9f540bf2ffea49cd63578fe473/egnn_pytorch/egnn_pytorch.py#L413

A small problem, this Tensor is not defined.

Thanks for your work.

Hi, thanks for your great work!

I have a question on how aggregations are computed for node embedding and coordinate embedding. In the paper, the aggregation for node embedding is computed over its neighbors, while the aggregation for coordinate embedding is computed over is computed over all others. However, in EGNN_sparse, I didn't notice such difference in aggregations.

I guess it is because computing all-pair messages for coordinate embedding makes 'sparse' meaningless, but I would like to double-check to see if I get this correctly. So anyway, did you do this intentionally? Or did I miss something?

My appreciation.

Hi - fast work coding these things up, as usual! Looking at the paper and your code, you're not using squared distance for the edge weighting. Is that intentional? Also, it looks like you are adding the old feature vectors to the new ones rather than taking the new vectors directly from the fully connected net - is that also an intentional change from the paper?

• Fixed some tiny errors in data flows for the PyG layers (dimensions and slices mainly)
• fixed the EGNN_Sparse_Network so now it works
• provides example for point cloud denoising (from gaussian masked coordinates), and showcases potential issues:
  • unstable (could be due to nature of data, not sure, but gvp does well on it)
  • not able to beat baseline (in contrast, gvp gets to 0.8 RMSD while this gets to the baseline 1 RMSD but not below it)

Hi, many thanks for the implementation!

I was quickly checking the code for the pytorch geometric implementation of the EGNN_sparse layer, and I noticed that it expects the first 3 columns in the features to be the coordinates. However, in the update method, features and coordinates are passed in the wrong order.

https://github.com/lucidrains/egnn-pytorch/blob/375d686c749a685886874baba8c9e0752db5f5be/egnn_pytorch/egnn_pytorch.py#L192

This may cause problems during learning (think of concatenating several of these layers), as they expect coordinate and feature order to be consistent.

One can reproduce this behaviour in the following snippet:

layer = EGNN_sparse(feats_dim=1, pos_dim=3, m_dim=16, fourier_features=0)

R = rot(*torch.rand(3))
T = torch.randn(1, 1, 3)

feats = torch.randn(16, 1)
coors = torch.randn(16, 3)
x1 = torch.cat([coors, feats], dim=-1)
x2 = torch.cat([(coors @ R + T).squeeze() , feats], dim=-1)
edge_idxs = (torch.rand(2, 20) * 16).long()

out1 = layer(x=x1, edge_index=edge_idxs)
out2 = layer(x=x2, edge_index=edge_idxs)

After fixing the order of these arguments in the update method then the layer behaves as expected (output features are equivariant, and coordinate features are equivariant upon se(3) transformation)

Hi Lucid!!

I find that when stacking multiple layers the output from the model rapidly goes to Nan. I suspect it may be related to the weights used for initialization.

Here is a minimal working example:

Make some data:

    import numpy as np
    import torch
    from egnn_pytorch import EGNN
    
    torch.set_default_dtype(torch.double)

    zline = np.arange(0, 2, 0.05)
    xline = np.sin(zline * 2 * np.pi) 
    yline = np.cos(zline * 2 * np.pi)
    points = np.array([xline, yline, zline])
    geom = torch.tensor(points.transpose())[None,:]
    feat = torch.randint(0, 20, (1, geom.shape[1],1))

Make a model:

    class ResEGNN(torch.nn.Module):
        def __init__(self, depth = 2, dims_in = 1):
            super().__init__()
            self.layers = torch.nn.ModuleList([EGNN(dim = dims_in) for i in range(depth)])
        
        def forward(self, geom, feat):
            for layer in self.layers:
                feat, geom = layer(feat, geom)
            return geom

Run model for varying depths:

    for i in range(10):
        model = ResEGNN(depth = i)
        pred = model(geom, feat)
        mean_absolute_value  = torch.abs(pred).mean()
        print("Order of predictions {:.2f}".format(np.log(mean_absolute_value.detach().numpy())))

Output : Order of predictions -0.29 Order of predictions 0.05 Order of predictions 6.65 Order of predictions 21.38 Order of predictions 78.25 Order of predictions 302.71 Order of predictions 277.38 Order of predictions nan Order of predictions nan Order of predictions nan

Hi, thanks for this implementation!

I was wondering if the pytorch-geometric implementation of this architecture is throwing the edge features out by mistake, as seen here

https://github.com/lucidrains/egnn-pytorch/blob/1b8320ade1a89748e4042ae448626652f1c659a1/egnn_pytorch/egnn_pytorch.py#L148-L151

Or maybe my understanding is wrong? Cheers,

I don't recommend normalizing the weights nor the coords.

• The weights are the coefficient that multiplies the delta in the i->j direction
• the coords are the deltas in the i->j direction Can't see the advantage of normalizing them beyond a naive stabilization that might affect the convergence properties by needing more layers due to the limited transformation that a layer will be able to do.

It works fine for denoising without normalization (the unstability might come from huge outliers, but then tuning the learning rate or clipping the gradients might be of help.)

Recently, I've tried to read EGNN paper and study your EGNN code. Actually, I had hard time to understand both paper and code because my major is not computer science. When studying your code, I realize that the shape of hidden_out and the shape of kwargs["x"] must be same to perform add operation (becaus of residual connection) in the class EGNN_sparse forward method. How can I increase or decrease the hidden dimension size of x?

I would like to get some advice.

Thanks for your consideration in this regard.

Hi, I found there may be a little mistake. In the input hint of class EGNN_Sparse of pyg version, the size of edge_index is (n_edges, 2). However, it should be (2, n_edges). Otherwise, the distance calculation will be not correct. """ Inputs: * x: (n_points, d) where d is pos_dims + feat_dims * edge_index: (n_edges, 2) * edge_attr: tensor (n_edges, n_feats) excluding basic distance feats. * batch: (n_points,) long tensor. specifies xloud belonging for each point * angle_data: list of tensors (levels, n_edges_i, n_length_path) long tensor. * size: None """

I'm using EGNN with 4 layers (where I also do global attention after each layer), and I'm seeing exploding gradients after 90 epochs or so. I'm using techniques discussed earlier (sparse attention matrix, coor_weights_clamp_value, norm_coors), but I'm not sure if there's anything else I should be doing. I'm also not updating the coordinates, so the fix in the pull request doesn't apply.

This removes the NaN bug completely (must also use norm_coors otherwise performance dies)

The NaN bug comes from the coors_mlp exploding, so forcing values between -1 and 1 prevents this. If coordinates are normalised then performance should not be adversely affected.