Equivariant Message Passing for Tensorial Molecular Properties

The paper introduced PaiNN, a molecular message-passing architecture that maintains scalar and vector atom features so it can predict rotationally equivariant tensorial properties.

Metadata

Field	Value
Paper	Equivariant message passing for the prediction of tensorial properties and molecular spectra
Authors	Kristof T. Schütt, Oliver T. Unke, Michael Gastegger
Year	2021
Venue	ICML 2021
arXiv	2102.03150
PMLR PDF	v139/schutt21a
Status	full note started

One-Line Takeaway

PaiNN keeps both invariant scalar features and equivariant vector features for each atom, letting molecular message passing propagate directional information without using heavy tensor product machinery at every layer.

Question

Distance-based molecular models such as SchNet are naturally invariant to translation and rotation:

$$d_{ij}=\Vert r_i-r_j\Vert .$$

This is good for scalar energy prediction, but some molecular properties are vectors or tensors. Examples include dipole moments, polarizability, and direction-sensitive spectral quantities.

The architecture question is:

$$\text{How can a molecular GNN represent directional information while preserving E(3) symmetry?}$$

PaiNN answers by giving each atom two feature channels:

$$s_i\in {\mathbb{R}}^F,V_i\in {\mathbb{R}}^{3\times F}.$$

The scalar feature $s_i$ is invariant. The vector feature $V_i$ rotates with the molecule.

Main Claim

Rotationally equivariant atom-wise vector representations can improve molecular property prediction and enable tensorial property prediction.

The durable architecture pattern is:

$$\text{atom scalar state}+\text{atom vector state}+\text{equivariant message passing}\to \text{scalar and tensor molecular properties}.$$

Compared with distance-only message passing:

$$h_i\to h_i^{\prime },$$

PaiNN updates coupled scalar and vector states:

$$(s_i,V_i)\to (s_i^{\prime },V_i^{\prime }).$$

Architecture Contract

Item	Contract
Input	atomic numbers $z_i$ and coordinates $r_i\in {\mathbb{R}}^3$
Node unit	atom
Edge geometry	relative direction $r_{ij}=r_j-r_i$ and distance $\Vert r_{ij}\Vert$
Hidden state	scalar atom features $s_i$ and vector atom features $V_i$
Main operator	equivariant message passing with scalar-vector coupling
Output	scalar molecular properties, vector properties, tensorial properties, spectra-related quantities
Symmetry target	translation invariance, rotational invariance for scalar outputs, rotational equivariance for vector/tensor outputs
Domain	molecular property prediction and molecular spectra

The contract differs from invariant molecular GNNs:

$$\text{scalar-only hidden state}\Rightarrow \text{good for invariant scalars}$$ $$\text{scalar + vector hidden state}\Rightarrow \text{can carry directional information}.$$

E(3) Symmetry

Let a global Euclidean transformation be:

$$r_i^{\prime }=R{r}_i+t,R\in O(3),t\in {\mathbb{R}}^3.$$

Scalar atom features should stay unchanged:

$$s_i^{\prime }=s_i.$$

Vector atom features should rotate:

$$V_i^{\prime }=R{V}_i.$$

This is the central equivariance condition. If the molecule rotates, the internal vector features rotate in the same way. A scalar energy prediction should remain invariant:

$$E(\{R{r}_i+t\})=E(\{r_i\}),$$

while a vector property should transform equivariantly:

$$\mu (\{R{r}_i+t\})=R\mu (\{r_i\}).$$

Scalar and Vector Features

PaiNN represents each atom with:

$$s_i\in {\mathbb{R}}^F,V_i=\left[\begin{array}{cccc}{v}_{i,1}& v_{i,2}& \cdots & v_{i,F}\end{array}\right]\in {\mathbb{R}}^{3\times F}.$$

Each vector channel $v_{i,f}\in {\mathbb{R}}^3$ can rotate with the molecule. The scalar channel can store invariant chemical context, while the vector channel stores orientation-aware information.

This is less expressive than a full hierarchy of high-order irreducible representations, but it is easier to implement and efficient for atomistic prediction.

Message Block

For an edge $j\to i$, define:

$$r_{ij}=r_j-r_i,d_{ij}=\Vert r_{ij}\Vert ,{\hat{r}}_{ij}=\frac{r_{ij}}{d_{ij}}.$$

Radial functions use the distance:

$$W(d_{ij})={\mathrm{MLP}}_{\theta }\left(\mathrm{RBF}(d_{ij})\right).$$

A simplified scalar message has the form:

$$\Delta s_i=\sum _{j\in \mathcal{N}(i)}{\varphi }_s(s_j,d_{ij}).$$

The vector message can use both the neighbor vector feature and the normalized direction:

$$\Delta V_i=\sum _{j\in \mathcal{N}(i)}\left[{\varphi }_{vv}(s_j,d_{ij})\odot V_j+{\varphi }_{vs}(s_j,d_{ij})\odot {\hat{r}}_{ij}\right].$$

This equation captures the architecture idea without tying the note to one implementation detail:

• scalar gates are invariant because they depend on scalar features and distances;
• $V_j$ rotates equivariantly;
• ${\hat{r}}_{ij}$ rotates equivariantly;
• sums of equivariant vectors remain equivariant.

Therefore:

$$\Delta V_i^{\prime }=R\Delta V_i.$$

Update Block

The update block mixes scalar and vector channels while preserving symmetry.

A common pattern is to derive scalar information from vector norms or dot products:

$$\Vert V_i\Vert =\left(\Vert v_{i,1}\Vert ,\dots ,\Vert v_{i,F}\Vert \right),$$

then use scalar gates to modulate vector channels:

$$V_i^{\prime }=V_i+\gamma (s_i,\Vert V_i\Vert )\odot V_i.$$

Scalar updates can depend on invariant vector summaries:

$$s_i^{\prime }=s_i+\psi (s_i,\Vert V_i\Vert ).$$

The design rule is:

$$\text{scalar output}\leftarrow \text{invariant combinations},$$ $$\text{vector output}\leftarrow \text{scalar gates}\times \text{equivariant vectors}.$$

This lets the model exchange information between scalar and directional channels without breaking equivariance.

Relation To SchNet And DimeNet

Paper	Geometry Carrier	Main Limitation Addressed By PaiNN
SchNet	pairwise distances	distance-only invariant hidden states do not carry vector direction directly
DimeNet	directed messages, distances, angles	direction is encoded through angular invariant features rather than persistent vector atom states
PaiNN	scalar and vector atom features	direction is explicitly carried as equivariant vector channels
EGNN	coordinate updates	learns equivariant coordinate dynamics but has a different scalar-coordinate update contract
SE(3)-Transformer	equivariant attention with representation structure	more general equivariant machinery, usually heavier

PaiNN is useful to read after SchNet and DimeNet because it makes the hidden representation itself equivariant, not only the geometric input features.

Tensorial Properties

For a vector molecular property such as dipole moment, the output must rotate with the molecule:

$${\mu }^{\prime }=R\mu .$$

An atom-wise vector readout can be written as:

$$\mu =\sum _i{\mu }_i(V_i,s_i,r_i).$$

For rank-2 tensor properties such as polarizability, a simplified equivariant construction can combine invariant scalar coefficients with equivariant vector outer products:

$$T=\sum _i\sum _k{\lambda }_{ik}{u}_{ik}{u}_{ik}^{\top },$$

where $u_{ik}$ is an equivariant vector derived from $V_i$ or coordinate-dependent vector features, and ${\lambda }_{ik}$ is invariant.

If coordinates rotate:

$$u_{ik}^{\prime }=R{u}_{ik},$$

then:

$$T^{\prime }=\sum _i\sum _k{\lambda }_{ik}(R{u}_{ik})(R{u}_{ik}{)}^{\top }=RT{R}^{\top }.$$

This is why tensorial molecular prediction is a stronger test than scalar energy alone.

Evidence To Read

Evidence	What It Supports
QM9 molecular property results	scalar molecular property prediction benefits from equivariant vector features
MD17 energy and force results	equivariant message passing can support molecular dynamics-style targets
tensorial property prediction	vector channels enable direction-sensitive outputs
spectra simulation case	learned potentials/tensorial properties can accelerate expensive electronic-structure workflows
ablations on equivariant features	vector channels are not only architectural decoration

The paper reports that PaiNN improves common molecular benchmarks while reducing model size and inference time relative to some previous networks, and applies the architecture to molecular spectra simulation.

Why This Matters

PaiNN is a central architecture for reading modern molecular geometric models because it separates three ideas cleanly:

1. invariant scalar chemical context;
2. equivariant vector directional context;
3. symmetry-preserving coupling between them.

This makes it a practical middle point between simple invariant GNNs and more general equivariant tensor/representation architectures.

Limitations

Limitation	Why It Matters
Mostly vector features	higher-order tensor information may require richer irreducible representations or tensor channels
Local cutoff dependence	long-range electrostatics and global conformational effects may be missed without additional mechanisms
Benchmark sensitivity	molecular benchmarks can be sensitive to split, conformer, and target preprocessing
Force consistency	if forces are derived from energy gradients, the energy architecture and training target must be handled carefully
Geometry availability	the model assumes 3D coordinates are available or generated elsewhere
Domain scope	atomistic property prediction differs from protein-ligand docking or sequence-only protein modeling

For docking or protein-ligand modeling, PaiNN-style equivariant representations are useful building blocks, but the full task also needs conformer generation, pocket context, interaction labels, and domain-specific evaluation.

What To Remember

The core memory is:

$$\text{PaiNN}=\text{scalar atom features}+\text{vector atom features}+\text{equivariant message passing}.$$

Read it as the point where molecular message passing stops treating geometry only as invariant edge features and starts carrying direction as part of the hidden state.

Links

• Equivariance
• E(3)
• Equivariant feature
• Invariant feature
• Equivariant GNNs
• Coordinate modeling contract
• Molecular property prediction
• Molecular dynamics
• Geometry for Structure Modeling
• SchNet
• DimeNet
• E(n) Equivariant GNN
• SE(3)-Transformer