Equivariant MLP block: `equivariant_glu`

Role

Point-wise blocks transform the features of each entity without looking at its neighbours: they mix channels, change the irrep layout and add expressiveness between message-passing steps. In MANDALA they appear as the pre/post maps and refinements of every update and as the necks and projections of the readout (E3MLP and its variants).

Built-in members of this kind: the E3MLP variants basic, normact, gate, gatemagnitudes, film, resnormact, resgatemagnitudes and bilinear; irrep_linear as a stage.

The new layer

Gated linear units (GLUs) are among the most effective feed-forward blocks in sequence models: a value path is multiplied elementwise by a sigmoid gate computed from the same input. The equivariant GLU keeps that structure but lets the gate see only invariants and act on whole copies:

\[\bm y=\bm x+\big(W_v\bm x\big)\odot_{\text{copy}}\sigma\big(\mathrm{MLP}(\iota(\bm x))\big),\]

where \(W_v\) is an equivariant linear map (o3.Linear), \(\iota(\bm x)\) are the per-copy invariants of \(\bm x\) and the MLP outputs one gate per output copy. The residual is used when input and output irreps coincide.

The gate sees the invariants of all input copies, so a scalar channel can open or close a vector channel, and the norm of an 2e copy can modulate a 1o output — couplings that a norm-based activation of each copy alone cannot express, at the cost of a single linear layer and a small MLP.

Symmetry. \(W_v\) mixes only copies of equal irreps (checklist 3a); gates are functions of invariants (2) applied to whole copies (3c).

Implementation

class EquivariantGLU(nn.Module):
    r"""``y = x + Linear_v(x) * sigma(MLP(inv(x)))`` (residual when shapes match).

    The value path is an equivariant linear map. The gate path sees only
    invariants of ``x`` and outputs one gate per output copy, which multiplies
    whole copies and therefore preserves equivariance.
    """

    def __init__(self, irreps_in, irreps_out=None, hidden=32):
        super().__init__()
        self.irreps_in = o3.Irreps(irreps_in)
        self.irreps_out = o3.Irreps(irreps_out or irreps_in)
        self.value = o3.Linear(self.irreps_in, self.irreps_out)
        self.gate = nn.Sequential(
            nn.Linear(self.irreps_in.num_irreps, hidden),
            nn.SiLU(),
            nn.Linear(hidden, self.irreps_out.num_irreps),
        )
        self.residual = self.irreps_in == self.irreps_out

    def forward(self, x):
        gates = torch.sigmoid(self.gate(copy_invariants(x, self.irreps_in)))
        y = self.value(x) * expand_per_copy(gates, self.irreps_out)
        return x + y if self.residual else y

@register_transform("equivariant_glu")
def _equivariant_glu(options, inputs, mapper):
    (features,) = inputs
    module = EquivariantGLU(global_irreps(features), **options)
    return module, [replace(features, irreps=str(module.irreps_out), width=None)]

Usage

- id: mlp_1
  transform: equivariant_glu
  inputs: [aggregated_1]
  outputs: [mixed_1]
  options: {irreps_out: 32x0e+16x1o+16x1e+8x2e+8x2o+8x3o+8x3e+4x4e}

o3.Linear maps input irreps that are absent from the output to nothing and output irreps that are absent from the input to zero; the factory reports the requested irreps_out as the output contract.

Variations

Stack two GLUs with a hidden layout wider than the input, as in transformer feed-forward blocks.
Add a bilinear term \(\mathrm{TP}(W_1\bm x,W_2\bm x)\) to the value path for self-interaction between degrees (the bilinear E3MLP variant does this with a residual).

Equivariant MLP block: equivariant_glu