Conceptual Framework

This document establishes a unified conceptual framework for equivariant atomistic machine-learning models.

We restrict the discussion to models whose features are organized as direct sums of spherical irreducible representations of \(SO(3)\) or \(O(3)\). Their mathematical foundations are developed in Equivariance.

We use a graph-based description of the physical system and the model. Atoms are nodes, selected ordered atomic pairs are directed edges, and features may be attached to nodes, edges, or the complete structure. This language makes message passing and periodicity explicit, but it does not require every model to be implemented as a conventional graph neural network. The graph notation is fixed in Atomistic Graph.

The framework is organized as a nested decomposition:

Model
├── Pipeline
└── Predictor
    ├── Physical Interface
    └── Backbone
        ├── Topology
        └── Geometry

1. Model = Predictor + Pipeline

A Model is the complete computation that maps its inputs to the output used for training or inference. Conceptually,

\[ \mathrm{Model} = \mathrm{Predictor} + \mathrm{Pipeline}. \]

Within this decomposition, the Predictor is the learnable part of the Model. The Pipeline comprises everything else required to turn one or more Predictor evaluations into the final result.

For a direct model, the Pipeline calls the Predictor once. For an iterative or generative model, it may repeatedly call the Predictor while updating a state, testing convergence, integrating a trajectory, or following a sampling schedule.

Examples:

Model family	Physical map represented by the Predictor	Pipeline
Direct prediction	Map the input directly to the desired output quantity.	Call the Predictor once.
DEQ	Map the current state \(y_t\) to the next state \(y_{t+1}=f_\theta(y_t,x)\).	Repeatedly update the state until reaching a fixed point.
Flow matching	Map an intermediate state to its velocity.	Integrate the Predictor-defined vector field.
Diffusion	Map the current noisy state to noise, score, or a denoised estimate.	Repeatedly denoise or sample along a noise schedule.

2. Predictor = Physical Interface + Backbone

A Predictor is a learnable mapping, usually a neural network,

\[ z_{\mathrm{out}} = f_\theta(z_{\mathrm{in}}), \]

where \(\theta\) denotes the learned parameters. Conceptually,

\[ \mathrm{Predictor} = \mathrm{Physical\ Interface} + \mathrm{Backbone}. \]

In the common case where the Physical Interface surrounds the Backbone,

\[ \mathrm{Predictor} = \mathrm{Decoder} \circ \mathrm{Backbone} \circ \mathrm{Encoder}. \]

Physical Interface

The Physical Interface connects task-specific physical quantities to generic node and edge irrep features.

It is usually realized primarily through an Encoder and a Decoder. The Encoder maps the physical input into generic node and edge irrep features, while the Decoder maps the resulting latent features to the target physical quantity.

Typical responsibilities include:

encoding atomic numbers as categorical node attributes and embedding basis identities and task conditioning;
converting a current physical state into node or edge features;
defining the output irreps required by the target;
mapping node features to atom-centered quantities;
mapping edge features to two-center operator blocks;
applying fixed basis transformations, masks, and physical reconstruction rules.

3. Backbone = Topology + Geometry

A Backbone is the main feature-learning component that operates on latent features.

We describe a Backbone through two relatively decoupled aspects:

\[ \mathrm{Backbone} = \mathrm{Topology} + \mathrm{Geometry}. \]

Topology determines how messages are passed, while Geometry determines the local equivariant operations applied at each step.

3.1 Topology: Generalized message passing

Topology describes graph-level message passing: where features live, which features communicate, and how messages are aggregated and used to update them.

At the highest level, we distinguish two families according to where the predicted physical quantity lives:

Node-targeted models construct features for quantities associated with individual atoms.
Edge-targeted models construct pair features for two-center quantities associated with atomic pairs.

For example, MACE is a node-targeted model whose Topology is organized around many-body message passing at each center atom.

Within both families, most models propagate information through the graph, so greater depth expands the spatial receptive field. Some models are instead strictly local: depth builds richer correlations inside a fixed local environment without expanding its spatial extent.

3.2 Geometry: Equivariant modules

Geometry describes what happens inside each local interaction or update step: how geometric information and irrep features are combined while preserving the required transformation laws. In a conventional graph network, it is the equivariant computation performed by each message-passing layer.

It includes:

distance and direction embeddings;
radial functions and spherical harmonics;
equivariant linear maps;
tensor products and angular-momentum coupling;
global-to-bond-frame rotations;
\(SO(2)\) mixing in local frames;
equivariant attention;
gates and spherical nonlinearities;
algebraic factorization and optimized execution of equivariant operations.

Representative geometric designs include:

NequIP: tensor-product convolution;
eSCN: bond-frame \(SO(2)\) computation;
Equiformer: equivariant attention and spherical nonlinearities.