SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Although it is possible to directly train the likelihood reconstruction term in the auto-regressive case, such as using the noise conditional maximum likelihood loss  $\mathbb{E}_{t\sim [0,T]}\mathbb{E}_{x_t\sim p_t}\log q_\theta (x_t)$  proposed in [50], we instead adopt trajectory score matching [28], which is applicable for non-autoregressive graph data. The score matching loss is defined as follows:
$$
\mathcal {L} _ {s c} := \mathbb {E} _ {t \sim [ 0, T ]} \mathbb {E} _ {p \left(x _ {0}\right) p \left(x _ {t} \mid x _ {0}\right)} \left[ \| \nabla \log p \left(x _ {t} \mid x _ {0}\right) - s _ {\theta} \left(x _ {t}, t\right) \| ^ {2} \right]. \tag {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  $t = 0$ , the score function for 3D structures can be interpreted as a "pseudo" 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 10 serves as a variational lower bound for the marginal log-likelihood  $\log q_{\theta}(x_0)$ . 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, the SE(3)-invariant density  $p_{3D}$  defined by Eq. 5 implies that the corresponding 3D score function  $\nabla_x p_{3D}(x)$  (represented as  $\mathrm{score}(P_t)$  in Fig. 2) is equivariant under  $SE(3)$ :
$$
\nabla_ {x} p _ {3 D} (\mathbf {R} x) = \mathbf {R} \nabla_ {x} p _ {3 D} (x), \forall \mathbf {R} \in S E (3), x \in \mathcal {C}.
$$
In conclusion, the symmetry principle mandates that the score neural network takes an 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.  $\nabla_{x}p_{3D}(x)$ : SE(3)-equivariant; 2.  $\nabla_{y}p_{2D}(y)$ : SE(3)-invariant.
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]):
$$
d y _ {t} = \left[ f \left(y _ {t}, t\right) - \frac {1 + \lambda^ {2}}{2} g ^ {2} (t) s _ {\theta} \left(y _ {t}, t\right) \right] d t + \lambda g (t) d B _ {t}, \tag {11}
$$
for  $\lambda > 0$ . Hence, it is theoretically necessary to optimize the second term  $\tilde{\mathbb{E}}_{p(x_0,x_t)}\left[\frac{\partial\log f_\theta(x_0,x_t)}{\partial\theta}\right]$  of Eq. 16. By employing the conditional probability formula, we have:
$$
f _ {\theta} \left(x _ {0}, x _ {t}\right) = \frac {p _ {\theta} \left(x _ {0} , x _ {t}\right)}{\int p _ {\theta} \left(x _ {0} , y\right) d y}. \tag {12}
$$
Figure 2: Architecture of MoleculeJAE. The inputs are ground-truth molecules with both 2D and 3D structures. MoleculeJAE 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 1.
It is important to note that the troublesome normalizing constant  $Z_{\theta}(t)$  is cancelled out in Eq. 12. In practice, the integral  $\int p_{\theta}(x_0,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,x^{+}))$  of the same data and simultaneously contrast the augmented views of different data  $(x^{-})$ . By treating the joint distribution as a similarity measure, Eq. 12 can be seen as implicitly imposing two regularization conditions on  $f_{\theta}(x_0,x_t)$ : ensuring a high probability for  $p(x_0,x_t)$  and a low probability for  $p(x_0,y)$ . This notion of similarity motivates us to refer to maximizing the second term of Eq. 16 as a contrastive task.
Contrastive surrogate However, estimating  $p_{\theta}(x_0,x_t)$  is challenging due to the intractability of closed-form joint distributions. To overcome this difficulty, we propose a black-box surrogate  $\varphi$  that represents the mapping from the latent representation  $h_\theta (x_t)$  to  $p_{\theta}(x_0,x_t)$  (following the notation in figure 2). Specifically, let  $(h_{\theta}(x_0),h_{\theta}(x_t))$  denote the representation pair obtained from the input molecule data  $(x_0,x_t)$ . Then, the surrogate of  $p_{\theta}(x_0,x_t)$  is defined by  $p_{\theta}(x_0,x_t) = \frac{1}{Z(t)}\exp \{-\frac{\|\varphi(h_{\theta}(x_0)) - \varphi(h_{\theta}(x_t))\|^2}{\tau^2(t)} +C(x_0)\}$ . Here,  $\tau (t)$  is a monotone annealing function with respect to  $t$ . By using Eq. 12, the unknown normalization constant  $Z(t)$  and  $C(x_0)$  cancel out, resulting in the following approximation:
$$
f _ {\theta} \left(x _ {0}, x _ {t}\right) \approx \frac {\exp \left\{- \| \varphi \left(h _ {\theta} \left(x _ {0}\right)\right) - \varphi \left(h _ {\theta} \left(x _ {t}\right)\right) \| ^ {2} \right\}}{\int \exp \left\{- \| \varphi \left(h _ {\theta} \left(x _ {0}\right)\right) - \varphi \left(h _ {\theta} (y)\right) \| ^ {2} \right\} d y}. \tag {13}
$$
Our surrogate is reasonable because our modeling target  $p(x_0, x_t)$  is derived from the joint distributions of a (continuous) Markov process. When  $x_t$  lies along the augmented trajectory of a specific data sample  $x_0$  and approaches  $x_0$  as  $t \to 0$ , the log-likelihood  $\log p_{\theta}(x_0, x_t)$  is also achieved a local maximum. Therefore, based on the Taylor expansion, the leading non-constant term takes the form of a quadratic expression.
Final Objective By combining the reconstruction and contrastive regularization tasks, we define the final multi-task training objective of MoleculeJAE as a weighted sum of the two tasks:
$$
\lambda_ {1} \mathcal {L} _ {s c} + \lambda_ {2} \mathcal {L} _ {c o}, \tag {14}
$$
where  $\lambda_1, \lambda_2$  are weighting coefficients, and  $\mathcal{L}_{co} \coloneqq \mathbb{E}_{t \sim [0,T]} \mathbb{E}_{(x_0,x_t) \sim p(x_0,x_t)} f_\theta(x_o,x_t)$ . In Fig. 2, the noise input  $x_t$  undergoes an encoding process to obtain the latent representation  $h_\theta(x_t)$ , which is then split into two branches:
1. The first branch passes through a score neural network  $\phi$  for reconstruction:  $s_{\theta}(x_t) = \phi(h_{\theta}(x_t))$ ;
2. The second branch incorporates the original data representation  $h_{\theta}(x_0)$  and further projects the pair  $(h_{\theta}(x_0), h_{\theta}(x_t))$  through a non-linear projection head  $\varphi$  for contrastive learning.
Note that the black-box projection function  $g$  (although not used in downstream tasks) also participates in the pretraining optimization, following the conventional contrastive learning frameworks [51, 52]. See Fig. 1 for a graphical illustration of MoleculeJAE's pipeline.
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 is mean absolute error, and the best and the second best results are marked in bold and bold, respectively.
Pretraining	α↓	∇E↓	EHOMO↓	ELUMO↓	μ↓	Cv↓	G↓	H↓	R²↓	U↓	U0↓	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	24.83	0.036	0.032	16.66	16.28	0.275	15.56	14.66	2.094
Distance Prediction	0.065	45.87	27.61	23.34	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	1.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	41.99	25.75	21.58	0.027	0.029	13.43	13.31	0.136	13.03	13.07	1.609
GeoSSL-DDM-1L	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
MoleculeJAE	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 2: 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 ↓	Benzene ↓	Ethanol ↓	Malonaldehyde ↓	Naphthalene ↓	Salicylic ↓	Toluene ↓	Uracil ↓
-(random init)	1.203	0.380	0.386	0.794	0.587	0.826	0.568	0.773
Type Prediction	1.383	0.402	0.450	0.879	0.622	1.028	0.662	0.840
Distance Prediction	1.427	0.396	0.434	0.818	0.793	0.952	0.509	1.567
Angle Prediction	1.542	0.447	0.669	1.022	0.680	1.032	0.623	0.768
3D InfoGraph	1.610	0.415	0.560	0.900	0.788	1.278	0.768	1.110
GeoSSL-RR	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.457	0.829	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
GeoSSL-DDM-1L	1.364	0.391	0.432	0.830	0.599	0.817	0.628	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, 39, 53]. In our model, the condition does not come from labels, but rather an encoding of the ground-truth molecule. Specifically, to obtain the latent representation  $h_{\theta}$  shown in Fig. 2, we implement two equivariant encoders that satisfy the  $SE(3)$  symmetry proposed in Section 3.1. One encoder takes the original molecule as input, and its output is used as a 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  $SE(3)$  equivariant GNN [54] that outputs invariant scalars will suffice.
Equivariant decoder. With our representation  $h_\theta$  which depends on  $(x_0, x_t, t)$ , the decoder part of MoleculeJAE is divided into two heads. One head is paired with node-wise SE(3) frames to match the  $SE(3)$  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 in Appendix A.
4.1 MoleculeJAE Pretraining
Dataset. For pretraining, we use PCQM4Mv2 [55]. It extracts 3.4 million molecules from PubChemQC [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, we take a simple version by mainly modeling the bond information (details in Appendix D). For the SDE models [28] for generating the joint trajectories, we consider both VE and VP, as in Eq. 6.
Baselines for 3D conformation pretraining. Recently, few works have started to explore 3D conformation pretraining. For instance, GeoSSL [5] provides comprehensive baselines. The initial
Table 3: Ablation studies of contrastive loss term in MoleculeJAE. The ablation results are on QM9.
Pretraining	α↓	∇E↓	EHOMO↓	ELUMO↓	μ↓	Cv↓	G↓	H↓	R2↓	U↓	U0↓	ZPVE↓
λ2=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
λ2=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
λ2=1	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 loss term in MoleculeJAE. The ablation results are on MD17.
Pretraining	Aspirin ↓	Benzene ↓	Ethanol ↓	Malonaldehyde ↓	Naphthalene ↓	Salicylic ↓	Toluene ↓	Uracil ↓
λ2=0	1.380	0.359	0.363	0.744	0.482	0.902	0.548	0.590
λ2=0.01	1.289	0.345	0.365	0.613	0.498	0.712	0.480	0.463
λ2=1	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 a 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 a 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 [7] 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 a 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  $U_{0}$  are the internal energies at 0K and 298.15K, respectively. The other two energies,  $H$  and  $G$  can be obtained from  $U$  accordingly. The main results are in Table 1. 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 [46] 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 C for details. The main results are in Table 2, and MoleculeJAE can outperform 12 baselines on 6 out of 8 tasks and reach the second-best for one of the remaining tasks.
4.4 Ablation Study on the Effectiveness of Contrastive Loss
As discussed in Equation (14), there is one important hyperparameter  $\lambda_{2}$  controlling the contrastive loss in MoleculeJAE. We want to conduct an ablation study on the effect of this contrastive term. As shown in Tables 3 and 4, we consider three values for  $\lambda_{2}$ : 0, 0.01, and 1.  $\lambda_{2} = 0$  simply means that we only consider the reconstructive task, and its performance is very close to  $\lambda_{2} = 0.01$ , i.e., the optimal results reported in Tables 1 and 2. However, as we increase the  $\lambda_{2} = 1$ , 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.
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 E for an in-depth discussion.
Jiujiu Chen would like to thank for her Associate Researcher Ruidong Chen, and Professor Xiaosong Zhang in UESTC for their special support.
[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 https://openreview.net/forum?id=CjTHVo1dvR.
[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 arXiv:2206.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:2110.07728, 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, Zhiming 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 S 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 arXiv:2102.10757, 2021.
[12] Lirong Wu, Haitao Lin, Zhangyang Gao, Cheng Tan, Stan Li, et al. Self-supervised on graphs: Contrastive, generative, or predictive. arXiv preprint arXiv:2105.07342, 2021.
[13] Xiao Liu, Fanjin Zhang, Zhenyu Hou, Li Mian, Zhaoyu Wang, Jing Zhang, and Jie Tang. Self-supervised learning: Generative or contrastive. IEEE Transactions on Knowledge and Data Engineering, 2021.
[68] Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. arXiv preprint arXiv:2210.08344, 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 on Learning Representations, 2023. URL https://openreview.net/forum?id=CjTHVo1dvR.
[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, Maurizio Delbracio, Hossein Talebi, Alexandros G Dimakis, and Peyman Milanfar. Soft diffusion: Score matching for general corruptions. arXiv preprint arXiv:2209.05442, 2022.
[73] Joseph Lemley, Shabab Bazrafkan, and Peter Corcoran. Smart augmentation learning an optimal data augmentation strategy. *IEEE Access*, 5:5858–5869, 2017.
[74] Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 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 conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022. doi: 10.1038/s41597-022-01288-4. URL https://doi.org/10.1038/s41597-022-01288-4.
A Methodology Detail
A.1 Details for Section 3.3
Following Eq. 8, we consider the 'marginal' distribution
$$
q _ {\theta} (x _ {0}) := \int p _ {\theta} (x _ {0}, x _ {t}) d x _ {t}
$$
by marginalizing out  $x_{t}$ . We define the corresponding energy function with respect to  $x_0$  as
$$
\bar {\mathbf {E}} _ {\theta} (x _ {0}) = - \log \int \exp (- \mathbf {E} _ {\theta} (x _ {0}, x _ {t})) d x _ {t}.
$$
From the definition of  $\bar{\mathbf{E}}_{\theta}(x_0)$ , the distribution  $q_{\theta}(x_0)$  shares the same normalization constant with  $q_{\theta}(x_0, x_t)$ . Therefore we have
$$
\frac {\partial \log f _ {\theta} (x _ {0} , x _ {t})}{\partial \theta} = - \frac {\partial \bar {\mathbf {E}} _ {\theta} (x _ {0} , x _ {t})}{\partial \theta} + \frac {\partial \bar {\mathbf {E}} _ {\theta} (x _ {0})}{\partial \theta},
$$
where  $f_{\theta}(x_0,x_t)$  denotes the normalized conditional probability  $q_{\theta}(x_t|x_0)$ . By the conditional probability formula, it is obvious that: it's obvious that
$$
\frac {\partial \log q _ {\theta} \left(x _ {0} , x _ {t}\right)}{\partial \theta} = \frac {\partial \log f _ {\theta} \left(x _ {0} , x _ {t}\right)}{\partial \theta} + \frac {\partial \log q _ {\theta} \left(x _ {0}\right)}{\partial \theta}. \tag {15}
$$
By taking the empirical expectation with respect to the sampled pair  $(x_0, x_t) \sim p(x_0, x_t)$ , we find that the gradient of the Maximum Likelihood allows the following decomposition:
$$
\tilde {\mathbb {E}} _ {p \left(x _ {0}, x _ {t}\right)} \left[ \frac {\partial \log p _ {\theta} \left(x _ {0} , x _ {t}\right)}{\partial \theta} \right] = \underbrace {\tilde {\mathbb {E}} _ {p \left(x _ {0}\right)} \left[ \frac {\partial \log q _ {\theta} \left(x _ {0}\right)}{\partial \theta} \right]} _ {\text {G r a d i e n t o f r e c o n s t r u c t i o n}} + \underbrace {\tilde {\mathbb {E}} _ {p \left(x _ {0} , x _ {t}\right)} \left[ \frac {\partial \log f _ {\theta} \left(x _ {0} , x _ {t}\right)}{\partial \theta} \right]} _ {\text {G r a d i e n t o f c o n t r a s t i v e l e a n i n g}}. \tag {16}
$$
Here, we use  $\tilde{\mathbb{E}}$  to denote the expectation with respect to the empirical expectation. In this way, we find that maximizing the likelihood  $p_{\theta}(x_0,x_t)$  is equivalent to solving:
$$
\operatorname {a r g m a x} _ {\theta} \left\{\tilde {\mathbb {E}} _ {p (x _ {0})} \left[ \frac {\partial \log q _ {\theta} (x _ {0})}{\partial \theta} \right] + \tilde {\mathbb {E}} _ {p (x _ {0}, x _ {t})} \left[ \frac {\partial \log f _ {\theta} (x _ {0} , x _ {t})}{\partial \theta} \right] \right\}.
$$
A.2 Details for Section 3.4