Molecule Joint Auto-Encoding: Trajectory Pretraining with 2D and 3D Diffusion
Authors: weitao Du
Jiujiu Chen
Xuecang Zhang
Zhi-Ming Ma
Shengchao Liu
Venue: NeurIPS 2023 poster (2023)
Keywords: Molecule Joint Auto-encoding
Molecule Joint Self-supervised Learning
Markov processes
contrastive learning
molecule representation learning
TL;DR: MoleculeJAE, a pretraining method for joint auto-encoding of molecule trajectories，effectively captures both 2D bond topology and 3D conformation information。

Abstract: 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.

Molecule Joint Auto-Encoding: Trajectory Pretraining with 2D and 3D Diffusion
Weitao D3
Jiujiu Chen
Xuecang Zhang
Zhiming Ma! Shengchao Liu
Department of Mathematics, Chinese Academy of Sciences 2 Shenzhen Institute for Advanced Study, University of Electronic Science and Technology of China 3 Huawei Technologies
4 Department of Computer Science and Operations Research, Université de Montréal
Abstract
Recently, artificial intelligence for drug discovery has raised increasing interest both machine learning and chemistry domains. The fundamental building block for drug discovery is molecule geometry and thus, the molecule's 's geometrical representation is the main bottleneck to better utilize machine learning techniques for drug discovery. In this work, we propose pretraining method for molecule joint auto-encoding (MoleculeJAE). MoleculeJAE and 3D conformation (geometry) information, and a diffusion process model is applied to mimic the augmented trajectories of such two modalities, based which, MoleculeIAE will learn the inherent chemical structure in a self-supervised manner. Thus, the pretrained geometrical representation in MoleculeJAA is expected benefit downstream geometry-related tasks. Empirically, MoleculeJAE proves its effectiveness by reaching state-of-the-art performance 15 out of 20 tasks by comparing it with 12 competitive baselines. The code is available on this website
learn both the 2D bond (topology)
Introduction
The remarkable progress self-supervised learning has revolutionized the fields molecule property prediction and molecule generation through the learning of expressive representations from large- scale unlabeled datasets [199] Unsupervised representation learning can generally be categorized into types 7 10 K3
generative-based methods, encouraging the model 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][5] trained using an autoregressive generation approach learning robust representations from molecular data remains significant challenge due to their complex graph structures. Compared to natural language and computer vision data [6] 17), molecular representations exhibit more complex graph structures and symmetries [8]] molecular dynamics follows the principle of coneequilibrumm statistical mechanics 19] diffusion models, as generative method
21). are a natural for 3D conformation generation. Previous researches have demonstrated the effectiveness of diffusion models, specifically 2D molecular graphs (22 23 or 3D molecular conformers 274 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?
inspired by -quiiuuilirrium statistical mechanies (20
duveitao@mass shengchao. .i Liu@umontreal. ca
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Gaussian Distribution
Score Matching
Positive
2D 3D Structure Refinement
Nogativee
Figure 1: Pipeline MoleculeJAE. For each individual molecule, MoleculeJAE utilizes reconstructive task to perform denoising. For pairwise molccules, it conducts a contrastive learning paradigm to the trajectory.
Answering this question is challenging, given that diffusion models are the first generative models based trajectory fitting. Specifically, diffusion models generate family of forward random trajectories by either solving stochastic differential equations [28 29 or discrete Markov chains 3.. and the generative model is obtained by learning to reverse the random trajectories remains unclear how the informative representation manifests itself in diffusion models, compared other generative models that explicitly contain semantically meaningful latent representation, such as GAN [4] 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 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.
23
30
32
33
Our approach. This work presents MoleculeJAE (Molccule Joint Auto Encoding), 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 respect the SE(3) symmetry of molecule data and is trained by fitting the joint distribution 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 corelation between the real data and its noised counterparts. Under certain approximations, this trajectory distribution modeling decomposes into a marginal distribution estimation and trajectory contrastive regularization task. This multi-task approach yields an effective and flexible framework that can simultancously handle various types of molecular data, ranging from SE(3) -uuuivaquiaa 3D conformers discrete 2D bond comnections. Furthermore, in contrast to diffusion models used for generation tasks that only accept noise as input, most downstream tasks have access to ground-trute trrth molecular structures. To leverage this advantage better, we incorporate an equivariant graph neural network (GNN) block into the architecture, inspired by 38 33, to efficiently encode crucial information from the ground-truff data. Analogous to conditional diffusion models 38 40), the encoder's output serves as meaningful guidence to help the diffused data accurately fitting the trajectory. summary, our self-supervised learning framework unifies both contrastive and generative learning approaches from a trajectory perspective, providing versatile and powerful molecular representation that can be applicd to various downstream applications.
Regarding experiments, we evaluate MoleculeJAE on 20 well-established tasks drawn from the geometric pretraining literature 66, 77, including the energy prediction stable conformation and force prediction along the molecule dynamics. Our empirical results support that MoleculeJAF can outperform 12 competitive baselines on 15 tasks. We further conduct ablation studies to verify the effectiveness of key modules in MoleculeJAE.
Background
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 trajectory augmentation
The concept of diffusion is widely used in various fields such as physics, mathematics, and computer science. In this paper, we take broad view of diffusion by defining it as a Markov process with discrete or continuous time index i, starting from a given data point 2g and producing a series of random transformations that lead to 24. If E {0, 1, ......... then 24 is discrete Markov chain where the probability distribution p(ax)) only depends on the distribution of zz-1. We further refer to the conditional probability p(1r)(ra-)) 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 [4 421, as sources for defining p(x,|21,-1. For instance, Gaussian blurring generates a specific transition probability by setting p(z, xT-1) = (x c), where the scaling factor с may depend 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 estab- lish unified framework for the trajectory augmentation examples mentioned carlier. Following the general approach in [33), each (cold) diffusion model is associated with a degradation Markov process t): indexed by time variable E [0, T), which introduces perturbations the initial data zo:
D
20
It.
(1)
Given that To ~ Pdata- the transformed variable It == D(Eg, t), is also a random variable. To represent the marginal distribution of D(-,t) at time t, is clear, we will use the shorthand notation D, to refer to D(..t) t). According to Equation define the pair (Io, Xt) as a multi-view of the sample xo a specific time (. It is important note
the notation (+()). When the context we
that, under certain conditions, a Markov process converges to a stationary distribution that effectively erases any memory of the initial data zig. Nevertheless, for small values of I, there a high probability that 22 remains within ball centered around zo.
In the context of the standard diffusion-basee generative model 228 32 process R3 is necessary to undo the effects of the degradation D, In 28 the degradation D t) is identified 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 (2 4 can reverse the degradation process D(. t). Specifically, the inversion property implies that Rr- shares the same marginal distribution as pr: RT- ~ Pr. Therefore, as varies from 0 to T, the reverse process R& gradually restores the data distribution. On the other hand, for cold diffusions, although the model is also trained by reconstructing the data from D., it does not explicitly assume the existence of a reverse process as rigorously as in the continuous diffusion- based generative model.
the presence of reverse
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):
dx. plz,, tt + (t)ddu., E [0,
(2)
Here, wy represents the standard Brownian motion. It is worth noting that while the solution 2; of 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:
(3) where 2 is sampled from (O, I). and both a(t) and (() are positive scalar functions. Using the forward SDE ( we approach the task of data synthesis by gradually denoising the noisy observations I, to recover 2i. which is also accomplished through solving 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
Zt = a(taro + 8())z,
SDE fundamentally relies on the score function V log Pe (x), which serves as the parameterization
objective in the score matching framework 28
3

---

Example of discrete (cold) diffusion In contrast to the continuous heat diffusion framework, we now introduce 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 & defined as:
(<=)( =
Multinomial(ze, p = 2x-1Qt),
(4)
where Qi represents a Markov transition matrix. An illustrative example is the absorbing diffusion. as introduced by [3), where
(1-
B,e
and Em is one-hot vector with value of at the absorbing state and zeros elsewhere. Typically, the state m is chosen as the 'masked' token. Hence, absorbing diffusion can be regarded as a masking process with dynamical masking ratio 3.. 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(
2.2 Graph-based molecular representation
In this paper, our primary focus on the graph representation of molecules. - denote the integrated molecular graph, consisting of n := [V] atoms and |E| bonds. The matrix P E R 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 E R=*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.
3 Method
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 3.22, we will introduce theoretical self-supervised learning framework for these trajectories. Our hypothesis is that good representation should encode the information from the distribution of augmented trajectories. This hypothesis forms the practical basis of 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 '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 2; := (t), (t), P(t)) from G. The challenging aspect lies in preserving the $ (3) symmetry of the position matrix P, which we will describe in detail below.
Given two 3D point clouds (molecular conformers) P1 and Pz, we say they are SE(3)-isometric if there exists an R E SE(3) such that Pi = RP2. In probabilistic terms, let pan pan (Tan) be the probability density of all 3D conformers denoted by C. Isometric conformers necessarily have the same density, i.e.,
(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(O) + P(t) should also be (Eqiivequi,,,,ant, satisfying
(aD(= PaD PaD (ReD)), A R E SE(3), 33D E C.
RP = P(Reap), V E SE(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, to Eq. 3 the F3)) equivariance automatically satisfied due to the SF)(3 equivariance of both the Gaussian random variable z and the original data zo. We leave the formal proof in Appendix A.

---

Regarding the 2D component, let (t) == (t), E(t)), where T2D consists of invariant scalars that remain unchanged under SE(3) transformations. Although z2p 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 22D (t) and P()) to undergo heat diffusion. By combining the 2D and 3D) parts, we introduce the system of SDEs for Zt:
(t) = - P(t)dt + a (t)ddui)
dH (t), t)dt + az(t)dus, dE(()= μs (E(1), ()dt az (1)duu..
(6)
It is worth mentioning that although the components of Eq. aare disentangled, the corresponding reverse diffusion process and its score function are entangled. common choice of p(z, t) is u(z,t) =2, then will utilize the explicit solution of equations (6) to generate the molecular augmentation trajectory. Additionally, it is also possible to perform equivariant diffusion the 2D and 3D joint representation through equivariant cold diffusion, following the approach described in (4). The detailed framework for this approach is provided 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 denoising from a trajectory perspective, According to the Kolmogorov extension theorem [46. every stochastic process defines probabilistic measure on the trajectory space, uniquely determined by set of finite-dimensional joint distributions that satisfy consistency conditions. Therefore, estimating the probability of the estimating the joint probabilities p(.',....,, Tt,) for each finite time sequence f1....,fy E [0. T].
trajectory space is equivalent to
From the standpoint of data augmentation, our specific focus is on learning the joint distribution of trajectory augmentation pairs: p(zn, .4) (solution of Eq. 6) which determines how the noisy dt corresponds to the original (denoised) 20-
To differentiate our notion of " Auto-Enccoiin (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 diffusion models (28). where multiple denoising processes from 2; to To 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. [T]) due 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.
from the "denoising" method utilized in DDPM 32
For given time & [0, J. let us assume that the probability of the random pair (zo, defined in Eq. [ is determined by joint density function p(zo, 2t). Our objective is to approximate p(zo, xx)) within a variational class po (rr, xt). Maximizing the likelihood of this joint distribution leads to the following optimization problem:
argmax II ps (wg, xf),
where {(r)) = represents the collection of n augmented samples from the training dataset. Following the tradition of Bayesian inference [7], we parameterize po by defining a joint energy func- tion (2 2,) such that: Po( xI,) = 2, Fo( where Zo(t) is the intractable normalization constant depending on t. Therefore, the maximal likelihood framework reduces to solving
argmin,EEpzo..z.)
[E xx)]
However, directly optimize po(xo. 2x) by taking the gradient with respect to 0 is challenging because 8 log pa (3t.aa) contains intractable term: @Za(U) To circumvent this issue, we borrow ideas from
by treating the transformed &; as an infinitely dimensional "label" of vo. More precisely, we consider the parameterized marginal density qo (20)
48
go (=o) = f po (20, as)da..
(8)
5

---

which involves integrating out Zt. We define the marginalized energy function with respect to as: Eo(.) := log
((): Now, let fo x+) denote the normalized conditional density function pe (z (zu), we have Olog fo(ze|zo]
aaa (wo. [aa(zo) By taking the empirical
80
expectation with respect to the finitely sampled pair (Io, Fe) ~ p( Ie), we can decompose the gradient of the maximum likelihood as follows (see Appendix A for the full derivation):
Ø Blog (oa,
[ol (z(o
log fo(mo. xz)
Bp(xo.Fe)
Ep(so)
+ E Ep(oat)))
(9)
09
00
00
here we use E denote the expectation with respect to the empirical expectation. Note that this decomposition holds for (z) for any two different time steps (0 S s, S 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 ho(2) (z) (the equivariant model for building he will be discussed in the next section). Based on Eq. 16 we introduce two tasks for training (x) that stem from this decomposition.
in Eq. 16aims to reconstruct the distribution of data samples. Therefore, we refer this term as the reconstruction task, which involves modeling the marginal distribution Pdota using qo (27)).
Reconstructive task. The first term Ep(wo)
a go
89
Although it is possible to directly train the likelihood reconstruction term in the auto-regressive case, such using the noise conditional maximum likelihood loss E.,.p, jo., log (a(a))) proposed in [50 we instead adopt trajectory score matching (28), which is applicable for non-autoregrressive graph data. The score matching loss is defined follows:
Lse =
V| log p(z)(or)))) -000-s0 (x(x,))))²
(10)
We choose this approach for two reasons: First, training the score function enables the model to generate new samples by solving the reverse process, thus facilitating generative downstream tasks. Second, when = O, the score function for 3D structures can be interpreted as a "pscudo" force field for the molecular system (6) containing essential information about the molecular dynamics. Furthermore, [44] provided a rigorous proof demonstrating that score matching formula T0 serves as a variational lower bound for the marginal log- likelihood log (zo). This theoretical guarantee solidifies the effectiveness of score matching as a training objective.
For molecule trajectories, the score function encompasses both the 2D and 3D components. Moreover, 1o.55inn Simplies that the corresponding 3D score function V (p)sD(r) (represented as score (i)) in Fig. 2 equivariant under SE()):
the SE(3)-invariant density pan defined by Eg.
VPPAD R.) RV ,PPD(()), VR E SF((3, x E C.
In conclusion, the symmetry principle mandates that the score neural network takes equivariant vector field (representing the positions of all atoms) and invariant atom features as input. It then produces two score functions that adhere to different transformation rules: 1. V=Pan(z): SE(3)- equivariant; VyP2n()) (y): SE(3)-invrriant
Contrastive task. We have demonstrated that optimizing the score matching task allows us to capture information about the marginal distributions of the trajectories. However, the joint distribution contains additional valuable information. As an example, consider the following parameterized random processes, all of which share the same marginal distributions but exhibit varying joint distributions (as proven in 44)):
1+X²
dy: (f), 1) (1)o O]d Ag(I)dB..
(11)
for X > 0. Hence, is theoretically necessary to optimize the second term E (x( aolog (at
of Eq. 16 By employing the conditional probability formula, we have:
x1)) = pa(aa, xx)
(12)
[2 v)dy

---

Ground-truth Molecule
Noised Molecule
Node Frame
Time!
Equhvriant
Equinarant
Representation By
Projectio
e RM,, kzone Fij
ER
Frejectio
Graph Representation
Nade Fooms
Projectin D
Figure 2: Architecture of MoleculeJAE. The inputs are ground-truth molecules with both 2D and 3D structures. MolcculeJAA. also adopts the noised inputs denoising, so as to model the trajectory distribution. The outputs are three score functions for conformer and bond representations, which flow into the general pipeline in Figure [
important to note that the troublesome normalizing constant Ze (4) cancelled out in Eq. 12 In practice, the integral f (2o. y)dy is empirically approximated using Monte Carlo sampling. To make a connection with contrastive learning (CL., recall that in CL, a common approach is to align the augmented views ((x, xx+) of the same data and simultancously contrast the augmented views of different data (x ). By treating the joint distribution as a similarity measure, Eq. 12 can be seen implicitly imposing two regularization conditions fe (2 Tz): ensuring high probability for p( 22) and a low probability for p(ro, y). This notion of similarity motivates us to refer to maximizing the second term of Eq. 16 as contrastive task.
Contrastive surrogate However, estimating po (200 .4) is challenging due to the intractability of closed-form joint distributions. To overcome this difficulty, we propose black-box surrogate i that represents the mapping from the latent representation ha(x;) to pe pe xx) (following the notation in figure 2.. Specifically, let (ha(2o), hx(x)) denote the representation pair obtained from the input molecule data (Fo, z1). Then, the surrogate of (o xx) defined by po zz)
Z(0 exp( [io (wo)) (==) + C(xo)). Here, r() is monotone annealing function with +T((]
respect to By using Eq. 12 the unknown normalization constant (t) and C(ro) cancel out,
resulting in the following approximation:
fe x1))
(13)
expl-1lyphoo@ee - (hh(
Our surrogate is reasonable because our modeling target p( 2ye) is derived from the joint distribu- tions of a (continuous) Markov process. When lies along the augmented trajectory of a specific data sample zy and approaches zo as t + O, the log- likelihoood log ps 2x) )) is also achieves local maximum. Therefore, based on the Taylor expansion, the leading non con-constant term takes the form of quadratic expression.
Final Objective By combining the reconstruction and contrastive regularization tasks, we define the final multi-task training objective of MoleculeJAF. as a weighted sum of the two tasks:
Aa + AzCcos
(14)
where Åι, X2 are weighting coefficients, and Lco := E2~ ) (No, wr). In Fig. 2 the noised input x, undergoes an encoding process obtain the latent representation (ax))), which is then split into two branches:
1. The first branch passes through a score neural network & for reconstruction: so(xz) = (hg(xz)));
2. The second branch incorporates the original data representation hy (zo) and further projects the pair (ho( he (xx)) through non-linear projection head y for contrastive learning.
Note that the black-box projection function g (although not used in downstream tasks) also participates the pretraining optimization, following the conventional contrastive learning frameworks 51 See Fig. [ for graphical illustration of MoleculeJAE's pipeline.
7

---

Table 1: Results on 12 quantum mechanics prediction tasks from QM9. We take 110K for training, 10K for validation, and 11K for testing. The evaluation mean absolute enror, and the best and the second best results are marked in bold and bold, respectively.
Pretraining
33 +
LLUMO 4
C+ 4
GL
H
R²
U 4
Uo 4
ZPVE
(random init)
0.060)
44.13
27 64
22.55
0.028
0.031
14.19
14.05
0.133
13,93
13.27
1.749
Type Prediction
0.073
45.38
28.76 76
24.83
0.036
0.052
16.66
16.28
0.275
15.56
14,66
2.094
Distance Prediction
0.065
45.87
0.031
0.033
14.83
15.81
0.248
15.07
15.01
1.837
Angle Prediction
0,066
48.45
29.02
24.40
0.034
0.031
14.13
13,77
0.214
13,50
13.47
1.861
3D InfoGraph
0.062
45.96
29.29
24.60
0.028
0.030
13.93
13.97
0.133
13.55
13.47
L.644
GeoSSL-RR
0.060
43.71
27.71
22.84
0.028
0.031
14.54
13.70
0.122
13.81
13.75
1.694
GeoSSL-InfoNCE
0.061
44.38
27.67
22.85
0.027
0,030
13.38
13.36
0.116
13.05
13.00
1.643
GeOSSL-EBM-NCE
0.057
43.75
27.05
22.75
0.028
0,030
12.87
12.65
0.123
13.44
12.64
1.652
3D InfoMax
0.057
42.09
25.90
21.60
0.028
0.030
13.73
13.62
0.141
13.81
13.30
1.670
GraphMVP
0.056
25.75
13.43
13.31
13.07
1.609
GeoSSL-DDM-ID IL
0,058
42.64
26.32
21.87
0.028
0.030
12.61
12.81
0.173
12.45
12.12
1.696
GeoSSL-DDM
0.056
42.29
25.61
21.88
0.027
0.029
11.54
11.14
0.168
11,06
10.96
1.660
MoleculeIAE
0.056
42.73
25.95
21.55
0.027
0.029
11.22
10.70
0.141
10.81
10.70
1.559
Table Results on eight force prediction tasks from MD17.
We take 1K for training. 1K for validation, and
48K to 991K molecules for the test concerning different tasks, The evaluation is mean absolute error, and the best results are marked in bold and bold, respectively.
Pretraining
Aspirin L
Benzene
Ethanol
Malonukdehyde
Naphthalene
Salicylic
Toluene
Uracil
(random init)
1.201
0.380
0,386
0.794
0.587
0.826
0.568
0.773
Type Prediction
1388
0.402
0,450
0.879
0.622
1,028
0.662
0.840
Distance Prediction
14477
0.396
0.434
0.818
0.793
0,952
0.509
1.567
Angle Prediction
1. 1.542
0.447
0.669
1.022
0,680
1.032
0.621
0.768
3D InfoGraph
0.415
0.560
0.788
1.278
0.768
GeoSS
1.215
0.393
0.514
1.092
0.596
0.847
0.570
0.711
GeoSSL-InfoNCE
1.132
0.395
0.466
0,888
0,542
0.831
0.554
0.664
GeOSSL-EBM-NCE
1.251
0.373
0.512
0.990
0.560
0.742
3D InfoMax
1.142
0.388
0.469
0.731
0.785
0.798
0.516
0.640
GraphMVP
1.126
0.377
0.430
0.726
0.498
0.740
0.508
0.620
Geo -D--III.
0.391
0.817
0.607
GeoSSL-DDM
1.107
0.360
0.357
0.737
0.568
0.902
0.484
0,502
MoleculeJAE
1,289
0.345
0.365
0.613
0.498
0.712
0,480
0.463
3.4 Model Architecture of MoleculeJAE
In pursuit of a meaningful latent molecular representation. we design our joint auto-encoder model as a conditional diffusion based model, inspired by [38] (3]. In our model, the condition does come from labels, but rather an encoding of the ground-trute molecule. Specifically, to obtain the latent representation hg shown in Fig. 2 we implement two equivariant encoders that satisfy the SE(3) symmetry proposed Section 3.1 One encoder takes the original molecule as input, and its output is used as condition to help the other encoder that encodes the noisy molecule for reconstruction. The only requirement for the architecture of the two encoders is that the output should be invariant. Therefore, any S S3) equivariant GNN 54 that outputs invariant scalars will suffice.
Equivariant decoder. With our representation ho which depends (ar0 Te. l), the decoder part of MoleculeIAE is divided into two heads. One head is paired with node-wise SE(3) frames to match the (E))) equivariant score function, while the other head generates an SE(3) invariant representation that is used for contrastive learning (see Fig. 2 for a complete illustration). Further details on the model design are left Appendix
4 Experiment
4.1 Molecule.AEE Pretraining
Dataset. For pretraining, we use PCQM4Mv2 (55) It extracts 3.4 million molecules from Pub- ChemQC [56] with both 2D topology and 3D geometry.
Backbone models. We want to highlight that MoleculeJAE is agnostic to backbone geometric GNNs. In this work, we follow previous works in using SchNet model [57 for 3D conformation. For the 2D GNN representation,
take simple version by mainly modeling the bond information (details in Appendix D). For the SDE models (8 for generating the joint trajectories, we consider both VE and VP, in Eq. 6
Baselines for 3D conformation pretraining. Recently, few works have started to explore 3D conformation pretraining. For instance, GeoSSL (5) provides comprehensive basclines. The initial

---

Table 3:
Ablation studies of contrastive loss term in MoleculeJAE. The ablation results are on QM9.
Pretraining
a
VE 4
Eposo 4
ELUMO 4
p |4
Ce
G
4
R² 4
U4
Uo
ZPVB +
λ 0
0.057
43.15
26.05
21.42
0.027
0.030
12.23
11.95
0.162
12.20
11.42
1.594
A2 0.01
0.056
42.73
25.95
21.55
0.027
0.029
11.22
10.70
0.141
10.81
10.70
1.559
λ
0.066
45.45
28.23
23.67
0.028
0.030
14.67
14,42
0.204
13.30
13.25
1.797
Table 4:
Ablation studies of contrastive Ioss term in MoleculeJAE. The ablation results are on MD17.
Pretraining
Aspirin 4 Benzene
Dihanol 4 Malonaldehyde Naphthalene + Salicylic
Toluene
Uracil
1.380
0.359
0.363
0.744
0.482
0.902
0.548
0.590
A2 0.01
1.289
0.345
0.365
0,613
0.498
0.712
0,480
0,463
A2
1.561
0.547
0.781
0.735
0.918
1.160
1.052
0.809
three baselines involve type prediction, distance prediction, and angle prediction, respectively aiming to predict the masked atom type, pairwise distance, and triplet angle. The next baseline is 3D InfoGraph. It is a contrastive SSL method and predicts whether the node- and graph-level 3D representation are for the same molecule. Last but not least, GeoSSL proposes new SSL family on geometry: GeoSSL-RR, GeoSSL-InfoNCE, and GeoSSL-EBM-NCE are to maximize the MI between the conformation and augmented conformation using different objective functions, respectively. GeoSSL-DDM optimizes the same objective function using denoising distance matching. GeoSSL-
DDM-1L 6 is special case of GeoSSL-DDM with one layer of denoising.
Baselines for 2D and 3D multi-modal pretraining- Additionally, several works have utilized both 2D topology and 3D geometry modalities for molecule pretraining. Vanilla GraphMVP 77) utilizes both the contrastive and generative SSL, and 3D InfoMax (58) only uses the contrastive learning part.
In the following, we provide the downstream tasks for applying our pre-trained MoleculeJAE. The experiment on joint generation of the 2D and 3D structures of molecules are provided in Appendix F
4.2 Quantum Property Prediction
QM9 [59 is dataset of 134K molecules, consisting of nine heavy atoms. It contains 12 tasks, which are related to the quantum properties. Among 12 tasks, the tasks related to the energies are very important, e.g., U and Un are the internal energies at OK and 298.15K, respectively. The other two energies, H and G can be obtained from U accordingly. The main results are in Table [ We can observe that MoleculeJAE can outperform 12 baselines on 9 out of 12 tasks. We want to highlight that these baselines are very competitive, and certain works (e.g., GeoSSL) also model the 3D trajectory. Noticeably, MoleculeJAE can reach the best performance on four energy-related tasks.
4.3 Molecular Dynamics Prediction
MD17 [6 is a dataset on molecular dynamics simulation. It contains eight tasks corresponding to eight organic molecules, and the goal is to predict the forces at different 3D positions. The size of each task ranges from 48K to 991K, and please check Appendix forr details. The main results in Table 2 and MoleculeJAF can outperform 12 baselines on out of 8 tasks and reach the second-best for one of the remaining tasks.
4.4 A Ablation Study on the Effectiveness of Contrastive Loss
As discussed in Equation (1)), there is one important hyperparameter A controlling the contrastive loss in MoleculeJAE. We want to conduct an ablation study on the effect of this contrastive term. As shown in Tables and 4 we consider three value for A2: 0, 0.01, and 1. A simply means that we only consider the reconstructive task, and its performance is very close to X2 = 0.01, i.e., the performance degrades by a large margin. Thus, we would like to claim that the contrastive term in MoleculeJAE is comparatively sensitive, and we need to tune this hyperparameter carefully to obtain optimal results.
the optimal results reported in Tables and 2 However, as we increase the X
5 Conclusion
In this work, we introduce a novel joint self-supervised learning framework called MoleculeJAE, which is based on augmented trajectory modeling. The term "joint" in our framework has two implications: Firstly, it signifies that our method is designed to model the joint distribution of trajectories rather than solely focusing on the marginal distribution. Secondly, it indicates that the

---

augmented molecule trajectory incorporates both 2D and 3D information, providing insights into different aspects of molecule representation. While our proposed method has demonstrated the best empirical results on 15 out of 20 geometry-related property prediction tasks, there are still areas left for improvement, such as architecture design. Please refer to Appendix F for an in-depth discussion.
Acknowledgement
Jiujiu Chen would like to thank for her Associate Researcher Ruidong Chen, and Professor Xiaosong Zhang in UESTC for their special support.
References
[1] Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint arXiv: 1905. 12265, 2019.
[2] Shengchao Liu, Mehmet F Demirel, and Yingyu Liang. N-gram graph: Simple unsupervised representation for graphs, with applications to molecules. Advances in neural information processing systems, 32, 2019.
[3] Ziniu Hu, Yuxiao Dong, Kuansan Wang. Kai- Wei Chang, and Yizhou Sun. Gpt-gnn: Generative pre-training of graph neural networks. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 1857-1867, 2020.
[4] Fan- Yun Sun, Jordan Hoffmann, Vikas Verma, and Jian Tang. Infograph: Unsupervised and semi-supervised graph-level representation learning via mutual information maximization. In International Conference on Learning Representations, ICLR, 2020.
[5] Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)- invariant denoising distance matching, In The Eleventh International Conference on Learning Representations, 2023. URL htttps://oppnrevieer..e wwt.:eeetoooooee wtt.:e/eeeeeeoooooodddCCCC//.o
[6] Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint ar 21V:22206.00133, 2022.
[7] Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre-training molecular graph representation with 3d geometry. arXiv preprint arXiv:2/ 10.0772, 2021.
[8] Shengchao Liu, Weitao Du Zhi-Ming Ma. Hongyu Guo, and Jian Tang. A group symmetric stochastic differential equation model for molecule multi-modal pretraining, In International Conference on Machine Learning, pages 21497-21526. PMLR, 2023.
[9] Shengchao Liu, Weitao Du, Yanjing Li, Zhuoxinran Li, Zhiling Zheng, Chenru Duan, Zhim- ing Ma, Omar Yaghi, Anima Anandkumar, Christian Borgs, et al. Symmetry-informed geometric representation for molecules, proteins, and crystalline materials. arXiv preprint arXiv:2306.09375, 2023.
[10] Yixin Liu, Shirui Pan, Ming Jin, Chuan Zhou, Feng Xia, and Philip Yu. Graph self-supervised learning: A survey, arXiv preprint arXiv:2103.00111, 2021.
[11] Yaochen Xie, Zhao Xu, Jingtun Zhang, Zhengyang Wang, and Shuiwang Ji. Self-supervised learning of graph neural networks: A unified review. arXiv preprint arX1v:2102.10757, 2021.
[12] Lirong Wu, Haitao Lin, Zhangyang Gao, Cheng Tan, Stan Li, et al. Self-supervieed on graphs: Contrastive, generative, or predictive. arXiv preprint arXiv:2 W0Niv2205.00334 073443 2021.
[13] Xiao Liu, Fanjin Zhang, Zhenyu Hou, Li Mian, Zhaoyu Wang, Jing Zhang, and Jie Tang. Self-supervised learning: Generative or contrastive. IEEE Transuctions on Knowledge and Data Engineering, 2021.
10

---

[14] Andrea Madotto, Zihan Liu, Zhaojiang Lin, and Pascale Fung. Language models as few-shot learner for trsk-oriented dialogue systems. arXiv preprint arXiv:2008. 06239, 2020.
[15] Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. A survey of large language models. arXiv preprint arXiv:2303. 18223, 2023.
[16] Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:1 1287-11302, 2021.
[17] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hicrarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.0612) 2022.
[18) Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.082119 2018.
[19] William W Wood and Jerome J Erpenbeck. Molecular dynamics and monte carlo calculations in statistical mechanics. Annual Review of Physical Chemistry, 27(1))319-3444 1976.
[20] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages
PMLR, 2015.
[21] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-damped langevin diffusion. arXiv preprint arXiv:21 2021.
[22] Jachyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In International Conference on Machine Learning, pages 10362-10383 PMLR, 2022.
[23] Clement Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pas- cal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209.1473: 14734, 2022.
[24] Weitao Du He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Namning Zheng, Bin Shao, and Tie- Yan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning, pages 5583-5608. PMLR, 2022.
[25] Shengchao Liu, Weitao Du, Zhiming Ma, Hongyu Guo, and Jian Tang. A group symmetric stochastic differential equation model for molecule multi-modal pretraining. 2023.
[26] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning, pages 9558-9568. PMLR, 2021.
[27] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning. pages 8867-8887. PMLR, 2022.
[28] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 2rXiv:20111.33555 2020.
[29] Weitao Du, Tao Yang. He Zhang, and Yuanqi Du. A flexible diffusion model. arXiv preprint arXiv:2206. 10365, 2022.
[30] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 119999999999, 2021.
[31] Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. Advances in Neural Information Processing 2022.
Systems,
11

---

[32] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 13:6840-6851, 2020.
[33] Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold diffusion: Inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.09333 2022.
[34] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu. David Warde-Farles Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-144 2020.
[35] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint urXiv:1312.61 14, 2013.
[36] Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning? Advances in neural information processing systems, 33:6827-6839 2020.
[37] Abhishek Sinha, Kumar Ayush, Jiaming Song, Burak Uzkent, Hongxia Jin, and Stefano Ermon. Negative data augmentation. arXiv preprint urXiv:2102.05113, 2021.
[38] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[39] Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakom. Diffusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages (0619-1062)) 2022.
[40] Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models, arXiv preprint a3X1v:2302.05543, 2023.
[41] Suorong Yang, Weikang Xiao, Mengcheng Zhang, Suhan Guo, Jian Zhao, and Furao Shen. Image data augmentation for deep learning: A survey. arXiv preprint arXiv:2204.086100 2022.
[42] Brandon Trabucco, Kyle Doherty. Max Gurinas, and Ruslan Salakhutdinov. Effective data augmentation with diffusion models. arXiv preprint arXlv:2302.0944, 2023.
[43] Qinsheng Zhang, Molei Tao, and Yongxin Chen, gddim: Generalized denoising diffusion implicit models. arXiv preprint 22Xiv:2206.05564, 2022.
[44] Chin-Was Sein-Weila Huang, Jae Hyun Lim, and Aaron C Courville. variational perspective on diffusion- based generative models and score matching. Advances in Neural Information Processing Systems, 34:22863-2288 22876, 2021.
[45] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. International conference on machine learning, pages
PMLR, 2021.
[46] Stefan Chmicla, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5):e 1000115, 2017.
[47] Ethan Goan and Clinton Fookes. Bayesian neural networks: An introduction and survey. Case Studies in Applied Bayesian Data Science: CIRM Jean-Morlet Chair, Fall 2018, pages 45-87, 2020.
[48] Will Grathwohl, Kuan- Chich Wang, Jömn-Henrt Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly energy based model and you should treat it like one. arXiv preprint arXiv: / 2.03223,, 2019.
[49] Stephen Zhao, Jorn-Herrik Jacobsen, and Will Grathwohl. Joint energy- models for semi-supervised classification. In ICML. 2020 Workshop on Uncertainty and Robustness in Deep Learning, volume 1, 2020.
12

