MEAGraph: Multi-Kernel Edge Attention Autoencoder

• The paper shows that MEAGraph learns compact atomic embeddings through unsupervised multi-kernel edge attention and selective edge pooling.
• The methodology integrates geometric sensitivity with parallel kernel transformations to build latent graphs that reveal distinct clusters of atomic environments.
• Empirical results demonstrate that the approach effectively prunes redundant data while reducing ML force-field RMSE in niobium, tantalum, and iron datasets.

Searching arXiv for the specified paper and closely related context papers to ground the article. Multi-kernel Edge Attention-based Graph Autoencoder (MEAGraph) is an unsupervised graph neural network autoencoder introduced for atomic data mining in computational chemistry and materials science, with the specific aim of learning compact, information-rich embeddings of atomic environments and using those embeddings for clustering, outlier detection, and dataset pruning without labels such as energies or forces (Sun et al., 15 Sep 2025). In the reported formulation, MEAGraph combines multiple linear kernel transformations with attention-based message passing to capture geometric sensitivity, then progressively removes edges judged less informative for reconstruction of the original per-atom feature matrix. The resulting latent graph is used to identify disjoint subgraphs of similar environments, providing a mechanism for reducing sampling bias in machine-learning force-field training sets (Sun et al., 15 Sep 2025).

1. Definition, objective, and problem setting

MEAGraph is described as an unsupervised GNN autoencoder with a two-fold objective. First, it learns compact, information-rich embeddings of atomic environments by selectively retaining only the most informative edges via multi-kernel edge attention. Second, it uses the learned embeddings, together with the induced edge-reduced latent graph, for downstream tasks such as clustering, outlier detection, and dataset pruning, without ever using labels in the representation learning stage (Sun et al., 15 Sep 2025).

The motivating problem is the construction of chemically diverse datasets while avoiding sampling bias during force-field development. Common dataset generation techniques are stated to be prone to oversampling regions of the potential energy surface, and such regions may be difficult to identify, isolate, or align with human intuition. Traditional clustering and pruning approaches can be useful, but the paper states that they often lead to information loss or fail to properly identify distinct regions of the potential energy surface because of the high dimensionality of atomic descriptors (Sun et al., 15 Sep 2025).

In practical use, MEAGraph takes as input a large atomic-descriptor dataset, explicitly including SOAP, ACSF, or bispectrum descriptors, and builds a similarity graph of atoms. Training then seeks to determine which similarities between atomic environments must be preserved to reconstruct the original feature matrix faithfully. This framing makes the method a representation-learning and graph-structuring procedure rather than a supervised force-field model (Sun et al., 15 Sep 2025).

2. Input graph construction and atomic-environment representation

Let $N$ be the total number of atoms in a dataset or sub-dataset. Each atom $i$ is represented by a handcrafted or learned feature vector

$x_i \in \mathbb{R}^F,$

and the feature matrix is written as

$$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$

The initial graph is denoted $G_0 = (V, E_0)$, where the node set is $V = \{1,\ldots,N\}$ and the edge set is built from pairwise distances in feature space (Sun et al., 15 Sep 2025).

The graph-construction procedure computes Euclidean distances

$d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$

converts them to similarities

$s_{ij} = \exp(-d_{ij}),$

and then applies min-max normalization to obtain

$$\hat s_{ij} = \frac{s_{ij} - \min_{p,q} s_{pq}}{\max_{p,q} s_{pq} - \min_{p,q} s_{pq}}.$$

Edges are retained according to a threshold $r_l$:

$i$0

The edges are unweighted, but are accompanied by the raw distances $i$1, or optionally $i$2, as geometric descriptors (Sun et al., 15 Sep 2025).

This construction establishes the initial relational substrate on which attention-based message passing operates. A plausible implication is that the model begins from descriptor-space proximity rather than from a physically predefined neighborhood graph, which is consistent with its role as a general-purpose analyzer of high-dimensional atomic datasets rather than a fixed atomistic simulator.

3. Multi-kernel edge attention and geometric sensitivity

A defining feature of MEAGraph is the use of $i$3 parallel linear kernels on node features, and implicitly on edge features, to capture multiple geometric length scales. At layer $i$4 and kernel index $i$5, the hidden representation of node $i$6 is

$i$7

To incorporate distance-based or angle-based information, the formulation allows a scalar geometric attribute $i$8 to be precomputed for each edge, with a kernel-specific transform

