Tensor-product convolution: `depthwise_convolution`

Role

The convolution creates messages: it couples the state of a neighbour with the geometry of the bond through a Clebsch–Gordan tensor product, weighted by a learned function of distance. It is the only place where angular momentum is created from geometry during message passing, and it dominates the cost of the network.

Built-in members of this kind: EquiConv inside the MANDALA update blocks (separate-weight or fully connected product, gate, radial channel weights), radial_tensor_product (factorized product, gate, radial scaling) and the neighbour-density product of neighbor_density.

The new layer

The depthwise (“uvu”) convolution couples every channel \(u\) of the sender with the spherical harmonics, keeping channels separate, and generates one weight per channel and per coupling path from the distance:

\[\big(\bm m_{ij}\big)^{L}_{u}=\sum_{\ell_1,\ell_2} w_{u}^{(\ell_1\ell_2\to L)}\!\big(\bm\varphi(r_{ij})\big)\, \Big(\bm h_{i,u}^{\ell_1}\otimes Y^{\ell_2}(\hat{\bm r}_{ij})\Big)^{L}, \qquad L\le L_{\max}.\]

This is the continuous-filter convolution of NequIP and MACE. Compared with the MANDALA default it trades the dense channel mixing of the separate-weight product for a fully distance-dependent filter on every path, at a cost that grows only linearly in the multiplicity. Channel mixing is then delegated to a cheap equivariant linear layer (here: Equivariant MLP block: equivariant_glu).

The output irreps are assembled from all reachable paths and sorted; the instruction list is re-indexed with the sorting permutation so that each path writes into its sorted slot.

Symmetry. The weights are functions of \(r\) (invariant), and the coupling is a Clebsch–Gordan product (checklist 3b). Sender features are gathered with edge_index[0]; messages live on the edge support.

Implementation

class DepthwiseConvolution(nn.Module):
    r"""``m_ij = sum_paths w_path(phi(r_ij)) * (h_j (x) Y(r_hat_ij))_path``.

    Every sender channel couples only with itself ("uvu" instructions), so the
    cost grows linearly with the multiplicity. The radial network emits one
    weight per channel and coupling path, i.e. a distance-dependent filter.
    """

    def __init__(self, node_irreps, angular_irreps, radial_dim, lmax=2, hidden=32):
        super().__init__()
        node_irreps, angular_irreps = o3.Irreps(node_irreps), o3.Irreps(angular_irreps)
        outputs, instructions = [], []
        for i, (mul, ir_in) in enumerate(node_irreps):
            for j, (_, ir_sh) in enumerate(angular_irreps):
                for ir_out in ir_in * ir_sh:
                    if ir_out.l <= lmax:
                        instructions.append((i, j, len(outputs), "uvu", True))
                        outputs.append((mul, ir_out))
        irreps_mid = o3.Irreps(outputs)
        irreps_out, permutation, _ = irreps_mid.sort()
        instructions = [
            (i, j, permutation[k], mode, train) for i, j, k, mode, train in instructions
        ]
        self.tp = o3.TensorProduct(
            node_irreps,
            angular_irreps,
            irreps_out,
            instructions,
            internal_weights=False,
            shared_weights=False,
        )
        self.radial = nn.Sequential(
            nn.Linear(radial_dim, hidden),
            nn.SiLU(),
            nn.Linear(hidden, self.tp.weight_numel),
        )
        self.irreps_out = irreps_out

    def forward(self, node, angular, radial, edge_index):
        sender = edge_index[0]
        return self.tp(node[sender], angular, self.radial(radial))

internal_weights=False, shared_weights=False makes e3nn accept one weight vector per edge; self.tp.weight_numel is the number of weights the radial network must emit.

@register_transform("depthwise_convolution")
def _depthwise_convolution(options, inputs, mapper):
    node, angular, radial, edge_index = inputs
    require_index(edge_index, "relation")
    support = require_same_support(angular, radial, edge_index)
    radial_irreps = global_irreps(radial, "edge")
    if not invariant_scalars(radial_irreps):
        raise ValueError("Radial features must be invariant scalars")
    module = DepthwiseConvolution(
        global_irreps(node, "node"),
        global_irreps(angular, "edge"),
        radial_irreps.dim,
        **options,
    )
    return module, [
        FeatureSpec("edge", "irreps", str(module.irreps_out), support=support)
    ]

Usage

- id: convolution_1
  role: message
  transform: depthwise_convolution
  inputs: [node_0, angular, radial, edge_index]
  outputs: [messages_1]
  options: {lmax: 4}

With node_0 = 32x0e and harmonics up to \(\ell=4\) this produces 32x0e+32x1o+32x2e+32x3o+32x4e; in the second layer, with a node state containing both parities, all parities up to \(L_{\max}\) appear.

Variations

Feed the receiver state into the radial network as well, making the filter depend on the chemical environment of both atoms.
Use the message support (message_edge_index, message_vectors, message_lengths) to pass messages on a different, e.g. shorter-range, graph than the one on which operators are predicted.

Tensor-product convolution: depthwise_convolution