Question about Message Construction

[lsnty5190]

Hi all! Thanks for your work on Multi-ACE.

The key component of MACE is the construction of A-basis:

and B-basis:

But I didn't find any scatter sum operation for the basis function in the code (the summation operation to construct A-basis), or the multiplication of A-basis (to construct the B-basis). The implementation is the direct tensor product between node_feats and edge_feats, which are augmented by spherical harmonics.
Am I missing something else? Looking forward to your reply. Thanks!

[davkovacs]

Hi! Thanks for your question!

The A basis is constructed slightly differently in the 1st and 2nd layers, but indeed that is an important component of MACE. The scatter sum operation is part of the interaction block. In the first layer:

mace/mace/modules/blocks.py

Line 530 in 44b6da4

message = scatter_sum(

In the subsequent layer(s):

mace/mace/modules/blocks.py

Line 603 in 44b6da4

message = scatter_sum(

The B basis is probably best thought of as being part of the update. It is formed on the nodes by taking symmetric tensor products of the A functions. You can find the operations in the product block of MACE, which indeed does not contain a scatter sum:

mace/mace/modules/blocks.py

Line 162 in 44b6da4

class EquivariantProductBasisBlock(torch.nn.Module):

Let me know if you have any further questions!

[lsnty5190]

Thanks for your answer. So the A basis is obtained by the message aggregation, and the B basis is computed by SymmetricContraction. Am I right :) ?

[davkovacs]

yes, that sounds right. Note, that the contraction computes the learnable weights times the B basis in one go, for details see the appendix of the paper.

[lsnty5190]

Hi davkovacs. Just another brainstorming though. What would happen when two A basis from two adjacent atomic clusters are multiplied, i.e., $\prod A_i A_j$ ? Would it still represent the many-body expansion or the interaction between two atomic clusters (maybe cluster merging)?

[ilyes319]

The product is over A basis on the same nodes, so these products do not occur. The many body expansion is implicitly computed by doing $(\sum_{j \in \mathcal{N}(i)} \phi_{ij})^{\nu}$. So you see, all the interactions are within clusters.

[davkovacs]

that would form a different set of many-body functions, that you would need to symmetrise slightly differently.
Another reason why we don't do that, is because we try to avoid having sum over ij, ie. N^2 scaling. In MACE the expensive operations are always done on the nodes (atoms) not on the edges.

[lsnty5190]

Right. I understand the "density trick" in MACE to reduce the computational overhead. But what I think is that maybe $\prod A_i A_j$ as mentioned could implicitly represent the 4-body interaction (still within O(N)), similar to $\prod A_{j1}A_{j2}A_{j3}$?

[davkovacs]

right, you get 4-body, but you still need to sum over ij, so not O(N).