PEANUT mini model architecture

To explore how little information is sufficient for meaningful energy predictions, a simplified PEANUT model was implemented that ignores the species of neighboring atoms. It still uses radial and spherical features, but instead of combining them via a tensor product, they are concatenated. The smoothing function is applied to both radial and spherical features to ensure smoothness. Messages are summed over neighbors, concatenated with the initial node embedding \(h_i^{(0)}\), and passed through an MLP to update the node vector, which is then used to predict atom-wise energy contributions. The full energy prediction is given by

\[\begin{split}E_{\mathit{full}} = \sum_i MLP_{\mathit{energy}}\Bigg(MLP_{\mathit{node}}\Big(h_i^{(0)} \,\big\|\, \sum_{j \in \mathcal{N}(i)} \begin{bmatrix} R_{ij} \\ Y_{ij}^{\mathit{smooth}} \end{bmatrix} \Big)\Bigg),\end{split}\]

where

\[\begin{split}\hat{h}_i = \Big(h_i^{(0)} \,\big\|\, \sum_{j \in \mathcal{N}(i)} \begin{bmatrix} R_{ij} \\ Y_{ij}^{\mathit{smooth}} \end{bmatrix}\Big) \in \mathbb{R}^{D + B + (L+1)^2}.\end{split}\]