Skip to content

Generalize Interaction GNN

Anthony Correia requested to merge anthonyc/generalise_interaction_gnn into main

This MR includes change that allow to add features to the Interaction GNN.

  • Be able to have a different # node and edge hidden units. There is no reason these 2 values should be equal in theory. In practice, reducing either of these numbers reduce performance, especially the # edge hidden units.
  • Define only_e property. If enabled, only the edge encodings e (and not the node encodings) are used by the edge and triplet classifiers. In practice, this does not reduce performance. This allows to reduce the number of parameters and to simplify the C++/CUDA implementation in Allen.
  • Define the use_xyz property. If enabled, the (x, y, z) coordinates are used instead of the cylindrical coordinates. In theory, it is easier to guess that 5 points are aligned using the (x, y, z) coordinates.
  • Define loss property that defaults to focal. Allow to use the cross_entropy loss. In practice, the focal loss leads to better performance.
  • Be able not to use any hidden activations. I wanted to try what is the performance without any hidden activations - it is worse. In a computer science standpoint, also think it is better that the default is no hidden activations rather than favouring any random hidden activations (which was ReLU).
  • Define output_activation. No output activations were used in the intermediate networks (node and edge encoders, node and edge networks). Setting it to SiLU improved performance.
Edited by Anthony Correia

Merge request reports