$i$9

or, more generally, a small MLP. In the simplest instantiation reported, geometric information enters the attention mechanism through a per-kernel distance-scale parameter $x_i \in \mathbb{R}^F,$0 (Sun et al., 15 Sep 2025).

For each kernel and encoder layer, MEAGraph first computes a linear transform followed by batch normalization:

$x_i \in \mathbb{R}^F,$1

It then evaluates latent-space pairwise distances

$x_i \in \mathbb{R}^F,$2

and constructs unnormalized attention logits

$x_i \in \mathbb{R}^F,$3

Kernel-wise normalization over the neighborhood of node $x_i \in \mathbb{R}^F,$4 yields

$x_i \in \mathbb{R}^F,$5

The message from $x_i \in \mathbb{R}^F,$6 for kernel $x_i \in \mathbb{R}^F,$7 is then

$x_i \in \mathbb{R}^F,$8

and the node update is

$x_i \in \mathbb{R}^F,$9

After all $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$0 kernels are computed in parallel, the outputs are fused by element-wise averaging:

$$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$1

The attention weights are likewise averaged,

$$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$2

and then min-max-normalized across edges at layer $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$3 to obtain $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$4 (Sun et al., 15 Sep 2025).

The paper characterizes this mechanism as learning multi-scale geometric attention. This suggests that kernel diversity is intended to prevent the representation from collapsing onto a single distance scale in descriptor space, while edge attention converts latent similarity into an explicit criterion for edge preservation.

4. Encoder, decoder, and training objective

The encoder has $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$5 layers. It takes node features $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$6 and the initial edge set $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$7 as input. For each layer $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$8, it applies the $$X_0 = [x_1; \ldots; x_N] \in \mathbb{R}^{N \times F}.$$9 parallel attention sub-layers and then performs edge pooling according to

$G_0 = (V, E_0)$0

where $G_0 = (V, E_0)$1 is a user-specified pooling rate (Sun et al., 15 Sep 2025).

The decoder is described as symmetric and runs in reverse order, from $G_0 = (V, E_0)$2 down to $G_0 = (V, E_0)$3. Starting from the latent features $G_0 = (V, E_0)$4 on the pruned graph $G_0 = (V, E_0)$5, it uses a symmetric GNN layer, with the same weights or learned decoder weights, and re-adds edges from the earlier graph levels. The final reconstruction of per-node features is

$G_0 = (V, E_0)$6

so that $G_0 = (V, E_0)$7 (Sun et al., 15 Sep 2025).

The activation $G_0 = (V, E_0)$8 is typically ReLU or ELU. If $G_0 = (V, E_0)$9, then each hidden state $V = \{1,\ldots,N\}$0, with $V = \{1,\ldots,N\}$1 (Sun et al., 15 Sep 2025).

Training is purely autoencoding on node features. The reconstruction objective for a batch $V = \{1,\ldots,N\}$2 is the mean squared error

$V = \{1,\ldots,N\}$3

An optional $V = \{1,\ldots,N\}$4 regularization term may be added:

$V = \{1,\ldots,N\}$5

The paper explicitly states that no KL or sparsity penalty is used, and that regularization arises implicitly via edge-pooling (Sun et al., 15 Sep 2025).

The training procedure uses Adam with a typical learning rate $V = \{1,\ldots,N\}$6, a batch size of 1–8 subgraphs per batch, and 50–600 iterations. The edge pooling rate $V = \{1,\ldots,N\}$7 is sampled uniformly from $V = \{1,\ldots,N\}$8 per batch during training and fixed at inference to a dataset-specific value. Early stopping monitors validation $V = \{1,\ldots,N\}$9 and stops if no improvement is observed for 10 epochs (Sun et al., 15 Sep 2025).

5. Latent graph structure, clustering, and pruning

After training, the full dataset or large subgraphs are forwarded through the encoder with fixed $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$0, and the final pruned edge set $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$1 is collected. The encoder’s pooling is reported typically to disconnect the graph into several disjoint subgraphs, and each connected component is treated as a cluster of atomic environments (Sun et al., 15 Sep 2025).

