Symmetric contractions which are element-dependent vs. not

[chaitjo]

mace/mace/modules/symmetric_contraction.py

Line 21 in 172c2d0

class SymmetricContraction(CodeGenMixin, torch.nn.Module):

In the symmetric contraction block, which implements Eq. 10 and 11 from the MACE paper, one of the key hyperparams is element_dependent which is by default set to True.

Based on my understanding, if we set this param to True, the contraction step will additionally consider the node_attrs (https://github.com/ACEsuit/mace/blob/main/mace/modules/blocks.py#L127) which are passed as the optional input y (https://github.com/ACEsuit/mace/blob/main/mace/modules/symmetric_contraction.py#L150).

My questions are:

• Does the paper's Eq. 10 describe the element dependent or non-dependent version of the symmetric contraction?
• Assuming that the initial node feature node_feat in the overall model (https://github.com/ACEsuit/mace/blob/main/mace/modules/models.py#L148) would already be incorporating information about the element type, why do we need the element dependent symmetric contraction?

[chaitjo]

And a similar query would apply to the implementation of the skip connection, e.g. https://github.com/ACEsuit/mace/blob/main/mace/modules/blocks.py#L513

Does this line of code follow Eq. 12 of the paper?

Is it something similar to the NequIP paper's residual connection: https://www.nature.com/articles/s41467-022-29939-5#Sec1

The weights of the Self - Interaction in the preceding formula are learned separately for each species.
(https://github.com/mir-group/nequip/blob/main/nequip/nn/_interaction_block.py#L138)

[chaitjo]

@ilyes319 tagging just ICYMI, no rush at all.

[ilyes319]

Sorry for the delay, forgot to reply! Exactly it corresponds to Eq.12. One subtility is that we do not use skip connection at the first layer as we dont want any shift in our prediction to get to true isolated atom energy. As the skip connection at the first layer would be a function of the species only (and not the positions), we dont use it.

[ilyes319]

Your understanding is completly correct! The equation to look at is actually Eq.11n You can see that the weights there have a z_i index. In the algo these are selected by the one hots as you mentionned.

While it is true that you would have the flexibility already, the full formulation is actually more flexible (the other would be a sort of tensor decomposition of what we are doing). Empirically it also gives the best results, for no extra cost!