Coordinating Cross-modal Distillation for Molecular Property Prediction ^††thanks: * These authors contributed equally.

Many works focus on applying GNN to 2D graphs. For example, Gilmer et al.[2] proposed Message Passing Neural Networks, which combine message passing and aggregating algorithms with GNN, forming a successful framework for MPP. Liu et al.[3] introduced a simple unsupervised representation for molecules that embeds the vertices in the molecule graph and constructs a compact feature by assembling the vertex embeddings. Ying et al.[10] proposed Graphormer, which utilizes a Transformer [11] in the GNN by effectively encoding the graph structure information. Although these methods have shown promising capabilities, their potential is limited as they do not use 3D coordinates that better represent molecular structure and energy.

3D molecular structures contain spatial information important to molecular property prediction, such as bond angles and lengths. Many recent methods try to answer how to utilize this 3D information in the GNN model. For instance, Klicpera et al.[5] proposed to let message embeddings interact based on the distance between atoms and the angle between directions to improve quantum mechanical property prediction. Lu et al.[12] proposed a Multilevel Graph Convolutional neural Network that extracts features from the conformation and spatial information with multilevel interactions. Liu et al.[7] proposed SphereNet, a 3D graph network framework with spherical message passing. 3D-based methods achieve significant performance improvements. But obtaining 3D structures is time-costing, which reduces the value of 3D-based methods in large-scale applications.

These methods use 3D conformers and 2D structures to construct positive and negative pairs while applying a contrastive loss during training. For example, Liu et al.[13] proposed the Graph Multi-View Pretraining framework, in which they perform self-supervised learning by leveraging the correspondence and consistency between 2D topological structures and 3D geometric views. Li et al.[14] developed GeomGCL, a dual-view geometric message passing network that utilizes the molecular geometry across 2D and 3D views. Hu et al.[15] proposed a strategy of pre-training GNN at the level of individual nodes as well as entire graph, so the GNN can learn useful local and global information. Stark et al.[16] proposed 3D Infomax, which pre-trains a GNN by maximizing the mutual information between its 2D and 3D embeddings. The drawbacks of contrastive methods lie in the requirements of multiple conformers. Large-scale conformer generation is costly, while one SMILES code usually generates only one 3D conformer with a specific force field.

Knowledge distilling (KD) is an intuitive solution for 3D to 2D cross-modality training. As our work relies on Transformer-based methods such as Graphormer, we only discuss Transformer related KD methods. Many works [17, 18, 19, 20, 21] have attempted to apply KD to the Transformer-based model. Upon molecular property prediction, Zhu et al.[22] proposed ST-KD, an end-to-end Transformer KD framework, bridging the knowledge transfer between graph-based and SMILES-based models. But ST-KD only KD the atom to atom attentions, neglecting the informative embeddings in atom tokens. Based on our experiment, gradient magnitudes of atom-token distillation are discrepant and unstable due to the variable molecular size. Distilling atom-token embeddings without correction will cause degradation problems.

To address this issue, we propose a theoretical insight to justify how to coordinate atom token embeddings and molecular features and formulate an auto-correcting distilling weight to dynamically tune and balance the gradients according to the number of atoms.

Following the previous works (Park et al. 2022; Ying et al. 2021), we integrate an absolute position encoding (APE) into graphormer [10] illustrated in Fig.2. The APE is the adjacent edge information of the node, which could distinguish some atoms with the same atom id easily in the input layer. Let $X^{0}$ denote the input tokens which are integrates with the APE, $X^{0}$ is calculated as:

$$X^{0}=X\oplus E,$$

(1)

where $\oplus$ is used to gather the information from neighbors. In specific, the $i$-th input token is calculated as:

$$X_{i}^{0}=MLP(\sum_{i=1}^{N+1}e_{i,j}).$$

(2)

For the 2D view, ${e}_{i,j}$ is the feature of original Chemical bonds. For the 3D view, we utilize radial basis functions (RBF) [23] as the transformed function $\widetilde{e}_{i,j}=RBF(d_{ij})$ to obtain the edge information, where ${d}_{ij}$ is the distance between atom i and atom j. This 3D embedding can encode diverse geometric factors

