lgatr.primitives.normalization

Multivector normalization.

Functions

equi_layer_norm(x[, channel_dim, gain, epsilon])

Equivariant LayerNorm for multivectors.
lgatr.primitives.normalization.equi_layer_norm(x, channel_dim=-2, gain=1.0, epsilon=0.01)[source]

Equivariant LayerNorm for multivectors.

Rescales input such that mean_channels |inputs|^2 = 1, where the norm is the GA norm and the mean goes over the channel dimensions.

Using a factor gain > 1 makes up for the fact that the GP norm overestimates the actual standard deviation of the input data.

Parameters:

  • x (torch.Tensor) – Multivectors with shape (…, 16).

  • channel_dim (int) – Channel dimension index. Defaults to the second-last entry (last are the multivector components).

  • gain (float) – Target output scale.

  • epsilon (float) – Small numerical factor to avoid instabilities. By default, we use a reasonably large number to balance issues that arise from some multivector components not contributing to the norm.

Returns:

outputs – Normalized multivectors with shape (…, 16).

Return type:

torch.Tensor