Message aggregation: `tempered_softmax_aggregation`

Role

Aggregation reduces the messages arriving at an atom into one equivariant vector. It must be invariant to the order of neighbours and must handle a variable number of them. The paper’s silicon ablation shows that the choice matters: replacing averaging by attention lowered the validation Hamiltonian MAE by 10 %.

Built-in members of this kind: sum, average and multi-head attention in the MANDALA node update (node_update_message_agg), indexed_aggregation and projected_attention as separate stages.

The new layer

Tempered softmax aggregation computes one invariant score per message and normalizes it over the neighbourhood with a learned temperature:

\[s_{ij}=\mathrm{MLP}\big(\bm\varphi(r_{ij})\oplus\iota(\bm m_{ij})\big),\qquad \alpha_{ij}=\frac{\exp(s_{ij}/\tau)}{\sum_{k\in\mathcal N(j)}\exp(s_{kj}/\tau)},\qquad \bm h_j=\sum_{i\in\mathcal N(j)}\alpha_{ij}\,\bm m_{ij},\]

where \(\iota(\bm m)\) collects one invariant per message copy (the value of 0e copies, the square of 0o copies and the norm of all others). The temperature interpolates between familiar reductions:

\[\tau\to\infty:\ \alpha_{ij}\to\frac1{|\mathcal N(j)|}\ \ (\text{mean}),\qquad \tau\to0:\ \alpha_{ij}\to\delta_{i,\arg\max_k s_{kj}}\ \ (\text{hard selection}).\]

Learning \(\log\tau\) lets each layer choose how selective it is: an early layer may average over a coordination shell, a later one may focus on the single most relevant bond.

The maximum score per receiver is subtracted before exponentiation, which keeps the softmax numerically stable for any number of neighbours.

Symmetry. Scores are functions of invariants only (checklist 2), so each \(\alpha_{ij}\) is invariant; a sum of equivariant messages weighted by invariants is equivariant (3c, 3d). The softmax runs over the incoming edges of each receiver, so it is also permutation invariant.

Implementation

def copy_invariants(x: torch.Tensor, irreps: o3.Irreps) -> torch.Tensor:
    """One rotation- and inversion-invariant number per irrep copy.

    Even scalars are invariant as they are. Odd scalars change sign under
    inversion, so their square is used. Every other copy contributes its norm.
    """
    values = []
    for (mul, ir), sl in zip(irreps, irreps.slices()):
        block = x[..., sl].reshape(*x.shape[:-1], mul, ir.dim)
        if ir.l == 0 and ir.p == 1:
            values.append(block[..., 0])
        elif ir.l == 0:
            values.append(block[..., 0].square())
        else:
            values.append(zero_safe_norm(block, dim=-1))
    return torch.cat(values, dim=-1)

class TemperedSoftmaxAggregation(nn.Module):
    r"""``h_i = sum_j alpha_ij m_ij``, ``alpha_ij = softmax_j(s_ij / tau)``.

    Scores ``s_ij`` are computed from invariants only (radial features and the
    norm of each message copy), so the weights are rotation invariant and the
    weighted sum stays equivariant. ``tau`` is learned; ``tau -> infinity``
    recovers mean aggregation, ``tau -> 0`` selects the strongest neighbor.
    """

    def __init__(self, message_irreps, radial_dim, hidden=32, temperature=1.0):
        super().__init__()
        self.irreps = o3.Irreps(message_irreps)
        self.score = nn.Sequential(
            nn.Linear(radial_dim + self.irreps.num_irreps, hidden),
            nn.SiLU(),
            nn.Linear(hidden, 1),
        )
        self.log_temperature = nn.Parameter(torch.tensor(math.log(temperature)))

    def forward(self, messages, radial, edge_index, node_reference):
        receiver, count = edge_index[1], len(node_reference)
        features = torch.cat([radial, copy_invariants(messages, self.irreps)], dim=-1)
        scores = self.score(features).squeeze(-1) / self.log_temperature.exp()
        scores = (
            scores - scatter(scores, receiver, dim_size=count, reduce="max")[receiver]
        )
        weights = scores.exp()
        weights = weights / scatter(weights, receiver, dim_size=count)[
            receiver
        ].clamp_min(1e-12)
        return scatter(weights[:, None] * messages, receiver, dim=0, dim_size=count)

The fourth input, node_reference, is any node feature; it fixes the number of receivers (atoms without incoming edges receive zeros) and the support of the output.

@register_transform("tempered_softmax_aggregation")
def _tempered_softmax_aggregation(options, inputs, mapper):
    messages, radial, edge_index, node_reference = inputs
    require_index(edge_index, "relation")
    require_same_support(messages, radial, edge_index)
    radial_irreps = global_irreps(radial, "edge")
    if not invariant_scalars(radial_irreps):
        raise ValueError("Attention scores must use invariant radial features")
    global_irreps(node_reference, "node")
    module = TemperedSoftmaxAggregation(
        global_irreps(messages, "edge"), radial_irreps.dim, **options
    )
    return module, [replace(messages, domain="node", support=node_reference.support)]

Usage

- id: aggregation_1
  role: message
  transform: tempered_softmax_aggregation
  inputs: [messages_1, radial, edge_index, node_0]
  outputs: [aggregated_1]
  options: {hidden: 32, temperature: 1.0}

Variations

Multiple heads with separate temperatures, concatenated or averaged.
Add the receiver’s own invariants to the score (a query, as in the built-in attention) to make the selection depend on the atom being updated.

Message aggregation: tempered_softmax_aggregation