Usually, for most of the existing distillation methods, the total loss is defined as:

$$\mathcal{L}(R,Y,\widetilde{R})={L}_{2D}(R,Y)+\mathcal{W}*L_{1}Loss({X},\widetilde{X}),$$

(8)

${L}_{2D}(R,Y)$ is the task-related loss, and $L_{1}Loss({X},\widetilde{X})$ is the distilling loss on all tokens and ${W}$ is the hyper-parameter. However, the virtual token ${X}_{0}$ is the target, which represents the global representation of the molecular, and the other tokens ${X}_{i}(i>0)$ are the atom representations that are auxiliary to the molecular representation. So the strategy of distillation should be divided into two modules: global molecular distillation and local atom distillation. In addition, We should select a proper weight to balance the global molecular distillation and local atom distillation.

The virtual token ${X}_{0}$ represents the molecular embedding and usually be used to predict molecular property, so distilling the virtual token embedding provides a global view of learning on 3D geometry structural knowledge. Considering that 2D graphs are seriously distinct from 3D conformers, we conduct the global molecular distilling from 3D to 2D on all layers which could learn more geometric information at different levels step by step. At last, the target is conducted by minimizing the following loss function given any pairs of $\widetilde{X}_{l,0}$ and $X_{l,0}$ in the training batch:

$$\mathcal{L}_{m}=\sum_{l=1}^{L}L_{1}Loss(\widetilde{X}_{0}^{l},X_{0}^{l})$$

(9)

The virtual token ${X}_{0}$ is the global view of molecular, and the token ${X}_{i}(i>0)$ is the local view of molecular. The virtual token is calculated from atom tokens by self-attention with the updating process. Considering the relationship between the virtual token and atom tokens, we use the local atom distilling to further boost the performance of the virtual token. One possible way is to calculate the similarity (distance) of each pairwise atom embedding, then the local atom calibration distillation loss is defined as:

$$\mathcal{L}_{a}=\sum_{l=1}^{L}\sum_{j=1}^{N+1}L_{1}Loss(\widetilde{X}_{j}^{l},X_{j}^{l})$$

(10)

For boosting the performance of distilling with global and local information, we formulate a novel coordinating weight according to the number of molecular size to scale gradient magnitudes. The analyzing and reducing process from two perspectives is as follows.

