On Training Derivative-Constrained Neural Networks
We refer to the setting where the (partial) derivatives of a neural network's
(NN's) predictions with respect to its inputs are used as additional training
signal as a derivative-constrained (DC) NN. This situation is common in
physics-informed settings in the natural sciences. We propose an integrated
RELU (IReLU) activation function to improve training of DC NNs. We also
investigate denormalization and label rescaling to help stabilize DC training.
We evaluate our methods on physics-informed settings including quantum
chemistry and Scientific Machine Learning (SciML) tasks. We demonstrate that
existing architectures with IReLU activations combined with denormalization and
label rescaling better incorporate training signal provided by derivative
constraints.
KaiChieh Lo, Daniel Huang
# On Training Derivative-Constrained Neural Networks

###### Abstract

We refer to the setting where the (partial) derivatives of a neural network's (NN's) predictions with respect to its inputs are used as additional training signal as a _derivative-constrained_ (DC) NN. This situation is common in physics-informed settings in the natural sciences. We propose an integrated RELU (IReLU) activation function to improve training of DC NNs. We also investigate _denormalization_ and _label rescaling_ to help stabilize DC training. We evaluate our methods on physics-informed settings including quantum chemistry and Scientific Machine Learning (SciML) tasks. We demonstrate that existing architectures with IReLU activations combined with denormalization/label rescaling better incorporate training signal provided by derivative constraints.

## 1 Introduction

Deep learning is increasingly being applied to physics-informed settings in the natural sciences. By _physics-informed_, we mean any situation where inputs and/or outputs in a dataset involve relationships based on physics (_e.g._, forces). The field of Scientific Machine Learning (SciML) (Karniadakis et al., 2021) is emerging to address the issues of applying machine learning (ML) to the physical sciences (_e.g._, physics-informed neural networks or PINNs (Raissi et al., 2019)). Domains include fluid dynamics (Sun et al., 2020; Sun & Wang, 2020), geo-physics (Zhu et al., 2021), fusion (Mathews et al., 2020), and materials science (Shukla et al., 2020; Lu et al., 2020). In the realm of quantum chemistry, there are promising results (Hermann et al., 2022; 2020) that attempt to solve the electronic-structure problem (i.e., predict energy from structure) or model an atomic system's energy surface (Hu et al., 2021; Gasteiger et al., 2021)

In the physics-informed setting, it is common to express constraints on the neural network's (NN's) predictions in terms of the NN's (partial) derivatives with respect to (w.r.t.) its inputs to express physical constraints. We call this a _derivative-constrained_ (DC) NN. Thus, training a DC NN addresses a subset of issues considered in physics-informed settings such as SciML. We emphasize that most settings do not use derivatives of the model w.r.t. its inputs to supply additional training signal even though most NN models are optimized with gradient-based methods.

One strategy for incorporating derivative constraints is to add additional terms containing the derivative constraints to a loss function so that multi-objective optimization can be performed. While it is possible to construct NNs with high predictive accuracy using this strategy, the resulting models may not incorporate derivative constraints efficiently. In the physics-informed setting, this translates into capturing less of the physics. We demonstrate that this occurs in many existing works in quantum chemistry and SciML where we obtain high predictive accuracy but lower accuracy on derivative constraints (see experiments, Sec. 5). This presents an opportunity to reevaluate aspects of training DC NNs and revisit best practices. We make the following contributions.

1. We propose a new activation function called an _integrated ReLU_ (IReLU) obtained by integrating a standard ReLU activation (Agarap, 2018) (Sec. 4.1). We intend IReLU's as a drop in replacement for activations in an existing architectures. Our main motivation for doing so is that training a DC NN involves higher-order derivatives. Consequently, the choice of activation function will impact the propagation of additional derivative information.
2. We propose _denomalizing_ NNs, _i.e._, removing all normalization layers, and _label rescaling_ as a dataset preprocessing method to stabilize training (Sec. 4.2). Our primary motivation for doing so is because we hypothesize that DC training of NNs is sensitive to _units_. Consequently, unit-insensitive normalization procedures (_e.g._, batch normalization (Ioffe and Szegedy, 2015)) that help stabilize training in standard settings may introduce artifacts in the DC training case.

We benchmark the performance of our proposed methods on a variety of datasets and tasks including quantum chemistry NNs (Schutt et al., 2017; Xie and Grossman, 2018; Gasteiger et al., 2020; Ma et al., 2021; Gasteiger et al., 2021) and PINNs (Raissi et al., 2019) (Sec. 5) used in SciML. We show that IReLUs combined with denormalization/label rescaling improve the learning of gradient constraints while retaining predictive accuracy.

