Equivariant normalization: `degree_rms_norm`

Role

Normalization keeps feature magnitudes in a range where nonlinearities and optimizers work well, decoupling the scale of each layer’s output from the scale of its input. As with nonlinearities, the statistics used for normalizing must be invariant, and means may only be subtracted from even scalars.

Built-in members of this kind: e3layernorm (graph-level equivariant layer normalization), graph_irrep_norm (per-graph normalization as in DeepH-E3), and the per-copy EquivariantRMSNorm used by e3mlp_pre_norm.

The new layer

degree_rms_norm normalizes each irrep term — all \(m\) copies of one \((\ell,p)\) — of every entity by its own root-mean-square value:

\[\bm y_c=\gamma_c\,\frac{\bm x_c}{\sqrt{\frac{1}{m(2\ell+1)}\sum_{c'\in\text{term}}\|\bm x_{c'}\|^2+\epsilon}}\;(+\beta_c\ \text{for }\texttt{0e}).\]

It sits between the existing variants: per-copy RMS normalization fixes the norm of every copy individually and erases the relative strength of copies, while graph-level normalization couples atoms of the same structure. Per-term normalization preserves the relative magnitudes of copies within a degree and treats each atom independently, so predictions do not depend on which other structures share a batch.

Symmetry. \(\sum_{c'}\|\bm x_{c'}\|^2\) is invariant, so dividing by it is multiplication by an invariant (checklist 6). The bias is restricted to even scalars (checklist 4).

Implementation

class DegreeRMSNorm(nn.Module):
    r"""``y = gamma_c x_c / sqrt(mean_{c in term} |x_c|^2 / d + eps) (+ beta_c)``.

    Each irrep term (all copies of one ``(l, p)``) of every entity is divided by
    its own root-mean-square magnitude. The statistic is invariant, so the
    normalization is equivariant; a bias is only added to even scalars.
    """

    def __init__(self, irreps, eps=1e-6):
        super().__init__()
        self.irreps, self.eps = o3.Irreps(irreps), float(eps)
        self.gain = nn.Parameter(torch.ones(self.irreps.num_irreps))
        scalars = sum(mul for mul, ir in self.irreps if ir.l == 0 and ir.p == 1)
        self.bias = nn.Parameter(torch.zeros(scalars))

    def forward(self, x):
        blocks, bias_cursor = [], 0
        for (mul, ir), sl in zip(self.irreps, self.irreps.slices()):
            block = x[..., sl].reshape(*x.shape[:-1], mul, ir.dim)
            rms = block.square().mean(dim=(-1, -2), keepdim=True).add(self.eps).sqrt()
            blocks.append((block / rms).reshape(*x.shape[:-1], mul * ir.dim))
            if ir.l == 0 and ir.p == 1:
                blocks[-1] = blocks[-1] + self.bias[bias_cursor : bias_cursor + mul]
                bias_cursor += mul
        return torch.cat(blocks, dim=-1) * expand_per_copy(self.gain, self.irreps)

def expand_per_copy(values: torch.Tensor, irreps: o3.Irreps) -> torch.Tensor:
    """Broadcast one value per irrep copy to every magnetic component."""
    repeats = torch.tensor(
        [ir.dim for mul, ir in irreps for _ in range(mul)], device=values.device
    )
    return values.repeat_interleave(repeats, dim=-1)

@register_transform("degree_rms_norm")
def _degree_rms_norm(options, inputs, mapper):
    (features,) = inputs
    return DegreeRMSNorm(global_irreps(features), **options), [features]

Usage

- id: norm_1
  transform: degree_rms_norm
  inputs: [mixed_1]
  outputs: [normalized_1]
  options: {eps: 1.0e-6}

Variations

Normalize per degree \(\ell\) across both parities, or across all non-scalar irreps jointly, to control how much relative scale is preserved.
Track running statistics over the training set, as batch normalization does, and freeze them for inference.

Equivariant normalization: degree_rms_norm