Scale of the output matrix value is too small

[lamthuy]

I tried a simple MatrixMACE, but the predicted matrices have very small values, on the order of 1e-6. I checked the output node features from MACE, and their scale is normal (around 5e-1). However, after a few tensor product transformations involving edge and node operations, the values shrink to around 1e-6. This creates a problem for learning because optimization becomes more difficult. I’d like to ask if there is a way to insert something like a SiLU activation after each tensor product layer to help rescale the outputs.

[pfebrer]

Hmm ok, I have never investigated this. Is this for all the values, or only for some?

I don't know where exactly you want to apply the SiLU, but if it's directly on the values you can only do it on the scalar values (l=0). If you do it on the higher l=1,2,3... you will break equivariance. Maybe it is worth it to play with the normalization options of e3nn when calling the tensor products. See path_normalization here: https://docs.e3nn.org/en/stable/api/o3/o3_tp.html#e3nn.o3.TensorProduct

I don't know exactly how it works, but if you reach some conclusion we can change this in the code 👍

[pfebrer]

So these values are for an untrained MatrixMACE I imagine, right?

[lamthuy]

Yes, at the beginning in the MatrixMACE the node_feats given by mace is in the order of 5e-1, so after the tensor square and tensor product the given value get as small as e-6 and that require the network to do rescaling which make the optimization much harder. So I have tried a few option such as using BatchNorm on the node_features it seem to solve the problem but the operations are still linear without any non-linearity so I guess all the heavy tasks including non-linear transformation are dependent on the MACE layers. I guess sometime we need on-linear transformation after each operation as well. If that is supported that would be great.

[pfebrer]

The non-linear operations can only be done on scalars if the model needs to stay equivariant, so that's quite a limitation. But sure they can be done. That's why Graph2Mat allows passing your own node and edge operations, so that you can experiment with whatever you want 🙂 The default behavior of Graph2mat is indeed to let the previous step (in this case MACE) do the heavy lifting.

Could you share the code of the modified operation with the Batchnorm? I can test it and it can go in if it simply learns better.

[lamthuy]

Thank you pfeber, I only add e3nn BatchNorm after each node_feats, however I think it depends on data, in general it should be managed by the users and they should choose which non-linearity they need.

I have another question. So far I see that the size of the tensor product weight is an order of magnitude larger than the size of mace. This make the training very memory inefficient. In your implementation, for predicting the matrix, do you only need the scalar part of the output? As the predicted matrices are scalars I guess the network does not need to have the complete output representation with l> 0?

[pfebrer]

I'm sorry but I don't understand what you mean. The matrix is predicted with the equivariance of your basis. So if the basis for which you want to predict a matrix has more than l=0, the matrix is certainly not made of scalars.

Yes the number of weights in Graph2Mat can become very large, this is another reason to use basis grouping.