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

Incorporating vector node features #10

Closed patriksimurka closed 2 months ago

patriksimurka commented 8 months ago

Hi, nice work! I was wondering what it would take to accommodate for systems with nodes that have additional non-scalar features. Any hints or snippets would be greatly appreciated. Thanks.

yilunliao commented 8 months ago

Hi @patriksimurka

The node embeddings have the shape (num_nodes, (1 + L_{max}) ** 2, num_channels). To add scalars to the node embeddings, we can first expand the scalar properties to have num_channels channels, and then add those scalar vectors to the type-0 part:

num_channels = 128
lmax = 6
num_nodes = 128
node_embedding = torch.randn(num_nodes, (1 + lmax) ** 2, num_channels)
input_scalar = torch.randn(num_nodes, 1, num_channels)
node_embedding[:, 0:1, :] = node_embedding[:, 0:1, :] + input_scalar

For type-1 vectors, you can do the following:

type_1_vectors = torch.randn(num_nodes, 3, num_channels)
node_embedding[:, 1:4, :] = node_embedding[:, 1:4, :] + type_1_vectors

Basically, the second dimension of node_embedding in the above examples corresponds to [one type-0 vector, one type-1 vector, ... one type-lmax vector]. So you can just add the type-L vector to the corresponding slice.