Normalization

Equivariant normalization rescales spherical features without changing their rotation rules. EquivariantRMSNorm combines a shared RMS scale for each sample with optional centering and learnable affine parameters that preserve equivariance.

EquivariantRMSNorm

EquivariantRMSNorm(
    l_max: int,
    channels: int,
    *,
    eps: float = 1e-05,
    learnable: bool = True,
    center: bool = True,
    weighting: Literal["l_balanced", "component"] = "l_balanced",
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Normalize complete SO(3) tensors with one RMS per sample.

Inputs contain one complete real SO(3) irrep for every l = 0, ..., l_max and use the l-major shape (..., (l_max + 1) ** 2, channels). When center is true, the scalar channels are first centered as

\[ \mu=\frac{1}{C}\sum_c x_{00c}, \qquad \widetilde{x}_{lmc}=x_{lmc}-\delta_{l0}\mu. \]

The default l_balanced weighting gives every l block equal weight:

\[ q=\frac{1}{C}\sum_c\frac{1}{l_{\max}+1} \sum_{l=0}^{l_{\max}}\frac{1}{2l+1} \sum_{m=-l}^{l}\widetilde{x}_{lmc}^2. \]

component weighting instead averages all (l, m, c) entries uniformly. The normalized tensor shares the sample-dependent scale \((q+\varepsilon)^{-1/2}\). When learnable is true, the result is multiplied by a learned weight[l, c] shared across m; centered scalar channels also receive a learned bias[c]:

\[ y_{lmc}=\gamma_{lc}\frac{\widetilde{x}_{lmc}} {\sqrt{q+\varepsilon}}+\delta_{l0}\beta_c. \]

Sharing parameters across m and restricting the bias to l = 0 makes the operation commute with rotations. Contiguous CUDA float16, bfloat16, and float32 inputs use fused Triton forward and first-order backward kernels; higher-order gradients and other inputs use differentiable PyTorch expressions.

Parameters:

l_max (int) –

Highest included non-negative l.
channels (int) –

Number of channels shared by every spherical component.
eps (float, default: 1e-05 ) –

Positive value added before the reciprocal square root.
learnable (bool, default: True ) –

Learn the equivariant scale and, when centering, scalar bias.
center (bool, default: True ) –

Subtract the mean of the l = 0 channels for every sample.
weighting (Literal['l_balanced', 'component'], default: 'l_balanced' ) –

Give each l block or each spherical component equal weight.
device (device | str | None, default: None ) –

Initial device of parameters and buffers.
dtype (dtype | None, default: None ) –

Initial floating dtype of parameters and weights.

Attributes:

weight (Tensor | None) –

Learned scale with shape (l_max + 1, channels), or None.
bias (Tensor | None) –

Learned scalar bias with shape (channels,), or None.

forward

forward(so3_tensor: Tensor) -> Tensor

Parameters:

so3_tensor (Tensor) –

[..., (l_max+1)**2, C].

Returns:

normalized_so3_tensor ( Tensor ) –

[..., (l_max+1)**2, C].