## 2 Related Work

There are at least two paths to improving training of DC NNs: (1) improving the loss function and (2) developing new architectures. In the first direction, loss functions used in training DC models often involve multiple terms so they are multi-objective optimization (MOO) problems. One solution to the MOO problem is to weigh each term in the loss function (Sener and Koltun, 2018; van der Meer et al., 2022; Bischof and Kraus, 2021), potentially in an adaptive manner (Li and Feng, 2022; Fernando and Tsokos, 2021; Xiang et al., 2022; Chen et al., 2018; Malkiel and Wolf, 2020; Heydari et al., 2019; Kendall et al., 2018; Lin et al., 2017). This is helpful in the SciML context since the different loss terms may use different units of measurements, and thus, have imbalanced label magnitudes (Wang et al., 2021). Weighing loss terms is also common in training quantum chemistry networks. We will demonstrate in Sec. 3.2 that it is difficult to control to motivate our methods.

In the second direction, we can also develop novel architectures that better incorporate domain knowledge. Domains such as quantum chemistry have custom designed NN architectures (Schutt et al., 2017; Xie and Grossman, 2018; Gasteiger et al., 2020; Ma et al., 2021; Gasteiger et al., 2021; Gasteiger et al., 2021; Gasteiger et al., 2021; Gasteiger et al., 2022; Schmitt et al., 2021; Zitnick et al., 2022; Passaro and Zitnick, 2023; Liao et al., 2023). The architectural improvements in these works focus on re-arranging interaction patterns (_e.g._, convolution layers), leveraging graph properties of atoms (_e.g._, molecular bond), and encoding invariances/equivariances.

We propose an activation function in Sec. 4.1 which we intend as a drop-in replacement for activations in existing architectures. Thus, we intend our activation to be applied to a wide range of architectures.

## 3 Training with Derivative-Constraints

We review an example of training with derivative constraints in the setting of quantum chemistry (Sec. 3.1). Then, we motivate our proposed methods with an experiment demonstrating the difficulty of incorporating gradient constraint information with traditional approaches (Sec. 3.2).

### Example: Potential Energy Surface Modeling

We use potential energy surface (PES) modeling from quantum chemistry as a concrete example to introduce DC training. A PES \(U:\mathbb{R}^{3A}\rightarrow\mathbb{R}\) gives the energy of a system with \(A\) atoms as a function of its atomic coordinates.1 A PES \(U(\mathbf{x})\) and its force field \(\mathbf{F}(\mathbf{x})\) can be evaluated by quantum mechanical simulation software such as Gaussian (Frisch et al., 2016) given the 3D Cartesian coordinates \(\mathbf{x}\) of the \(A\) atoms, _i.e._, its structure. The force field is the negative gradient of the PES and can be used to simulate the dynamics of the \(A\) atoms. In particular, when \(\mathbf{F}(\mathbf{x})=0\), there are no forces acting on \(\mathbf{x}\) which means that the configuration of \(\mathbf{x}\) is stable.

Footnote 1: We are ignoring symmetries. Technically, \(U:\mathbb{R}^{3A-5}\rightarrow\mathbb{R}\) for general atomistic systems and \(U:\mathbb{R}^{3A-6}\rightarrow\mathbb{R}\) when it is planar.

Conservation of energy is expressed with the following derivative constraint

\[-\nabla_{\mathbf{x}}U(\mathbf{x})=\mathbf{F}(\mathbf{x}) \tag{1}\]
that relates the gradient of the PES with the negative force. This connects simulation of a system with its changes in energy. As a result, this means that we can find stable configurations of atomistic systems in nature by finding local minima on the system's PES, since \(\mathbf{F}(\mathbf{x})=-\nabla_{\mathbf{x}}U(\mathbf{x})=0\). Consequently, we can study molecules and materials _in silico_ if we can model a PES efficiently and accurately enough.

We can construct a surrogate model of a system's PES by fitting a NN \(f_{\theta}\) with parameters \(\theta\) to a dataset \(\mathcal{D}=\{(\mathbf{x}^{i},E^{i},F^{i})_{i}:1\leq i\leq N\}\) where \(\mathbf{x}^{i}\) are atomic coordinates, \(E^{i}\) is an energy, and \(F^{i}\) are forces - the negative gradient of the energy w.r.t. \(\mathbf{x}^{i}\). This dataset can be created from quantum mechanical simulation software. A surrogate model can be used to accelerate the computation of a PES since quantum mechanical simulation software can be compute-intensive to run. To train a NN on this dataset, we can use the following multi-objective loss function

