Orbital Mixer: Using Atomic Orbital Features for Basis Dependent Prediction of Molecular Wavefunctions

Leveraging ab initio data at scale has enabled the development of machine learning models capable of extremely accurate and fast molecular property prediction. A central paradigm of many previous works focuses on generating predictions for only a fixed set of properties. Recent lines of research instead aim to explicitly learn the electronic structure via molecular wavefunctions from which other quantum chemical properties can directly be derived. While previous methods generate predictions as a function of only the atomic configuration, in this work we present an alternate approach that directly purposes basis dependent information to predict molecular electronic structure. The backbone of our model, Orbital Mixer, uses MLP Mixer layers within a simple, intuitive, and scalable architecture and achieves competitive Hamiltonian and molecular orbital energy and coefficient prediction accuracies compared to the state-of-the-art.


Introduction
An explosion of interest has surrounded applying machine learning (ML) methods to quantum chemistry with a plethora of interesting application areas such as learning interatomic potentials (Behler & Parrinello, 2007;Unke et al., 2021c;Bartók et al., 2010;Smith et al., 2017;Chmiela et al., 2017;2018;Schütt et al., 2018;Unke & Meuwly, 2019;Unke et al., 2021b;Batzner et al., 2021;Klicpera et al., 2020;Liu et al., 2021a;Schütt et al., 2021), constructing density functionals (Snyder et al., 2012;Brockherde et al., 2017;Ryczko et al., 2019;Kalita et al., 2021;Li et al., 2021), predicting spectroscopic properties (Gastegger et al., 2017;Westermayr & Marquetand, 2020), optoelectronic properties Mazouin et al., 2021;Lu et al., 2020;Gladkikh et al., 2020), activation energies (Lewis-Atwell et al., 2021;Grambow et al., 2020), and a variety of physical properties throughout chemical compound space (Montavon et al., 2013;De et al., 2016;von Lilienfeld et al., 2020;Keith et al., 2021;Liu et al., 2021b;Tielker et al., 2021;Bratholm et al., 2021). Quantum chemistry workflows can obtain such chemical and physical information by modelling the electronic Schrodinger equation in a chosen basis set of localized atomic orbitals that is then used to derive the ground-state molecular wavefunction. Using ML, we can rather directly predict the molecular electronic structure which then provides access to a plethora of these derived properties without needing to train specialized models for each property of interest. Previous works of Schütt et al.  (SchNOrb) and most recently Unke et al. (Unke et al., 2021a) (PhiSNet) present deep learning architectures for predicting molecular wavefunctions and electronic densities by purposing only information of the atomic coordinates and molecular composition. Though inputs to these models rely only on the raw features of the molecule, they are trained on molecular wavefunctions from real quantum chemistry calculations, which necessarily associates the model's predictions with a prescribed basis.
Here, we present an alternate approach for predicting molecular orbitals that explicitly supplies basis set-specific information as input to a deep learning architecture that models interactions of atomic orbitals representations. Compared to only atomic coordinates and molecular composition, a complete basis set can provide a much higher dimensional and information rich representation of a molecular configuration (Qiao et al., 2020b). Our model, called Orbital Mixer, purposes characteristics of the atomic orbitals jointly with their spatial overlap to predict the orbital coefficients that define the molecular electron density. The backbone of the Orbital Mixer architecture uses MLP mixer (Tolstikhin et al., 2021) layers to efficiently model interactions between atomic orbital representations, and to ultimately predict the electronic Hamiltonian F for a molecular configuration, which can be diagonalized to obtain molecular orbitals. Directly operating on atomic orbital representations provides a strong inductive bias for Orbital Mixer when predicting the Hamiltonian, F which is represented in the same atomic orbital basis.
We evaluate Orbital Mixer against SchNOrb and PhiSNet on three separate molecular configuration datasets. Similar to SchNOrb, but unlike PhiSNet, Orbital Mixer is not explicitly covariant with respect to rigid molecular rotations but rather is trained using data augmentation to learn this equivariance, in line with previous interatomic-potential work, although explicitly-covariant schemes can also be developed (Mailoa et al., 2019;Park et al., 2021). We report improved accuracy and data and parameter efficiency when predicting the electronic Hamiltonian F, molecular orbital coefficients and energies compared to SchOrb while also performing competitively compared to PhiSNet. We demonstrate how integrating Orbital Mixer into quantum chemistry workflows using the predicted Hamiltonian F as an initial guess to DFT calculations achieves improved convergence speeds compared to default methods. Lastly, directly predicting the electronic structure gives us access to a variety of derivable physical chemical properties without needing to train separate ML models for each property of interest. We show excellent agreement between Orbital Mixer predicted and reference calculations for HOMO-LUMO gap energies and electronic dipole moments. The Orbital Mixer architecture benefits from simple and intuitive construction while leveraging strong inductive biases operating directly on atomic orbital representations to achieve competitive prediction accuracies.

