Model Architecture

Graph Representation of Migration Hops

Each lithium migration hop is represented as a graph constructed from the atomic supercell.

Graph construction:

• Nodes: All atoms in the supercell.

• Special node: A centroid atom labeled “X” marks the diffusion bottleneck.

• Edges: Atoms connected within a radial cutoff distance.

• Node features:

  • Atomic number

  • Distance to centroid

• Edge features:

  • Interatomic distance
    
    

This representation allows the model to focus on local geometric environments surrounding the migration pathway.

Geometry-Aware Graph Neural Network

We implemented a message-passing graph neural network (GNN) using PyTorch Geometric.

Architecture components:

1. Edge-conditioned convolution layers (NNConv)

2. Nonlinear activation (ReLU)

3. Global pooling (mean aggregation)

4. Fully connected regression head

5. Scalar output representing migration barrier

The model learns structural interactions directly from atomic geometry without manual feature engineering.

Output:

Predicted migration barrier in electron volts (eV).