\[\text{Loss}(f_{\theta},\mathcal{D})=\sum_{i=1}^{N}\alpha\|f_{\theta}(\mathbf{x }^{i})-E^{i}\|^{2}+\beta\|-\nabla_{\mathbf{x}}f_{\theta}(\mathbf{x}^{i})-F^{ i}\|^{2}\,. \tag{2}\]

The terms \(\alpha\) and \(\beta\) are hyper-parameters that weigh the contributions of the first term involving \(f_{\theta}\)'s predictions and the second term involving \(f_{\theta}\)'s gradients. Training \(f_{\theta}\) with a gradient-based method will thus involve second-order derivatives. The second term enforces the conservation of energy since it constrains the force prediction of the model to be the observed force. Thus, conservation of energy is violated when training signal from derivative constraints is not efficiently incorporated.

More generally, we can have arbitrary derivatives, constraints, and datasets involving additional supervised signal that contains these constraints for different situations (_e.g._, thermodynamics, fluid flow). These constraints can come from the underlying partial differential equation (PDE) that describes the natural process. As before, we can add these constraints to a loss term and optimize as before. We refer the reader to the SciML literature (Raissi et al., 2019; Karniadakis et al., 2021; Meng et al., 2022) for more examples.

### Motivation

We hypothesize that one fundamental challenge with training a DC model \(f_{\theta}\) is that it is more difficult to incorporate derivative constraint information in the loss term compared to model prediction information in the loss term. As a quick test of our hypothesis and motivation for our methods, we report the following experiment in the setting of quantum chemistry.

ExperimentWe train a selected NN designed for quantum chemistry on a training set consisting of \(\{(\mathbf{x}^{i},E^{i},F^{i})_{i}:1\leq i\leq N\}\) tuples for varying \(\beta\) values while holding \(\alpha=1\) constant to compare the relative difficulty of predicting the energy versus predicting the force (_i.e._, involving

Figure 1: Comparing relative difficulty of learning energy (prediction) versus force (derivative constraint) with SchNet (Schütt et al., 2017) on Asprin molecule in MD17 dataset by varying \(\beta\) in loss function (Eq. 2).
the gradient). As a reminder, the terms \(\alpha\) and \(\beta\) control the relative importance of each term in the loss function in Eq. 2. Thus, whether \(\alpha>\beta\) or \(\alpha<\beta\) can be used as a proxy to determine which term in the loss is more difficult to learn. In particular, we consider learning of the derivative signal to be more difficult if \(\beta>\alpha\) gives lower force loss on the test set compared to energy loss on the test set on a relative basis. We use the relative basis since energies and forces are given in related units, \(\frac{\text{real}}{\text{mol}}\) versus \(\frac{\text{real}}{\text{mol}\,\Lambda}\) respectively, and so direct comparison is not possible.

DetailsIn this experiment, we choose a classic NN, Schnet (Schutt et al., 2017). The implementation and hyper-parameters of SchNet are taken from the Open Catalyst Project (OCP) (Chanussot* et al., 2021), a joint effort to between computer scientists and chemistry/material scientists to solve the PES modeling problem. We select the MD17 (Chmilela et al., 2017) dataset which contains \((\mathbf{x}\cdot E^{i},F^{i})_{i}\) tuples for 8 different small molecules. As terminology, each \(\mathbf{x}^{i}\) is also called a _conformation_. We train Schnet on \(50,000\) conformations of the Asprin molecule for \(50\) epochs.2 Instead of normalizing the data before training, we train the networks directly on \(\mathbf{x}^{i}\) so that a surrogate PES can directly predict energies and forces with the same units. The mean energy of Asprin in our training set \(-406,737.28\frac{\text{real}}{\text{mol}}\) and the variance is \(35.36\frac{\text{real}^{2}}{\text{mol}^{2}}\). The mean force is \(423.87\frac{\text{real}}{\text{mol}\,\Lambda}\) and the variance \(779.99\frac{\text{real}^{2}}{\text{mol}^{2}\,\Lambda^{2}}\). Thus, the absolute value of the mean energy is roughly \(3\) orders of magnitude larger than the mean force. We will comment more on this in Sec. 4.2. We use \(\beta=\{0.01,0.1,1,10,30,50,100,200\}\).

Footnote 2: We have also trained Schnet for \(300\) epochs, but have observed fast convergence.

