Interpretation of intermediate variables in equiformer_v2 model

1

hello，OCP team!
I would like to ask about the meaning behind the intermediate variable x in the equiformer_v2 model and the information it contains. In the forward part of EquiformerV2_OC20 model, after the structural data goes through “Initialize data structures”, “Initialize node embeddings” and “Update spherical node embeddings” part, the information about x.embedding will be generated. Finally, x is processed in the “Energy estimation” part to predict the energy.
The parts of the code where I have questions are as follows:

# Update spherical node embeddings
for i in range(self.num_layers):
    x = self.blocks[i](
        x,  # SO3_Embedding
        atomic_numbers,
        edge_distance,
        edge_index,
        batch=data.batch,  # for GraphDropPath
    )

# Final layer norm
x.embedding = self.norm(x.embedding)

If my struct data is a 5-atom struct, and the “atomic_numbers” part of the struct data read from the lmdb file is a tensor([1,2,3,4,5]), this means that the five atoms are H, He, Li, B, Be.The size of x.embedding is [5×49×128]. Can I interpret this [5×49×128] vector to mean that each [49×128] vector contains information about one of the atoms? Do these five [49×128] vectors correspond to each of the previous five atoms? For example, the first [49×128] vector corresponds to information about the H atom, and the last [49×128] vector corresponds to information about the B atom？
I would really appreciate it if you could help me answer this question.

2

Correct - each [49x128] vector corresponds to a different atom. data.atomic_numbers will tell you the ordering of your atoms. i.e.

data.atomic_numbers = [1, 2, 3, 5, 4]
x.embedding.shape = [5x49x128]

x[0].shape = [49x128] # this corresponds to data.atomic_numbers[0] == 1 (H)

Hope this clarifies it!