atomicarchitects / equiformer_v2

[ICLR 2024] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
https://arxiv.org/abs/2306.12059
MIT License
213 stars 26 forks source link

Inquiry About rescale_factor in SO3.py #18

Open QuantumLab-ZY opened 2 months ago

QuantumLab-ZY commented 2 months ago

Dear Developer,

I am writing to inquire about the rescale_factor = math.sqrt(length / (2 * mmax + 1)) defined in the SO3_Grid and CoefficientMappingModule classes within SO3.py. Could you please explain its purpose? Additionally, would omitting rescale_factor affect the network's equivariance?

Thank you for your assistance.

Best regards,

Yang Zhong

yilunliao commented 2 months ago

Hi @QuantumLab-ZY

Please see my responses below:

  1. The purpose of rescale_factor.

The reason is that we only have a subset of m components (from 0 to mmax) when projecting to grids instead of having all the m components (from 0 to lmax). Since we want the grid features to be well-normalized (i.e., std = 1 and mean = 0), we rescale the remaining m components with rescale_factor.

  1. Removing rescale_factor.

Removing it would not affect equivariance. It is similar to rescaling a feature uniformly, so equivariance will still be preserved.