Rotation-equivariant Graph Neural Networks for Learning Glassy Liquids Representations

TL;DR

This work tackles the challenge of connecting static structure to dynamics in glassy liquids by enforcing SE(3) roto-translation equivariance in a Graph Neural Network. The authors design an SE(3)-equivariant GNN using spherical harmonics, Clebsch–Gordan tensor products, and a radial- SH kernel to build multi-layer, steerable representations of local structure, paired with a decoder that predicts particle mobility across multiple timescales. They systematically study the impact of input choices (thermal vs quenched IS positions, local potential energy) and network depth, demonstrating superior performance with fewer parameters and strong generalization across temperatures, culminating in evidence that the learned representation acts as a robust structural descriptor or order parameter. The approach also delivers interpretable insights, linking early layers to local density fields and enabling transfer-learning analyses that reveal the stability and transferability of the learned structural representation across state points.

Abstract

The difficult problem of relating the static structure of glassy liquids and their dynamics is a good target for Machine Learning, an approach which excels at finding complex patterns hidden in data. Indeed, this approach is currently a hot topic in the glassy liquids community, where the state of the art consists in Graph Neural Networks (GNNs), which have great expressive power but are heavy models and lack interpretability. Inspired by recent advances in the field of Machine Learning group-equivariant representations, we build a GNN that learns a robust representation of the glass' static structure by constraining it to preserve the roto-translation (SE(3)) equivariance. We show that this constraint significantly improves the predictive power at comparable or reduced number of parameters but most importantly, improves the ability to generalize to unseen temperatures. While remaining a Deep network, our model has improved interpretability compared to other GNNs, as the action of our basic convolution layer relates directly to well-known rotation-invariant expert features. Through transfer-learning experiments displaying unprecedented performance, we demonstrate that our network learns a robust representation, which allows us to push forward the idea of a learned structural order parameter for glasses.

Figures (15)

• Figure 1: Input Graph with its input features. Node features are the one-hot encoded particle types (invariant features, $l=0$), and edge attributes $\mathbf{a_{ij}}$ are split: the direction is embedded in Spherical Harmonics $Y(\mathbf{\hat{a}_{ij}})$ and the norm is retained separately. Throughout this paper, we depict each rotational order with a given color: $l=0$ (red), $l=1$ (green), $l=2$ (blue). The relative length of each is a reminder that each requires $2l+1$ real values to be stored on the machine.
• Figure 2: Equivariance to 2D rotations. Simple case in which the input and output fields have the same dimension, $2$. $\pi(r)$ represents the action of the rotation operator on the input field, $\pi'(r)$ on the output one. In general they can be different, here since input and output are in the same space, they are equal. The mapping $\mathcal{K}$ acts in an equivariant way, indeed it commutes with the rotation. In practice, $\mathcal{K}$ corresponds to the action of one of our neural network's layers. It represents a mapping from one internal representation, later denoted by $\mathbf{h}$ to the updated one $\mathbf{h}'$. Each representation is a vector field that spans the 3D simulation box and consists of 3D vectors (one such field per channel). Here we depict a more generic and readable case.
• Figure 3: Overview of the convolution layer, summarizing Eqs. (\ref{['eq:convolution']},\ref{['eq:self_int']}). For each neighboring node, the node and edge features are combined (with C-G product) and multiplied by the learned radial filter $\varphi$. Before performing this operation, the one-hot encoded particle type is concatenated to $\mathbf{h_i}$ by adding 2 $l=0$ channels (not shown, for simplicity). Because multiple triplets come out of the C-G product, we obtain a much larger representation (left part of inset). This intermediate representation is narrowed down using a linear layer (one for each $l_O$ and each channel).
• Figure 4: Overall Architecture. Top: node and edge features are fed to each convolution layer. Each SE-convolution layer $L=0,\ldots, 7$ refines the output $\mathbf{h}_i^{(L)}$.
• Figure 5: Multi-variate vs uni-variate and influence of inputs. (top) Correlation $\rho$ between the true and the predicted propensity for A particles at temperature $T=0.44$ as function of timescale. Marker shapes distinguish multi-variate and uni-variate approaches. Colors picture the input type: red for thermal positions ($\{\mathbf{x}_i^{th}\}$), blue for quenched (Inherent Structures, IS) positions ($\{\mathbf{x}_i^{IS}\}$) and light-blue for combined ($\{\mathbf{x}_i^{th}\} + E_{pot}^{IS}$). Error-bars represent the best and the worst $\rho$ for ten identical models trained independently with different random seed initialisation, and are comparable with marker's sizes. (bottom) Correlation $\rho$ as function of training (and testing) temperature. Two timescales are shown: $\tau = \tau_\alpha$ (full markers) and $\tau = 0.5 \tau_{LJ}$ (empty markers). Color code and marker code identical to that of the top plot. The multi-variate, thermal positions + $E_{pot}(IS)$ choice is a good compromise to maintain high performance across timescales.