Writing new layers and objectives

An E3GNN for operators is assembled from a small number of layer kinds. This section implements a new member of every kind and combines all of them into a complete, trainable network. Each layer is a self-contained example of the extension mechanism; together they form the layer cookbook in examples/layer_cookbook:

Kind	New layer	Idea
Radial basis: bessel_basis	`bessel_basis`	Orthogonal Bessel functions with a smooth polynomial cutoff
Node encoder: periodic_table_embedding	`periodic_table_embedding`	Species embeddings that share rows and columns of the periodic table
Edge encoder: radial_angular_encoder	`radial_angular_encoder`	Radially modulated spherical harmonics per output copy
Tensor-product convolution: depthwise_convolution	`depthwise_convolution`	“uvu” tensor product with distance-dependent path weights
Message aggregation: tempered_softmax_aggregation	`tempered_softmax_aggregation`	Attention with a learned softmax temperature
Equivariant nonlinearity: squash_activation	`squash_activation`	Capsule-style norm squashing with a learned knee
Equivariant normalization: degree_rms_norm	`degree_rms_norm`	Per-entity RMS normalization of each irrep term
Equivariant MLP block: equivariant_glu	`equivariant_glu`	Gated linear unit with invariant gates
Residual update: rezero_residual	`rezero_residual`	Learned per-copy skip weights initialized at zero
Edge update: endpoint_bilinear_edge_update	`endpoint_bilinear_edge_update`	Bilinear coupling of both endpoint states
Operator readout: pair_mixture_readout	`pair_mixture_readout`	Low-rank mixture of shared readouts conditioned on the atom pair
Output envelope: learned_pair_envelope	`learned_pair_envelope`	Learned Slater-type radial prior per atom pair

Assembling a network from new layers builds the complete network from these layers, verifies its equivariance on a periodic structure with s, p and d orbitals, and trains it on the bundled silicon data.

Adding an observable objective extends the objective instead: define a differentiable scalar from the predicted matrices, register its loss, and enable it in YAML.

The extension mechanism

A layer consists of a module and a factory. The factory is registered under a name that recipes use in transform:.

import torch
from torch import nn
from e3nn import o3
from net.pipeline import FeatureSpec, register_transform

class ScaleByLength(nn.Module):
    """Multiply every component of an edge feature by a learned f(r)."""

    def __init__(self, hidden=16):
        super().__init__()
        self.f = nn.Sequential(nn.Linear(1, hidden), nn.SiLU(), nn.Linear(hidden, 1))

    def forward(self, values, lengths):
        return values * self.f(lengths.reshape(-1, 1))

@register_transform("scale_by_length")
def build(options, inputs, mapper):
    values, lengths = inputs                       # positional FeatureSpecs
    if values.domain != "edge" or values.irreps is None:
        raise ValueError("scale_by_length expects irrep-valued edge features")
    if lengths.dim != 1 or lengths.support != values.support:
        raise ValueError("expects one length per edge of the same support")
    return ScaleByLength(**options), [values]      # output has the input's contract

The factory is the place to encode the layer’s contract:

which domain each input must have (node, edge, relation…);
which inputs must be invariant (only 0e channels), e.g. everything fed to an MLP or used as a weight;
that edge-like inputs share the same support;
the exact irreps of every output.

Once registered, the layer is available to every recipe:

- id: damp_edges
  transform: scale_by_length
  inputs: [edge_1, edge_lengths]
  outputs: [edge_1_damped]
  options: {hidden: 16}

Registration happens when the defining module is imported. The command-line trainer becomes extension-aware with a one-line import (examples/layer_cookbook/train.py):

from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).resolve().parent))
import layers  # noqa: E402,F401  (registers the cookbook transformations)

from scripts.train import main  # noqa: E402

if __name__ == "__main__":
    main()

The symmetry checklist

A layer is E(3)-equivariant if each of the following holds; every cookbook page explains how its layer satisfies them.

Checklist

Geometry enters only through relative vectors (edge_vectors, edge_lengths, displacements), never through absolute positions.
Learnable functions act only on invariants. MLPs, softmaxes, gates and embeddings may read 0e channels, norms of copies and squares of 0o channels — nothing else.
Non-scalar copies are only (a) mixed linearly with copies of the same irrep, (b) coupled through Clebsch–Gordan tensor products, (c) multiplied by invariant numbers, or (d) gathered/summed over graph indices.
Biases only on even scalars. An additive constant on a 1o or a 0o copy breaks rotation or inversion symmetry.
Pointwise nonlinearities only on even scalars; odd scalars admit odd functions only.
Normalization statistics are invariant (norms, means of norms), and means are subtracted only from even scalars.
Undefined directions are zero. Onsite edges have no direction; any quantity derived from it must vanish (non-scalar harmonics) or be finite by construction (e.g. \(\sin x/x\)).

Certifying a layer

The cookbook’s verification script checks each layer numerically with the identity \(f(D(g)\,\bm x)=D(g)\,f(\bm x)\) for a random rotation, an improper rotation and the inversion. The helper is short enough to reuse for any new layer:

def assert_covariant(name, apply, inputs_and_irreps, output_irreps, atol=1e-9):
    """``apply(*inputs)`` must commute with the action of every group element."""
    for label, g in group_elements().items():
        transformed = [
            x @ o3.Irreps(irreps).D_from_matrix(g).T if irreps is not None else x
            for x, irreps in inputs_and_irreps
        ]
        expected = apply(*(x for x, _ in inputs_and_irreps))
        expected = expected @ o3.Irreps(output_irreps).D_from_matrix(g).T
        observed = apply(*transformed)
        error = (observed - expected).abs().max().item()
        if error > atol:
            raise AssertionError(f"{name}: {label} error {error:.2e}")
    print(f"  {name:<34s} covariant")

D(g) is o3.Irreps(irreps).D_from_matrix(g); inputs that are invariant (irreps=None) are left unchanged. Perturb the parameters away from their initial values before testing: zero-initialized weights hide many mistakes. The script also contains a model-level test that rotates a whole periodic structure; Assembling a network from new layers shows it at work.

Running the certificate:

PYTHONPATH=src:. python examples/layer_cookbook/verify.py

Layer covariance
  radial_angular_encoder             covariant
  depthwise_convolution              covariant
  tempered_softmax_aggregation       covariant
  squash_activation                  covariant
  degree_rms_norm                    covariant
  equivariant_glu                    covariant
  rezero_residual                    covariant
  endpoint_bilinear_edge_update      covariant
  pair_mixture_readout               covariant
  learned_pair_envelope              covariant
  bessel_basis                       invariant, finite at r = 0
  periodic_table_embedding           invariant by construction
Model covariance
  initial  rotation           max |dH| = 2.9e-10
  initial  improper rotation  max |dH| = 2.9e-10
  initial  inversion          max |dH| = 0.0e+00
  trained  rotation           max |dH| = 2.1e-10
  trained  improper rotation  max |dH| = 2.1e-10
  trained  inversion          max |dH| = 0.0e+00
  parameters: 289,311
New observable objective
  hamiltonian_trace: configured weight and nonzero finite gradients verified
All cookbook layers and the cookbook model are E(3)-equivariant.

Replacing SquashActivation.forward with an elementwise tanh — a classic mistake — makes the model test fail immediately with a rotation error of order \(10^{-3}\), five orders of magnitude above the tolerance.