Coordinating the total loss Comparing the global molecular loss $\mathcal{L}_{m}$ and local atom loss $\mathcal{L}_{a}$, we can draw the conclusion that $\mathcal{L}_{a}$ is unstable due to the sum operation on all tokens. This means that when the molecular is very large(with big N), the final loss ${L}$ is exploding. So, we should add the function $f(N)$ according to the variable number of atoms to scale and calibrate the local atom loss $\mathcal{L}_{a}$. The formulation is as follows:

	$\displaystyle\mathcal{L}$	$\displaystyle=\mathcal{L}_{2D}+(\sum_{l=1}^{L}L_{1}Loss(\widetilde{X}_{0}^{l},X_{0}^{l})$		(11)
		$\displaystyle+f(N)*\sum_{l=1}^{L}\sum_{j=1}^{N+1}L_{1}Loss(\widetilde{X}_{j}^{l},X_{j}^{l}),$

$L$ is fixed and $N$ is variable which represents the number of atoms. In order to eliminate the effect of sum operation on all tokens, it’s intuitive to set $f(N)=\frac{1}{N}$. In addition, we couldn’t add $\frac{1}{N}$ to the molecular loss, which will dissipate the loss of the virtual token as N increases, so it’s necessary to separate global and local loss, and only scale the atom loss ${{L}_{a}}$. Generally, to make the total loss not explode, we should add a function $f(N)\propto{O}(\frac{1}{N})$ or $f(N)\propto{o}(\frac{1}{N})$. ${O}$ means “is of the same order as” and ${o}$ means “is ultimately smaller than”.

Coordinating the molecular gradient Since the virtual token is the target embedding, its gradient is much important. So we will calculate the relationship between its gradient and N. To simplify the forward of transformer, $X_{i}^{l+1}$ in $(l+1)$ layer is formed as:

$$X_{i}^{l+1}=MHA(X^{l})\\ \&=\sum_{j=0}^{N+1}({W}^{l}_{ji}*{X}^{l}_{j})$$

(12)

The ${W}^{l}_{ji}$ is the attention weight. In the back-propagation, the gradient of the virtual token in $l$th layer is as formulated:

	$\displaystyle{grad(X_{0}^{l})}$	$\displaystyle=\nabla_{X_{0}^{l}}{L}$		(13)
		$\displaystyle=\sum_{k=0}^{N+1}((\nabla_{X_{k}^{l+1}}{L})*(\nabla_{X_{0}^{l}}{X_{k}^{l+1}}))$

Finally, for transformer, the virtual token gradient $grad(X_{0}^{l+1})$ in ${l}$ layer can be described in follow:

$\displaystyle{grad(X_{0}^{l})}$

$\displaystyle\propto{{O}({N}^{2})},$

(14)

Different from multi-head attention in transformer, a simple GNN always use ${MLP}$ to calculate ${X}^{l+1}_{i}$ with ${X}^{l}_{0}$ . So $(\nabla_{X_{0}^{l}}{X_{k}^{l+1}})$ is constant, then the result is:

$\displaystyle{grad(X_{0}^{l})}$

$\displaystyle\propto{{O}({N})},$

(15)

For detailed proof, please refer to the Section1 in Supplementary Material. For transformer, the formula indicates that the gradient of virtual token $grad(X_{0}^{l})$ will be scaled by ${N}^{2}$, and N is the number of atoms. So scaled with $\frac{1}{N}$ couldn’t solve the unstable gradient of the virtual token. However, for a simple GNN, $\frac{1}{N}$ is enough.
Conclusion To this end, for transformer, we proved the weight $f(N)=\frac{1}{{N}^{2}}$ is suitable to scale the losses and balance the gradient magnitudes of the global molecular information and local atom information. So the total objective function is formulated as:

$$\mathcal{L}=\mathcal{L}_{2D}+(\mathcal{L}_{m}+\frac{1}{{N}^{2}}*\mathcal{L}_{a}).$$

(16)

PCQM4Mv2 dataset [9] served as an ”ImageNet Large Scale Visual Recognition Challenge” in the field of graph ML is a quantum chemistry dataset originally curated under the PubChemQC project. It contains 3.8 million molecular graphs. The 3D structure for training molecules is calculated by DFT, which is equipped with only one conformer for each SMILES code.

MolHIV dataset [24] is adopted from the MOLECULENET [25] with 41K molecular graphs for molecular property prediction. All the molecules are pre-processed using RDKIT [26]. Each graph represents a molecule, where nodes are atoms, and edges are chemical bonds. Input node features are 9-dimensional, containing atomic number and chirality, as well as other additional atom features such as formal charge and whether the atom is in the ring. Input edge features are 3-dimensional, containing bond type, bond stereochemistry as well as an additional bond feature indicating whether the bond is conjugated.

Table I presents the results for molecular property prediction task on the OGB-LSC PCQM4Mv2 datasets [9]. What’s worth mentioning is that contrastive methods are hard applicable because the PCQM4Mv2 is equipped with only one conformer for each SMILES code. Compared with established methods, our model is currently the best on the PCQM4Mv2 leaderboard. As shown in Table I, our approach achieves the significant MAE improvement by 6% compared with the state-of-the-art method, indicating the effectiveness of our proposed method.

We further investigate our proposed method on the MOLHIV dataset. The comparison results are shown in Table II. From Table II, we can see that our method outperforms other methods, achieving prominent improvement by 1.39% on AUC compared with baseline. In addition, the origin MOLHIV dataset is not attached with 3D information, so we use the rdkit to generate the 3D position as the 3D coordinates of atoms. It should be noted that the generated 3D coordinates with randomly selected 3D conformers are not matched the real coordinates, so the generated force fields are different from the real force fields. The weakness of the 3D teacher model limits the performance of the student. Our approach can still obtain improvement with biased conformers, outperforming state-of-the-art methods and software (e.g.rdkit) generated 3D coordinates-based results. These phenomenon demonstrate the effectiveness of our distilling framework.