DIRECTIONAL MESSAGE PASSING FOR MOLECULAR GRAPHS
DIRECTIONAL MESSAGE PASSING FOR MOLECULAR GRAPHS
3 REQUIREMENTS FOR MOLECULAR PREDICTIONS
Molecular Dynamics: Molecular Dynamics (MD) predicts the forces $$Fi$$ acting on each atom. We can predict a potential instead of forces and then obtaining the forces via backpropagation to the atom coordinates. $$F_i(X,z) = - \frac{\partial}{\partial x_i} f{\theta}(X,z)$$.
Further, we can add this into the loss function:
$$ L{MD}(X,z) = |f{\theta}(X,z) - \hat t(X,z) | + \frac{\rho}{3N} \sum{i=1}^{N} \sum{\alpha=1}^{3} \bigg| -\frac{\partial f\theta(X,z)}{\partial x{i \alpha} } - \hat F _{i \alpha}(X,z) \bigg|
$$
$$\hat t = \hat E$$ is the ground-truth energy. $$\hat F$$ is the ground-truth force.
4 DIRECTIONAL MESSAGE PASSING
$$hi^{l+1} = f{update}(hi^{l}, \sum{j \in Ni} f{int}(hj^{l} , e{ij}^{l} ) )$$
$$f{update}$$ is the update function $$f{int}$$ is the interaction function $$e_{ij}^{l}$$ is the edge embedding
Directionality.
Directional embeddings. each neighborhood breaks up to one rotational invariance, which also introduces additional degrees of freedom generate a separate embedding $$m_{ij}$$ is embedding from atom $$i$$ to atom $$j$$ by applying the same learned filter in the direction of each neighboring atom
Interaction via cosine basis. use the directional information associated with each embedding by leveraging the angle $$\alpha(kj,ji) = x_k x_j x_i$$ then use the cosine basis representation empirical findings that this basis representation provides a better inductive bias than the raw angle
Message embeddings. $$hi = \sum{j \in \mathcal{N}i} m{ji}$$ and update the message $$m{ji}$$ based on the incoming messages $$m{kj}$$
$$ m{ji} = f{update}(m{ji}, \sum{k \in \mathcal {N}-{i} } f{int} (m{kj}, e{RBF}, a{CBF} ) )
$$
5 BESSEL FUNCTIONS AS A RADIAL BASIS
interaction function depends on both the angle and pairwise distance earlier works used a set of Gaussian radial basis functions, with means distributed uniformly or exponentially we propose to use orthogonal basis, which can provide a helpful inductive bias
From Schrodinger to Bessel.
6 DIRECTIONAL MESSAGE PASSING NEURAL NETWORK (DIMENET)
Directional Mesage Passing Neural Network (DimNet) is based on a streamlined version of PhysNet, where we integrated directional message passing and spherical Bessel functions.
Embedding block the atom type embeddings are genrated ramdonly $$h_i$$ the message embeddings is initialized as
$$ m{ji} = \sigma([h_j || h_i || e{RBF} ]W + b)
$$
Interaction block resembles the Residual block and skip connection
Output block messages are summed up per atom to get $$hi = \sum_j m{ji}$$ then transformed using multiple dense layer to generate the atom-wise output $$t_i^l$$ the final output is $$t = \sum_i \sum_l t_i^l$$
Continuous Differentiability multiple choices were made to reach twice continuous differentiability. 1. make the activation function $$\sigma(x) = x$$ instead of ReLU 2. multiply the radial basis functions $$e_{RBF}(d)$$ with an envelope function $$u(d)$$, the first and second derivatives go to 0 at the cutoff c. 3. DimNet doesn't use auxiliary data, but atom types and positions alone
Experiments
Models SchNet, PhysNet, PPGN, MEGNet-simple, Cormorant, sGDML
QM9 molecular properties prediction
MD17 molecular dynamics
Ablation studies to study the effect of directional message passing and radial Bessel asis, we ablate them individually and compare the standaridized MAE and logMAE on QM9