search

atomicarchitects / equiformer_v2

[ICLR 2024] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
https://arxiv.org/abs/2306.12059
MIT License
218 stars 27 forks source link

Question about Output Head #19

Closed feyhong1112 closed 1 month ago

feyhong1112 commented 3 months ago

I read the paper, but I'm not quite sure what 'Output Head' means. For example, if I have a 3D molecular graph of my protein pocket, with each node embedded with features from the BLOSUM62 matrix, and Does 'scalar' mean the combination of node features and the adjacency matrix into a 1D representation? Similar to Graphormer output?

Thank you for your fantastic project and to reply my simple question in advance.

yilunliao commented 3 months ago

Hi @feyhong1112

Please see my responses below:

  1. What does 'Output Head' mean?

Output heads here refer to how to predict labels from the node embeddings and here have energy and force heads. The energy head is a feedforward network (here), and the force head is an attention module (here).

  1. Does 'scalar' mean the combination of node features and the adjacency matrix into a 1D representation?

No. 'Scalar' means the type-0 vectors of node/edge embeddings (a subset of the equivariant features). The adjacency matrix is how we represent graph structures (e.g., (1) dense arrays with each graph having the same number of nodes/edges after padding as in Graphormer or (2) a flattened 1D array with different numbers of nodes/edges for each graph as in this work) and is not part of the node embeddings.

  1. Similar to Graphormer output?

For energy prediction, I think how we predict energy is similar to how Graphormer does -- we only use the scalar part of node embeddings. For force prediction, it is different. More specifically, we use all type-L vectors in node embeddings to obtain the output type-1 vectors while Graphormer uses some scalars to weigh the edge vectors of relative positions.