Self-Supervised Learning for Ordered Three-Dimensional Structures

TL;DR

This work forms a set of geometric tasks suitable for the large-scale study of ordered three-dimensional structures, without requiring any human intervention in data labeling, and builds deep rotation- and permutation-equivariant neural networks based on geometric algebra.

Abstract

Recent work has proven that training large language models with self-supervised tasks and fine-tuning these models to complete new tasks in a transfer learning setting is a powerful idea, enabling the creation of models with many parameters, even with little labeled data; however, the number of domains that have harnessed these advancements has been limited. In this work, we formulate a set of geometric tasks suitable for the large-scale study of ordered three-dimensional structures, without requiring any human intervention in data labeling. We build deep rotation- and permutation-equivariant neural networks based on geometric algebra and use them to solve these tasks on both idealized and simulated three-dimensional structures. Quantifying order in complex-structured assemblies remains a long-standing challenge in materials physics; these models can elucidate the behavior of real self-assembling systems in a variety of ways, from distilling insights from learned tasks without further modification to solving new tasks with smaller amounts of labeled data via transfer learning.

Figures (5)

• Figure 1: Two types of ordered structural data used in this work. (a) Periodic three-dimensional systems can be represented by their basic repeat unit, or unit cell. (b) Unit cells can be tiled to fill space; however, non-periodic systems including liquids and (c) quasicrystals have surfaces and may exhibit nontrivial symmetries. (d) Ordered structures often have different numbers of neighbors at different radial distances. We select the 20 nearest neighbors of each particle to feed into our deep learning models for the tasks presented here. In this diagram, we show three possible permutations of the 10 nearest neighbors around a particle in a CsCl-type crystal structure.
• Figure 2: Architectures for equivariant, self-supervised transfer learning tasks. (a) The core shared by all architectures consists of a series of learned, rotation- and permutation-equivariant Geometric Algebra Attention (GAlA) blocks propagating one rotation-invariant and one rotation-equivariant signal for each input point. Architectures used in this work stack $N=3$ sets of layers inside the core. (b) The head layers for the autoencoding task produce two vectors using geometric algebra attention, which is converted into a bottleneck value specifying an orientation of the point cloud via SVD and a cross product. The rotation-invariant values are passed through a standard variational autoencoder kingma_auto-encoding_2014. Finally, geometric algebra attention between the per-input embeddings and the orientation vectors generates the output value. (c) The architecture for denoising simply adds a layer of rotation-equivariant, permutation-equivariant attention. For shift identification and nearest bond regression, we instead use a permutation-invariant final attention layer. (d) For classification tasks (frame classification and noisy bonds), a simple MLP is applied to generate logits for classification. Frame classification uses a final permutation-invariant reduction over the point cloud before the MLP, while the noisy bond architecture uses permutation-equivariant attention.
• Figure 3: (a) Two-dimensional projection of embeddings for a set of self-assembled structures using a self-supervised model trained on bulk unit cell data. The generated embeddings hold useful information about three-dimensional structure, as they show distinct signatures for the different structures and have a reasonable distribution even for crystalline surfaces and the icosahedral quasicrystal, which were not included in the training set: a capability which is necessary for models to generalize well for out-of-distribution observations. (b) Representations learned by self-supervised models can perform better than classical methods: here, embeddings learned for a frame classification model on solids can distinguish between pre-crystallization liquids significantly better than the Steinhardt $Q$ order parameters are capable of.
• Figure 4: Transfer learning performance of rotation-equivariant deep neural networks. Models are first trained on all idealized crystal structure unit cell data available for a source task, then fine-tuned on a limited amount of data for a target task. Blue lines indicate performance for cases where the "core" weights of the network are not updated during retraining; red lines indicate performance for updating all weights, and black lines indicate direct performance of pure (non-transfer learning) models with limited training data. For ease of comparison to transfer learning results, results from direct training on limited data are reproduced in off-diagonal plots as dashed lines. Gray-shaded areas indicate one standard error of the mean over 10 independent training replicas. Data are available in numerical form in Appendix \ref{['app:xfer_tables']}.
• Figure 5: (a) Histogram of predicted frames over a simulated cooling trajectory of particles forming the $cI2$-W ("body-centered cubic") crystal structure. Models trained on simple self-supervised tasks can help automate the application of human knowledge, such as identifying distinct gas, liquid, and solid phases here. (b--d) System snapshots and bond-orientational order diagrams dzugutov_formation_1993 (BOODs) showing the development of order for the gas (b), liquid (c), and solid (d) phases at the times indicated by red arrows. Peaks in BOODs indicate global alignment of bonds and can serve as a helpful fingerprint to identify crystallization.