ResultsWe report training loss curves in Fig. 2a and test loss for various \(\beta\) values in Fig. 2b. We observe that the energy loss divided by \(1000\) (since the energies are roughly \(3\) orders of magnitude larger) is typically much lower than the force loss. This gives evidence that it is more difficult to incorporate derivative constraint information in the loss term compared to model prediction information in the loss term. Moreover, it is not easy to improve the relative difference, even for large values of \(\beta\) which may make the energy loss worse. Finally, we observe that the training losses for each choice of \(\beta\) converges to roughly the same level so that there is a trade-off between learning energies and learning forces.

## 4 Methods

In this section, we propose two ideas to be used in conjunction to improve learning of derivative constraints. First, we propose a new activation function called an integrated ReLU (IReLU) activation function (Sec. 4.1). Second, we introduce _denormalization_ and _label rescaling_ to help stabilize training of DC NNs with IReLU activations (Sec. 4.2).

### Integrated ReLU Activation

The simple observation that we make concerning training a DC NN is that it involves higher-order derivatives. Consequently, higher-order derivatives of activation functions will also be used during the training process. This motivates us to revisit the choice of activation function for DC training since ordinary activation functions have been designed in the setting where only the first-order derivative is used.

We use a simple idea to construct an activation for DC NN training: use an integrated form of an ordinary activation function such as a ReLU when we are performing DC training of NNs. The intuition for doing so is the following: if we only fit derivative constraints, then we should use a ordinary activation (_e.g._, ReLU) after we have taken the derivative of the original activation in the model. We will discuss this intuition more in Appx. A.

Define the _integrated ReLU_ (IReLU) to be the activation

\[\mathrm{IReLU}(x)=\int_{0}^{x}\mathrm{ReLU}(y)dy=\max(0,0.5*x^{2})\,. \tag{3}\]

We focus on the IReLU activation in this work since the ReLU is a popular activation. Naturally, the idea of an IReLU can be applied to other activations (Maas et al., 2013; Clevert et al., 2015; Ramachandran et al., 2017; Hendrycks & Gimpel, 2016) as well.
### De-Normalization and Label Rescaling

Conventional wisdom is that normalization techniques such as batch normalization may help accelerate the training of NNs as well as improve the stability of training. Notably, centering and rescaling the internal values in a NN might be crucial when using IReLU's since these activations produce higher responses compared to traditional activations. Dataset normalization is also common practice to ensure that input features are set on equal footing. Intuitively, normalization techniques remove the _units_ of the features and dataset.

We hypothesize that DC NNs are more sensitive to _units_ compared to typical training without derivative constraints because of the linearity of derivatives, _i.e._, \(\nabla f(c\mathbf{x})=c\nabla f(\mathbf{x})\). We can interpret the constant \(c\) as determining the _units_ of \(\mathbf{x}\), which also determines the units of the derivative \(c\nabla f(\mathbf{x})\). In particular, this constant \(c\) will appear in the loss term in while training a DC NN and so the loss function is sensitive to the choice of units on the inputs \(\mathbf{x}\). We emphasize that a typical setting does not have derivatives of the NN w.r.t. its inputs \(\mathbf{x}\), and so these units will not appear in the loss. If our hypothesis holds, then we will need to develop alternative approaches to stabilizing training than the typical unit-insensitive approaches.

Towards this end, we propose two techniques. First, we propose _denormalization_, _i.e._, the removal of all normalization techniques in a NN architecture. Second, we propose a simple _label rescaling_ procedure where we scale the labels in a dataset \(\mathcal{D}=(\mathbf{x}^{i},\ell_{1}^{i},\ldots\ell_{n}^{i})_{i}\) by a suitable constant \(C\) defined as

\[C=\max_{\ell_{j}^{i}}\{C:0\leq\frac{\ell_{j}^{i}}{C}\leq 1,C\text{ is power of 10 }\}\;. \tag{4}\]

In the PES modeling example, this means we use the same constant \(C\) for both energy and force labels. Intuitively, what label rescaling does is set the units of the model's predictions and derivatives. The loss function in the PES modeling example becomes

\[\sum_{i}\alpha\|f_{\theta}(\mathbf{x}^{i})-\frac{E^{i}}{C}\|^{2}+\beta\|-\nabla _{\mathbf{x}}f_{\theta}(\mathbf{x}^{i})-\frac{F^{i}}{C}\|^{2} \tag{5}\]

with label rescaling. Thus, label rescaling plays a similar role to \(\alpha\) and \(\beta\) in setting units, the difference being that the units are set on the model as opposed to the loss. We emphasize that in label rescaling, we do not normalize the dataset inputs.

## 5 Experiments

