SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we provide the details for the model design described in Section 3.4. Specifically, we illustrate how to obtain the latent representation  $h_{\theta}$  from  $x_0 = (P_0, H_0, E_0)$  and  $x_t = (P_t, H_t, E_t)$  (the solution of Eq. 6 at time  $t$ ). It is important to note that in standard diffusion generative models [28, 32], only  $x_t$  is fed into the score neural network.
The 3D part  $(H_0, E_0)$  and  $(H_t, E_t)$  are transformed by two equivariant 3D GNNs, and we denote the invariant node-wise output representation by  $f_0$  and  $f_t$ . Additionally, we embed the time  $t$  (which is essential for trajectory learning) using Fourier embedding, following [28]:
$$
\operatorname {E m d} (t) := \operatorname {F o o r i e r} (t).
$$
We obtain the 3D node-wise representation by concatenating the Fourier embedding,  $f_{0}$ , and  $f_{t}$ :
$$
\mathbf {N o d e} _ {3 D} = \mathbf {M L P} [ \operatorname {E m d} (t) \| f _ {0} \| f _ {t} ].
$$
Up to this point, we haven't utilized the 2D information  $(E_0,E_t)$ . Unlike the node-wise 3D representation, we encode  $(E_0,E_t)$  as a weighted adjacency matrix  $W$ :
$$
W = \mathbf {M L P} [ \operatorname {E m d} (t) \| E _ {0} \| E _ {t} ].
$$
Combing  $\mathbf{Node}_{3D}$  and the adjacency matrix  $W$ , we implement a dense graph convolution neural network (GCN) to obtain the final (node-wise) latent representation  $h_{\theta}$ , following [22]:
$$
h _ {\theta} = \mathbf {G C N} \left(\operatorname {N o d e} _ {3 D}, W\right).
$$
Figure 3: Modular Design of MoleculeJAE.
3D output  $\mathbf{Score}(P_t)$ . To obtain the  $SE(3)$  equivariant and reflection anti-equivariant vector field [60]  $\mathrm{Score}(P_t)$  from  $h_\theta$ , we implement the node-wise equivariant frame and perform tensorization, following [24, 54]. Consider node  $i$  with 3D position  $\mathbf{x}_i$ , let  $\bar{\mathbf{x}}_i \coloneqq \frac{1}{N} \sum_{\mathbf{x}_j \in \mathcal{N}(\mathbf{x}_i)} \mathbf{x}_j$  be the center of mass around  $\mathbf{x}_i$ 's neighborhood. Then the orthonormal equivariant frame  $\mathcal{F}_i \coloneqq (\mathbf{e}_1^i, \mathbf{e}_2^i, \mathbf{e}_3^i)$  with respect to  $\mathbf{x}_i$  is defined as:
$$
\left(\frac {\mathbf {X} _ {i} - \bar {\mathbf {X}} _ {i}}{\| \mathbf {X} _ {i} - \bar {\mathbf {X}} _ {i} \|}, \frac {\bar {\mathbf {X}} _ {i} \times \mathbf {X} _ {i}}{\| \bar {\mathbf {X}} _ {i} \times \mathbf {X} _ {i} \|}, \frac {\mathbf {X} _ {i} - \bar {\mathbf {X}} _ {i}}{\| \mathbf {X} _ {i} - \bar {\mathbf {X}} _ {i} \|} \times \frac {\bar {\mathbf {X}} _ {i} \times \mathbf {X} _ {i}}{\| \bar {\mathbf {X}} _ {i} \times \mathbf {X} _ {i} \|}\right), \tag {17}
$$
where  $\times$  denotes the cross product, which is  $SE(3)$  equivariant and reflection anti-equivariant. Then, we transform  $h_\theta \in \mathbf{R}^{N\times L}$  to  $h_{3D} \in \mathbf{R}^{N\times 3}$ :
$$
h _ {3 D} = \left(h _ {1}, h _ {2}, h _ {3}\right) := \mathbf {M L P} \left(h _ {\theta}\right).
$$
Finally, for node  $i$
$$
\mathbf {S c o r e} ^ {i} (P _ {t}) = h _ {1} \cdot \mathbf {e} _ {1} ^ {i} + h _ {2} \cdot \mathbf {e} _ {2} ^ {i} + h _ {3} \cdot \mathbf {e} _ {3} ^ {i}.
$$
2D output  $\mathbf{Score}(E_t)$ .  $\mathbf{Score}(E_t)$  represents the probability gradient with respect to the molecule bonds. Therefore,  $\mathbf{Score}(E_t)$  has the same shape as the dense adjacency matrix and is  $SE(3)$  invariant. To obtain  $\mathbf{Score}(E_t)$ , we leverage graph multi-head attention [61] to obtain a dense attention matrix  $A_{t}$  (see [61] for the explicit formula):
$$
A _ {t} = \mathbf {A t t} (h _ {\theta}).
$$
Then, we apply an MLP to the edge-wise  $A_{t}$  to obtain the final 2D score function:
$$
\operatorname {S c o r e} \left(E _ {t}\right) = \mathbf {M L P} \left(A _ {t}\right).
$$
See Figure 3 for a graphical representation.
A.3 A Cold Alternative for Section 3.1
Unlike continuous data that can take values ranging from  $(- \infty, \infty)$ , 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  $Q_{t}$  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  $q(z_{t}|x)$ , defined as:
$$
q \left(z _ {t} | x\right) := x Q _ {1} \dots Q _ {t}
$$
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  $t \to \infty$  should not depend on the original data  $x$ , so that we can use it as a prior distribution for inference. A common choice is a uniform transition over the classes, where the corresponding  $Q_{t}$  is defined as:
$$
Q _ {t} := \alpha_ {t} \mathbf {I} + (1 - \alpha_ {t}) \mathbf {1} _ {d} \mathbf {1} _ {d} ^ {T} / d
$$
with  $\alpha_{t}$  transitioning from 1 to 0. By letting  $Q_{t}$  act on each node and each edge, we have defined the discrete diffusion for  $H(t)$  and  $E(t)$ :
$$
q \left(H _ {t} \mid H _ {t - 1}\right) = H _ {t - 1} \cdot Q _ {t}, \quad \text {a n d} \quad q \left(E _ {t} \mid E _ {t - 1}\right) = E _ {t - 1} \cdot Q _ {t}.
$$
On the other hand, although the 3D structure is represented as a continuous point cloud, the molecule's conformer structures at 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\}_{i=1}^c$  be the collection of these rough 3D conformers, and denote uniform sampling on set  $\{P_i\}_{i=1}^c$  as  $\mathbf{Uniform}(\{P_i\}_{i=1}^c)$ . Then, we can construct a discrete diffusion for the 3D structure as follows, similar to Eq. 3:
$$
P (t) = \alpha (t) \mathcal {F} _ {0} ^ {- 1} \cdot P (0) + \beta (t) \mathbf {U n i f o r m} \left(\left\{\mathcal {F} _ {i} ^ {- 1} \cdot P _ {i} \right\} _ {i = 1} ^ {c}\right), \tag {18}
$$
where  $\mathcal{F}$  is a global equivariant frame obtained by averaging the node-wise equivariant frames  $\mathcal{F}_i$  defined by Eq. 17. It can be verified that by projecting the original 3D structure with the inverse of its global frame  $\mathcal{F}^{-1}$ , 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.
The primary objective of self-supervised learning (SSL) is to acquire meaningful and compressed representations from unlabeled data, often measured by their ability to estimate a specific data distribution  $p_{data}$ . 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  $p_{data}$  (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 a 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, are more subtle as they involve perturbations along specific directions of the
natural data that can be identified using the information provided by the score function  $\nabla \log p_{data}(x)$  and  $\nabla \log p_{data}(x|y)$  conditioned on labels  $y$ .
Once the augmented trajectories are established, our learning objective is to fit the joint distribution of the random trajectories. From a practical standpoint, our framework lies at 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 a 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 a similar phenomenon within a more general unsupervised framework, as illustrated in Eq. 4, 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 self-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 (AttrMask [1], GPT-GNN [3]) or detecting the positive and negative pairs in a contrastive manner (ContextPred [1], MolCLR [69]). Meanwhile, several works have studied the molecular geometry pretraining. GeoSSL [70] presents a comprehensive benchmark, including distance prediction, angle prediction, and an 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 in a contrastive and generative manner, and 3D InfoMax [58] is a special case of GraphMVP 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 to the score matching training objective tailored for continuous probability distributions, an alternative training objective suitable for categorical probability distributions is the restoration loss proposed in [33]:[33]:
$$
\min  _ {\theta} \mathbb {E} _ {t} \mathbb {E} _ {x _ {0} \sim p _ {\mathrm {d a t a}}} \| R _ {\theta} (D (x _ {0}, t) - x _ {0} \|. \tag {19}
$$
It 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  $x_0$  and the score function are equivalent up to scaling constants, as the distribution of  $x_{t-1}$  can be obtained in closed form through the forward process of DDPM using  $x_0$  and the one-step noise.
Moreover, [72] proposed a reparameterization of the score matching loss, where the backbone neural network models the residual:  $r_t = x_t - x_0$ . Under the notation of Formula 3, the Soft Score Matching (SSM) in [72] is defined as:
$$
\mathcal {L} _ {s s m} := \mathbb {E} _ {t \sim [ 0, T ]} \mathbb {E} _ {p (x _ {0}) p (x _ {t} | x _ {0})} [ \| \alpha (t) (s _ {\theta} (x _ {t}, t) - r _ {t}) \| ^ {2} ].
$$
Note that this reparameterization only works for the case when the closed form solution of the forward process has the form of Eq. 3.
C Dataset Specification
Table 5: Some basic statistics on MD17.
Pretraining	Aspirin ↓	Benzene ↓	Ethanol ↓	Malonaldehyde ↓	Naphthalene ↓	Salicylic ↓	Toluene ↓	Uracil ↓
Train	1K	1K	1K	1K	1K	1K	1K	1K
Validation	1K	1K	1K	1K	1K	1K	1K	1K
Test	209,762	47,863	553,092	991,237	324,250	318,231	440,790	131,770
D Implementation and Hyperparameter
In this section, we provide the detailed hyperparameters for implementing the experiments.
Table 6: Hyperparameter specifications for MoleculeJAE of property prediction.
Hyperparameter	Value
epochs	{50, 100}
learning rate	{5e-4, 1e-4}
β	{[0.1, 10]}
number of steps	{1000}
λ1	{1}
λ2	{0, 0.01, 1}
Table 7: Hyperparameter specifications for MoleculeJAE of molecule generation.
Hyperparameter	Value
epochs	{3000, 10000}
learning rate	{2e-4, 3e-4}
number of layers	12
number of diffusion steps	{500}
diffusion noise schedule	cosine
mlp hidden dimensions	{X: 256, E: 128, y: 128, pos: 64}
λtrain	{X: 0.4, E: 2, y: 0, pos: 3, charges: 1}
From an algorithmic standpoint, as briefly mentioned in 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_{\theta}(x_0,x_t)$  and subsequently applying automatic differentiation with respect to the data variable  $x_0$ . This methodology simplifies our two-head decoder framework described in Section 3.4 to a 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 a form of trajectory augmentation in our current framework remains unparameterized. As the next step, it is important to parameterize the trajectory augmentation process similar to automatic data augmentation techniques [73, 74]. 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 a parameterization, we can enhance the flexibility and effectiveness of our approach in modeling realistic molecule dynamics. This extension would allow for a more fine-grained control over the trajectory augmentation process and enable the optimization of molecule dynamics in a principled manner. Exploring these possibilities represents an exciting direction for future research.
Table 8: Results on GEOM.
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, 569 epoch)	84.5	99.6	79.7	100.0	0.059	0.021	0.008	0.003	2.16
Table 9: Results on QM9.
Model	Mol stable ↑	Atom stable ↑	Validity ↑	Unique ↑	AtomTV ↓	BondTV ↓	ValW1 ↓	Bond Lengths W1 ↓	Bond Angles W1 ↓
EDM	90.7	99.2	91.2	98.5	0.021	0.002	0.011	0.001	0.44
MiDi (2D+3D, epoch 10000)	94.83	99.66	95.38	97.4	0.023	0.006	0.008	0.007	1.407
MoleculeJAE (2D+3D, Pretrained on PCQM4Mv2, epoch 3995)	97.42	99.80	97.65	97.35	0.006	0.001	0.002	0.004	0.561
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
framework3.1. The pretraining method still follows Figure 1 and 3, with the backbone 3D GNN of  $(P(t),H(t))$  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 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 of [75]. Besides the MiDi model in [75], we also use EDM [27] as our baseline. For our pretrained model, we finetune MoleculeJAE without the contrastive loss for 600 epoch in GEOM, and 4000 epoch in QM9.
Results The experimental results are provided in Table 8 and 9, with our hyperparameters setting given in Table 7. We achieve state-of-art performance for both datasets.