Deep Learning for Orbital Prediction
Machine learning for molecular inference has experienced impressive success in recent years, showcasing spectacular predictive accuracy enabled by large quantities of ab initio data, the incorporation of prior physical and chemical knowledge, and invariant and/or equivariant architectures (Keith et al., 2021). A common paradigm of these works interprets molecules as connected graphs and uses message passing to model interactions as a function of single-particle contributions. A variety of increasingly complex graph convolutional operations have been proposed for these purposes, such as the pioneering work of SchNet (Schütt et al., 2018) introducing continuous filter convolutions, the message passing designed based on physical principles and attention of PhysNet (Unke & Meuwly, 2019), the explicitly covariant network operations of Cormorant (Anderson et al., 2019), Tensor field networks (Thomas et al., 2018) and NequIP (Batzner et al., 2021), among others (Schütt et al., 2021;Klicpera et al., 2020;Schütt et al., 2017;Klicpera et al., 2021;Liu et al., 2021a;Haghighatlari et al., 2021;Unke et al., 2021b). While these works have demonstrated excellent expressively and accuracy for molecular property prediction, each network is trained to predict only a predetermined set of scalar, vector, or sometimes tensor quantities. This fundamental design of these networks therefore requires training separate bespoke models for each molecular property of interest.
In contrast, a recent line of research strives instead to ascertain molecular wavefunctions by predicting the Hamiltonian matrix that satisfies the electronic Schrodinger equation from which physical and chemical properties can be derived. The short history of these methods begins with Hedge and Bowen (Hegde & Bowen, 2017), where they predict the Hamiltonian for two simple Copper and Carbon (diamond) systems using kernel ridge regression. Schütt et al.  then proposed a deep learning architecture called SchNet for Orbital (SchNOrb) that uses the SchNet architecture and pair-wise features to predict the Hamiltonian block-wise, establishing baselines for molecule configurations from the popular MD17 dataset . Follow-up work from Gastegger et al. (Gastegger et al., 2020) reports improved accuracy on select molecules by applying SchNOrb trained on a minimal basis set representation of molecular wavefunctions. More recently, Unke et al. (Unke et al., 2021a) propose PhiSNet, which draws upon insights of SE(3)-equivariant models to maintain that Hamiltonian predictions remain explicitly covariant with respect to rigid rotations or translations while also reporting significantly improved prediction accuracies. Notably, Nigam et al. (Nigam et al., 2021) devise similarly equivariant Hamiltonian representations for uses in other applications such as kernel machines.

Modeling Atomic Interactions
Architectural choices around modeling interactions between atoms and other molecular are of central interest in the design of neural networks for molecular inference. The aforementioned common choices of graph neural network structure or convolution/mixing operation dependent on (pairwise) atomic distances have the benefit of incorporating geometric information into the structure of the network and prioritizing local interactions. These approaches come at the cost of additional hyperparameters or increased network complexity, and frequently requires imposing a distance cutoff that prevents the network from directly modeling long-range atomic interactions (Unke & Meuwly, 2019).
In lieu of a graph neural network architecture, we draw on approaches for mixing spatial information from other deep learning domains, namely those from computer vision, where vision transformers (ViTs) have prompted rethinking of the standard approach of convolutional networks (CNNs) (Dosovitskiy et al., 2020). Orbital Mixer is directly based on the recently-proposed MLP Mixer vision architecture (Tolstikhin et al., 2021), a competitive but dramatically simpler alternative to both ViTs and CNNs. The MLP Mixer architecture splits an input image into patches and alternates between patch-wise and channel-wise mixing operations via simple multi-layer perception (MLP) layers. Though developed for vision, the MLP Mixer can be trivially adapted to other domains, and by replacing image patches with atomic orbital and overlap matrix information, we obtain a model that learns atomic interactions at any range without the need for complex or hand-engineered mixing operations. We further find that the favorable complexity of MLP layers, particularly as compared to graph transformers, to be of benefit in the molecular inference domain, where scalability to larger systems is particularly desirable.

