eLayerNorm#

class eLayerNorm(in_rep, eps=1e-06, equiv_affine=True, bias=True, device=None, dtype=None, init_scheme='identity')[source]#

Bases: Module

Equivariant Layer Normalization.

Given \(\mathbf{x}\in\mathcal{X}\), we first move to the irrep-spectral basis \(\hat{\mathbf{x}} = \mathbf{Q}^{-1}\mathbf{x}\), compute one variance scalar per irreducible block via irrep_radii(), and normalize each block uniformly:

\[\hat{\mathbf{y}} = \frac{\hat{\mathbf{x}}}{\sqrt{\mathbf{\sigma}^{2} + \varepsilon}}, \qquad \mathbf{y} = Q\hat{\mathbf{y}}.\]

The layer is equivariant:

\[\rho_{\text{in}}(g) \mathbf{y} = \text{LayerNorm}(\rho_{\text{in}}(g) \mathbf{x})\]

since the statistics are computed per irreducible subspace (which are preserved by the group action).

When equiv_affine=True the learnable affine step is performed directly in the spectral basis using the per-irrep scale/bias provided by eAffine.

Parameters:

in_rep (Representation) – description of the feature space \(\rho_{\text{in}}\).
eps (float) – numerical stabilizer added to each variance.
equiv_affine (bool) – if True, applies an eAffine in spectral space.
bias (bool) – whether the affine term includes invariant biases (only used if equiv_affine).
device – optional tensor factory kwargs.
dtype – optional tensor factory kwargs.
init_scheme (Literal['identity', 'random'] | None)

Note

This layer appears to generate numerical instability when used in equivariant transformer blocks. Use eRMSNorm instead in such cases.

forward(input)[source]#

Normalize per irreducible block and (optionally) apply the spectral affine transform.

Return type:: Tensor
Parameters:: input (Tensor)