Equivariant nonlinearity: `squash_activation`

Role

Without nonlinearities a stack of equivariant layers collapses to a polynomial of fixed degree in the input. Elementwise activations, however, break equivariance for non-scalar irreps: \(\tanh\) applied to each component of a vector does not commute with rotations. Equivariant nonlinearities therefore act on vectors only through their norms, through invariant gates, or through functions on the sphere.

Built-in members of this kind: normact, gate_magnitudes, gate_scalars_mlp and s2act (nonlin_kind), the e3nn gate inside EquiConv, and irrep_activation as a stage.

The new layer

The squash function of capsule networks rescales each copy of an irrep according to its own norm:

\[\bm x_c\ \longmapsto\ \frac{\|\bm x_c\|}{\kappa_c^2+\|\bm x_c\|^2}\,\bm x_c, \qquad \kappa_c=\operatorname{softplus}(\theta_c)+10^{-3}.\]

The output norm \(\|\bm x_c\|^2/(\kappa_c^2+\|\bm x_c\|^2)\) grows quadratically for small inputs and saturates at one for large ones: the layer suppresses weak, noisy directions and caps large ones, with a learned knee \(\kappa_c\) per copy. Even scalars, which admit any pointwise function, use SiLU.

Symmetry.

Rotations. The scale factor depends on \(\|\bm x_c\|\), which is invariant, so the map multiplies each copy by an invariant (checklist 3c).
Inversion. An odd-parity copy changes sign, \(\bm x_c\to-\bm x_c\). The squash is an odd function of \(\bm x_c\), so its output changes sign too. This makes the same formula valid for pseudoscalars 0o, where a function like SiLU would be wrong (checklist 5).

Implementation

class SquashActivation(nn.Module):
    r"""``v -> v |v| / (kappa^2 + |v|^2)`` for every copy except even scalars.

    The map rescales each copy by a function of its norm, so it commutes with
    rotations; it is odd in ``v``, so it also commutes with inversion for odd
    parity copies. Even scalars use SiLU. ``kappa`` is learned per copy.
    """

    def __init__(self, irreps):
        super().__init__()
        self.irreps = o3.Irreps(irreps)
        self.raw_knee = nn.Parameter(torch.zeros(self.irreps.num_irreps))
        self.register_buffer(
            "even_scalar",
            torch.tensor(
                [ir.l == 0 and ir.p == 1 for mul, ir in self.irreps for _ in range(mul)]
            ),
        )

    def forward(self, x):
        norms = []
        for (mul, ir), sl in zip(self.irreps, self.irreps.slices()):
            block = x[..., sl].reshape(*x.shape[:-1], mul, ir.dim)
            norms.append(zero_safe_norm(block, dim=-1))
        norm = torch.cat(norms, dim=-1)
        knee = nn.functional.softplus(self.raw_knee) + 1e-3
        squash = norm / (knee.square() + norm.square())
        scale = expand_per_copy(squash, self.irreps)
        even = expand_per_copy(self.even_scalar.to(x.dtype), self.irreps).bool()
        return torch.where(even, nn.functional.silu(x), x * scale)

zero_safe_norm returns zero with a zero gradient at the origin, so the layer has well-defined gradients even for vanishing copies.

@register_transform("squash_activation")
def _squash_activation(options, inputs, mapper):
    (features,) = inputs
    module = SquashActivation(global_irreps(features), **options)
    return module, [features]

Usage

- id: activation_1
  transform: squash_activation
  inputs: [normalized_1]
  outputs: [node_1]

Variations

Use a gate that depends on all invariants of the entity (as gate_scalars_mlp does) instead of each copy’s own norm, which lets scalars switch vectors on and off.
Replace the rational squash by \(\tanh(\|\bm x\|/\kappa)\,\bm x/\|\bm x\|\) for a bounded, norm-preserving-at-small-scale variant.

Equivariant nonlinearity: squash_activation