Mace parities for SO3-restricted symmetry / Use mace-torch weights

[bhcao]

Here p=-1 (by default of e3nn.Irreps.spherical_harmonics), however in mace-torch p=1 (interaction_irreps is called sh_irreps_inter in mace-torch). It seems that in mace-jax, the author forgot to add p=1 as well.

https://github.com/instadeepai/mlip/blob/main/mlip/models/mace/models.py#L213

if not self.include_pseudotensors:
    interaction_irreps = e3nn.Irreps.spherical_harmonics(self.l_max) # so3: 0e+1o+2e+3o+...
else:
    interaction_irreps = e3nn.Irreps(e3nn.Irrep.iterator(self.l_max)) # not so3: 0e+0o+1e+1o+2e+2o+3e+3o+...

https://github.com/ACEsuit/mace/blob/main/mace/modules/models.py#L145

if not use_so3:
    sh_irreps = o3.Irreps.spherical_harmonics(max_ell) # not so3: 0e+1o+2e+3o+...
else:
    sh_irreps = o3.Irreps.spherical_harmonics(max_ell, p=1) # so3: 0e+1e+2e+3e+...

def generate_irreps(l):
    str_irrep = "+".join([f"1x{i}e+1x{i}o" for i in range(l + 1)])
    return o3.Irreps(str_irrep)

sh_irreps_inter = sh_irreps
if hidden_irreps.count(o3.Irrep(0, -1)) > 0:
     sh_irreps_inter = generate_irreps(max_ell) # not so3: 0e+0o+1e+1o+2e+2o+3e+3o+...

[olivier-peltre]

Thanks a lot for reporting and documenting this discrepancy with the MACE-torch repo @bhcao!

Here is how we understand the current situation:

• mlip models all enforce full $E_3$ / $O_3$ equivariance:

  • with include_pseudotensors=True, all representations below the l_max degree are included.
  • with include_pseudotensors=False, only representations with the parity of harmonic polynomials are kept, i.e. p = (-1)**l which may be $\pm 1$.

  In the snippet you shared, interaction_irreps contains tensor products after the message-passing block, where the mixing of source node features with spherical activations may create pseudo-tensor parities. We have not noticed any significant impact on accuracy when disabling pseudo tensors, although this option is beneficial on runtime and memory. In both cases, we have both p=-1 and p=1 parities.

• mace-torch models have the option to enforce only $SO_3$ equivariance (i.e., drop equivariance with respect to reflections)

  • with use_so3=False, the spherical activation block generates the same representation as we always return from our spherical activation block, i.e. "0e+1o+2e+3o+..."
  • with use_so3=True, the spherical activation block overrides parities to p=1, thus creating pseudotensors artificially and breaking the equivariance under reflections.
  • if hidden_irreps contains "0o", then it seems the interaction block returns the same irreps as we do with include_pseudotensors=True. I guess here that hidden_irreps should coincide with our node_irreps variable (i.e. target of the EquivariantProductBasisBlock) but I have a hard time keeping track of irreps variable names.

Please let me know if you think I got something wrong here! I understand that discarding reflections to enforce only $SO_3$ equivariance might allow for more compact and efficient models, but predictions on enantiomers could then differ in principle. This is why it should not be confused with our include_pseudotensors flag that serves also improves the model's compactness, but without breaking $O3$ equivariance.

We may look into your issue as part of a future release: I labelled it a feature request for now as we didn't see anything broken or incorrect on our side, but please correct me if you think otherwise.

[bhcao]

Got it. It means there is no direct correspondence between your implementation and mace-torch. So, how can I convert the model trained by mace-torch to the jax version?

Explain:

• The issue above was found when I tried to write a code for model conversion. Those who only want to perform inference and not training need to seamlessly convert official torch models, such as mace-off, to jax version, thereby leveraging the acceleration capabilities of jax. Training from scratch is costly for them. See Feature Request: Example for Using Pre-trained MACE-MP Models in MACE-JAX ACEsuit/mace-jax#10.
• However, I found many inconsistencies during the implementation process. While inconsistencies in layer order are relatively easy to resolve, I am still struggling with mismatches in parameter dimensions, such as those in Contraction layers. I was wondering if you could provide some advice.

[olivier-peltre]

Thanks for providing more context!

While some people in our team have tried to match them as close as possible, the upstream version from MACE-JAX and the MACE-torch do not numerically agree to our knowledge. Differences may reside in tiny details (e.g. how RBFs are normalised -- we have reported those we spotted to the MACE group).

Although I understand your motivation, I doubt that the issue you linked will be addressed by the MACE-JAX maintainer(s) nor that we will aim for a 1-to-1 correspondence between MACE-torch and our codebase, as full transferability is too much of a challenging task.

If you cannot afford retraining a MACE-JAX model from scratch, I suggest that you have a look at our models on HuggingFace. If you're specifically interested in MACE-MP, note this one has only been trained on molecules, though we may release MACE-MP like models trained on materials in the future.