We benchmark the performance of our proposed methods on a variety of architectures, datasets, and tasks including quantum chemistry NNs (Sec. 5.1) and PINNs (Sec. 5.2) used in SciML. Scripts we used for the experiments can befound in Github: [https://github.com/lbai-lab/dcnn-training](https://github.com/lbai-lab/dcnn-training)

### Quantum Chemistry

We separate our experiments in quantum chemistry by dataset since different atomistic systems can have different properties. We use the MD17 dataset (Sec. 5.1.1), which has small organic molecules, and OC22 (Chanussoot* et al., 2021) (Sec. 5.1.2), which contains large atomistic systems and metals.

#### 5.1.1 Experiments on MD17

Our first experiment tests the efficacy of our methods across different architectures for task of potential energy surface modeling. We select SchNet (Schutt et al., 2017), CGCNN Xie & Grossman (2018), ForceNet Hu et al. (2021), DimeNet++Gasteiger et al. (2020), and GemNet Gasteiger et al. (2021). SchNet and CGCNN are based on a convolutional NN architectures. ForceNet, DimeNet++, and GemNet are based on graph NNs.

We select the MD17 dataset. For each molecule in MD17, we randomly select \(50000\), \(6250\), and \(6250\) conformations from the dataset as training set, validation set and testing set. The molecules include Asprine (Asp.), Benzene (Ben.), Ethanol (Eth.), Malonaldehyde (Mal.), Naphthalene (Nap.), Salicylic acid (Sal.), Toulene (Tol.), and Uracil (Ura.). We present more details in Appx. B.
Baseline models are trained with the same training configuration and model hyperparameters given by OCP (Chanussot* et al., 2021). We note that Schnet was originally benchmarked on MD17 whereas the other NNs have been tested on other datasets. We train for 50 epochs on the MD17 dataset. Given the fast convergence of the training loss in MD17 (Fig. 2a), we consider 50 epochs is sufficient for models to fully learn the energy and forces. We use a batch size of \(20\) as recommended in the literature (Chanussot* et al., 2021). We use the Adam optimizer with a learning rate of \(0.0001\).

Tab. 1 compares the performance between models trained with original settings and models trained with our proposed methods. We denormalize all networks with normalization layers. For architectures which consist of multiple interaction-output blocks (_e.g._, DimeNet++ and GemNet), we were only able to replace activation layers in output blocks with IReLU as training with all activations replaced proved to be unstable. For label rescaling, we use the constant \(C=1000000\) since the energies in MD17 are on the order of \(400000\). The results of our proposed methods are noted with * in the table. To investigate the individual contribution of IReLU, denormalization, and label rescaling, we also conduct ablation studies (Appx. C). We also experiment on different dataset sizes (Appx. D).

For each architecture, we report the energy loss (upper row) and the force loss (bottom row) separately. We use the units of the original dataset, \(\frac{\text{real}}{\text{mol}}\) for energy and \(\frac{\text{real}}{\text{mol}A}\) for the force respectively. In general, our methods improve upon force loss across most molecules and most architectures. In particular, there is significant improvement in learning forces for CGCNN, DimeNet++ and ForceNet. We observe cases where our method performs worse on forces (_e.g._, SchNet and GemNet) but provides competitive performance. Perhaps surprisingly, our methods also improve the energy loss (38 out of 40 cases). We might reason that better incorporating force information would lead to improved learning of the physics. Nevertheless, it would be an interesting direction of future work to study this in more detail.

#### 5.1.2 Experiments on OC22 Dataset

To validate our methods on more datasets, we also compare the performance with and without our proposed methods on OC22 (Chanussot* et al., 2021). OC22 contains \(62331\) relaxations of oxides calculated at the DFT level. It contains a wide range of crystal structures (_e.g._, monoclinic,

\begin{table}
\begin{tabular}{|l|c|c|c|c|c|c|c|c|} \hline Model & Asp. & Ben. & Eth. & Mal. & Nap. & Sal. & Tol. & Ura. \\ \hline SchNet & 53.00 & 125.43 & **5.97** & 37.54 & 197.06 & 65.46 & 129.80 & 54.88 \\  & **1.21** & 0.39 & **0.68** & **1.00** & **0.72** & **1.02** & **0.67** & **0.86** \\ \hline SchNet* & **28.31** & **0.56** & 13.39 & **11.27** & **5.53** & **27.57** & **0.34** & **20.36** \\  & 1.67 & **0.36** & 1.53 & 1.67 & 1.01 & 1.12 & 1.04 & 1.23 \\ \hline \hline CGCNN & 239.07 & 98.86 & 38.36 & 104.45 & 130.98 & 197.09 & 124.29 & 133.51 \\  & 14.20 & 6.11 & 8.25 & 14.27 & 8.46 & 8.50 & 9.73 & 9.08 \\ \hline CGCNN* & **137.12** & **16.63** & **4.63** & **30.56** & **66.68** & **78.29** & **36.73** & **81.02** \\  & **7.13** & **0.61** & **2.96** & **4.99** & **2.38** & **4.77** & **2.91** & **4.93** \\ \hline \hline DimeNet++ & 47.43 & 133.70 & 236.94 & 856.01 & 1096.34 & 580.41 & 301.06 & 669.06 \\  & 13.82 & 12.45 & 6.41 & 8.87 & 7.82 & 10.41 & 8.64 & 6.15 \\ \hline DimeNet++* & **2.59** & **1.68** & **0.58** & **20.30** & **5.18** & **3.65** & **0.77** & **5.96** \\  & **2.52** & **0.25** & **0.29** & **0.78** & **0.70** & **1.79** & **0.46** & **0.71** \\ \hline \hline ForceNet & 755.49 & 239.53 & 1048.44 & 874.22 & 1677.71 & 165.17 & 369.33 & 1964.51 \\  & 21.95 & 117.79 & 17.61 & 31.52 & 18.16 & 22.36 & 18.33 & 46.47 \\ \hline ForceNet* & **13.80** & **19.09** & **3.18** & **4.52** & **5.98** & **41.43** & **5.63** & **14.54** \\  & **0.89** & **0.33** & **1.15** & **1.36** & **0.35** & **0.34** & **0.50** & **0.26** \\ \hline \hline GemNet & 2201.94 & 352.19 & 470.32 & **5.23** & 564.06 & 1238.93 & 193.50 & 228.10 \\  & **0.34** & **0.22** & 0.28 & **0.27** & **0.17** & **0.30** & **0.12** & **0.20** \\ \hline GemNet* & **13.16** & **6.31** & **7.10** & 15.36 & **5.87** & **5.67** & **8.13** & **21.98** \\  & 0.93 & 1.27 & **0.27** & 0.70 & 0.47 & 0.73 & 0.47 & 7.07 \\ \hline \end{tabular}
\end{table}
Table 1: Comparison of model performance trained with original settings and our proposed methods (*). Mean energy loss (kcal/mol) on the upper row and mean force loss (kcal/mol/A) on the bottom row of the eight molecules in MD17 trained on state-of-the-art models.
tetragonal) that contain heavier elements (_e.g._, metalloids, transition metals). Compared to MD17, OC22 consists of various large structure/metals (more than \(100\) atoms) mixed together in the training, validation, and testing set. In this dataset, the absolute value of the mean energy is only 1 order of magnitude larger than the mean force (compared to 3 in MD17). We randomly select \(200000\) molecules in OC22 training split as our training set and \(25000\) molecules in OC22 validation (out of domain) split as our testing set.

Baseline models are trained with the same training configuration and model hyperparameters given by OCP (Chanussoot* et al., 2021). We train all models for \(50\) epochs. We use a batch size of \(20\) as recommended in the literature (Chanussoot* et al., 2021) for SchNet and CGCNN. Due to hardware memory limitations we tested DimeNet++ and ForceNet with a batch size \(10\). We were not able to test GemNet due to hardware limitations. We use the Adam optimizer with a learning rate of \(0.00001\) for ForceNet and \(0.0001\) for others.

Tab. 2 shows that our methods produce better force predictions for CGCNN, DimeNet++, and ForceNet. In those models where we improved force losses, CGCNN and DimeNet++ also improve energy losses. It would be interesting to investigate why SchNet with our methods performs worse on OC22. As a reminder, both SchNet and CGCNN are based on convolutional architectures, and CGCNN's performance is improved with our method. We emphasize again that we do not modify the given architectures beyond replacing the activation functions (when appropriate).

### Physics-informed Neural Networks

To validate the generalization ability of our proposed methods in other domains aside from quantum chemistry, we also experiment on physics-informed neural networks (PINNs) (Raissi et al., 2019; Karniadakis et al., 2021; Wu et al., 2018). PINNs are a general family of models that use a NN to predict the solution of a partial differential equation (PDE). The solution of a PDE is a latent function \(\psi(x,t)\) that describe physical measurements (_e.g._, temperature and velocity) as a function of spatial coordinates \(x\) and time \(t\). PINNs enforce physical constraints on the solution \(\psi(x,t)\) by enforcing that the solution satisfies the governing PDE in the loss function.

In general the loss function for a PINNs takes the form below

\[\mathcal{L}(\psi,\mathcal{D})=\mathcal{L}_{f}\left(\psi,\mathcal{D}_{f}\right) +\mathcal{L}_{ICBC}\left(\psi,\mathcal{D}_{IC},\mathcal{D}_{BC}\right) \tag{6}\]

where \(\psi\) is a learned PDE solution (_e.g._, a NN) and \(\mathcal{D}=\left(\mathcal{D}_{f},\mathcal{D}_{IC},\mathcal{D}_{BC}\right)\) is a dataset consisting of several components containing additional constraints. The first term

\[\mathcal{L}_{f}\left(\psi,\mathcal{D}_{f}\right)=\frac{1}{\left|\mathcal{D}_{ f}\right|}\sum_{(\mathbf{x},t,\mathbf{y})\in\mathcal{D}_{f}}\mathcal{F}\left( \frac{\partial\psi(\mathbf{x},t)}{\partial\mathbf{x}},\frac{\partial\psi( \mathbf{x},t)}{\partial t},\frac{\partial^{2}\psi(\mathbf{x},t)}{\partial \mathbf{x}^{2}},\frac{\partial^{2}\psi(\mathbf{x},t)}{\partial\mathbf{x} \partial t},\ldots,\mathbf{y}\right) \tag{7}\]

gives the predicted solution's loss evaluated on a spatial-temporal grid \((\mathbf{x},t)\in\mathcal{D}_{f}\) using a function \(\mathcal{F}\). The loss is a function of additional derivatives of the predicted solution \(\psi\) w.r.t. its inputs according to the governing PDE. Thus, the loss function for a PINN may contain many higher-order derivatives.

The second term

\[\mathcal{L}_{ICBC}\left(\psi,\mathcal{D}_{IC},\mathcal{D}_{BC}\right)=\frac{1} {\left|\mathcal{D}_{IC}\right|}\sum_{(\mathbf{x},\mathbf{i})\in\mathcal{D}_{IC }}\mathcal{I}(\psi(\mathbf{x},0),\mathbf{i})+\frac{1}{\left|\mathcal{D}_{BC} \right|}\sum_{(\mathbf{x},t,\mathbf{b})\in\mathcal{D}_{BC}}\mathcal{B}(\psi( \mathbf{x},t),\mathbf{b}) \tag{8}\]

\begin{table}
\begin{tabular}{|l|c|c|} \hline Model & Energy MAE Loss (eV) & Force MAE Loss (eV/A) \\ \hline SchNet & **6.94** & **0.10** \\ \hline SchNet* & 22.4621 & 0.12 \\ \hline \hline CGCNN & 233.90 & 0.32 \\ \hline CGCNN* & **71.33** & **0.08** \\ \hline \hline DimeNet++ & 5.10 & 0.09 \\ \hline DimeNet++ & **0.57** & **0.01** \\ \hline \hline ForceNet & **4.48** & 0.33 \\ \hline ForceNet* & 5.90 & **0.10** \\ \hline \end{tabular}
\end{table}
Table 2: Comparison of performance on OC22 with original settings and our proposed methods.
enforces constraints on the PDE solution given by initial conditions (\(\mathcal{D}_{IC}\)) and boundary conditions (\(\mathcal{D}_{BC}\)). \(\mathcal{I}\) and \(\mathcal{B}\) are the respective loss functions for the initial conditions and boundary conditions. These conditions further constrain the solution of a PDE. There can be multiple IC and BC loss terms, Moreover, the IC and BC loss terms (not shown) can also involve higher-order derivatives in certain PINNs.

For our experiments with PINNs, we adapt baseline architectures, datasets and training configurations from PDEBench (Takamoto et al., 2022; 2022). PDEBench provides implementations and benchmarks of SciML models including PINNs for learning (1) the Advection equation (Sec. 5.2.1), (2) the compressible fluid dynamic equation (Sec. 5.2.2), and (3) the diffusion-reaction equation (Sec. 5.2.3). In the baseline architecture, all PINNs use the same MLP (multi-layer perceptron) architecture with \(6\) hidden layers of \(40\) neurons in each to simulate their latent function \(\psi\). The MLP uses Tanh activation function for every hidden layer. We use the same library DeepXDE (Lu et al., 2021) to construct and train the backbone MLP.

For all PDES, we compare training a PINN to learn the PDE with activations replaced with IReLU activations. We did not find a need for label rescaling. For comparison, we also add in batch normalization (BN) to study its impact. Following PINN convention, we measure the model performance by using the mean square error (MSE) of the PINN's latent function prediction (_i.e._, \(\psi(x,t)\)). Thus, the loss terms which involve higher-order derivatives are taken purely as constraints. Loss terms labeled with \(\dagger\) are the terms which involve derivatives of the model w.r.t its inputs. To give more fine-grained information about how our methods impact each component of the loss term, we provide the MSE value of each loss term evaluated in the last epoch of training along side the predictive MSE evaluated in testing set.

#### 5.2.1 Advection Equation

The Advection equation has a simple PDE that involves first-order partial derivatives of the model w.r.t. its input in \(\mathcal{L}_{f}\), one initial condition in \(\mathcal{L}_{ICBC}\), and one boundary condition in \(\mathcal{L}_{ICBC}\). Thus, this task tests our method's performance on loss functions in DC training with 3 terms. The full loss function associated with the Advection equation is presented in the Appx. E.1. Both standard training and training with our methods use \(15000\) epochs on the 1D Advection dataset with the Adam optimizer and an initial learning rate of \(0.001\) following PDEBench.

Tab. 3 presents the predictive loss evaluated on the test set and the training loss of each term in the loss function evaluated in the last epoch of training since these terms act as constraints. The model trained with our method achieves the best predictive loss (\(\mathcal{L}_{f}\)) on the test set. We also improve two of the training loss terms. It is also interesting to observe that adding batch normalization (BN) decreases performance.

#### 5.2.2 Compressible Fluid Dynamics Equation

The compressible fluid dynamics (CFD) equation contains 6 total initial and boundary conditions in \(\mathcal{L}_{ICBC}\). Like the Advection equation, it also involves first-order derivatives of the model w.r.t. its input in \(\mathcal{L}_{f}\). Thus, this task tests our method's ability to handle many terms in the loss function. The full loss function associated with the CFD equation is presented in the Appx. E.2. Both standard training and training with our methods use \(15000\) epochs on the 1D CFD dataset from PDEBench with the Adam optimizer and an initial learning rate of \(0.001\) following PDE bench.

\begin{table}
\begin{tabular}{|l|c|c|c|c|} \hline Method & MSE & \(\mathcal{L}_{f}^{\dagger}\) & \(\mathcal{L}_{IC}\) & \(\mathcal{L}_{BC}\) \\ \hline Tanh + BN & 1.70 & 8e-16 & 5e-5 & 3e-5 \\ \hline IReLU + BN & 2.42 & 9e-12 & 8e-6 & 2e-6 \\ \hline Tanh (original) & 1.59 & 1e-5 & **3e-6** & 6e-6 \\ \hline IReLU (ours) & **0.99** & \(<\)**e-45** & 0.08 & \(<\)**e-45** \\ \hline \end{tabular}
\end{table}
Table 3: MSE and loss terms of the Advection equation. \(\mathcal{L}_{f}\) is the loss of the PDE which governs the Advection equation, \(\mathcal{L}_{IC}\) is the loss on the initial conditions, and \(\mathcal{L}_{BC}\) is the loss on the boundary conditions.
Tab. 4 presents the predictive loss \(\mathcal{L}_{f}\) evaluated on the test set. We also include the training loss of each term in \(\mathcal{L}_{ICBC}\) evaluated in the last epoch of training. The model trained with IReLU and denormalization performs the best on all training terms. Our methods perform exceptionally well on the term involving derivatives (\(\mathcal{L}_{f}\) loss). As before, we also observe that adding BN decreases performance.

#### 5.2.3 Diffusion Reaction equation

The diffusion-reaction (DR) equation involves the most complex PDE in terms of derivative constraints. The loss on the PDE solution's prediction contains second-order derivatives, which means that third-order derivative information is used during NN training. Additionally, there are boundary conditions that utilize gradient information to describe the solution's boundary state at each given time step. Thus, this task tests our method's ability to cope with higher-order derivatives and derivative constraints in multiple terms. The full loss function associated with the diffusion reaction equation is presented in the Appx. E.3. For the DR equation, both standard training and training with our methods use \(100\) epochs on the 2D DR dataset (\(1000\) samples) from PDEBench with the Adam optimizer and an initial learning rate of \(0.001\).

Tab. 5 presents the MSE loss on the testing set and training loss of each term in \(\mathcal{L}_{ICBC}\) in the last epoch of training. The models with IReLU and denormalization have the best predictive loss, which involves second-order derivative information. Additionally, we also achieve competitive performance on the boundary conditions that involve derivatives. This experiment demonstrates that our proposed methods has ability to fit gradients of NN that have order more than two. Note that the IReLU has a vanishing third-order derivative so other integrated activations may perform better.