Question on EquivariantLayerNorm and Noise AccumulatedNormalization

[lsnty5190]

Hi, thanks for your impressive work. Here I have two questions about the details in the pretrained denoising process:

1. The EquivariantLayerNorm. By running the test_scalar_invariance in the tests/test_equivariance.py and set layernorm_on_vec: whitened, I found that the model is not equivariant. The problem exists in the addition of regularized matrix self.eps * reg_matrix. The model would be equivariant if I remove this part.
2. The Noise AccumulatedNormalization. In your implementation, the noises are accumulated through all batches and normalized. Then they are treated as the supervised signal to compute the loss. I guess such an operation (AccumulatedNormalization) is conducted to ensure the mean=0, std=1 after noise_scale. But you can directly set noise = torch.randn_like(data.pos), data.pos_target = noise and data.pos = data.pos + noise * self.hparams['position_noise_scale'] without the AccumulatedNormalization. Could I ask the reason? Thanks.

[shehzaidi]

Hi, thanks for your questions!

Regarding EquivariantLayerNorm, the additional self.eps * reg_matrix term is essentially a precaution there to avoid inverting an ill-conditioned matrix. Ideally, self.eps would be small enough to not affect the equivariance. Without that, I noticed a bad covar matrix occasionally disturbed training. Feel free to play around with it, and let me know if you have other ideas!

Regarding AccumulatedNormalization, indeed you can also do it directly to have normalized noise targets. We inherited this from the implementation of denoising for noisy nodes, where sometimes, it's easier to have an automatic normalizer instead of manually preprocessing the noise to have unit variance.

Hope that helps!