---

[50] Henry Li and Yuval Kluger. Autoregressive generative modeling with noise conditional maxi- mum likelihood estimation. arXiv preprint arXiv:2210.10715, 2022.
[51] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
[52] Kartik Gupta, Thalaiyasingam Ajanthan, Anton van den Hengel, and Stephen Gould. Under- standing and improving the role of projection head in self-supervised learning. arXiv preprint arXiv:2212.11490 2022.
[53] Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2021.
[54] Weitao Du, Yuanqi Du, Limei Wang, Dieqiao Feng, Guifeng Wang, Shuiwang Ji, Carla Gomes, and Zhi-Ming Ma. new perspective on building efficient and expressive 3d equivariant graph neural networks. arXiv preprint arXiv:2304.0477
2023.
[55] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint wwXiv:2005.00663 2020.
[56] Maho Nakata and Tomomi Shimazaki. Pubchemqe project: a large-scale first pest-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 17(6):1300-1333 2017.
[57] Kristof T Schüit,, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24))241722, 2018
[58] Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning. pages 20502.20502. PMLR, 2022.
[59] Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):17, 2014.
[60] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. arXiv preprint arXiv:2206.01729. 2022.
[61] Jinheon Baek, Minki Kang, and Sung Ju Hwang. Accurate learning of graph representations with graph multiset pooling. arXiv preprint arXiv:2102. 2021.
[62] Lei Huang, Hengtong Zhang, Tingyang Xu, and Ka-Chun Wong. Mdm: Molecular diffusion model for 3d molecule generation, 2022.
[63] Thomas A Halgren. Merck molecular force field. i. basis, form, scope, parameterization and performance of mmff94. Journal of computational chemistry, 17(5-6):40-519, 1996.
[64] Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning recover global and local spectral embedding methods. arXiv preprint arXiv:2205./15500 2022.
[65] Michael F Mathieu, Junbo Jake Zhao, Junbo Zhao, Aditya Ramesh, Pablo Sprechmann, and Yann LeCun. Disentangling factors of variation in deep representation using adversarial training. Advances in neural information processing systems, 29, 2016
[66] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 2023.
[67] Yao Zhu, Jiacheng Sun, and Zhenguo Li. Rethinking adversarial transferability from a data distribution perspective. In International Conference on Learning Representations, 2021.
13

