A question about MACE descriptors

[shoubhikraj]

I have been trying to use MACE, especially the foundation models and thinking about what information they learn. I know that we can get node features from MACE layers with the calculator.get_descriptors() function. As I understand, these descriptors are for each node i.e. each atom(?)

Would it be possible to calculate a global descriptor for the entire molecule from the MACE architecture? If so, how would you recommend constructing that? Does it make sense to construct similarity metrics e.g. dot product or euclidean distance on the node features? I would appreciate any thoughts and suggestions.

[gabor1]

Hello!

This is a generic task, some old literature for something similar is here: https://pubs.rsc.org/en/content/articlelanding/2016/cp/c6cp00415f (part of the paper is on the SOAP descriptors, but the problem is local -> global descriptors is central and the discussion there applies to MACE invariant descriptors too).

Other points to make are that the full receptive field of MACE is large, (2*the layer cutoff), so straight averaging over the atoms (or in a similar vein to Deep Sets, applying a nonlinear function to the invariants before averaging them) might work well for most systems which are not much bigger than the full receptive field. What you are missing with this is still global relative orientational descriptors. One partial solution is to take the non-scalar parts of the MACE descriptors on different atoms and dot product them (or in general tensor-contract them, in the case of the invariant parts you could multiply them), but make sure you do it in a permutationally invariant way (i.e. eventually average over all pairs). It would be interesting to try to do this in a systematic (i.e. improvable) way.

To go further, you could put a global transformer readout on the equivariant descriptors, which would increase the flexibility and offer you full range - but needs extra training obviously.