Rotation-Invariant Vector Normalization for Molecular Force Learning

Rotating a molecule should leave its energy unchanged while rotating each force vector by the same rotation, as required in $SO(3)$-equivariant molecular learning. However, conventional normalization layers in neural networks are designed for scalar activations; when applied to vector channels through component-wise centering or scaling, they can break this symmetry. We study normalization modules for equivariant molecular networks under force supervision, where invariant scalar channels and covariant vector channels require different geometric treatment. We formalize an $SO(3)$-equivariance-preserving recipe that computes vector scale statistics from rotation-invariant squared norms and applies the resulting rescaling identically across Cartesian coordinates, which preserves $SO(3)$-equivariance without vector shifts. Building on this recipe, we introduce grouped vector RMS (GroupRMS), which shares vector RMS estimates across channel groups to reduce estimator variance and stabilize vector scaling in small-batch training. The proposed variant is lightweight and compatible with standard scalar normalization branches. We instantiate these ideas in practical scalar-vector normalization blocks for PaiNN-style backbones and evaluate them across several molecular benchmarks. On the paired three-trajectory subset, the strongest GroupRMS-family variant, GroupNorm-GroupRMS, improves average force MAE from 0.1190 to 0.0383 and average energy MAE from 0.0675 to 0.0424 relative to BatchNorm. These results support the proposed rotation-compatible vector scaling and grouped scale estimation as effective tools for robust molecular force learning.