DFT fundamentals
Electronic structure calculations typically represent electrons in a basis-set of atomic orbitals (AO) {|φ i } i=N orbs i=1 meant to describe the available electron orbitals of the system. (Here we use the braket notation |x to represent a quantum state in the complex Hilbert space.) In the Hartree-Fock (HF) model, the electron energies are given by the Hartree-Fock equations, represented in matrix form as which determine the molecular orbital (MO) wavefunctions C im |φ i and their associated MO energies m = E mm , where E is a diagonal matrix (Lehtola et al., 2020). Electrons populate the lowest energy MOs in accordance with the Pauli exclusion principle, which then define the electronic and chemical properties of the system. The Hamiltonian F in Eqn. 1 approximates the singleelectron energy operatorF within the set of basis functions {|φ i } with matrix elements defined as F ij = φ i |F |φ j , while the overlap matrix S describes the relationships between the different basis functions via the inner product S ij = φ i |φ j . The formulation for density-functional theory (DFT) is almost identical, with the replacement of the Kohn-Sham Hamiltonian for the Fock matrix and Kohn-Sham orbitals for molecular wavefunctions; we will use the F matrix or term "Hamiltonian" for either.
The generalized eigenvalue problem in Eqn. 1 can be solved to determine the electron density D ij = k C ik C jk , where the summation is carried over the k indexing the lowest energy MOs which are occupied. The remaining unoccupied orbitals that do not enter the density matrix calculations are called virtual orbitals and are only defined up to an arbitrary unitary transformation . However, the matrix elements F ij themselves actually depend on the electron density D requiring that Eqn. 1 be solved in a selfconsistent manner. Typically, DFT and HF begin with an initial guess for the electron density D and/or orbital occupations C, followed by an iterative procedure that alternatingly 1) uses the density estimate D to calculate the Hamiltonian F and 2) solves Eqn. 1 to yield an improved estimate for the electron density, until some convergence criteria is met. The major source of computational expense in DFT comes from the number of these self-consistent iterations that must be performed to obtain converged electron density estimates that may in turn be used in downstream quantum chemical calculations.
While in principle the eigenvectors C and eigenvalues E contain the same information as the Hamiltonian F and the overlap matrix S, using ML to directly predict C and E is complicated by possible state degeneracies and the coefficients being defined only up to an arbitrary phase (Gastegger et al., 2020;Schütt et al., 2019). In contrast, the Hamiltonian F is better behaved as a smooth function of the atomic coordinates, and combined with S and Eqn. 1, can be used to determine C and E. These properties make the Hamiltonian F a more suitable target for ML-enabled prediction. The goal of this work is therefore to learn to reliably predict the Hamiltonian F for a given molecular configuration and thereby alleviating some computational expense required in self-consistently solving Eqn. 1.