This cluster structure is central to the model’s use in data mining. Because the representation-learning stage does not use energies or forces, cluster discovery is induced by the autoencoder’s effort to preserve those node-feature relationships that are most relevant to reconstruction. In the paper’s framing, the latent graphs therefore naturally decompose into clusters of similar environments suitable for post hoc analysis and reduction of redundancy (Sun et al., 15 Sep 2025).

The pruning rule operates within this cluster decomposition. Within each cluster $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$2, one may randomly remove up to a fraction $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$3 of atoms, or remove entire clusters whose size exceeds a threshold. The resulting reduced set is intended to retain diversity while removing redundant sampling (Sun et al., 15 Sep 2025).

A common misconception would be to treat MEAGraph as a direct predictor of force-field labels. The reported workflow does not do this: it uses label-free graph autoencoding to produce embeddings and latent connectivity, and then applies basic pruning techniques on the discovered clusters. The paper’s contribution is therefore methodological in representation learning and unsupervised dataset optimization rather than in supervised regression itself (Sun et al., 15 Sep 2025).

6. Empirical findings and reported behavior

The reported experiments cover niobium, tantalum, and iron datasets, using SOAP or bispectrum inputs depending on the case (Sun et al., 15 Sep 2025).

Dataset/task	Input	Reported finding
Nb dislocation clustering	SOAP	Latent clusters recover bulk vs. two dislocation-core environments in a split screw dislocation, matching DFT-force-based ground truth coloring
Ta dataset pruning	bispectrum	Up to 47% of atoms pruned yields a smaller training set whose ML force-field test RMSE is lower than the full dataset, approximately $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$4
Fe cluster-wise ablation	bispectrum	Approximately 66% of clusters can be removed fractionally without harming RMSE; a few small clusters sharply degrade performance upon removal

For niobium dislocation clustering, MEAGraph latent clusters are reported to recover bulk and two dislocation-core environments in a split screw dislocation, matching DFT-force-based “ground truth” coloring. In that same setting, all other unsupervised methods considered—kmeans, DBSCAN, spectral, and AP—are reported to fail to align with DFT force magnitudes (Sun et al., 15 Sep 2025).

For the tantalum dataset, MEAGraph-based clustering and pruning permit removal of up to 47% of atoms while producing a smaller training set whose ML force-field test RMSE is lower than that of the full dataset, with the paper reporting an improvement of approximately $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$5. Random pruning or DBSCAN-based pruning either increases RMSE or shows high variance, while affinity-propagation is described as more stable but without monotonic improvement (Sun et al., 15 Sep 2025).

For the iron dataset, cluster-wise ablation shows that approximately 66% of clusters can be removed fractionally without harming RMSE, while removal of a few small clusters sharply degrades performance. The paper interprets this as evidence that some small clusters have unique importance (Sun et al., 15 Sep 2025).

The hyperparameter study reports that $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$6 kernels and $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$7 layers suffice for all three elemental datasets, and that the pooling rate $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$8 is learned adaptively during training and then fixed per dataset at inference (Sun et al., 15 Sep 2025). This suggests that the architecture can remain relatively shallow while still producing a useful latent graph decomposition.

7. Position within unsupervised atomic data mining

MEAGraph is presented as a lightweight, label-free graph autoencoder that learns multi-scale geometric attention, produces a compact latent graph whose connected components are natural clusters of atomic environments, and delivers superior clustering and dataset-pruning performance compared with classical unsupervised methods in the reported studies (Sun et al., 15 Sep 2025). Within the scope of the paper, its principal significance lies in coupling representation learning, edge selection, and graph decomposition into a single unsupervised pipeline for high-dimensional atomic-descriptor datasets.

The method also makes explicit a distinction between feature reconstruction and label supervision. Because the loss is defined on reconstruction of $d_{ij} = \|x_i - x_j\|_2,\qquad i \neq j,$9 rather than on energies or forces, any downstream improvement in force-field training is mediated through the quality of the learned clustering and pruning strategy. A plausible implication is that MEAGraph is most naturally interpreted as a data-centric pre-processing or dataset-analysis method for machine-learning force fields.

The reported formulations and pseudocode are stated to be sufficient to reproduce or extend MEAGraph for other high-dimensional atomic-descriptor datasets (Sun et al., 15 Sep 2025). Within that framing, the model occupies a specific role: unsupervised atomic data mining for representation learning, clustering, outlier detection, and dataset optimization under conditions where oversampling and descriptor-space complexity make conventional pruning strategies unreliable.