SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weitao Du $^{1,3}$  * Jiujiu Chen $^{2,3}$  Xuecang Zhang $^{3}$  Zhiming Ma $^{1}$  Shengchao Liu $^{4}$ †
$^{1}$  Department of Mathematics, Chinese Academy of Sciences
2 Shenzhen Institute for Advanced Study, University of Electronic Science and Technology of China
<sup>3</sup> Huawei Technologies Ltd
$^{4}$  Department of Computer Science and Operations Research, Université de Montréal
Recently, artificial intelligence for drug discovery has raised increasing interest in both machine learning and chemistry domains. The fundamental building block for drug discovery is molecule geometry and thus, the molecule's geometrical representation is the main bottleneck to better utilize machine learning techniques for drug discovery. In this work, we propose a pretraining method for molecule joint auto-encoding (MoleculeJAE). MoleculeJAE can learn both the 2D bond (topology) and 3D conformation (geometry) information, and a diffusion process model is applied to mimic the augmented trajectories of such two modalities, based on which, MoleculeJAE will learn the inherent chemical structure in a self-supervised manner. Thus, the pretrained geometrical representation in MoleculeJAE is expected to benefit downstream geometry-related tasks. Empirically, MoleculeJAE proves its effectiveness by reaching state-of-the-art performance on 15 out of 20 tasks by comparing it with 12 competitive baselines. The code is available on this website.
The remarkable progress in self-supervised learning has revolutionized the fields of molecule property prediction and molecule generation through the learning of expressive representations from large-scale unlabeled datasets [1-9]. Unsupervised representation learning can generally be categorized into two types [7, 10-13]: generative-based methods, encouraging the model to encode information for recovering the data distribution; and contrastive-based methods, encouraging the model to learn invariant features from multiple views of the same data. However, despite the success of large language models like GPT-3 [14, 15] trained using an autoregressive generation approach, learning robust representations from molecular data remains a significant challenge due to their complex graph structures. Compared to natural language and computer vision data [16, 17], molecular representations exhibit more complex graph structures and symmetries [18]. As molecular dynamics follows the principle of non-equilibrium statistical mechanics [19], diffusion models, as a generative method inspired by non-equilibrium statistical mechanics [20, 21], are a natural fit for 3D conformation generation. Previous researches have demonstrated the effectiveness of diffusion models, specifically 2D molecular graphs [22, 23] or 3D molecular conformers [24-27], for molecular structure generation tasks. However, a crucial question remains: Can diffusion models be effectively utilized for jointly learning 2D and 3D latent molecular representations?
*duweitao@mass.ac.cn
†shengchao.liu@umontreal.ca
Figure 1: Pipeline of MoleculeJAE. For each individual molecule, MoleculeJAE utilizes the reconstructive task to perform denoising. For pairwise molecules, it conducts a contrastive learning paradigm to fit the trajectory.
Answering this question is challenging, given that diffusion models are the first generative models based on trajectory fitting. Specifically, diffusion models generate a family of forward random trajectories by either solving stochastic differential equations [28, 29] or discrete Markov chains [23, 30, 31], and the generative model is obtained by learning to reverse the random trajectories [32, 33]. It remains unclear how the informative representation manifests itself in diffusion models, as compared to other generative models that explicitly contain a semantically meaningful latent representation, such as GAN [34] and VAE [35]. In addition to the generative power of diffusion models, we aim to demonstrate the deep relationship between the forward process of diffusion models and data augmentation [36, 37], a crucial factor in contrastive learning. As a result, we ask whether a trajectory-based self-supervised learning paradigm that leverages both the generative and contrastive benefits of diffusion models can be used to obtain a powerful molecular representation.
Our approach. This work presents MoleculeJAE (Molecule Joint Auto-Encoding), a novel trajectory learning framework for molecular representation that captures both 2D (chemical bond structure) and 3D (conformer structure) information of molecules. Our proposed method is designed to respect the  $SE(3)$  symmetry of molecule data and is trained by fitting the joint distribution of the data's augmented trajectories extracted from the forward process of the diffusion model. By training the representation in this manner, our framework not only captures the information of real data distribution but also accounts for the correlation between the real data and its noised counterparts. Under certain approximations, this trajectory distribution modeling decomposes into a marginal distribution estimation and a trajectory contrastive regularization task. This multi-task approach yields an effective and flexible framework that can simultaneously handle various types of molecular data, ranging from SE(3)-equivariant 3D conformers to discrete 2D bond connections. Furthermore, in contrast to diffusion models used for generation tasks that only accept noise as input, most downstream tasks have access to ground-truth molecular structures. To leverage this advantage better, we incorporate an equivariant graph neural network (GNN) block into the architecture, inspired by [38, 39], to efficiently encode crucial information from the ground-truth data. Analogous to conditional diffusion models [38, 40], the encoder's output serves as a meaningful guidance to help the diffused data accurately fitting the trajectory. In summary, our self-supervised learning framework unifies both contrastive and generative learning approaches from a trajectory perspective, providing a versatile and powerful molecular representation that can be applied to various downstream applications.
Regarding experiments, we evaluate MoleculeJAE on 20 well-established tasks drawn from the geometric pretraining literature [6, 7], including the energy prediction at stable conformation and force prediction along the molecule dynamics. Our empirical results support that MoleculeJAE can outperform 12 competitive baselines on 15 tasks. We further conduct ablation studies to verify the effectiveness of key modules in MoleculeJAE.
In this section, we introduce the diffusion mechanism and relevant notations as a powerful data augmentation framework and elaborate on its instantiation in the context of molecular graph data.
2.1 Diffusion Mechanism as a trajectory augmentation
The concept of diffusion is widely used in various fields such as physics, mathematics, and computer science. In this paper, we take a broad view of diffusion by defining it as a Markov process with a discrete or continuous time index  $t$ , starting from a given data point  $x_0$  and producing a series of random transformations that lead to  $x_t$ . If  $t \in \{0, 1, 2, \ldots\}$ , then  $x_t$  is a discrete Markov chain where the probability distribution  $p(x_t)$  only depends on the distribution of  $x_{t-1}$ . We further refer to the conditional probability  $p(x_t | x_{t-1})$  as the transition probability of the Markov chain. This definition enables us to model the evolution of data over time, allowing us to generate a sequence of augmented data points that span different temporal regimes.
In fact, we can adapt classical data augmentations, such as spatial rotation, color distortion, and Gaussian blurring [41, 42], as sources for defining  $p(x_{t}|x_{t - 1})$ . For instance, Gaussian blurring generates a specific transition probability by setting  $p(x_{t}|x_{t - 1}) = \mathcal{N}(x_{t - 1},\epsilon)$ , where the scaling factor  $\epsilon$  may depend on time  $t$ . This transition probability rule perturbatively transforms the data for each step, and connecting all steps of transformed data produces a trajectory. In a similar vein, random masking can be seen as a discrete Markov chain, while continuous rotation can be viewed as a deterministic Markov process. By employing this approach, we can effectively expand the scope of one-step data augmentation operations to encompass trajectory-based scenarios.
In this paper, we build upon the concepts of generalized diffusion processes proposed in [33] to establish a unified framework for the trajectory augmentation examples mentioned earlier. Following the general approach in [33], each (cold) diffusion model is associated with a degradation Markov process  $D(\cdot ,t)$  indexed by a time variable  $t\in [0,T]$ , which introduces perturbations to the initial data  $x_0$ :
$$
x _ {0} \xrightarrow {D (\cdot , t)} x _ {t}. \tag {1}
$$
Given that  $x_0 \sim p_{data}$ , the transformed variable  $x_t \coloneqq D(x_0, t)$ , is also a random variable. To represent the marginal distribution of  $D(\cdot, t)$  at time  $t$ , we use the notation  $p_t(\cdot)$ . When the context is clear, we will use the shorthand notation  $D_t$  to refer to  $D(\cdot, t)$ . According to Equation (1), we define the pair  $(x_0, x_t)$  as a multi-view of the sample  $x_0$  at a specific time  $t$ . It is important to note that, under certain conditions, a Markov process converges to a stationary distribution that effectively erases any memory of the initial data  $x_0$ . Nevertheless, for small values of  $t$ , there is a high probability that  $x_t$  remains within a ball centered around  $x_0$ .
In the context of the standard diffusion-based generative model [28, 32], the presence of a reverse process  $R_{t}$  is necessary to undo the effects of the degradation  $D_{t}$ . In [28], the degradation  $D(\cdot ,t)$  is identified as the solution of a stochastic differential equation (SDE), and it has been demonstrated that a family of SDEs, including the probability flow ordinary differential equation (ODE) proposed in [28, 43], can reverse the degradation process  $D(\cdot ,t)$ . Specifically, the inversion property implies that  $R_{T - t}$  shares the same marginal distribution as  $p_t$ :  $R_{T - t}\sim p_t$ . Therefore, as  $t$  varies from 0 to  $T$ , the reverse process  $R_{t}$  gradually restores the data distribution. On the other hand, for cold diffusions, although the model is also trained by reconstructing the data from  $D_{t}$ , it does not explicitly assume the existence of a reverse process as rigorously as in the continuous diffusion-based generative model.
Example of heat diffusion We use the term 'heat' to describe a diffusion process that involves injecting Gaussian noise. Under this definition, the continuous limit of heat diffusion can be expressed by the solution of stochastic differential equations (SDEs):
$$
d x _ {t} = \mu \left(x _ {t}, t\right) d t + \sigma (t) d w _ {t}, \quad t \in [ 0, T ] \tag {2}
$$
Here,  $w_{t}$  represents the standard Brownian motion. It is worth noting that while the solution  $x_{t}$  of a general SDE may not follow the Gaussian distribution, for the commonly used SDEs such as the Variance Exploding SDE (VE) and Variance Preserving SDE (VP), the solution can be explicitly expressed as:
$$
x _ {t} = \alpha (t) x _ {0} + \beta (t) z, \tag {3}
$$
where  $z$  is sampled from  $\mathcal{N}(0, I)$ , and both  $\alpha(t)$  and  $\beta(t)$  are positive scalar functions. Using the forward SDE (2), we approach the task of data synthesis by gradually denoising the noisy observations  $x_{t}$  to recover  $x_{0}$ , which is also accomplished through solving a reverse SDE. For the general formula and further details, readers can refer to [28, 32]. One fact we will use later is that the reverse SDE fundamentally relies on the score function  $\nabla_{x} \log p_{t}(x)$ , which serves as the parameterization objective in the score matching framework [28, 44].
Example of discrete (cold) diffusion In contrast to the continuous heat diffusion framework, we now introduce a family of discrete Markov Chains as an instance of discrete (cold) diffusion. This example is particularly suitable for modeling categorical data. In the unconditional setting, the transition probability of the discrete diffusion at time step  $t$  is defined as:
$$
q \left(x _ {t} \mid x _ {t - 1}\right) = \text {M u l t i n o m i a l} \left(x _ {t}, p = x _ {t - 1} Q _ {t}\right), \tag {4}
$$
where  $Q_{t}$  represents a Markov transition matrix. An illustrative example is the absorbing diffusion, as introduced by [30], where
$$
Q _ {t} = \left(1 - \beta_ {t}\right) I + \beta_ {t} e _ {m} ^ {T},
$$
and  $e_m$  is a one-hot vector with a value of 1 at the absorbing state  $m$  and zeros elsewhere. Typically, the state  $m$  is chosen as the 'masked' token. Hence, absorbing diffusion can be regarded as a masking process with a dynamical masking ratio  $\beta_t$ . Similarly to the heat diffusion, the corresponding reverse process can also be formulated as a discrete diffusion. However, the training objective shifts from fitting the score function (which only exists in the continuous case) to directly fitting the 'reconstruction probability'  $p(x_0|x_t)$ .
2.2 Graph-based molecular representation
In this paper, our primary focus is on the graph representation of molecules. Let  $G = (V,E,P,H)$  denote the integrated molecular graph, consisting of  $n := |V|$  atoms and  $|E|$  bonds. The matrix  $P \in \mathbf{R}^{n \times 3}$  is utilized to represent the 3D conformer of the molecule, containing the positions of each node. Moreover, within the graph  $G$ , the edge collection  $E$  represents the graphical connections and chemical bond types between atoms. Additionally, we have the node-wise feature matrix  $H \in \mathbf{R}^{n \times h}$  where, for the purposes of this article, we consider the formal charges and atom types as the components of  $H$ . In the method section, we will demonstrate how to build trajectories based on  $G$ .
Now we illustrate the process of obtaining molecule augmented trajectories, based on which, MoleculeJAE is proposed to estimate the trajectory distribution. In Section 3.1, we will outline the construction of equivariant molecular trajectories. Subsequently, in Section 3.2, we will introduce our theoretical self-supervised learning framework for these trajectories. Our hypothesis is that a good representation should encode the information from the distribution of augmented trajectories. This hypothesis forms the practical basis of our core reconstructive and contrastive loss, which will be discussed in Section 3.3. Finally, in Section 3.4, we will present the key architectures. A comprehensive discussion of the related works is in Appendix B.
3.1 Equivariant Molecule Trajectory Construction
In the field of molecular representation learning, our objective is to jointly estimate the distribution of a molecule's 2D topology (including atom types and chemical bonds) and its 3D geometries (conformers). Building upon the notations in Section 2.2, our goal is to construct augmented trajectories  $x_{t} \coloneqq (H(t), E(t), P(t))$  from  $G$ . The challenging aspect lies in preserving the  $SE(3)$  symmetry of the position matrix  $P$ , which we will describe in detail below.
Given two 3D point clouds (molecular conformers)  $P_{1}$  and  $P_{2}$ , we say they are SE(3)-isometric if there exists an  $R \in SE(3)$  such that  $P_{1} = RP_{2}$ . In probabilistic terms, let  $p_{3D}(x_{3D})$  be the probability density of all 3D conformers denoted by  $\mathcal{C}$ . Isometric conformers necessarily have the same density, i.e.,
$$
p _ {3 D} \left(x _ {3 D}\right) = p _ {3 D} \left(\mathbf {R} x _ {3 D}\right), \quad \forall \mathbf {R} \in S E (3), x _ {3 D} \in \mathcal {C}. \tag {5}
$$
Utilizing the symmetry inductive bias has been shown [24, 45] to greatly reduce the complexity of data augmentation. Additionally, in order to respect the  $SE(3)$  symmetry, the augmented trajectory  $P(0) \to P(t)$  should also be SE(3)-equivariant, satisfying
$$
\mathbf {R} P (x _ {3 D}) = P (\mathbf {R} x _ {3 D}), \forall \mathbf {R} \in S E (3).
$$
This condition imposes a rigorous restriction on the form of the forward SDE (2) when it is applied to generate the augmented trajectory of conformers. However, if we constrain the form of  $x_{t}$  to Eq. 3, the  $SE(3)$  equivariance is automatically satisfied due to the  $SE(3)$  equivariance of both the Gaussian random variable  $z$  and the original data  $x_{0}$ . We leave the formal proof in Appendix A.
Regarding the 2D component, let  $x_{2D}(t) \coloneqq (H(t), E(t))$ , where  $x_{2D}$  consists of invariant scalars that remain unchanged under  $SE(3)$  transformations. Although  $x_{2D}$  is categorical in nature, we can treat them as continuous scalars that can be perturbed by Gaussian noise. During inference (generation), these continuous scalars are quantized back into categorical integers (see [22] for details). This allows both  $x_{2D}(t)$  and  $P(t)$  to undergo heat diffusion. By combining the 2D and 3D parts, we introduce the system of SDEs for  $x_t$ :
$$
\left\{ \begin{array}{l} d P (t) = - P (t) d t + \sigma_ {1} (t) d w _ {t} ^ {1}, \\ d H (t) = \mu_ {2} (H (t), t) d t + \sigma_ {2} (t) d w _ {t} ^ {2}, \\ d E (t) = \mu_ {3} (E (t), t) d t + \sigma_ {3} (t) d w _ {t} ^ {3}. \end{array} \right. \tag {6}
$$
It is worth mentioning that although the components of Eq. 6 are disentangled, the corresponding reverse diffusion process and its score function are entangled. A common choice of  $\mu (x,t)$  is  $\mu (x,t)\coloneqq -x$ , then will utilize the explicit solution of equations (6) to generate the molecular augmentation trajectory. Additionally, it is also possible to perform equivariant diffusion of the 2D and 3D joint representation through equivariant cold diffusion, following the approach described in (4). The detailed framework for this approach is provided in the Appendix A.
3.2 Equivariant Molecule Trajectory Learning with Density Auto-Encoding
To provide a formal foundation for self-supervised learning from augmented trajectories, we revisit the concept of denoising from a trajectory perspective. According to the Kolmogorov extension theorem [46], every stochastic process defines a probabilistic measure on the trajectory space, uniquely determined by a set of finite-dimensional joint distributions that satisfy consistency conditions. Therefore, estimating the probability of the infinite-dimensional trajectory space is equivalent to estimating the joint probabilities  $p(x_{t_1},\ldots ,x_{t_k})$  for each finite time sequence  $t_1,\dots ,t_k\in [0,T]$ . From the standpoint of data augmentation, our specific focus is on learning the joint distribution of trajectory augmentation pairs:  $p(x_0,x_t)$  (solution of Eq. 6), which determines how the noisy  $x_{t}$  corresponds to the original (denoised)  $x_0$ .
To differentiate our notion of "Auto-Encoding" from the "denoising" method utilized in DDPM [32] (denoising diffusion probabilistic models), it is important to highlight that traditional diffusion-based generative models are based on marginal distribution modeling. In this framework, joint distributions induce marginal distributions, but not vice versa. This distinction becomes more apparent in the continuous SDE formalism of diffusion models [28], where multiple denoising processes from  $x_{t}$  to  $x_{0}$  can be derived while still sharing the same marginal distributions. However, the joint distributions of a probabilistic ODE flow significantly differ from those of the reverse SDE (as shown in Eq. 11) due to the deterministic nature of state transitions in an ODE between timesteps. In the following, we formally demonstrate this observation from the perspective of maximizing the trajectories' joint probabilistic log-likelihood.
For a given time  $t \in [0, T]$ , let us assume that the probability of the random pair  $(x_0, x_t)$  defined in Eq. 1 is determined by a joint density function  $p(x_0, x_t)$ . Our objective is to approximate  $p(x_0, x_t)$  within a variational class  $p_\theta(x_0, x_t)$ . Maximizing the likelihood of this joint distribution leads to the following optimization problem:
$$
\operatorname {a r g m a x} _ {\theta} \prod_ {i = 1} ^ {n} p _ {\theta} \left(x _ {0} ^ {i}, x _ {t} ^ {i}\right), \tag {7}
$$
where  $\{(x_0^i,x_t^i)\}_{i = 1}^n$  represents the collection of  $n$  augmented samples from the training dataset. Following the tradition of Bayesian inference [47], we parameterize  $p_{\theta}$  by defining a joint energy function  $\mathbf{E}_{\theta}(x_0,x_t)$  such that:  $p_{\theta}(x_0,x_t) = \frac{1}{Z_\theta} e^{-\mathbf{E}_{\theta}(x_0,x_t)}$ , where  $Z_{\theta}(t)$  is the intractable normalization constant depending on  $t$ . Therefore, the maximal likelihood framework reduces to solving
$$
\operatorname {a r g m i n} _ {\theta} \mathbb {E} _ {p (x _ {0}, x _ {t})} \left[ \mathbf {E} _ {\theta} (x _ {0}, x _ {t}) \right].
$$
However, directly optimize  $p_{\theta}(x_0,x_t)$  by taking the gradient with respect to  $\theta$  is challenging because  $\frac{\partial\log p_{\theta}(x_0,x_t)}{\partial\theta}$  contains an intractable term:  $\frac{\partial Z_{\theta}(t)}{\partial\theta}$ . To circumvent this issue, we borrow ideas from [48, 49] by treating the transformed  $x_{t}$  as an infinitely dimensional "label" of  $x_0$ . More precisely, we consider the parameterized marginal density  $q_{\theta}(x_0)$  as:
$$
q _ {\theta} \left(x _ {0}\right) := \int p _ {\theta} \left(x _ {0}, x _ {t}\right) d x _ {t}, \tag {8}
$$
which involves integrating out  $x_{t}$ . We define the marginalized energy function with respect to  $x_{0}$  as:  $\bar{\mathbf{E}}_{\theta}(\cdot)\coloneqq -\log \int \exp (-\mathbf{E}_{\theta}(\cdot ,x_{t}))dx_{t}$ . Now, let  $f_{\theta}(x_0,x_t)$  denote the normalized conditional density function  $p_{\theta}(x_{t}|x_{0})$ , we have  $\frac{\partial\log f_{\theta}(x_t|x_0)}{\partial\theta} = -\frac{\partial\mathbf{E}_{\theta}(x_0,x_t)}{\partial\theta} +\frac{\partial\bar{\mathbf{E}}_{\theta}(x_0)}{\partial\theta}$ . By taking the empirical expectation with respect to the finitely sampled pair  $(x_0,x_t)\sim p(x_0,x_t)$ , we can decompose the gradient of the maximum likelihood as follows (see Appendix A for the full derivation):
$$
\tilde {\mathbb {E}} _ {p \left(x _ {0}, x _ {t}\right)} \left[ \frac {\partial \log p _ {\theta} \left(x _ {0} , x _ {t}\right)}{\partial \theta} \right] = \tilde {\mathbb {E}} _ {p \left(x _ {0}\right)} \left[ \frac {\partial \log q _ {\theta} \left(x _ {0}\right)}{\partial \theta} \right] + \tilde {\mathbb {E}} _ {p \left(x _ {0}, x _ {t}\right)} \left[ \frac {\partial \log f _ {\theta} \left(x _ {0} , x _ {t}\right)}{\partial \theta} \right], \tag {9}
$$
here we use  $\tilde{\mathbb{E}}$  to denote the expectation with respect to the empirical expectation. Note that this decomposition holds for  $(x_{s}, x_{t})$  for any two different time steps  $(0 \leq s, t \leq T)$  of the trajectories. In the next section, we decompose Equation (16) into two sub-tasks, leading to our goal of optimizing these two parts simultaneously, as discussed below.
3.3 Reconstructive and Contrastive Tasks
In what follows, we denote the latent representation of data  $x$  by  $h_\theta(x)$  (the equivariant model for building  $h_\theta$  will be discussed in the next section). Based on Eq. 16, we introduce two tasks for training  $h_\theta(x)$  that stem from this decomposition.
Reconstructive task. The first term  $\tilde{\mathbb{E}}_{p(x_0)}\left[\frac{\partial\log q_\theta(x_0)}{\partial\theta}\right]$  in Eq. 16 aims to reconstruct the distribution of data samples. Therefore, we refer this term as the reconstruction task, which involves modeling the marginal distribution  $p_{data}(x_0)$  using  $q_{\theta}(x_0)$ .