vgsatorras / egnn

MIT License
420 stars 75 forks source link

Question about velocity updates in N-body system experiment #8

Closed ndkhanh360 closed 2 years ago

ndkhanh360 commented 2 years ago

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: image

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!

vgsatorras commented 2 years ago

Hi,

This is a good point which I am explaining in more detail in an upcoming version of the paper. Both ways become similar in practice. When unrolling the recursive equation 7 through different layers you obtain an equation of the form:

x_i^{l+1} = \sum_i (x_i^l - x_j^l) \phix^l(m{ij}^l)\phi_v^l + sum_i (x_i^{l-1} - x_j^{l-1}) \phix^{l-1}(m{ij}^{l-1})\phi_v^{l} \phi_v^{l-1} + ... + sum_i (x_i^0 - x_j^0) \phix^0(m{ij}^0)\prod_l \phi_v^l(h_l) + \prod_l \phi_v^l(·)*v_0

Notice the update in every layer becomes the sum of all previous coordinate updates (x_i^l - x_j^l) for all layers "l" plus the initial velocity v^0 multiplied by a scalar value depending on the given layer "l". Therefore, in practice, we can just sum the original initial velocity v^0 multiplied by an inferred scalar value in every layer. I am adding an explanation of this in the Appendix.

Best, Victor