Basis set-specific molecular encodings
Many existing deep learning architectures for molecular inference use exclusively basis-independent information, even when predicting basis-dependent information (e.g., the Hamiltonian F). In contrast, we explicitly choose a basis that the network will model, and represent each molecular configuration through basis-dependent quantities. This featurization approach is similar OrbNet (Qiao et al., 2020b;a), which uses symmetry-adapted atomic orbital features within a GNN architecture for molecular property prediction. In particular, we use the overlap matrix S to capture geometric information together with categorical features corresponding to each atomic orbital. Elements of the overlap matrix S ij measure the spatial overlap of localized atomic orbital basis functions, and provide a detailed description of the molecular geometry specific to the choice of basis set. Unlike the Hamiltonian F, overlap matrix elements are independent of electron density, and can be quickly computed given only a choice of basis set and the atomic positions.
Though supplying the model with atomic coordinates directly as inputs provides benefits -most notably, native access to analytic derivatives with respect to the atomic positions, which may be later used for e.g., electronic property optimization (Schütt et al., Figure 1. Orbital Mixer neural network architecture for predicting molecular electronic structure. Orbital Mixer leverages properties of atomic orbitals that comprise the DFT basis set to predict the Hamiltonian F in Eqn. 1, from which the molecular orbital coefficients defining the electron density are derived. stronger inductive bias than basis-independent features like atomic coordinates, and gives a more direct link when predicting the Hamiltonian F that is expressed within the same basis set and carries the same symmetries as the overlap matrix S. These representations are then combined with categorical encodings of atomic orbital features that comprise the basis set. Together, the overlap matrix S and categorical atomic orbital features provide detailed geometric and basis set-specific descriptors for us to effectively predict the molecular electronic structure.

The architecture
Orbital Mixer presents a novel MLP mixer-based architecture for modelling molecular electronic structure in a basis of localized atomic orbitals. A schematic illustration of the Orbital Mixer deep learning architecture and data flow is presented in Fig. 1. We begin with enumerating a list of atomic orbitals comprising the set of basis functions used to model the electronic Hartree-Fock or Kohn-Sham wavefunction of the system. The size and complexity of the basis set determines the achievable accuracy of the electronic structure calculation at the cost of computational expense, and depends on the particular system and properties of interest. Each atomic orbital is characterized with five categorical features that uniquely specify each orbital within the basis set: the index and element of the atom at which the orbital is centered, the principal quantum number n, the azimuthal quantum number l, and the magnetic quantum number m l . Each feature for each orbital is transformed into a d hidden -dimensional vector using a separate learned embedding layer and assembled into a tensor of shape N orbs × d hidden × N features . Finally performing a summation over the N features dimensions yields the complete multi-feature embedding for the atomic orbital basis of shape N orbs × d hidden .
While the multi-feature embedding describes properties of each atomic orbital independent of the molecular geometry, the overlap matrix S ∈ R Norbs×Norbs provides a representation of the molecular geometry by measuring the integrable overlap of the atomic orbital functions basis which depends on the spatial arrangement of the atoms where the orbitals are localized. Each row i of S corresponds to the overlap of an atomic orbital |φ i with the other N orbs orbitals; these rows are processed with an initial MLP to yield d hidden -dimensional atomic orbital representations, and their interactions are modeled with subsequent MLP Mixer layers. As the overlap matrix S is covariant with respect to rigid rotations of the atomic coordinates, this processing step is important to capture the interdependence between the atomic orbitals and learn a globally aware overlap matrix embedding of shape N orbs × d hidden . The multi-feature embedding is then added to this overlap matrix embedding yielding a complete atomic orbital embedding that captures both categorical properties specific to each atomic orbital along with global spatial information of the molecular geometry.
The atomic orbital representations are further refined with another series of MLP mixers, followed by a row-wise MLP to reshape the N orbs × d hidden representations into the target N orbs × N orbs dimensionality of the Hamiltonian F. As the off-diagonal elements of the Hamiltonian F i,j are indicative of coupling between the atomic orbitals, the MLP mixer layer effectively captures these interactions between atomic orbitals useful for predicting the off-diagonal matrix elements. On the other hand, the diagonal elements of the Hamiltonian F i,i correspond to energies of each of the atomic orbital basis functions, and in practice are often significantly larger than their off-diagonal counterparts in magnitude. To account for these physical and numerical differences, we apply a separate MLP row-wise to the atomic orbital embeddings generating a single scalar value for each of the N orbs atomic orbitals. This N orbs -dimensional vector is assembled into a diagonal matrix and added to the N orbs × N orbs dimensional output of the interaction branch, constituting a diagonal correction to the interaction branch and providing an inductive bias delineating the on-and offdiagonal elements of the Hamiltonian F. Finally, as the Hamiltonian F is always symmetric, the diagonally corrected representationsF are symmetrized to obtain the com-plete predicted Hamiltonian F = 1 2 (F +F T ). From the predicted Hamiltonian F we can obtain the molecular orbital coefficients, energies and electron density by solving Eqn. 1, giving us access to the molecular electronic structure and a plethora of chemical and physical properties for the system.

Training procedures
Orbital Mixer is trained end-to-end with mini-batch stochastic gradient descent and the ADAM optimizer (Kingma & Ba, 2014) using a simple mean squared error (MSE) between the true F (true) and predicted F (pred) Hamiltonians, Importantly, both the input overlap matrix S and the Hamiltonian F are covariant with respect to rigid rotations of the atomic coordinates r ∈ R Natoms×3 . We perform data augmentation to ensure Orbital Mixer leans this covariance where during training a random rotation matrix R performs a rigid rotation of the atomic coordinates r = rR T for each training sample. The corresponding covariant change to the overlap S and Hamiltonian F matrices due to this rigid rotation R is then accounted for using Wigner D R ∈ R Norbs×Norbs rotation matrices (Wigner, 1931) via a unitary transformation, We also incorporate during training a separately maintained exponential moving average of the model parameters that are then used at inference time, as we find this leads to improved generalizability. Orbital Mixer in total contains ∼38M parameters, which is ∼2.5× fewer than the ∼93M parameters in the SchNOrb architecture, but ∼2× more than the ∼17M parameters of PhiSNet. A full accounting of all neural network, optimizer and scheduler hyperparameters is provided in the Appendix along with a complete PyTorch (

Predicting molecular electronic structure
The neural network architecture proposed in this work, Orbital Mixer, is capable of accurately predicting the Hamiltonian F used to determine the electronic density and other derivable physical and chemical properties for a variety of molecular systems. We evaluate Orbital Mixer on conformational geometries of small molecule molecular dynamics trajectories taken from the MD17 dataset . Namely, we select the three molecules (Ethanol, In Fig. 2 we present a detailed comparison of Hamiltonian F and MO coefficients and energies predictions generated by Orbital Mixer trained with 950 configurations and SchNOrb trained with 25K configurations. We visualize in Fig. 2A the matrix element-wise test set MAE of the 1 Reported PhiSNet cosine similarities are rounded to fewer significant digits than we report in this work. 2 Results for PhiSNet are taken from Ref. (Unke et al., 2021a) and SchNOrb models are trained in this work using the publically available implementation at: https://github.com/atomistic-machinelearning/SchNOrb predicted Uracil Hamiltonian F. Beyond the concentration of both models' prediction errors along the diagonal (as expected, due to the diagonal entries' significantly larger magnitude), SchNOrb model predictions produce significantly larger errors in select off-diagonal blocks, while errors from Orbital Mixer predictions are comparatively lower throughout the off-diagonal elements. In Fig. 2B we show cosine similarity between SchNOrb and Orbital Mixer predicted MO coefficients delineated for the 29 occupied orbitals of Uracil. Orbital Mixer never performs worse than ∼0.983 cosine similarity for any particular occupied orbital, while SchNOrb only performs better than our worst performer for 6/29 occupied orbitals. We also show in Fig. 2C the MAE between the ground truth and predicted occupied MO energies for both Orbital Mixer and SchNOrb. Here we notice our largest error orbital (orbital 2 at ∼0.012 eV) performs better than all the predicted SchNOrb orbitals and ∼52% better than the best SchNOrb orbital. We emphasize Orbital Mixer achieves these predictions accuracies using ∼10× fewer training samples than SchNOrb, but nevertheless on average performs worse than PhiSNet (Table 1).

Improved SCF convergence
The ability to accurately predict the Hamiltonian F for a molecular system enables integration of Orbital Mixer with electronic structure calculations. Namely, the Hamiltonian predictions generated by Orbital Mixer can be used as initial guess in Eqn. 1, which is then solved using the self-consistent field (SCF) method to arrive at a converged estimate of the Hamiltonian F, and therefore the electron density. The speed of these DFT calculations is in large part determined by the number SCF iterations required to reach convergence. High quality initial guesses for the Hamiltonian F can greatly expedite the SCF procedure and enable higher-throughput DFT calculations. We perform experiments testing speed-up with respect to number of SCF iterations by initializing DFT calculations using Orbital Mixer predicted Hamiltonian F. Reported in Fig. 3A is the distribution of the SCF iterations to convergence for DFT calculations performed on 250 test set Uracil configurations using both the default PySCF and Orbital Mixer predicted Hamiltonian F for initialization. The Orbital Mixer initialization achieves an impressive ∼44% improvement in the number of SCF iterations required to reach convergence compared to the default PySCF initialization strategy. Fig. 3B visualizes this improvement by tracking the difference in the total energy estimated after each SCF iteration as compared to the terminal converged energy estimate. Although both initialization schemes eventually reach the same energy difference criterion of 10 −9 Ha, all configurations using the Orbital Mixer initialization reliably converge after only at most 9 SCF iterations. In similar SCF initialization experiments, SchNorb reports a speedup of ∼15% and PhiSNet a speedup of ∼47% SCF iterations for Uracil. Similar comparisons for SCF convergence improvement of Orbital Mixer for the Ethanol and Malondialdehyde datasets are provided in the Appendix. We expect similar performance for other selfconsistent field applications, from higher-accuracy quantum chemistry methods to inhomogeneous copolymers and nanoparticles (Arora et al., 2016).

Derived physical chemical properties
We interrogate the ability of Orbital Mixer to directly calculate physical and chemical properties by using Orbital Mixer predicted Hamiltonian F and molecular orbital coefficients to calculate HOMO-LUMO gaps and electronic dipole moments, both of which are physically meaningful and often measurable quantities. Reported in Table. 2 is a comparison of predicted HOMO-LUMO gap energies and dipole mo- ments between Orbital Mixer and SchNOrb; we note that dipole moments or HOMO-LUMO gap energies were not reported in the original PhiSNet paper. The results for Orbital Mixer show overall excellent accuracies with respect to the reference DFT calculations for both HOMO-LUMO gap energies and dipole moments, at both training sizes MAEs below than 0.017 eV and 0.034 D, respectively. We again outperform SchNOrb using ∼10× fewer training data on all three benchmark molecules. Orbital Mixer performs particularly well on the most challenging molecule Uracil: using only 950 training samples, we generate HOMO-LUMO gap energy and dipole moment predictions with errors orders of magnitude smaller than those from a SchNOrb model trained with 25K configurations. These results highlight the ability of Orbital Mixer to capture chemically and physically meaningful molecular properties using a single neural network architecture, circumventing the need to develop separate specialized models for each property of interest.

Conclusions
We present in this work a deep learning workflow for predicting molecular electronic structure directly in a basis of localized atomic orbitals. Compared to SchNOrb and PhiSNet, which generate predictions as a function of only atomic coordinates and molecular composition alone, Orbital Mixer leverages strong inductive biases by operating jointly on basis set-specific atomic orbital representations

B. Neural network architecture and training settings
Throughout the Orbital Mixer architecture we use a hidden dimension of d hidden = 1024 and GELU activation functions (Hendrycks & Gimpel, 2016). The initial row-wise MLP used to process the rows of the overlap matrix S includes two dense layers with an expansion factor of 2 for the intermediate hidden representation, such that the complete action of the MLP involves the following sequence of transformations onto the shape of the atomic orbital representations: N orbs → 2d hidden GELU − −− → d hidden . The MLPs within the MLP Mixer layers for the overlap matrix and interaction branch similarly use an expansion factor of 2 with no dropout. The initial MLP Mixer applied to the overlap matrix S consists of n layers = 2 Mixer layers, while the second MLP Mixer applied to the atomic orbital embeddings used for predicting the Hamiltonian F within the interaction branch uses n layers = 6 Mixer layer. The row-wise MLP used to reshape the N orbs × d hidden representation processed by the MLP Mixer in the interaction branch into the target N orbs × N orbs dimensionality consists of a GELU non-linearity followed by a single dense layer. The row-wise MLP in the diagonal correction branch similarly uses a GELU non-linearity followed by a single dense layer. A complete PyTorch implementation of Orbital Mixer with accompanying training scripts is publicly available at DOI:10.18126/cu4h-d2mm (Shmilovich et al., 2022).
Orbital Mixer is trained using the ADAM optimizer with the default PyTorch parameters and a mini-batch size of 32 samples. Evaluation and testing are performed using the model obtained with an exponential moving average over all parameters during training time employing a decay rate of 0.999 per step. Gradient clipping is applied to clip gradient norms to a maximum value of 0.001, as we find this helps to stabilize training. We use an initial learning rate of 3 × 10 −4 which we decay by a factor of γ = 0.8 every n decay training steps. All Orbital Mixer models are trained for 120 hours on Nvidia Tesla V100 32GB GPUs, after which time we observe the training and validation losses to plateau. When training reference SchNOrb models we follow training procedures outlined in the SchNOrb paper for the 25K dataset originally handled in the paper. We use all the same training settings when training SchNOrb models on the 950 sample training set size, except we modify the patience to decay the learning rate from the original 15 epochs when training with 25K samples to 150 epochs when training with 950 samples. Training of all SchNOrb models was still stopped in each case once the learning rate dropped below the 5 × 10 −6 threshold. Provided in Table 3 is a breakdown of all these training settings and parameters for both Orbital Mixer and retrained SchNOrb models.