---

[68] Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. arXiv preprint arXiv:2210.083:4, 2022.
[69] Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence, 4(3): 279-287, 2022.
[70] Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)- invariant denoising distance matching. In The Eleventh International Conference Learning Representations, 2023,
[71] Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. 2023.
[72] Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alexandros G Dimakis, and Peyman Milan- far. Soft diffusion: Score matching for general corruptions. arXiv preprint arXiv:2209.06 054424444 2022.
[73] Joseph Lemley, Shabab Bazrafkan, and Peter Corcoran. Smart augmentation learning optimal data augmentation strategy. leee Access, 5:5858-5888 5569, 2017.
[74] Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1 113-123, 2019.
[75] Clement Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for molecule generation, 2023.
[76] Simon Axelrod and Rafael Gómez-Bombarelli. Geom, energy-annotated molecular conforma- tions for property prediction and molecular generation. Scientific Data, 9(1):185, 2022. doi: 10. 1911038/415970022211288 URL http::://oo....... org/10..00000. 1038/ 0038/441597-0 0011441722220222222212
14

---

A Methodology Detail
A.1 Details for Section 3.3
Following Eq. we consider the 'marginal' distribution
go (=())) = / po (2o, Iz dxz
by marginalizing out zg. We define the corresponding energy function with respect to zo as
Ez(zo) -log log / exp(-EEs 24) dr))....
From the definition of (xo). the distribution (00) shares the same normalization constant with qe (zo, Zz). Therefore we have
i log (x))
DE (zo, vs)
(xo
00
00
00
where fo(2o, 2z) denotes the normalized conditional probability g0(xx(z0). By the conditional probability formula, it is obvious that: it's obvious that
i log (o xτ) ① log fo 24) a llg go
(15)
00
00
By taking the empirical expectation with respect to the sampled pair (Io, 21) ~ p( x²), we find
that the gradient of the Maximum Likelihood allows the following decomposition:
Ö log (o xe) Ep(xo)
Ö log (o()))
i log fo x4)
Ep(zo.at))
+ Bptro.re)
(16)
00
00
00
Gradient reconstruction
Gradient of contrustive learning
Here, we use È. to denote the expectation with respect to the empirical expectation. In this way, we find that maximizing the likelihood po(2o, (e)) is equivalent to solving:
T log (=o)
T Olg fo( xz
argmax,
E Ep(zo)
E p(zo..))
00
00
A.2 Details for Section 3.4
In this section, we provide the details for the model design described in Section 3.4 Specifically, illustrate how to obtain the latent representation he from 2o = (Po, Ho, Eo) and 2, = (P,, Hy, E,) (the solution of Eq. at time t). It is important to note that in standard diffusion generative models 28 32), only xe is fed into the score neural network.
The 3D part (Ho. Fo) and (H,, E, Jare transformed by two equivariant 3D GNNs, and we denote the invariant node-wise output representation by fo and ft. Additionally, we embed the time (which is essential for trajectory learning) using Fourier embedding, following (28
Emd(L) = Fourier(!).
We obtain the 3D node-wise representation by concatenating the Fourier embedding, fo. and fú:
NodeaD MLP Emd(t) || fo ft].
Up this point, haven't utilized the 2D information (Ea, Ee). Unlike the node-wise 3D represen- tation, we encode (Fo, Fe) as a weighted adjacency matrix W:
W MLP [Emd(t) Eo Er].
Combing Nodesp and the adjacency matrix W, we implement dense graph convolution neural network (GCN) to obtain the final (node-wise) latent representation he, following [22]:
hg GCN (NodeaD. W).
15

---

Score (E(t))
Score (H(t))
Score (P(t))
3D GNN
Phbaads Cijj
Node Frame
GCN
Nodeup
3D GNN
3D GNN
Fourler
E(O) 11 E(t)
H(0)
P(t), H(t)
Figure Modular Design MolcculeJAF.
3D output Score(P.) Pt). To obtain the SE(3) equivariant and reflection anti-equivariant vector field 60 Score Pa) from hg. we implement the node-wise equivariant frame and perform tensorization, following 24 [4]. Consider node with 3D position x;, let x, ; Σx,eN(s) Xj be the center
of mass around x, neighhorhood. Then the orthonormal equivariant frame F, == (ef e e%) with respect x; is defined as:
X; [x; X
8 [|x;
X; x;|
(17)
X x
[xx
x,|
X
where x denotes the cross product, which is S3)(3 equivariant and reflection anti-equivariant, Then, we transform he R NxL. haD E Nx3.
h&D (11,, hb) MLP (he).
Finally, for node
Score' = h1 ef + h2 e + ha e
2D output (E).. Score(E:) represents the probability gradient with respect the molecule bonds. Therefore, Score (E, has the same shape as the dense adjacency matrix and is SE(() invariant. To obtain Score [E,,, we leverage graph multi-head attention [6]] obtain dense attention matrix A, (see [61] for the explicit formula):
A = Att (ho).
Then, we apply MIP to the edge-wise A, to obtain the final 2D score function:
Score E.) MLP (As).
See Figure 3 for graphical representation.
16

---

A.3 A Cold Alternative for Section 3.1
Unlike continuous data that can take values ranging from (-o0, oo). categorical data only takes a finite number of values. For example, in this paper, the atom type takes its value from the periodic table of elements, while the bond type takes its value from the set {0, 1, 2, 3, 4]. For these finite state spaces, we can replace our continuous diffusion framework with a discrete (cold) diffusion framework.
To define a discrete diffusion, we need to specify the transition matrix Qr for each time t, as stated in Eq. 4 Following [23 29 62], to effectively estimate the likelihood of the diffusion model, we further require that the noised state g(2r (2), defined as:
(<((2=) := xQ
to be equivariant under permutation. This requirement can be satisfied if we diffuse separately on each node and edge feature. Additionally, we require that the limiting distribution as + co should not depend on the original data so that we can use it as prior distribution for inference. A common choice is uniform transition over the classes, where the corresponding Qr is defined as:
Q = az (1 0: 1a1//
with ot transitioning from to O. By letting Qt act each node and each edge, we have defined the discrete diffusion for (1) and E(t):
q (H, He-1 Qu and ((E,)E,-1) E Q+
On the other hand, although the 3D structure is represented as a continuous point cloud, the molecule's conformer structures the stationary states are discrete (5) and labeled by the energy levels of molecules. Through classical force field simulation, [63 provides a discrete set of rough 3D conformers (with their corresponding energies) for each molecule. Let {P:}i=1 be the collection of these rough 3D conformers, and denote uniform sampling on set {P:) as Uniform( (P.)) Then, we can construct a discrete diffusion for the 3D structure as follows, similar to Eq.
P(t) a())FF P(O) + ((t)Uniformi (aarrm(((5))
(18)
where F is global equivariant frame obtained by averaging the node-wise equivariant frames F; defined by Eq. 17 It be verified that by projecting the original 3D structure with the inverse of its global frame F
the above equation is SE((3) invariant.
Different from Eq. (3), where the randomness comes from Gaussian noise, the randomness in Eq. (18) arises from finite uniform sampling. In conclusion, we have constructed an SE)(3 invariant discrete 3D diffusion based on the ground truth and rough 3D conformers.
B Related Works
The primary objective of self-supervised learning (SSL) is to acquire meaningful and compressed representations from unlabeled data, often measured by their ability estimate specific data distribution Pdata- This leads to the hypothesis that a well-trained pretrained representation can enhance the robustness and generalization of downstream tasks. However, a challenging gap exists between the ground-truth distribution Pdota (if available) and our finite set of natural data samples. To address this gap, self- supervised learning techniques such as contrastive learning and manifold learning [64] often require additional data beyond the natural samples. Contrastive learning, for example, employs data augmentation to introduce positive and negative samples, while adversarial training (65] introduces adversarial samples through data perturbation for improving the robustness of downstream classification tasks.
Adversarial and contrastive samples In this work, we propose diffusion framework for trajectory data augmentation based on the forward process of diffusion. A related concept that emphasizes the trajectory viewpoint is consistency, introduced in [66], which requires that the points on the backward path map to the same endpoint, corresponding to the condition of the trajectories' joint distribution. Additionally, the learned reverse (generative) process of a pretrained diffusion model has been employed to generate synchronized augmentation samples 42 Adversarial samples [67], compared to data augmentation, more subtle they involve perturbations along specific directions of the
17

---

natural data that can be identified using the information provided by the score function V log Pdata (x) and V log Pattta (a y) conditioned on labels y.
Once the augmented trajectories are established, our learning objective to fit the joint distribution of the random trajectories. From practical standpoint, framework lies the intersection of two unsupervised learning paradigms: contrastive learning and generative learning. In fact, [68] proposes an intriguing contrastive lower bound (positive sample part) for the reconstruction loss of masked auto-encoder (AE). The authors further demonstrate that explicitly incorporating the negative sample part into the masked AE enhances the performance of the learned representation. Similarly, we observe similar phenomenon within more general unsupervised framework, as illustrated in Eq. where the masking process is realized through cold diffusion.
Molecular representation pretraining The labels for molecules are scarce due to the laborious and expensive cost, and sulf-supervised pretraining has been widely adopted to alleviate this issue. For the molecular topology pretraining, existing unsupervised pretraining methods utilize either the reconstructing the masked subgraphs in molecules (AtrMaak (.. GFT-GNN (3) or detecting the pos-
itive and negative pairs in contrastive manner (ContextPred M. MolCLR (69 Meanwhile, several works have studied the molecular geometry pretraining. GeoSSL 70) presents a comprehensive benchmark, including distance prediction, angle prediction, and MI-based coordinate denoising framework specifically designed for conformations. Recent progress has started to combine these two research directions, i.e., molecule topology and geometry joint pretraining. GraphMVP [7] first proposes to maximize the MI between these two modalities a contrastive and generative manner, and 3D InfoMax (58) is a special case of GraphMMVP by only considering the contrastive part. On the other hand, Unimol 71 introduced a novel 3D transformer for molecular pretraining, where their pretraining tasks involve reconstructing masked topological and geometric subgraphs in combination. The proposed MoleculeJAE in this work follows this research line.
Alternative training objectives of diffusion models In addition the score matching training objective tailored for continuous probability distributions, an alternative training objective suitable for categorical probability distributions is the restoration loss proposed
33
min E. Exoo
|| Ro (D(oo, 1) -ool--
(19)
is worth noting that compared to the score matching loss, Equation 19 is closer to the classical auto- decoder reconstruction loss if we disregard the time variable t. Furthermore, [31] has demonstrated that for discrete DDPM models, predicting zig and the score function are equivalent up to scaling constants, as the distribution of zz- can be obtained in closed form through the forward process of DDPM using Zy and the one-step noise.
Moreover, [72) proposed reparameterization of the score matching loss, where the backbone neural network models the residual: T; x4 To- Under the notation Formula the Soft Score Matching (SSM) in 72 is defined as:
Cram == Em Sem (llp ()
r
Note that this reparameterization only works for the case when the closed form solution of the forward process has the form of Eq.
C Dataset Specification
Table 5:
Some basic statistics on MD17.
Pretraining
Aspirin
4
Benzene 1
Ethanol 1
Malonaldehyde
Naphthalene 4
Salicylic
Toluene J
Uracil 1
Train
1K
1K
1K
1K
Validation
1K
IK
1K
1K
1K
1K
IK
1K 131,770
Test
209.762
47,863
553,092
991,237
324,250
318,231
440,790
D Implementation and Hyperparameter
In this section, we provide the detailed hyperparameters for implementing the experiments.
18

---

Table 6: Hyperparameter specifications for MoleculeJAE of property prediction.
Hyperparameter
Value
epochs
[50, 100)
learning rate
[5c-4, le-4)
B
[00.1, 10]]
number of steps
1000]
[1)
[0. 0.01..1.11 1
Table 7: Hyperparameter specifications for MolcculcIAE of molecule generation.
Hyperparameter
Value
epochs
[3000, 10000)
learning rate number of layers
[2e-4, 3e-4)
12
number of diffusion steps
[500]
diffusion noise schedule
cosine
mip hidden dimensions
[X: 256, E: 128, 128, pos: 64)
Atrain
[X: 0.4, E: 2, y: 0, pos: 3, charges:
E
Limitations
From algorithmic standpoint, as briefly mentioned Section 3.3 the parameterization of the score function for the reconstructive task can be obtained from the contrastive surrogate by numerically marginalizing the joint distribution p& (Z) and subsequently applying automatic differentiation with respect to the data variable Zg. This methodology simplifies our two-head decoder framework described in Section 3.4 to single head. Despite the additional computational cost associated with this approach, it holds value in terms of its theoretical significance and is therefore worth exploring in the future.
Another limitation is the treatment of the forward (noising) process of diffusion models as form of trajectory augmentation our current framework remains unparameterized. As the next step, is important to parameterize the trajectory augmentation process similar to automatic data augmentation For instance, the flexible diffusion framework proposed in [29] could be adapted to the molecule graph setting, allowing for the joint optimization of molecule (forward + backward) trajectories. By incorporating such parameterization, we can enhance the flexibility and effectiveness of our approach in modeling realistic molecule dynamics. This extension would allow for a more fine- fine-grained control over the trajectory augmentation process and enable the optimization of molecule dynamics in principled manner. Exploring these possibilities represents exciting direction for future research.
techniques 73 74
More Results
Table 8: Results on GEOM.
Valdby
Unoget
AnoeTV
BoneTV
Lenghts Angles
EDM 3000 gooch)
55
92.9
34.8
100.8
0.12
1100)
0.112
B002
6,35
epestt
824
MoleoleeAE Pretnice
845
Table 9: Results on QM9.
Mocel
Mai sable 1
Aloss stilic T
Validte 1
Usaquc" AMeeTV
BeaaTY
Boad Leazio- W1
Bveal Augles WL.
EDM
95.2
$1.3
0:01
0.64
MIN 20 gueds 1000)
AL
99.0
95.5
0033
n036
0/02
0.00
400
PCOM4M-22com
2D + 3D molecule structure generation We evaluate MoleculeJAE's performance on unconditional molecule generation tasks. We generate both the graph structure and the conformer simultaneously, which is the joint distribution of the molecule's 2D and 3D structures. This task also gives us a chance to test if our MoleculeJAE pretraining framework is compatible with the (discrete) cold diffusion
19

---

framewor 3.1 The pretraining method still follows Figure I and 3 with the backbone 3D GNN of (P(t), Ht substituted by the MiDi transformer in 75] The discrete diffusion formulas are the same as 75
Dataset We still apply the PCQM4Mv2 dataset as the pretraining dataset. For the unconditional molecule generation downstream dataset, we fine-tune tune our model on the QM9 dataset and the GEOM- DRUGS dataset (76.. GEOM comprises 430,000 drug-sized molecules with an average of 44 atoms 181 atoms. We follow 75
to split the dataset.
Evaluation We test the generation performance by calculating the probability distance of 2D and 3D structures between our generated samples with the test set. The precise definition of these metrics are provided in section 5.1 [75]. Besides the MiDi model in [75], we also use EDM 22 as our baseline. For our pretrained model, we finetune MoleculeJAE without the contrastive loss for 600 epoch in GEOM, and 4000 epoch QM9.
The experimental results provided in Table and with hyperparameters setting given Table 7 We achieve state-of-art performance for both datasets.
Results
20


Reviews:
[{"review_id": "ankoQUIFyH", "reviewer": "Reviewer_LUAH", "summary": "The authors propose an auto-encoding method for learning molecular embeddings from both 3D and 2D information jointly. The method is loosely related to diffusion methods in that embeddings of data augmentation trajectories are learned via a score-based reconstruction loss and contrastive loss. The embeddings give SOTA performance on a number of QM and activity prediction tasks.", "questions": "- The unnumbered equation after Eq. 5 does not make sense - how does one apply a rigid transformation to time t?\n- How are the results in Tables 1 and 2 obtained? What model is trained on the learned embeddings?\n- What is the intuition behind this entire paper? Why does embedding a trajectory give better performance than embedding just the stationary data? How do these results depend on the choice of trajectories?", "limitations": "Some limitations are addressed in the appendix.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "### Originality \n\nThe method is original in that it encodes both 3D and 2D molecular information via learning to embed data augmentation trajectories.\n\n### Quality\n\nThe results indicate that the resulting embeddings are SOTA.\n\n### Clarity \n\nThe paper could be made much clearer (see Weaknesses).\n\n### Significance\n\nOptimal molecular embeddings are a perennial desire for all kinds of downstream molecular tasks. The SOTA performance of these embeddings suggests that this method represents a significant contribution.", "weaknesses": "- Section 2 Background is a meandering presentation of a number of related ideas, some of which are not immediately relevant for the main paper. Those should be moved to the appendix, and the remainder should be made much crisper so that it is clear to the reader how the remaining sections will proceed.\n  - the unnumbered equation after Eq. 4 is missing an $e_m$.\n- It is not at all clear what data augmentations are used in this paper. E.g. what are $\\mu_1$ and $\\mu_2$ in Eq. 6?\n- The presentation of the MLE in Section 3.2 is unnecessarily confusing. There is no need to introduce energy-based models, Eq. 9 can be derived without the energy functions. Note that the energy functions are never referred to again. The only thing is the surrogate gaussian KDE, but again that does not need to be motivated by energy functions.\n\nTypos:\n- The statement $E_θ(x0) = \\log q_θ(x_0)$ is incorrect since it is missing a $\\log Z_\\theta$.\n- Most of the score expressions in section 3.2 are missing a $\\log$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We express our sincere gratitude for your meticulous examination and insightful feedback. \n\n**Weaknesses: Organization of Section 2 and Section 3.2 and Equation Missing $e_m$** : To rectify this, we will enhance the clarity of Section 2 by trimming the general introduction of the diffusion mechanism . Our intention is to directly establish a linkage between our equivariant trajectory approach and the classical data augmentation technique, employing the heat and discrete (cold) diffusion equations (3) and (4). The terminology \"energy\" at line 195-196 is employed solely to represent the logarithm of the probability density function, a convention consistent with Bayesian optimization literature. \n\nTo response your concern on introducing the general energy based models,  We will streamline the MLE exposition in Section 3.2, omitting the energy-based models' introduction for Equation 9 derivation.\n\nFurthermore, to optimize space usage, we will introduce a comprehensive roadmap at the conclusion of the introduction. This addition will illuminate the forthcoming sections' structure, easing the reader's navigation through our paper.\n\n**Clarification on Data Augmentations in Eq. 6**: We deeply appreciate your scrutiny. We commit to providing unequivocal clarity by explicitly stating that $\\mu(x,t) = -x$ for both $x=E(t)$ and $H(t)$ in Eq. 6. We also use a discrete equation of $E(t)$ and $H(t)$ (formula is given in appendix) for an additional gneration task suggested by other reviewers. Please concult the general response. \n\n**Addressing Typos on missing $\\log$**:  We will delete the statement on $E_{\\theta}(x_0)$. We want to emphasis that the rest of the  argument remains valid due to the cancellation of $Z_{\\theta}$ by conditional probability definition. The assertion regarding missing \"score\" expressions in Section 3.2 seems a misunderstanding,  since no \"score' appeared in section 3.2. \n\n**Questions**:\n\nQ1: The unnumbered equation following Eq. 5 will be rectified: t will be replaced with $x_{\\text{3D}}(t)$.\n\nQ2: For results in Tables 1 & 2, they are obtained by re-running standard pretraining-and-finetuning results (the references are given in Line 295-304) in one V-100 GPU. Specifically:\n- In the pretraining stage, we pretrain an encoder (e.g., the equivariant encoder in Figure 1) using various pretraining algorithms.\n- In the finetuning stage, we add a linear prediction head on the pretrained encoder and then take an end-to-end finetuning.\n\nQ3: The paper's intuition is outlined in sections 2.1, 3.2, and 3.3. Essentially, we propose trajectory utilization (inspired by real molecule dynamics and diffusion generative models) for continuous data augmentation. Via standard unsupervised representation learning, we extract learning objectives from augmented data distributions. Our model captures joint trajectory distributions, culminating in a combination of contrastive and generative learning objectives in section 3.3, supported theoretically in section 3.2.  As to the question of  \"How do these results depend on the choice of trajectories\",  we recognize the merit in investigating this avenue further. It aligns intriguingly with the exploration of whether a **composition of diverse data augmentations (trajectory-based or not)** can enhance representation learning. Consequently, we intend to address this pertinent question in our upcoming \"Future Outlook\" section."}, {"review_id": "tTSoW7L95q", "reviewer": "Reviewer_fgxq", "summary": "This paper proposes a new representation learning method for molecules using 2D and 3D structures. The joint distribution between original molecules and augmented molecules is decomposed into reconstructive and contrastive tasks. The proposed model, MolecularJAE, simultaneously tackles both tasks with the help of SE(3) equivariant GNN models. MolecularJAE is evaluated on quantum property prediction and molecular dynamics prediction and achieves a competitive result against baseline models.\n", "questions": "- Why does the energy function formulation (Eq 7) reformulate the maximum likelihood estimator to the one above line 196, even though the normalizing constant depends on theta?\n- The reconstructive task with score matching is advertised as a tool for generating new samples (line 222). Are there any experiments on the generative perspective? Given that the gain from the contrastive task is not significant in the experiments, the score-matching model seems well modeling the underlying distribution of graphs.", "limitations": "I couldn't find any negative societal impact.\n", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Diffusion over two different modalities hasn't been explored as far as I know.\n- The decomposition of joint distribution into reconstructive and contrastive tasks is new and interesting.\n", "weaknesses": "- It is known that the 3D structure of a molecule follows certain physical rules. The diffusion process used in this paper does not account for the prior knowledge from the domain.\n- The performance of the proposed model is worse than some recently proposed models, which hasn't been included in the experiment section. For example, the performance of Equiformer[1] and PaiNN[2] on QM9 is better than the proposed model across almost all tasks.\n- The presentation of this paper can be improved further. For example, Figure 1 does not provide any meaningful information to readers. Figure 2 is also difficult to digest. Polishing these figures with additional details and proper explanations would improve the accessibility of the manuscript.\n- There are too many typos throughout the main text and appendix.\n\n[1] Liao, Yi-Lun, and Tess Smidt. \"Equiformer: Equivariant graph attention transformer for 3d atomistic graphs.\" ICLR 2023.\n\n[2] Schütt, Kristof, Oliver Unke, and Michael Gastegger. \"Equivariant message passing for the prediction of tensorial properties and molecular spectra.\" ICML 2021.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We express our sincere gratitude for your meticulous examination and insightful feedback. Now we want to address your concerns and clarify misunderstandings in detail.\n\n**Weaknesses**\n\n1. **Incorporating Physical Rules**:  We appreciate your suggestion to integrate such knowledge into our diffusion process. However, we must acknowledge the practical challenge in precisely determining the drift and diffusion coefficients for Eq.6. We've explored the avenue of parameterizing these coefficients, as per a flexible framework that paramterizing general physical diffusions [1], yet found the resultant improvement to be marginal. In response, we intend to incorporate the pertinent experimental outcomes into the appendix as a releted emperical study.\n\n2. **Performance Analysis**:  In light of your observation, we emphasize our flexibility in employing various geometric GNN backbone models, as articulated in Line 290. Our pretraining framework operates agnostically to the choice of backbone models, including PaiNN and Equiformer. We contend that **evaluating a pretraining method fairly necessitates maintaining a fixed backbone model while testing diverse pretraining frameworks**, as demonstrated in our experimental section. While we are indeed conducting experiments with a PAINN backbone model [2] as per your suggestion (see the general response), we are confident that our current experiments substantiate the efficacy of our approach.\n\n3. **Enhancing Clarity of figures and Correcting Typos**:  We recognize your concerns regarding Figure 1 and Figure 2. We perceive Figure 1 as illustrative of our novel trajectory modeling pipeline, encompassing both contrastive and denoising elements. Similarly, Figure 2 elucidates our novel joint 2D and 3D information encoding while preserving symmetry.  \n\nWe appreciate your feedback on typographical errors and will address them in line with other reviewers' suggestions. To expedite this process, we kindly request the specific locations of these errors.\n\n**Questions**\n\n1. We appreciate your astute observation. To address this, we will incorporate the expectation of the normalization constant $E_x(Z_{\\theta}(x))$ into Eq. 7 to accurately represent the energy function formulation. Note that this factor will disappear in Eq. 8, therefore doesn't affect the the correctness of the rest of the argument.\n\n2. We've indeed explored the generative perspective through experiments in our ablation study. Moreover, we perform an additional generative task to our pretrained representation. As can be seen in Fig.2, the right-hand side of the model's output is ready for downstream generation tasks. Please consult the experiment section of the general response for details.\n\n[1] \"A flexible diffusion model\"  ICML 2023\n\n[2] \"Equivariant message passing for the prediction of tensorial properties and molecular spectra\"  ICML 2021."}, {"review_id": "zL99PlgRA5", "reviewer": "Reviewer_s6Nt", "summary": "The paper proposes MoleculeJAE, an auto-encoder for both 2D and 3D molecule diffusion trajectories. The model learns the trajectories jointly in a self-supervised manner. Empirically, MoleculeJAE achieves competitive results on property and force prediction benchmarks.", "questions": "There is another work that introduced timestep embeddings in their diffusion process: https://ojs.aaai.org/index.php/AAAI/article/view/25639/25411 What are the advantages of contrastive learning?\n", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The joint diffusion of 2D and 3D is novel. It is widely concerned that point cloud diffusion is less aware of 2D information. I admire the authors' contribution to this problem. I would also thank the authors for providing a detailed theoretical analysis.", "weaknesses": "1. I did not really get the motivation of modeling diffusion trajectories. Personally, I think the intermediate states are less informative than x_0. Could you please further explain your idea, and provide some additional ablation results if possible, for it seems that the current ablation study shows that the role of the contrastive loss is not very significant in many cases? Thanks.\n2. I believe the baselines of QM9 experiments are not state-of-the-art. Please consider this work: https://openreview.net/forum?id=tYIMtogyee\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. Certain aspects of your concerns related to motivation and experiments have also been covered in the general response. Please also review the comprehensive response provided there for further clarifications.\n\n**Weaknesses: Motivation of Modeling Diffusion Trajectories and contrastive learning** \n\n1. Motivation:  In Sections 2.1, 3.2, and 3.3, we furnish an in-depth understanding of our motivation. To expound briefly, Section 2.1 introduces a novel trajectory-based approach, inspired by real molecule dynamics and diffusion generative models, to serve as continuous data augmentation. We subsequently adhere to the established protocol of unsupervised representation learning, aiming to deduce a learning objective from the augmented trajectory data distribution. Our focus centers on modeling the collective (joint) distribution of these equivariant trajectories, ultimately culminating in a fusion of a distinct contrastive element and generative learning objectives, as outlined in Section 3.3. The theoretical underpinning is furnished in Section 3.2. \n\nRegarding your concern about the **informativeness of intermediate states**, we have addressed this aspect in Lines 95-96, highlighting that for small t values, xt remains proximal to x_0. This phenomenon also informs our decision to abstain from inducing white noise typical in diffusion generative models (in experiments, we keep the noise schedule to be small). When the trajectory aligns closely with the original data, we consider it an augmentation of the initial dataset (just like Gaussian bluring and color distrotion are treated as augmentation of CV data).  On the other hand, scine the trajectory is continuous, each amount of transformation between two successive steps is small, and we care more on the correlation of successive steps.\n\n2. We acknowledge your observation regarding the incremental significance of the **contrastive loss** in our ablation study.  We believe adding more optimization tricks customized for contrastive learning  (like adding the momentum update, EMA) may improve the abaltion results. The relevant experiments are currently in progress. However, we want to emphsis that we have theoretically demonstrated that the **joint distribution** of the trajectories contain more information than **marginal distributions**, and we utilized contrastive learning to extract the additional information for achieving better molecular representation.\n\n**Weaknesses: State-of-the-Art Baselines** \n\nRes: We appreciate your recommendation to include state-of-the-art baselines in our QM9 experiments. In Line 290, we emphasize the versatility of our pretraining framework with respect to geometric GNN backbones. Our approach **encompasses a fair testing protocol where the backbone remains constant (Schnet), while diverse pretraining tasks are explored**, as evident in our experimental section. In response to your suggestion, we have already integrated the pretraining framework in https://openreview.net/forum?id=tYIMtogyee into our experimentation (marked as **GeoSSL-DDM-1L**). Please see the general response (point 2 of the experiments part) for more details.\n\n**Question on a related work**\n\nRes: We will add the related work (https://ojs.aaai.org/index.php/AAAI/article/view/25639/25411) involving timestep embeddings in the diffusion process. However, we wish to highlight a fundamental distinction between our representation learning context and the generative model focus of the referenced work. Our trajectory distribution is harnessed as an augmented data distribution, resulting in the derivation of a **combined contrastive and generative approach** from the joint trajectory distribution, as expounded in our learning objectives. While contrasting with prior works, we also emphasize the intrinsic value of contrastive learning as a way of extracting the additional information from the joint ditribution rather than the marginal distributions (**which cannot been down by just introducing timestep embeddings**), substantiating its necessity theoretically and through ablation studies. \n\nSidenote: please see the general response for an additional downstream task of jointly 2D and 3D generation.\n\n\n\n\n"}, {"review_id": "j3lUWf66TZ", "reviewer": "Reviewer_oitq", "summary": "The paper introduces a pretraining method for molecule joint auto-encoding (MoleculeJAE) for 2D molecular topology and 3D molecular geometry. Their approach adopts SE(3) symmetry and is trained by fitting the joint distribution of the trajectories from the forward process of the diffusion model.\n\nThe authors treat the 2D molecular structure and 3D geometry as continuous objects, which are perturbed by Gaussian noise in the diffusion process. Unlike traditional diffusion models that denoise data by marginal distribution modeling, their 'auto-encoding' approach focuses on learning the joint distribution of data pair $(x_0, x_t)$. The overall task is then optimized by two objective functions, including a reconstruction loss and a contrastive loss, as well as a surrogate model.\n\nExperimentally, MoleculeJAE is pre-trained on the PCQM4Mv2 dataset and is used to perform downstream tasks on QM9 and MD17 datasets. The experimental results are competitive on property prediction and dynamics prediction tasks.\n\nOverall, this work is well-motivated and novel, with a theoretical justification for their approach. This paper presents good quality work.", "questions": "1. Jointly learning 2D and 3D molecular data is not completely underdeveloped. Previous works on diffusion models, [1][2], have shown that diffusion models can be used to learn 2D and 3D latent molecular representations effectively jointly. They perform better on molecular generation tasks than single-modal diffusion models, and it is proper to cite these works and mention the difference.\n\n2. I wonder if the authors consider extending this auto-encoding framework to discrete graph topology, as a discrete representation and diffusion process is more reasonable for graph-like objects. If not, I would like to hear the reason.\n\n3. I would like to see the error bars (standard deviation), the number of model parameters, and the total training time compared with different methods.\n\n[1] Hua, Chenqing, et al. \"MUDiff: Unified Diffusion for Complete Molecule Generation.\" arXiv preprint arXiv:2304.14621 (2023).\n\n[2] Vignac, Clement, et al. \"Midi: Mixed graph and 3d denoising diffusion for molecule generation.\" arXiv preprint arXiv:2302.09048 (2023).\n\n\nTypo: \n1. Line 257, the log-likelihood logp is also achieves $\\rightarrow$ the log-likelihood logp also achieves", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "First of all, this work is well-motivated, as jointly learning 2D and 3D molecular representations by diffusion models is still at the stage of development. MoleculeJAE can learn both 2D bond topology and 3D conformation geometry information and 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.  \n\nIt is also novel that the joint 2D and 3D learning of diffusion models is trained in a self-supervised learning manner and optimized by a contrastive learning objective function. MoleculeJAE 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.\n\nUnlike traditional diffusion modeling, MoleculeJAE uses a contrastive learning paradigm to fit the trajectory of pairwise molecules, while also utilizing the reconstructive task to perform denoising for each individual molecule. By using contrastive learning, it allows the model to align the augmented views of the same data and simultaneously contrast the augmented views of different data. Also, ablation studies show the importance to have a contrastive loss in the task.\n", "weaknesses": "For methodology, the authors treat the 2D graph topology as a continuous object that is perturbed by Gaussian noise. I understand it is easier and reasonable to put the three things, atom features, bonds, and coordinates, under a single SDE framework. However, following previous works [1][2], describing the 2D graph topology as a discrete object with a discrete diffusion process is more reasonable, and their models show superior results. \nExperimentally, MoleculeJAE shows better results with only small margins. However, the authors do not compare the number of model parameters, the overall training time, and the error bars (standard deviation) with existing methods, so it is unclear to me how well MoleculeJAE surpasses those methods.\n\n[1] Vignac, Clement, et al. \"DiGress: Discrete Denoising diffusion for graph generation.\" arXiv preprint arXiv:2209.14734 (2022).\n\n[2] Hua, Chenqing, et al. \"MUDiff: Unified Diffusion for Complete Molecule Generation.\" arXiv preprint arXiv:2304.14621 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. Certain aspects of your concerns related to motivation and experiments have also been covered in the general response. Please also review the comprehensive response provided there for further clarifications.\n\n**Weaknesses**\n\nA：Thank you for your insightful comments regarding our methodology and the referenced works. As we have outlined in section 2.1 and the appendix, our approach is indeed compatible with discrete diffusions [1] [2]. Moreover, we have extended our 3D structure modeling to incorporate the discrete diffusion framework. For detailed information, kindly refer to Appendix A.3. We also want to emphsis that our framework is mainly for unsupervised representation learning, as a result, certain related works focused exclusively on generation tasks have been inadvertently excluded. We will add the two references as you suggested.\n\nIn response to your concerns about experiments, we have introduced a **new generation downstream task grounded in discrete diffusions**, which serves as a supplement to our property prediction downstream task. For further details, please consult the last part of the general response. \n\n**Question 1**\n\nA: We will discuss the relation between the mentioned papers and our method in the revised version. Notably, our approach introduces a novel unsupervised molecule representation learning framework centered around trajectory modeling. This contrasts with the focus of [1] and [2], which primarily concentrate on generation tasks.\n\n**Question 2**\n\nA: Your inquiry about extending our auto-encoding framework to discrete graph topology is astute. We have indeed considered this extension, and Appendix A.3 provides comprehensive details. From a methodology standpoint, our approach inherently encompasses discrete trajectories, akin to the formulation in [1] and [2].  Furthermore, an additional experiment addressing discrete generation has been performed in the **general response**.\n\n**Question 3**\n\nA: Thank you for raising these questions, and they are important details. We have illustrated the details below.\n1. On the standard deviations on geometric property prediction tasks, we want to kindly point out that the computational resources for running such standard deviations are huge. Concretely, SchNet is by far the fastest geometric model (in comparison to PaiNN and Equiformer), and it may take around 12 hours for one task. We have 13 (pretraining algorithms) * 12 (tasks) * 12 (GPU hours) * 5 (seeds) = 9.36K GPU hours or 290 GPU days. Such computational cost is huge and impossible for us to add during the rebuttal period. This is also why the existing geometric pretraining methods (as the baselines in this work) do not include the standard deviations.\n2. The model parameters depend on the backbone model.  Roughly speaking, the numer of parameters is the sum of backbone model, the projection head, and the MLP for the time-embedding. For the experiment in the paper,  the number of model parameters is 1.353M.\n3. The training time for different pretraining algorithms are listed below:\n| Pretraining Algorithm | min / epoch |\n| -- | -- |\n| Type Prediction | 8 |\n| Distance Prediction | 7 |\n| Angle Prediction | 8 |\n| GeoSSL-RR | 10 |\n| GeoSSL-InfoNCE | 10 |\n| GeoSSL-EBM-NCE | 11 |\n| GeoSSL-DDM-1L | 12 |\n| GeoSSL-DDM | 19 |\n| 3D InfoMax | 9 |\n| GraphMVP | 11 |\n| Molecule JAE- light | 17 |\n| Molecule JAE | 20 |\nHere, we take all possible ingredients of Molecule JAE (see Fig. 2) into account, and the diffusion steps are also taken to be the maximal.  **Reducing the noise schedule of the diffusion model in Molecule JAE will further decrease the pretraining time** (Molecule JAE- light for non-generative downstream tasks).\n\n**Typos**: Thank you for highlighting the typos. We will rectify the mentioned error in the revised version.\n\n[1] Hua, Chenqing, et al. \"MUDiff: Unified Diffusion for Complete Molecule Generation.\" \n\n[2] Vignac, Clement, et al. \"Midi: Mixed graph and 3d denoising diffusion for molecule generation.\" "}]

Author rebuttal:
## General Response
We thank all the reviewers for their time, and valuable feedback for improvements.  All relevant works and typos mentioned by reviewers will be discussed in the revised version. Now, we clarify and address some common issues that have been raised by the reviewers.

**Motivation and intuition** (Reviewer LUAH, s6nt, oitq): We appreciate the opportunity to clarify the motivation and intuition behind our approach. In **Section 2.1**, we introduced the concept of utilizing trajectories as continuous data augmentation, drawing inspiration from real molecule dynamics and diffusion generative models. Following the standard procedure of unsupervised representation learning, we derived a learning objective from the augmented data distribution. Our methodology involves modeling the joint distribution of trajectories (as elaborated in **Section 3.2**), which culminates in a novel combination of contrastive and generative learning objectives, outlined in **Section 3.3**.  In conclusion, **we learn a powerful representation by fitting the joint distribution of augmented equivariant trajectories which contain both 2D and 3D molecular information**. We hope this clarification better emphasizes the foundation and significance of our proposed approach.

**Paper organization** (LUAH, fgxq):  To enhance clarity, we will relocate certain content from the **background section** to the appendix, ensuring a smoother flow for readers. Additionally, we acknowledge the importance of providing a clear **roadmap** for our work. In response, we will incorporate a comprehensive roadmap in the final paragraph of our introduction. Furthermore, more detailed illustrations of the **figures** will be given within the main text.

**Experiments**  (fgxq, s6Nt, oitq): A fundamental misunderstanding is that our primary aim is to present a comprehensive unsupervised molecule representation learning framework rooted in equivariant symmetric trajectories. Therefore, the most straightforward way to demonstrate our method's effectiveness is to **fix an equivariant backbone neural network** (we choose Schnet), and test if the representation trained by our method can beat other representation learning methods.  In contrast, some **experiments suggested by reviewers (e.g., comparing our QM9 experimental results to PaiNN[1], GNS-TAT[2]) involve different research settings**, and we believe our core experiments have effectively showcased the efficacy of our approach.  To further address your suggestions, we have conducted three additional experiments (** only partial results have been shown due to severe GPU shortage in the rebuttal period**).

1. We performed experiments **replacing the Schnet backbone with PaiNN** on certain QM9 tasks, and our MoleculeJAE achieved state-of-the-art performance. **The table is given in the attached PDF**.

2. We want to highlight that this work [2] comprises a backbone model (GNS-TAT) and a pretraining algorithm (GeoSSL-DDM-1L). (1) The default GNS-TAT in [2] is not an equivariant neural network. (2) **We have already experimented GeoSSL-DDM-1L**. As shown in Tables 1 & 2 on page 8, we can observe that MoleculeJAE is better than GeoSSL-DDM-1L.

3. Reviewer oitq suggests us to test if our **pretrained representation performs well for generation task**. We apply our generalized (discrete) framework in Appendix A.3 for a joint 2D and 3D structure generation (the detailed setting is given in [3]), revealing significant enhancements over non-pretrained models: 

Table 1: Results on GEOM-Drugs with explicit hydrogens. We pretrained our model MoleculeJAE on PCQM4Mv2 for 24 epochs and then finetuned the generative head (see Figure 2) of MoleculeJAE on GEOM-Drugs for 211 epochs. As the table shows, MoleculeJAE achieves SOTA results on 9 out of 10 metrics.

|Model|Mol stable↑|Atom stable↑|Validity↑|Unique↑|AtomTV↓|BondTV↓|ValW1↓|Bond Lengths W1↓|Bond Angles W1↓|
| ------------ | ------------ | ------------ | ------------ | ------------ | ------------ | ------------ | ------------ | ------------ | ------------ |
|EDM (3000 epoch)|5.5|92.9|34.8|100.0|0.212|0.049|0.112|0.002|6.23|
|MiDi (2D+3D, 3000 epoch)|69.2|99.0|67.4|100.0|0.059|0.024|0.036|0.012|5.47|
|MoleculeJAE (Pretrained on PCQM4Mv2)|84.5|99.6|79.7|100.0|0.059|0.021|0.008|0.003|2.16|

We believe these supplementary experiments underscore the robustness and versatility of our approach.

[1] Equivariant message passing for the prediction of tensorial properties and molecular spectra

[2] Pre-training via Denoising for Molecular Property Prediction

[3] Midi: Mixed graph and 3d denoising diffusion for molecule generation


Meta-review: The authors propose an auto-encoding method for learning molecular embeddings from both 3D and 2D information jointly. It achieves strong performance on a number of QM and activity prediction tasks. The presentation of the paper can be improved (make the related work more concise, removing MLE and energy-based model). Additional experiment comparing with Liao et al 2023, Shutt et al 2021, can be added. 

Based on the reviews, I full support inclusion of this paper in the conference.
Decision: Accept (poster)