# Learning Hierarchical Representations for Explainable Chemical Reaction Prediction
## 3 Method
### Learning Strategy and Objectives  
#### 3.2.1 Coarse-Fine Level Contrastive Learning  
In order to improve the generalization ability of the proposed model, we apply the contrastive learning strategy. In addition, we design specific data augmentation strategies for coarse-level and fine-level contrastive learning, respectively. We first introduce the contrastive learning and then describe the coarse-fine level data augmentation as follows. Contrastive learning is effective for preserving the generalization of deep learning models by maximizing the lower bound of the mutual information (MI) of the input and the output representation. The lower bound is calculated by the MI of the output representations extracted by the deep models from two augmented views of the original input. To alleviate the collapse caused by large variations between the positive pairs that especially and frequently occur in coarse-level learning, we use an asymmetric mutual information neural estimation loss [30] to separately optimize our model at the coarse level and the fine level. Since the objective calculations at the two levels are the same, we take the coarse-level processing as an example and show the loss function of the \(j\)-th instance in the training set:  
\[L_{CL\_coarse}(j)=-\frac{1}{2}(h_{pred}^{\top}\big{(}\mathbf{c}_{j}\big{)}\hat\mathbf{{c}}_{p}-h_{pred}^{\top}\big{(}\mathbf{c}_{p}\big{)}\hat{\mathbf{c}}_{j})+\lambda_{CL} \log(\sum_{\mathbf{c}_{n}\in\text{Neg}(\mathbf{c}_{j})}\exp(\mathbf{c}_{j}^{\top}c_{n})), \tag{5}\]  
where \(h_{pred}\) is the prediction head following [30,31], \(\mathbf{c}_{p}\in\text{Pos}\big{(}\mathbf{c}_{j}\big{)}\) is the sequence encoding feature of the augmented view of the \(j\)-th instance, and variables with hat symbols denote that the stop gradient operations are applied to them. For the data augmentation, we take inspiration from the method [28] which manages to augment data by randomly changing atoms or bonds outside the reaction centers because these atoms and bonds are relatively not important in the reaction. In this paper, we design a probabilistic data augmentation for the SMILES representation. For the \(i\)-th token in the SMILE sequence, we define \(p_{i}\) to represent the probability of changing this token in the augmentation operation. The larger \(p_{i}\) means the farther token is away from the reaction center, and vice versa. We apply two kinds of changes to the SIMILES tokens. First, the rotation of the molecule SMILES randomization [32,33] is introduced and \(p_{i}\) is the probability of conducting the rotation operation on the left part of the \(i\)-th token in the molecule. Second, we randomly move the branched chain to other places with the same 1-step local context in the molecule, and \(\sum_{i\in\mathbf{S}}p_{i}\) is the probability of whether to move the branched chain, where \(\mathbf{S}\) denotes the set of the token indices of the corresponding branched chain. An example of the second kind can be found in the box of \"Fine-level augmentation\" in Figure 1b. In this work, the previously introduced attention coefficient vectors \(\mathbf{\alpha}_{coarse}\) and \(\mathbf{\alpha}_{fine}\) measure the importance of the token in the chemical reaction, and we believe that tokens with larger coefficient values are more likely to be the reaction centers. The probability \(p_{i}\) can be computed according to the attention coefficients. Specifically, for the data augmentation at the coarse level, \(p_{i}\) is calculated by  
\[p_{i}=\text{Softmax}(1-\mathbf{\alpha}_{coarse}[i],\tau), \tag{6}\]  
where \(\text{Softmax}(\cdot)\) is the Softmax operation and \(\tau\) denotes the temperature. During training, we gradually decrease the value of the temperature. Besides, we add a data augmentation strategy that directly regards two reactions in the same coarse-level but different fine-level categories as a positive pair. An example can be seen in the box of \"Coarse-level augmentation\" in Figure 1b.