Learning Regularized Positional Encoding for Molecular Prediction
Abstract
Machine learning has become a promising approach for molecular modeling.
Positional quantities, such as interatomic distances and bond angles, play a crucial role in molecule physics.
The existing works rely on careful manual design of their representation.
To model the complex nonlinearity in predicting molecular properties in an more end-to-end approach, we propose to encode the positional quantities with a learnable embedding that is continuous and differentiable.
A regularization technique is employed to encourage embedding smoothness along the physical dimension.
We experiment with a variety of molecular property and force field prediction tasks.
Improved performance is observed for three different model architectures after plugging in the proposed positional encoding method.
In addition, the learned positional encoding allows easier physics-based interpretation.
We observe that tasks of similar physics have the similar learned positional encoding.
Introduction
The prediction of molecular properties is crucial for the discovery of molecules with desired properties, with applications in drug, material, and energy industries.
Ab initio quantum chemical calculations, such as these based on density functional theory (DFT), give accurate predictions for the chemical properties with a high computational cost.
Semi-empirical methods [1] are faster but much less accurate.
Machine learning based molecular predictor is a promising alternative solution to balance accuracy and efficiency.
In recent years scientists have started leveraging machine learning methods in molecular prediction tasks [2, 3].
An effective machine learning model should represent the key physical quantities properly.
Positional quantities, such as interatomic distances and bond angles are among the most fundamental physical quantities for ML-based molecular predictors.
According to Born-Oppenheimer approximation, a molecular system is uniquely determined by the nucleus positions and charges.
Almost all molecular properties are thus a function of the positions and types of the atoms.
Moreover, interatomic distances and bond angles are invariant to transformations such as rotations and translations.
These symmetric properties make them preferred inputs to build machine learning models to exploit the symmetry of physical problems [4, 5, 6].
The design of an effective positional quantity representation faces challenges from the complex physics in molecules.
First, the interactions in molecules are often nonlinear or even non-monotonic with respect to the positional quantities.
Lennard-Jones potential, for instance, is a function of a repulsive term and an attractive term.
Second, multiple nonlinear effects co-exist
and they scales differently with positional quantities.
For example, Lennard-Jones potential involves two polynomial terms, while Morse potential contains two exponential terms.
Ideally, a good representation of positional quantities should help the machine learning models to effectively model these multi-modal non-linearity.
Existing methods of positional quantities representations can be grouped into three categories.
( i )
Directly using original scalar form as the model inputs [6].
Although neural networks can learn any function of the inputs according to the universal approximation theorem, it is practically challenging given the multiple complex physics in molecules mentioned above.
( ii ) Using a manually defined set of simple transformations, such as the reciprocal of interatomic distances used by Wang et.al.
Such representations maybe more effective than directly using original scalar forms, but they are only ideal for limited potential functions.
For example, Lennard-Jones potential contains polynomial terms of reciprocal of interatomic distances but Morse potential do not.
The generalizability of such a method is limited.
( iii ) Using the representation of positional quantities under a set of manually picked basis functions , such as the Gaussian kernels [8, 9] or Bessel functions [10].
This method also has limited generalizability and the efficiency depends on the choice of the basis functions.
Motivated by the success of learnable positional encoding in natural language processing (NLP) [11, 12], we propose to use learnable positional encoding in molecular properties prediction.
Compared to the existing methods, the learnable position embedding is not limited to a manually picked function format.
Empirically we demonstrate that our method improve the accuracy on multiple molecular property and force field tasks compared to state-of-the-art models.
Our method works as a plug-in module, combined with any models with positional quantities as inputs.
Moreover, we propose a smoothness regularization technique, which filters out unnecessary embedding fluctuations and makes physics-based interpretation possible.
We demonstrate that the learned embedding reflects the dependence of different tasks on short and long-range interactions, and observe that tasks of similar physical nature result in similar positional encoding.
Our contribution is two-folds.
We propose a learnable positional encoding method that is continuous and differentiable.
The method can be used to represent interatomic distance, bond angles and other positional quantities with a series of nonlinear transformations.
Experiments show that new state-of-the-art results are achieved with this technique.
We propose a regularization method for the proposed positional encoding.
Based on physical intuition, this method encourages the embedding to form a meaningful manifold for easier visualization and interpretation from a physical point of view.
The regularization reduces overfitting and can further increase model accuracy for certain tasks.
Related works 2.1 Positional encoding
There are generally two types of positional encoding, fixed and learned.
Fixed positional encoding represent positions with manually picked basis functions.
In the initial version of transformer [13], the discrete token positions are represented by a series of sine and cosine functions of a range of fixed frequencies.
Gaussian kernels are employed to represent some continuous physical quantities such as interatomic distances in works of molecule modeling [8, 9].
DimeNet [10, 14] and GemNet [15] and
SphereNet [16] instead use Spherical Bessel functions and spherical harmonics.
Learned positional encoding represent each position with a learned vector.
Although initially the fixed positional encoding was proposed for transformers [13], the learnable positional encoding has become more popular technique in natural language processing (NLP) community, as in BERT [11] and GPT models [17, 18, 12].
Besides transformers, learnable positional encoding is used with convolutional neural networks [19].
2.2 Molecular modeling
Convolution neural networks are equivariant to translations and this enables them to generalize better on image problems.
This characteristic inspires researcher to use them in molecular properties prediction, where translation equivalence is a useful inductive bias [20, 4, 5].
Graph neural networks are another popular family of models used in modeling molecules.
Method
In this section, we introduce our method to encode positional quantities such as interatomic distance or bond angle, to a fixed-length learnable vector, h ( x ) ∈ R s , where s is the embedding size.
3.1 Continuous and Differentiable Positional Encoding
In contrast to the application in NLP, where the positional embedding are designed for discrete tokens, the positional quantities are continuous for molecule modeling.
To represent these continuous physical quantities with embeddings on discrete points, we divide the space of x into n bin bins, and obtain the embedding of a arbitrary x by interpolating the embeddings of the nearest bin centers.
A simple implementation is via linear interpolation
[Equation]
where ⌊ x ⌋ is the largest bin center smaller than x , and ⌈ x ⌉ is the smallest bin center larger than x .
The linear interpolation method is simple but it makes h ( x ) not differentiable with respect to x , as the derivative g ( x ) ≡ dh ( x ) dx ∈ R s is not continuous at the bin centers.
g ( x ) is required if the trained model is used in calculation that involving the derivative with respect to x .
For example, the potential force acting on atoms is the derivative of potential energy with respect to their positions.
To resolve this issue, we consider higher order interpolation.
We choose the Cubic Hermite spline as follows,
[Equation]
where the interpolation coefficients are
[Equation]
The positional encoding h ( x ) defined in this way is continuous and differentiable with respect to x .
This implementation involves two sets of learnable parameters, { h ( x i ) } and { g ( x i ) } , where i = 1 , 2 , ..., n bin is the bin index, and { x i } are the bin centers.
The total number of additional trainable parameters is 2 sn bin .
The proposed encoding method provides a simple way to build a universal approximation of any continuous functions of x .
The accuracy of the approximation can be easily increased by adding the number of bins.
3.2 Smoothness Regularization
As the positional quantities are continuous, we assume the embedding should be locally smooth.
The difference between the embedding of two adjacent bins should not be large.
Following this intuition we consider a smoothness loss, defined as the average relative change of a embedding, h ( x i ) , compared to the next bin, h
( x i +1 ) .
[Equation]
where || .
|| is 2-norm.
During training, L smooth is combined with the original training loss L orig as the new training loss
[Equation]
where λ is a hyperparameter to control the contribution of the smoothness loss.
L smooth acts as a regularization term as it reduces the flexibility of the positional encoding.
This may help the model generalize in case of small amount of data.
Even if a bin is never directly trained during training (i.e., no x fall near this bin during training), its parameters are still updated because its difference with neighbor bins is regularized by the smoothness loss.
We will show that smoothness regularization can improve model accuracy and provide physical interpretability in the next section.
3.3 Implementation
The regularized positional encoding is proposed as a plug-in for existing models.
When a model uses positional quantities as their input, one may consider replacing these quantities by the proposed embedding, as illustrated in Figure ?? .
In this work [Table] we consider three backbond models: EGNN [6],
DimeNet++ [14] [Table] and a Transformer[13] based model 2 .
The experiments are conducted on a Nvidia
Tesla V100 GPU
[Table]
Experiments and discussion 4.1 Molecular force fields
[Table]
We experiment with a set of molecular force field datasets, MD17 [24].
The task is to predict the force acting on each atom given the 3D geometry (configuration) of the molecule.
Following [25], we use 1000 training and test configurations.
Surprisingly, as illustrated in Table [Table] both EGNN and the Transformer-based model shows a large error, although these two architectures previously achieve success in a various tasks [6, 26].
We hypothesise the data characteristics of the current task makes the relatively simple representation of the interatomic distance r by EGNN and the Transformer-based model performing not well on this task.
Molecular force fields are highly sensitive to the change of 3D geometries.
The 3D geometries of the same molecule in MD17 dataset only slightly different from each other, while the forces acting on atoms change significantly across molecular configurations
[Table] We then use the proposed positional encoding method (+PosEnc in Table 1) to represent the interatomic distance.
This significantly reduce the test error for both EGNN and the transformer-based model.
When we employ the proposed smoothness regularization technique (+Smooth in Table 1), the test error further decreases.
We compare the results with a few recent works, sGDML [27], FCHL19 [28], DimeNet [10],
SphereNet [29], NequIP [25], and PaiNN [30].
They either use kernels or carefully designed spatial basis functions to represent the positional quantities.
We show that with the proposed learnable positional encoding technique, similar or better results on force fields can be achieved, as illustrated in Table 1.
4.2 Molecular properties
We then experiment with the molecular property prediction tasks using the QM9 [31] dataset.
This dataset consists the DFT calculation results of 134 k stable small organic molecules.
The
[Table]
C v
G
H
R
U U
ZPVE
[Table]
geometries minimal in energy with a variety of corresponding chemical properties are included: polarizability ( α ), energy HOMO energy LUMO energy gap
∆
HOMO HOMO ) electronic
The results compared with NMP [21], SchNet [20], Cormorant [32], L1Net [4], and LieConv [5].
Compared with the original backbone models, using the proposed positional encoding method (PosEnc) can improve the test error for both EGNN and DimeNet++, as listed in Table 2.
The smoothness regularization technique (Smooth) can further increase the accuracy for several tasks.
We also conducted experiments using a 2-layer MLP with ReLU activation.
For DimeNet++, replacing the manually-designed representation by an MLP generally increases test error compared to the original version.
For EGNN, MLP increases error for ϵ HOMO and R 2 and reduces error for the other tasks compared to the original version.
The improvement however is not as significant as PosEnc.
4.3 Regularized Embedding
The smoothness regularization helps the learned embedding to be more smooth for easier interpretation.
The first 10 dimensions of the learned embedding are visualized in Figure 1 as a function of the normalized distance x as
[Equation]
Without the smoothness regularization, the learned embedding appear to have no obvious patterns.
The difference between the embeddings of two adjunct bins are large.
In contrast, the embedding learned with smoothness regularization appear to be more smooth and show a simpler pattern.
The embedding are generally nonlinear and non-monotonic.
For the embedding learned for this particular task, U 0 , we observe that the turning points concentrate at the small distance , and the embedding
do not change much at the large distance. We conduct PCA analysis and visualize the learned embedding on a 2D space in Figure 2.
The embedding learned without regularization does not show a clear manifold, while the embedding learned with regularization forms a low-dimensional manifold.
This 2-D manifold is not a straight line, consistent with the non-linearity and non-monotonicity observed above.
We further investigate the learned distance embedding in four aspects: non-linearity, nonmonotonicity, diversity, and smoothness.
• The non-linearity is measured based on Pearson’s correlation, ρ Pearson , between the distance and the embedding.
Non-linearity = 1 − 1 ls (cid:80) li =1 (cid:80) sj =1 ρ 2 Pearson ( h ij , x ) , where l is the number of model layers and h ij
is the j -th dimension of the embedding in the i -th layer.
• The non-monotonicity is quantified based on Spearman’s correlation, ρ Spearman , between the distance and the embedding
[Figure]
[Figure]
• The diversity indicates the average similarity between different dimensions of the learned embedding.
A high diversity indicate that the embedding consists of diverse transforms.
[Figure]
• The smoothness is defined based on the smoothness loss Smoothness = 1 − L smooth .
[Table]
These metrics are calculated for the embedding learned with smoothness regularization with EGNN model on QM9 dataset.
For comparison, we consider a toy embedding made of three polynomial functions as its three dimensions: h polynomial ( x ) = { x, x 2 , x 3 } .
As illustrated in Table 3, the learned embedding generally show higher non-monotocity and diversity compared to the polynomial embedding.
The smoothness of the learned embedding are close to the polynomial embedding for some tasks.
This indicates that, the proposed positional encoding method learn a embedding that contain multiple nonlinear or non-monotonic transformations, yet remain smooth.
More physical-based interpretation of the learned embeddings are relegated to the appendix.
4.4 Physics-based Interpretation
In this appendix, we present some physic-based interpretation of learned embedding in various tasks from Section 4.2.
Different tasks show different characteristics, as illustrated in Table 3.
U [-
-]0 and H show a high non-linearity, while µ and R 2 show a low non-linearity.
In this section, we conduct a few case studies to investigate the connection between embedding characteristics and the physical nature of the corresponding task.
4.4.1 Short vs. Long-Distance Physics
The smoothness regularization encourage the embedding to only change with x when such change increases the accuracy.
By analyzing how much the learned distance embedding changes with distance, we know which region of distance the model learned to focus on.
We quantify the change by the normalized derivative.
[Equation]
As illustrated in Figure 3, for both U 0 and H , ˆ g ( x ) is large at small ˆ x , and then quickly drop to a small value close to zero as ˆ x increases.
This indicates that the embedding values do not change much when the distance is greater than certain ”cutoff” value.
This is consistent with the physical nature of the problem.
U 0 and H of the system considered in QM9 are generally dominated by short-range interactions.
There is no significant long-range forces such as electrostatic force in these systems.
In contrast, for both µ and R 2 , ˆ g ( x ) show a much flatter profile.
This indicates that both long distance and short distance are important.
This observation is consistent with the physical natural of the tasks.
Dipole moment ( µ ) is roughly the charges timed by the distance.
By definition it is a quantity sensitive to both short and long distance.
Therefore the embedding learned for this task has different value at different distance, as shown by the non-zero ˆ g ( x ) over full ˆ x range.
Electronic spatial extent ( R 2 ) reports how far out the electronic density extends with significant probability.
This roughly measure the molecule size and is obviously not just sensitive to short interatomic distance.
These case studies illustrate that the learned embedding reflect the physical dependence on short and long distance for different tasks.
4.4.2 Monotonic vs. Non-monotonic Physical Dependence
The non-monotonicity of the learned distance embedding illustrated in Table 3 is rooted in the physical nature of the tasks.
For physical quantity governed by non-monotonic dependence on distance, the corresponding learned distance embedding shows a high non-monotonicity.
The internal energy is connected to the potential forces acting on the atoms.
The force can be non-monotonic w.r.t.
distance.
Short distance is often dominated by repulsive forces while attractive forces are more significant at larger distance.
Consistent with this non-monotonic physics, The embeddings learned for U 0 and U 0 have the highest non-monotonicity score, 0.33, listed in Table 3.
In contrast, if there is no strong non-monotonic physical dependence on distance, the learned distance embedding has a low non-monotonicity score.
For instance, R 2 roughly measure the molecule size, as mentioned above, and is monotonic with distance.
The distance embedding learned for Rconsequently shows the a relatively low non-monotonicity score, 0.18.
4.4.3 Similar Physics have Similar Distance Embedding
The tasks can be closely related by their similar physical nature.
Some relation is obvious.
Both U and U 0 are internal energy, and there is a high correlation between their values in the QM9 dataset.
Some physical connection is not straightforward.
For example, the energy gap ∆ ϵ have been used to predict the polarizability α [33, 34], but such relation can not be simply observed in the original data.
The Pearson correlation between α and ∆ ϵ is close to zero in QM9.
Can the learned distance embedding capture such physical relation?
We hypothesize that if two tasks are governed by the similar set of physical laws, the embedding learned for these two tasks should share certain similarity.
We quantify the similarity between the embedding for task a and task b based on Pearson correlation.
[Equation]
The pairwise similarity is visualized in Figure 4.
The energies U , U 0 , H , G and ZPVE show a high similarity ( ≥ 0 . 88 ) between each other.
Another group of physical quantities, α , µ , and ∆ ϵ are closely related to the reactivity of the molecules.
The pairwise embedding similarity for this group is high ( ≥ 0 . 90 ).
In comparison, the similarity cross these two groups are relatively low.
This observation indicates that tasks of the similar physical nature have the similar embedding.
Conclusion
We propose and demonstrate a regularized positional encoding method for molecular properties prediction tasks.
As a plug-in module, the proposed method improves the accuracy of three model architecures, over variety molecular property and force field prediction tasks.
The novel smoothness regularization technique encourages the learned embedding to form a simpler low-dimensional manifold for easier physics-based interpretation.
Key characteristics of the embedding are connected to the physical nature of the tasks.
Tasks of similar physics have the similar learned embedding.