lgatr.layers.linear.EquiLinear

class lgatr.layers.linear.EquiLinear(in_mv_channels, out_mv_channels, primitives, in_s_channels=0, out_s_channels=0, bias=True, initialization='default')[source]

Bases: Module

Linear layer.

The forward pass maps multivector inputs of shape (..., in_channels, 16) to multivector outputs of shape (..., out_channels, 16) as

outputs[..., j, y] = sum_{i, b, x} weights[j, i, b] basis_map[b, x, y] inputs[..., i, x]
= sum_i linear(inputs[..., i, :], weights[j, i, :])

plus an optional bias term for outputs[..., :, 0] (biases on other multivector components would break equivariance). Here basis_map are precomputed (see lgatr.primitives.linear) and weights are the learnable weights of this layer. The basis_map includes 5 elements if the full Lorentz group is considered, and 10 elements if only the connected subgroup is considered. See lgatr.primitives.config.PrimitivesConfig for the subgroup option.

If there are auxiliary input scalars, they transform under a linear layer and mix with the scalar components of the multivector data. The scalars argument to forward() must be provided when in_s_channels > 0 and must not carry scalar data when in_s_channels == 0; a mismatch raises. A zero-channel scalar tensor is the “no scalars” convention and is always accepted (unlike the other layers, in_s_channels may be 0).

This layer supports four initialization schemes:

  • "default": preserves (or slightly reduces) the variance of the data in the forward pass.

  • "small": variance of outputs is approximately one order of magnitude smaller than for "default".

  • "unit_scalar": outputs will be close to (1, 0, 0, ..., 0).

  • "almost_unit_scalar": similar to "unit_scalar", but with more stochasticity.

The "almost_unit_scalar" initialization is used for the second argument of GeometricBilinear; "small" is used to combine attention heads. Everything else uses "default".

Parameters:

  • in_mv_channels (int) – Input multivector channels.

  • out_mv_channels (int) – Output multivector channels.

  • primitives (PrimitivesConfig) – LGATr primitives configuration.

  • in_s_channels (int) – Input scalar channels. Use 0 for no scalar inputs.

  • out_s_channels (int) – Output scalar channels. Use 0 for no scalar outputs.

  • bias (bool) – Whether a bias term is added to the scalar component of the multivector outputs.

  • initialization (str) – Initialization scheme; one of "default", "small", "unit_scalar", "almost_unit_scalar".

forward(multivectors, scalars=None)[source]

Apply the most general equivariant linear map to the inputs.

Parameters:

  • multivectors (Tensor) – Input multivectors of shape (..., in_mv_channels, 16).

  • scalars (Tensor | None) – Optional input scalars of shape (..., in_s_channels).

Return type:

tuple[Tensor, Tensor | None]

Returns:

  • outputs_mv – Output multivectors of shape (..., out_mv_channels, 16).

  • outputs_s – Output scalars of shape (..., out_s_channels), or None if out_s_channels == 0.

reset_parameters(initialization, gain=1.0, additional_factor=None)[source]

Initialize the weights of the linear layer.

We follow the initialization philosophy of Kaiming et al., which preserves the variance of the activations during the forward pass. This implementation deviates from the torch.nn.Linear default to take the communication between scalar and multivector channels in our linear layer into account. See inline comments for details.

Parameters:

  • initialization (str) – Initialization scheme; one of "default", "small", "unit_scalar", "almost_unit_scalar". See EquiLinear for details.

  • gain (float) – Gain factor for the activations. Should be 1.0 if the previous layer has no activation, sqrt(2) if it has a ReLU activation, and so on. Can be computed with torch.nn.init.calculate_gain().

  • additional_factor (float | None) – Empirically, slightly decreasing the data variance at each layer gives better performance. The PyTorch default initialization uses an additional factor of 1/sqrt(3) (cancelling the sqrt(3) that naturally arises in uniform initialization). See https://github.com/pytorch/pytorch/issues/57109 and https://soumith.ch/files/20141213_gplus_nninit_discussion.htm for discussion. Defaults to 1/sqrt(3).

Return type:

None