Neural network embeddings based similarity search method for atomistic systems
Yilin Yang, Mingjie Liu and John R. Kitchin, “Neural network embeddings based similarity search method for atomistic systems”, Digital Discovery 1(5), 636–644 (2022). https://doi.org/10.1039/d2dd00055e
This summary was written by a language model from the paper and has not been checked line by line. Rely on the published paper for details.
Question
Large atomistic databases such as the Materials Project, PubChem and Open Catalyst 2020 (OC20) invite searches for structures that resemble a target with desirable properties. Molecular similarity search is well developed (string fingerprints like SMIfp, substructure keys like MACCS and BCI), but those fingerprints rely on predefined bonding patterns and cannot handle periodic systems such as adsorbates on surfaces, where bonding is not well defined. Periodic descriptors like ACSF and SOAP exist but scale poorly with the number of elements. The paper asks whether the learned atom embeddings of a graph neural network can serve as a general fingerprint for similarity search across molecules, bulk crystals and adsorption systems (Yang et al. 2022).
Approach
The workflow has a preparation stage and a query stage. A modified GemNet-dT is trained on the energies of a database; the modification keeps only the atomic embedding from the last interaction block (the vector used to predict each atomic energy) rather than combining all blocks, which limits dimensionality and cost at the price of some fine detail. The trained model computes an embedding for every atom, and these are indexed with the FAISS library using its IndexIVFPQ approximate nearest neighbor (ANN) method, which combines k-means inverted files with product quantization, using L2 distance. At query time the same model embeds each atom of the query structure and the index returns its k nearest atom embeddings. For whole-structure search, candidate structures are ranked by how many of their atoms were matched, with the summed embedding distances breaking ties; large k favors globally similar structures and small k favors strong local matches (k was about ten times the number n of results wanted).
Three datasets were used: QM9 (about 134k small organic molecules), the Materials Project (more than 126k bulk crystals) and the OC20 IS2RS relaxed adsorption configurations (about 460k for training). Separate GemNet models were trained on each. For QM9 the model was trained on raw rather than element-referenced energies, to push elemental information into the embeddings (Yang et al. 2022).
Key findings
On QM9 the energy MAEs were 4.57 eV (training) and 4.97 eV (validation), much worse than published benchmarks because raw energies were used. Even so, a benzene query returned five molecules that all contained a six-membered ring, sometimes with one nitrogen substituted or with small hydroxyl or amino groups attached. A Morgan fingerprint and Tanimoto similarity baseline instead returned some molecules with larger rings. As a test of energetic relevance, Gaussian process regression models (FLARE) were trained on retrieved molecules and used to predict the benzene energy: with molecules from the embedding search the error fell to 0.04 eV (standard deviation 0.02 eV), versus 3.64 eV (17.14 eV) for Morgan fingerprint neighbors and 5.76 eV (43.44 eV) for random QM9 molecules with up to 15 configurations. A joint substructure search for a hydroxyl and an amino group, combining results like a logical and, retrieved a molecule containing both in similar chain positions.
For the Materials Project the MAEs were 0.62 eV and 1.42 eV, a gap the authors attribute to extrapolation in the validation set. A query oxygen atom in an Al2Cu3O6 cell retrieved only oxygen atoms with similar Al–O–Cu geometry at embedding distances of 0.06 to 0.11, with the seventh hit jumping to 0.17 and looking less similar, even though the search pool was never restricted to oxygen.
For OC20 the MAEs were 0.82 eV and 0.92 eV. A query O atom on a tilted hollow site of two Pd and one Ag returned O atoms on hollow sites with two Pd atoms, and their p-projected densities of states had cosine similarities above 0.6 to the query, higher than four randomly chosen adsorbed O atoms. An acetylene query on a three-Cu hollow site (searching on the two carbon embeddings only) returned acetylene adsorbates, the top two on Cu3 hollows, with carbon p-DOS cosine similarities above 0.65, again higher than random examples (Yang et al. 2022).
Significance
The work shows that embeddings learned while fitting energies act as property-aware fingerprints that encode element identity, local geometry and, indirectly, electronic structure, and that they apply equally to molecules and periodic solids. Paired with ANN indexing, this gives a fast way to query very large catalyst databases, for example to find cheaper materials with similar adsorption environments. The authors stress that the framework is modular: any descriptor or deep learning model and any vector search engine could replace GemNet and FAISS, and the similarity found depends on the property the network was trained on. Data and indexes are archived at https://doi.org/10.1184/R1/19968323 (Yang et al. 2022).