Physical Consistency Bridges Heterogeneous Data in Molecular Multi-Task Learning
Authors: Yuxuan Ren
Dihan Zheng
Chang Liu
Peiran Jin
Yu Shi
Lin Huang
Jiyan He
Shengjie Luo
Tao Qin
Tie-Yan Liu
Venue: NeurIPS 2024 poster (2024)
Keywords: AI for science; molecule structure generation; diffusion model; physical prior;
Primary area: machine_learning_for_physical_sciences

Abstract: In recent years, machine learning has demonstrated impressive capability in handling molecular science tasks. To support various molecular properties at scale, machine learning models are trained in the multi-task learning paradigm. Nevertheless, data of different molecular properties are often not aligned: some quantities, e.g. equilibrium structure, demand more cost to compute than others, e.g. energy, so their data are often generated by cheaper computational methods at the cost of lower accuracy, which cannot be directly overcome through multi-task learning. Moreover, it is not straightforward to leverage abundant data of other tasks to benefit a particular task. To handle such data heterogeneity challenges, we exploit the specialty of molecular tasks that there are physical laws connecting them, and design consistency training approaches that allow different tasks to exchange information directly so as to improve one another. Particularly, we demonstrate that the more accurate energy data can improve the accuracy of structure prediction. We also find that consistency training can directly leverage force and off-equilibrium structure data to improve structure prediction, demonstrating a broad capability for integrating heterogeneous data.

Physical Consistency Bridges Heterogeneous Data in Molecular Multi-Task Learning
Yuxuan Ren'; Dihan Zheng",, Chang Liu; Peiran Jin, Yu Shi, Lin Huang, Jiyan He', Shengjie Luo', Tao Qin, Tie-Ya Liu
Microsoft Research AI for Science
Abstract
In recent years, machine learning has demonstrated impressive capability in han- dling molecular science tasks. To support various molecular properties at scale, machine learning models trained in the multi-task learning paradigm. Neverthe- less, data of different molecular properties are often not aligned: some quantities, e.g. equilibrium structure, demand more cost compute than others, e.g. energy, so their data are often generated by cheaper computational methods at the cost lower accuracy, which cannot be directly overcome through multi-task learning. Moreover, is not straightforward to leverage abundant data of other tasks to benefit particular task. To handle such data heterogeneity challenges, the specialty of molecular tasks that there are physical laws connecting them, and design consistency training approaches that allow different tasks exchange infor- mation directly as to improve one another. Particularly, we demonstrate that the more accurate energy data can improve the accuracy of structure prediction. We also find that consistency training can directly leverage force and rffiiiiuiibriumm structure data to improve structure prediction, demonstrating broad capability for integrating heterogeneous data.
exploit
1 Introduction
The field of machine learning has witnessed a blossom of progress in solving molecular science tasks in recent years, including molecular property prediction 2 3 machine-learning force field (energy/force prediction) 4
13 molecular structure 19 20] In molecular research, these tasks are often required jointly: for example, energy prediction is required for molecular stability and dynamics, and equilibrium structure (conformation) prediction offers the most probable and characteristic structure for understanding molecule interaction and functions. Multi-task learning is hence adopted, where a model trained to predict multiple properties using the same number of decoders (output heads) built
1
6
7 3 electronic structure 10
generation [14
15 16 and design 18
shared encoder (backbone model) (Fig. I 21 This paradigm is also used for pre-training a model by leveraging as much data as possible that are scattered over various tasks and domains 22 23 24).
Nevertheless, science tasks have some unique challenges beyond conventional machine learning tasks, which cannot be adequately addressed by multi-task learning alone. It is more costly curate dataset for molecular science tasks since it calls for running physics- theoretic computation algorithms, which come with stringent securacy-ffficinncy trade-off. Molecular-scient science datasets are hence generated each with a specific algorithmic portfolio that is economic for the particular purpose.
"qqual contribution.
Work done during an internship at Microsoft Research Al for Science. t Correspondence <changliu @ microsoft.com>
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

This incurs two challenges regarding data heterogeneity. (1) Different properties in a dataset may come from algorithms in different levels of theory, meaning different levels of accuracy. This limits the prediction accuracy on some tasks. For example, labeling the energy of molecular structure yet affordable for common molecules using algorithms in the density functional theory (DFT) level, but producing the equilibrium structure of molecule requires repeated energy evaluations hence is tens to hundreds of times more costly. As a result, DFT-level equilibrium structure data are available only in limited scale [25] 26 27 while larger-scale datasets [28] 29] have to resort to lower levels of theory to generate equilibrium structures at scale, which come with lower level of accuracy. Using such data, multi-task learning alone cannot predict structures in accuracy higher than the data-generee gettaaagenrtiinn method. (2) Different datasets focus on different tasks, so combining these datasets to enhance the performance a particular task is not straightforward. For example, there are datasets [30 31 32 that are concerned with force prediction and off-equilibrium structures, which do not provide direct supervision equilibrium structure prediction. Although including these additional tasks in multi-task learning could help learn a better encoder (or, representation), this is only based on an empirical observation from general machine learning perspective and does not directly exploit relevant physical information in the additional tasks.
In this work, we highlight that science tasks also provide fortunate specialties that come to the rescue. In contrast to conventional machine learning tasks which are primarily defined by data, science tasks originate in fundamental physical laws, and data are rather the demonstration of such laws. These laws im- pose explicit constraints between tasks, hence define the "physical consistency" between model predic- tions on these tasks. By enforcing such consistency, model predictions for different tasks are connected and can explicitly share the information in the data of one task to the prediction for other tasks, hence bridg- ing data heterogeneity. From another perspective, while sharing a common encoder connects various decoders the input end. physical consistency closes the loop by connecting the decoders at the output end (Fig. [T. With the additional information from the physical consistency, capabilities beyond conven- tional multi-task learning are enabled: (1) data from a higher-level of theory of one task can improve the accuracy of a physically related task, and (2) the abundant data of physically related task can directly improve the performance of the concerned task.
data consistency
E data consistency data
R argmin E(R')
R pred.
E pred.
X pred.
p(R) exp
Structure Decoder D (0g)
Energy Decoder 2D ()(0)
Task X Decoder De0
Encoder g(()
A
molccular input
Illustration of the idea of physi- cal consistency. To support multiple tasks ("Task X" represents a general task), the model (blue solid lines) builds multiple de- coders on a shared encoder, which are trained by multi-task learning with data of respective tasks (green dotted double arrows). Physical consistency losses enforce physical laws be- tween tasks (orange dashed double arrows), hence bridge data heterogeneity and directly improve one task from others.
Figure 1:
Concretely, we demonstrate the practical value of the physical consistency between energy prediction and equilibrium structure prediction, two central tasks in molecular science. The consistency can be constructed from two perspectives: the equilibrium structure of a molecule is the structure that attains the minimal energy of the molecule (3.2. and the equilibrium structure is a sample of the thermodynamic equilibrium distribution at low temperature (Sec. 3.3) When adopting the denoising diffusion formulation [33 34 35) for structure prediction, we show that the two physical
laws can be translated into consistency loss functions which connect the energy preegyypre model with the structure- predicciooiiiiionn model at large and small diffusion steps. respectively. We apply the consistency losses in the multi-task learning on the PubChemQC 13LYY/6-31G*//PM which is perhaps the largest public dataset with DFT-level (B3LYP/6-31**) energy labels thus highly relevant to model pre-training. But its equilibrium structures are generated in the semi-empirical level (PM6), which is in a lower level of accuracy. We use the consistency losses to transmit the information of the higher-level theory in the energy model to the structure model by only optimizing parameters of the structure decoder, which achieves the followings. (1) Consistency losses improve structure prediction accuracy beyond multi-task learning alone, when evaluated using DFT-level structures from the PCQM4Mv2 (27) and QM9 (25) datasets. The advantage persists even after both have undergone finetuning. (2) Datasets providing DFF-levell force labels on off-equilibrium structures are better leveraged to improve structure prediction accuracy beyond only including force prediction in multi-tas, task learning. Since force is the gradient of the energy, the additional data allow
dataset [29],
2

---

the energy model to learn a better energy landscape, which in turn leads to more accurate structure prediction through the consistency losses. This is a more direct pathway for data on an additional task to benefit structure prediction beyond learning a better representation in multi-task learning. We remark that the improvement not from any more accurate structure data, but from bridging data heterogeneity using the "free lunch" redeemed from physical laws.
1.1 Related Work
Structure prediction. In recent years, deep generative models have been used as a powerful tool to generate molecular structures. Due to the subtlety that structure being rotated as whole is essentially the same structure, early methods opt to generate intermediate geometric variables, such as
inter-atomic distances [6 37 or torsional angles [38 39]. Directly generating Cartesian coordinates of atoms has also been explored where the rotational equivalence is handled by alignment [40] or leveraging gradient of distances 42). More recently, diffuusin models have been used to generate torsional angles 43), or atom coordinates [6 44] using equivariant models. Given the generality and superior performance,
adopt equivariant diffusion model for structure prediction. While these works may generate multiple meta-stable structures, we aim at the single equilibrium structure for each molecule, and investigate the benefit from an energy model.
Leveraging physical laws between tasks. A noticeable example is leveraging the connection between energy and thermodynamic equilibrium distribution to compensate for potentially biased data to better learn the distribution. The equilibrium distribution in a canonical ensemble is the Boltzmann distribution, which is directly determined by the energy function. From the machine learning perspective, the energy function provides an unnormalized density function of the target distribution. Using a flow-based model 45). Boltzmann generators [4] inject the energy supervision via the evidence lower bound objective. Bootstrapped a-divergence objective is thereafter introduced to mitigate mode collapse [46). Zheng et al. [17 use a diffusion model where the energy supervises the score model at the start diffusion step and is propagated to intermediate steps by enforcing PDE. Similar techniques are also explored for conventional machine learning tasks [47(88) 48). More recently, Bose al.
directly connect the energy to intermediate-ster scores. In the opposite samples to learn a coarse- grained energy model. In parallel with these works, the current work is devoted to the investigation of leveraging the connection between energy and structure prediction, and is not using an oracle energy function considering the cost of DFT calculation.
49
direction, Arts et al. [50] leverage Boltzmann-distribure
coarse-grained
2 Technical Background
Before delving into details of the consistency between energy and structure prediction, first intro- duce the problem formulation, and diffusion- based generative formulation for structure prediction.
Chemically, a molecule is specified by the types (i.e., chemical elements) of atoms and bonds, jointly forming a molecular graph g. A molecule in physical reality can take different structures (conformations) R E Ax i.e., the collection of 3-dimensional coordinates of its A atoms. Many properties of molecule G depend on its specific structure R, e.g. the (inter-atomic potential) energy,
so energy prediction is in the form E (0) (R.) with model parameters e.
Among possible structures, the equilibrium structure is the that attains the minimal energy, and is the most representative structure for the molecule. As mentioned, the diffusion formulation has been preferred choice for equilibrium structure prediction, which samples from distribution PG (R) concentrated at the equilibrium structure. For this, a primitive structure is sampled from a simple distribution, e.g., the standard Gaussian, which undergoes diffusion process that transforms the simple distribution to the desired distribution. This is done by reversing the process in the opposite direction, which easier to construct. For example, the process can be taken as the Langevin diffusion that converges to standard Gaussian [35]: dR, = B, V log.(((;, 0, 1) dt + vB, dW,
& R. dt. + 1/3.d.. where 3, is time dilation factor [51. and W, denotes the Wiener process. and ends after a sufficiently long period I when the distribution converges, pr N I). The reverse process is known to
2
The process starts from the desired distribution Pç.0 PG
3

---

follow [52
dR;-
Br-i R Bo--Voogge.r-iiRieee - W
R,-
(I)
where T which transforms I) at 0 to the desired distribution at = T. hence generation process is constructed. To simulate the process, the only unknown is the (Fisher's)
score function V log Pg.. (R) at each diffusion instant, for which machine- model sg. (0)
(R)
is introduced. To learn to fit V log Pg.r (R)), practical approach is by optimizing the denoising score matching loss [3 (Ro) RR) ) g2 (R,) VR. log log Ro)) I2, which is convenient since from the Langevin diffusion, we can derive p (R, Ra) = (R, R, (1 ā II), where ā exp(- J 3, ds), which is known distribution. Leveraging the reparameterization trick, the
loss is reformed as 35
2
(0) Ennit(e; (1 ai Pg. ( EN(e....) sg, (
+ v1- (v
VI C'e
where the expectation w.r.t E (R.) can be estimated by averaging over data. Once the score model trained, it can generate structures by replacing PogPs.. and simulating Eq. [ only Pg.o (R)) ) is desired (instead of po Ro.T)), then an equivalent simulation approach can be adopted
known as probability- ODE 35]:
BT dR;
R logpg.r-:(R)) dē,
2
which is equivalent to the deterministic process in denoising diffusion implicit model (DDIM) (54
An alternative to the score-pree poreeprediction formulation is the denoising formulation [55 56 By defining
"denoising model" | (R.) which formulates the score model following:
va,s
9.l (R4) R.
(0) S.t
RR,)
(4)
a
the training loss function Eq. in terms of S (R.) becomes:
(8)
S ( (Va, + a;e;; Roo
E Ps. (
Ee.
(5)
which follows the intuition to denoise a perturbed structure by predicting the original structure. This formulation better aligns with the notion of "structure prediction" hence can be benefited from and matches structure pre-training strategies [57
successful model architectures [8
3 Method
We begin with the basic formulation of multi-task learning. We then present the two consistency training approaches between energy and structure prediction, based physical laws between the two tasks. The approach to directly leveraging pessically-relatee datasets is described at last.
3.1 Multi-Tash Learning for Energy and Structure Prediction
Both energy and structure prediction tasks require a comprehensive understanding of the input molecular graph G and structure R, so a shared encoder E/ (R) with parameters & is employed. The time step / is required by the diffusion formulation, which is taken as 0 for energy prediction indicating the input structure unperturbed and real. For energy and structure prediction, the
corresponding decoders D (00)) and ((2s) are introduced. For structure prediction, the denoising formulation adopted (end of Sec. 2. Under this formulation, the two tasks are handled by:
(h.)) 0n) (R) (0n) (R),
(#.0s) (R) = D (0s) (g( (R), R).
(6)
On one datapoint (G, R, E). the multi-task loss is (c.f. Eq. (5)): Lmutti- (,, OE, 0s G, R, E)
= by (p.00)) 0z) (R) El + E,
a C Ee.
s50 9.
(Va.R. VI- 2
(7)

---

A subtlety with the models is geometric invariance and equivariance. Indeed, if a structure R is translated and rotated as a whole, the resulting atom coordinates represent essentially the same structure. The energy and the probability density should keep invariant after the transformation. For energy prediction, this can be guaranteed by using an invariant encoder E (R) (no requirement
on D( ). For the probability density, it is known [59] that rotational invariance can be achieved by the invariance of pg.r, which is satisfied by (O I), and the equivariance of the denoising model G.t (())..))RR..............)))))))
s(4b.0)) (R). For this reason, the input structure R re-enters the structure decoder D.0a), which is implemented with an equivariant architecture. Translational invariance of density can be achieved by centering the structures; see ref. [60] for reasoning.
Nevertheless, multi-task learning is restricted by the level of accuracy of training data. As mentioned, structure data are often generated in a lower level of accuracy than energy data due to the more demanding nature. To alleviate this limitation, we exploit physical laws between molecular energy and structure, and propose consistency training losses accordingly to bridge the energy and structure models, thereby enhancing the accuracy of structure prediction from the more accurate energy model.
3.2 Optimality Consistency
One direct relationship between energy and equilibrium structure is that the equilibrium structure Rey minimizes the energy, i.e., Rey = argming (R.).. To enforce this optimality condition, we propose an optimality consistency loss Laptim-cons, in the form of "increase after perturbation" loss. It is based on the idea that the energy of the equilibrium structure should be lower than that of its perturbed version. Denoting R (d.0s) as the model-predicted equilibrium structure, the loss on a molecule G can be written as:
pred
I Trpptim-coos coms (0s) Φ, OF. 9) E, max fo. th 0&) (R pred
Fng (1.00))
(R (t.0s) + +m)]
(8) where n ~ (0, Diag(o")) is a small perturbation, and each element of 0² is independently sampled from (ni l Since the purpose is to improve structure prediction accuracy by leveraging the energy model which has seen more accurate labels, so the consistency loss only optimizes the parameters exclusively for the structure prediction utility, i.e., structure decoder parameters Os- Other parameters dp and 0g do not optimize this loss. In this way, the consistency loss would not contaminate the energy prediction model with the less accurate structure prediction model.
The standard way produce R ((t.0s) requires simulating the diffusion process (Eq. (T) or the equivalent ODE (DDIM) (Eq. D)) which calls the denoising model recursively. So optimizing Loptim-cons would involve backpropagation through the simulation process, which can be impractically costly and numerically unstable. Fortunately, we can exploit the intuition in the denoising formulation and find much cheaper way to generate structure. The intuition of the denoising model Sg. (R,) is to recover the original structure Ro from the perturbed structure R,, Although this is informationally impossible at the instance level, the denoising model still has a definite learning target at the distributional level, which is R. [R,] [56). Particularly, when = T, the correlation between Ro and Rr diminishes, so Sç. (Rr) learns to output E. RRo) the expectation of the target distribution, which is the equilibrium structure since the distribution concentrates at that structure. Under this perspective, the model-predicted structure can be generated by R (b.0g) ( prod
G.T (Rr), where
Rr (O, I). This only requires one evaluation of the denoising model.
Nevertheless, the rotational invariance of the structure distribution introduces more subtleties.
Let ((0) RAX3 + Ax3 be a rotationally equivariant function; that is, for any rotation matrix Q E SO(3) and structure R E RRA33 we have (RQ) ( Then, for any target structure R. E ARAxx3, the minimizer of the denoising loss function L( EN(( ()) -
Proposition 1.
R 112 is the zero map; that is, S(0) (R) = for any R.
See Appendix A for proof. This conclusion reveals that the learning target of the denoising model at is trivially zel-zero, which cannot serve to generate the equilibrium structure. To circumvent this, a simple choice is to denoise from a time step that is close but smaller than T:
(4) ($p00s) (E), E~ N(0,1) (for large T)
(9)
prod
G.r
5

---

The target of the denoising model at r is gg. ()))))())) R = e]. which is equivariant w.r.t R.,-
So the input R., provides an orientation information hence breaking the rotational symmetry of the corresponding distribution p( (Ro R, The resulting expectation then would not average a structure over orientations evenly, hence not zero. More explicitly, since p(Ro|Rr) p(Ro) (Ro) x
p(Ro) exp
RR Ro-R./va..1. we have pR,)(R,) & p(Ro) Rt (1/ā, - 1)1), 2(1/0--1)
where the Gaussian factor assigns larger probability along the direction of R., hence breaks the rotational symmetry from p(Ro). Under this choice, the optimality consistency loss in Eq. (8) is
specified as: Loptim-cons((; dp, Gg.
= EE max 0, Eqe (.0.0r)
0r) (s(00.0) (E)) - E G. (e) (e)+
(0.000))
(for large 7) (10)
0rr
3.3 Score Consistency
The optimality consistency loss only supervises the denoising model at large time steps. For small time steps, alternative perspective the physical law between equilibrium structure and energy can help. Physically, a molecule G in real system can take different structures with different probabilities. When the system is in thermodynamic equilibrium, the probability distribution of the structures can be determined from the energy function of the molecule. Particularly, in a system with fixed volume and temperature T, the structure distribution is the well-known Boltzmann distribution, Pn:c.T (R. x exp (-Fis (R)(R)((k77) approaches zero, the distribution becomes concentrated on the equilibrium structure. This aligns with the learning target of the diffusion model for equilibrium structure prediction. Through the expression of the Boltzmann distribution, the structure model can thus be connected to the energy model.
where kB is the Boltzmann constant. When temperature
To enforce this connection, ideally, the density funetion Pg should match that defined by the energy model. Since the latter only provides an unnormalized density, we enforce their scores to match: E V log Pg
(db.0s) (R)) modeled by the denoising model
th (x))ss) (R) + V (t (R)/(kB ()
where (R) a reference distribution. Nevertheless, is computationally costly to eval- uate the density function from the diffusion model: log (tb00s
(R) = log Pr (R ( +
(3AVa. V So & vār A
(0.0.) g.t (R. (R(000))) dt, where R Rejo,T] the solution to the ODE dR.
((..0s)
Y (V..R. 2
0s) Gt / (R.) (R.)) with initial condition R,- R B3 Significant computational
cost would be incurred from invoking and backpropagating through an ODE solver to evaluate and optimize the density.
We hence turn to another way to leverage this connection. Note from Eq. 4 the denoising model can be used to recover the score model, which targets the score function of the marginal distribution at the corresponding diffusion time instant. Particularly, the score model at f = 0 should approximate the score of the desired distribution, which is Pa:G,7 for small T. The energy model can hence provide supervision to the score model by enforcing this connection: Tuwore- cons
(Os db, 0r, 9)
Vā S (s) D.r
(R,) - R, VR.
E
( 5
(R,)
Ep, (R.)
(for small 7)
(11)
ār
KBT
In this score consistency loss, we have avoided taking the time step T exactly zero for numerical stability consideration. Indeed, when r - 0, the denoising model G.r approaches the identity map, and ā,- approaches 1, making the first term in Eq. (II) an indeterminate form of type 0/0, which may render numerical stability issues. For the reference distribution to generate data evaluate the loss, one can choose either the data distribution Ps.0 for which the energy model gives more confident results, or the perturbed distribution pg,r (can be sampled by adding noise to a data sample
Ro following p R, RR)) for relevance to how the denoising model S (().es) is invoked. Through some trials, we found the latter gives slightly better results.
9.r
Finally, as is the case for the optimality consistency loss, the score consistency loss also only optimizes the parameters 0s of the structure decoder, to ensure the energy model would not be misled by the less accurate structure data. To implement this unconventional optimization requirement, detailed algorithms for the two consistency losses in Appendix B.I in terms of the actual model (8s)
list
components e D (er)) and ()0))

---

3.4 Leveraging Physically-Related Data
Besides the difference in the level of theory to generate data, the heterogeneity of molecular-science datasets also lies in the difference of concerned quantitics. Compared to the enormous chemical space, available datasets for equilibrium structure are still not abundant, while there is vast amount of data generated for physically related but different tasks, for example, labels of atomic forces, and data on off-equilibrium structures. We highlight that the consistency losses can leverage such datasets in an explicit way to further improve equilibrium structure prediction. Note that the consistency losses Eqs. {TO II works by offering the information model, in the form of energy landscape on the structure space; i.e., ranking different structures in optimality consistency, and providing energy gradient in score consistency. For better learning the landscape, the force labels, which are negative gradients of the energy. provides first-order information of the landscape, and energy and force labels on multiple off off-equilibrium structures enable better exploration on the structure space. This approach provides a more direct and concrete information path to equilibrium structure prediction than helping learn a better representation in multi-task learning. For learning a better energy landscape, the force labels are used to directly supervise the gradient of the energy model. The loss term for a datapoint (G, R. F) is:
higher level of theory in the energy
Liuce ④b 0 R.
(R)
where F is the force label. Note that there may be multiple (R, F) data pairs for one molecule G. which provide even richer information on the energy landscape.
4
Experiments
In this section, we demonstrate the advantages of incorporating the proposed consistency losses into multi-task learning. Implementation details are provided in Appendix B
4.1
Setup
Datasets. We consider multi-task learning of energy and structure prediction on the PubChemQC
B3LYP/6-31G///PM dataset (29] (abbreviated as PM6), which is seemingly the largest (~86M molecules) public available dataset with DFT-level property labels, hence a preferred setting for pre-training molecular model. The energy labels are in the DFT (B3LYP/6-31C) G level, while the equilibrium structures are produced at the semi-empirical PM6 61 level, which is less accurate than DFT. Consistency training is hence considered to improve structure prediction accuracy using the more accurate energy data. To evaluate the effect of improved structure prediction accuracy beyond the PM6 level, the accuracy is evaluated against structures generated at the DFT level, which are available in the PCQM4Mv2 dataset 27 (abbreviated as PCQ) and the QM9 dataset 225
Evaluation. Each of the evaluation datasets of PCQ and QM9 is spilt into three disjoint sets for training. validation, and test. The training and validation sets are for optional fine-tuning (see Sec. (4.).
Following existing convention [1 44 40), each test set is prepared by uniformly randomly selecting 200 distinct molecules from PCQ or QM9 that do not appear in the training dataset (PM6). which already makes the test molecules sufficiently dissimilar from training molecules (Appendix C.6).
On each test molecule, we sample 200 structures using the model, calculate their rooted mean square deviations (RMSDs) against the equilibrium structure in the test set, and evaluate the mean and the minimum over these RMSDs. Due to the geometric invariance of the structure distribution, the RMSD is evaluated after translational and rotational alignment of two structures using the Kabsch algorithm [62]. We consider both the denoising (Eq. ) and the DDIM (Eq. (3)) approaches for structure sampling. We also provide coverage evaluation results in Appendix C.2
In each setting, we independently repeat the evaluation process for five times using different random seeds, and report the mean of the repeats in the following tables. The standard deviations and t-test p-values are collectively provided in Appendix C.5) settings using consistency training, both the optimality (Eq. (10)) and score consistency losses (Eq. (TI)) are added to the multi-task training loss (Eq. (7). Validation results for training terms of both energy prediction and structure generation provided in Appendix C.4
7

---

Caraistancy
Carnislancy
Carsisiency
-007
2011
-1900
-2200 11899 -2000
-2/0
2009
Phed medeeceed eneegyy
Phedtco peddccadd energy
20011
-2001
3
Predicted energy of
Predicted aner Rag
Prodic and wnergy Rag
Figure 2: Comparison of energy (eV) on the model-generated structure Ryred using the denoising method and the equilibrium structure Req in the PCQ dataset. Each point represents the model- predicted energy values on the two structures for one test molecule. Models are trained on (left) the PM6 dataset, (middle) the PM6 dataset and SPICE force dataset, and (right) the PM6 dataset with subset of force labels. The closer a point lies to the diagonal line, the closer the energy of the predicted structure is the minimum energy, indicating a closer prediction of equilibrium structure.
4.2 Structure Prediction Results
We first evaluate the effect of consistency training over multi-task training following the above settings. The results are shown in [ We see that consistency training enhances the accuracy of structure prediction beyond multi-task learning consistently, without using any more accurate structure data. The improvement is significant when compared to the standard deviations provided in Appendix Table C.8, which are as low as around 0.003 Á (Appendix Table C.10 verifies t-test significance). We note that this improvement not the cost of lower energy prediction accuracy. as indicated by the energy prediction results in Appendix Table (.6lef..
To further examine that the improvement is from the effect of consistency training, evaluate the energy of the predicted structure and the true DFT-level equilibrium structure in the PCQ dataset for each test molecule using the model, which is shown in the scatter plot of Fig. 2 left). We observe that when using consistency training, the overall energy is reduced, indicating that the predicted structures indeed have lower energy hence closer the true DFT-level equilibrium structure. To quantify this improvement, in Appendix C.3 we define an "energy gap" metric, and the results shown in Table C.5 consolidate the observation.
4.3 Results using Physically-Relateee Data
We next investigate the effect of consistency training for leveraging physically related data, as explained in (.4 For this, we consider the force data in the SPICE dataset (PuhChem subset) [22. The force data are also available on multiple off-equilibrium structures for each molecule. We note the subtlety that the setting of DFT in SPICE, WB97M- (09977-D3BB)/de222222VP
is different from that for generating energy labels in the PM6 dataset. (Up to our knowledge, there does not seem to exist a public force dataset that matches the DFT setting as the PM6 dataset.) On one hand, calculation on near-equilibrium structures of small molecules is not very sensitive to DFT settings, especially for force calculation (even less affected than energy), while the force labels on multiple structures could be more valuable to learning the energy landscape despite the mismatch. So we still consider it as
Table 1: Test RMSD (A; lower is better) of structure prediction by multi-task learning and consistency learning PM6 dataset.
Test Set
PCQ
QM9
Generated by
Denoising
DDIM
Denoising
DDIM
Struet. Stat.
Mean
Min
Mean
Min
Mean
Min
Mean
Min
Multi-Task Consistency
1.189
0.655
1.041
0.361
0.928
0.545
0.669
0.197
1.158
0.645
1.007
0.346
0.848
0.490
0.650
0.194
8

---

Table 2: Test RMSD (A; lower is better) of structure prediction by multi-task learning and consistency leamning the PM6 dataset with additional SPICE force dataset PM6 subset force data.
Test Set
PCQ
QM9
Additional
Training Data
Generated by
Denoising
DDIM
Denoising
DDIM
Struct. Stat.
Mean
Min
Mean
Min
Mean
Min
Mean
Min
SPICE force
Multi-Task Consistency
1.161
0.631
1.047
0.373
0.876
0.486
0.670
0.207
1.147
0.590
1.013
0.345
0.842
0.485
0.644
0.194
PM6 subset
Multi-Task Consistency
1.199
0.672
1.027
0.365
0.914
0.545
0.648
0.193
force
1.113
0.629
1.019
0.351
0.836
0.488
0.646
0.192
relevant investigation setting. Note energy labels in SPICE are not used, and additional structures are not used for training the structure decoder. On the other hand, to reduce the gap, we also generated in-house force labels subset of PM6 structures using the same DFT setting (B3LYP/6-31G" as for the PM6 dataset energy labels. The systematic error controlled, although for each molecule labeled structure.
there is only
The results after incorporating SPICE force data and PM6 subset force data in training are shown in Table 2 We first observe that consistency training still outperforms multi-task training in structure prediction in all cases (see Appendix Tables C.8 and C.10] for standard derivations and t-test p- values). We note that this improvement is not at the cost of a lower energy prediction accuracy. as indicated by Appendix Table T.blee middmid indicates energy prediction is also not compromised in consistency training in this case. When compared to Table [ data does not uniformly enhance multi-task learning performance, since the mechanism learn a better representation is still implicit and indirect and may require extensive tuning. In contrast, using consistency losses improves structure prediction more consistently when physically related data are available in training. These observations indicate that consistency loss training can potentially assist the model in more effectively utilizing data from different sources or modalities.
we see that the inclusion of force
Energy analysis is presented in Fig. middle) utilizing SPICE force data, and in Fig, righht with force labels on a subset of PM6 molecules. The data indicate that training with consistency loss results in lower predicted energies than multi-task training. Furthermore, we observe that the structures predicted by models trained with force datasets (Fig- 2 middle and right) have lower predicted energies compared to those trained exclusively on the PM6 dataset (Fig. 2 left), illustrating the advantage of incorporating force labels.
4.4 Fine-Tuning Results
In addition to the zero-shot prediction evaluations, we further fine-tune the pre-trained models investigated in Sec. 4.2 using DFT-level structures in PCQ or QM9 training datasets. The results presented in Table indicate that the inclusion of consistency loss in pre-training still enhances the accuracy of structure prediction even after fine-tuning See Appendix Tables C.Iaad statistical significance. This could be attributed to that using consistency losses in pre-training can already inform the model of more accurate structure information, leading to a state that more relevant to the underlying physics, which is a favored starting point for further improvements through
Table 3: Test RMSD (A; lower is better) after finetuning for structure prediction pre-trained by multi-task learning and consistency learning on the PM6 dataset.
Test Set
PCQ
QM9
Generated by
Denoising
DDIM
Denoising
DDIM
Struct. Stat.
Mean
Min
Mean
Min
Mean
Min
Mean
Min
Multi-Tase
1.158
0.614
0.921
0.220
0.889
0.467
0.501
0.090
Consistency
1.152
0.610
0.918
0.218
0.835
0.420
0.493
0.076

---

fine-tuning, Results of fine-tuning models that are pre-trained with SPICE force and PM6 subset force are shown in Appendix C.T which further demonstrates the advantages of the consistency loss.
5 Conclusion and Discussion
This work leverages physical laws between molecular tasks to bridge data heterogeneity in multi-task learning. Consistency losses are designed to enforce physical laws between inter-atomic potential energy prediction and equilibrium structure prediction. They have shown to improve structure prediction beyond the typical accuracy level of structure data by leveraging abundant energy data in higher level of accuracy, and can directly leverage force and off-equilibrium structure data to further improve the accuracy. The advantage still holds after finetuning. We would like to highlight additional data (e.g., more accurate structure data)
that the improvement comes "for free"" required, demonstrating the value of physical laws learning molecular tasks. The idea bears broader generality data heterogeneity is ubiquitous in the science domain, and data for specific task are often limited in either abundancy or accuracy.
The current work is limited to the consistency between energy and structure prediction, while more consistency laws can considered molecular science. Apart from mentioned works in connecting energy and thermodynamie distribution, more possibilities include electronic structure and molecular properties, and fine-grained and coarse-grained structures and macroscopic statistics. The significance of improvement in this work is still limited by the abundance of the data involved. Further improvement can be expected with more abundant/divers, but possibly less accurate structure data from, e.g., RDKit [63] or experimental measurements, or energy/force datasets in matching level of theory that more extensively explore the structural space.
References
[1] Oliver T Unke and Markus Meuwly. PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of chemical theory and computation, 15(6): 3678-3693, 2019.
[2] Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tin-Yan Liu. Do Transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877-28888 2021.
[3] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pages 9377-9388. PMLR, 2021.
[4] Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and E Weinan. Deep potential molecular dynamics: scalable model with the accuracy of quantum mechanics. Physical Review Letters, 120(14):14301,, 2018.
[5] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal directional graph neural networks for molecules. Advances Neural Information Processing Systems, 34: 7790-6802. 2021.
[6] Chi Chen and Shyue Ping Ong. universal graph deep learning interatomic potential for the periodic table. Nature Computational Science, 2(11):718-728 2022.
[7] Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:1 1423-11436, 2022.
[8] Albert Musaclian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
[9] He Li, Zun Wang, Nianlong Zou, Meng Ye, Runzhang Xu, Xiaoxun Gong, Wenhui Duan, and Yong Xu. Deep-learning density functional theory Hamiltonian for efficient ab initio electronic-struccure calculation. Nature Computational Science, 266) 2022.
10

---

[10] He Zhang, Chang Liu, Zun Wang. Xinran Wei, Siyuan Liu, Nanning Zheng, Bin Shao, and Tie- Yan Liu. Self-consistency training for ftunttonal-ttheee Hamiltonian prediction. In Forty-first International Conference on Machine Learning, 2024. htt.htttps:::orrreeeee
net/forum?id-Vs44Y
[11] James Kirkpatrick, Brendan McMorrow, David HP Turban, Alexander L Gaunt, James Spencer, Alexander GDG Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, et al. Pushing the frontiers of density functionals by solving the fractional electron problem. 2021.
Science, 374(6573)
[12] R. Remme, T. Kaczun, M. Scheurer, A. Dreuw, and F. A. Hamprecht. KineticNet: Deep learning a transferable kinetic energy functional for orbital-free density functional theory. The Journal of Chemical Physics, 159(14):1441) 10 2023. ISSN 0021-9606. doi: 10. (063//5.11822))
[13] He Zhang, Siyuan Liu, Jiacheng You, Chang Liu, Shuxin Zheng, Ziheng Lu, Tong Wang, Nanning Zheng, and Bin Shao. Overcoming the barrier of orbital-free density functional theory for molecular systems using deep learning. Nature Computational Science, Mar 2024. ISSN (662-8457, doi: 10. 10.1038/s43588 005588822400005---
[14] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 1147, 2019.
[15] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873):533-58, 2021.
[16] Emiel Hoogeboom, Víctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8887. PMLR, 17-23 Jul 2022. URL /ttppptss////ooooeeee hoogeboom222.. html
mlr press/v162/
[17] Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng, Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, He Zhang, Shidi Tang, Hongxia Hao, Peiran Jin, Chi Chen, Frank Noé, Haiguang Liu, and Yan Liu. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence, May 2024. ISSN 2522-5839, doi:
[18] Joseph L Watson, David Juergens, Nathaniel Bennett, Brian Trippe, Jason Yim, Helen Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with RFdiffusion. Nature,
10879977):09--1000, 2023.
[19] John B Ingraham, Max Baranov, Zak Costello, Karl Barber, Wujie Wang, Ahmed Ismail, Vin- cent Frappier, Dana M Lord, Christopher Ng-Thow-Hing, Erik R Van Vlack, et al. Illuminating protein space with programmable generative model. Nature, pages 1-9, 2023.
[20] Rafael Bombrrlo,,,,elli, Jennifer N Wei, David Duvenaud, José Miguel Hernández- Benjamín Sånchez-Lengeling Dennis Sheberla, Jorge Ryan Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science, 4(2):268-2766 2018.
Timothy D Hirzel,
[21] Rich Caruanaf. Multitask learning. Machine learning, 28:41-75, 1997.
[22] Gengmo Zhou, Zhifeng Gao. Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-Mol: A universal 3d molecular representation lcarning framework. In The Eleventh International Conference on Learning Representations, 2023. URL https
//
[23] Duo Zhang, Xinzijian Liu, Xiangyu Zhang, Chenggian Zhang, Chun Cai, Hangrui Bi, Yiming
Du, Xuejian Qin, Jiameng Huang, Bowen Li, Yifan Shan, Jinzhe Zeng, Yuzhi Zhang, Siyuan Liu, Yifan Li, Junhan Chang, Xinyan Wang, Shuo Zhou, Jianchuan Liu, Xiaoshan Luo, Zhenyu Wang,
11

---

Wanrun Jiang. Jing Wu, Yudi Yang. Jiyuan Yang, Manyi Yang, Fu- Fu-Qiang Gong, Linshuang Zhang, Mengchao Shi, Zui Dai, Darrin M. York, Shi Liu, Tong Zhu, Zhicheng Zhong, Jian Lv, Jun Cheng, Weile Jia, Mohan Chen, Guolin Ke, Weinan E, Linfeng Zhang, and Han Wang. DPA-2: Towards a universal large atomic model for molecular and material simulation. arXiv preprint arXiv:2312. 15492, 2023.
[24] Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. One transformer can understand both 2d & 3d molecular data. In The Eleventh International Conference Learning Representations, 2023.
[25] Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[26] Maho Nakata and Tomomi Shimazaki. PubChemQC project: a scale first- pritcprinciples electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 13000 1308, 2017.
[27] Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. OGB- LSC: large-scale challenge for machine learning on graphs. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
[28] Maho Nakata, Tomomi Shimazaki, Masatomo Hashimoto, and Toshiyuki Maeda. PubChemQC PM6: Data sets 221 million molecules with optimized molecular geometries and electronic properties. Journal of Chemical Information and Modeling, 60(12),5811-5899 2020.
[29] Maho Nakata and Toshiyuki Maeda. PubChemQC B3LYP//-31GG//PMG dataset: the elec- tronic structures of 86 million molecules using B3LYP/6-31G* calculations. arXiv preprint Sttt::3055188554 2023.
[30] Stefan Chmiela, Alexandre Tkatchenko, Huziel 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)=e663001,, 2017.
[31] Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T Unke, Adil Kabylda, Huziel E Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9())eadd00873, 2023.
[32] Peter Eastman, Pavan Kumar Behara, David L Dotson, Raimondas Galvelis, John E Herr, Josh T Horton, Yuezhi Mao, John Chodera, Benjamin P Pritchard, Yuanqing Wang, et al. SPICE, dataset of drug-like like molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11, 2023.
[33] Jascha Sohl--ickstten, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[34] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020.
[35] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[36] Gregor Simm and Jose Miguel Hernandez-Lobato. A generative model for molecular distance geometry. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 8949-8958. PMLR, 13-18 Jul 2020. httplattt://proooo 1imm20a.htll
.lr. halr.press/v19/
[37] Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng. and Jian Tang. Learning neural generative dynamies for molecular conformation generation. In International Conference on Learning Representations, 2021. URL htst:p/s//pppeeeeeeeeeee//rrrrra
12

---

[38] Bowen Jing. Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learn- ing from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.01411, 2020.
[39] Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen, William Green, and Tommi Jaakkola. Geomol: Torsional geometric generation of molecular 3D conformer ensembles. Advances Neural Information Processing Systems, 34:13757-13769, 2021.
[40] Jinhua Zhu, Yingee Xia, Chang Liu, Lijun Wu, Shufang Xie, Yusong Wang, Tong Wang, Tao Qin, Wengang Zhou, Houqiang Li, Haiguang Liu, and Tin Liu. Direct molecular conformation generation. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
[41] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9558-9566 PMLR, 18 24 Jul 2021. StkLLhtpps:/prrceeddii
mlr. wwt.press/v33/ www.hhhihh...... html
[42] Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predicting molecular con- formation via dynamic graph score matching. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Infor- mation Processing Systems, volume 34, pages 19784-19795. Curran Associates, Inc.,
2021. URL https://prooceedingg. neurips wttcccccpppooo wwwpccp/pppp.eeeeeeee paape_fflls/pppepppppppppppepppp ffiless/paeer/222//512ee eff.dd/eeebbb7111/77200111199999
pdf
[43] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, volume 35, pages 24240-2444 2425222233322 Curran Associates, Inc., 2022.
[44] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. GeoDiff: geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[45] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR, 2015.
[46] Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobal Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023. URL httpss:/ net/forum?idd-CTII XCTVF
[47] Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion- based generative modeling. In NewrIPS 2022 Workshop Score-Based Methods, 2022.
[48] Francisco Vargas, Will Sussman Grathwohl, and Amaud Doucet. Denoising diffusion samplers.
In The Eleventh International Conference on Learning Representations, 2023. URL https // /oopeneevew..eeeeee wtw//fouumii dadpvnffTAbuu
[49] Avishek Joey Bose, Tara Akhound -adeeh, Jarrid Rector-Brooks Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Niko- lay Malkin, and Alexander Tong. Iterated denoising energy matching for sampling from Boltzmann densities. In International Conference on Machine Learning. 2024.
[50] Marloes Arts, Victor Garcia Satorras, Chin-Wei Huang, Daniel Züggner, Marco Federici, Cecilia Clementi, Frank Noé, Robert Pinsler, and Rianne van den Berg. Two for one: Diffusion models and force fields for coarse- grained molecular dynamies. Journal of Chemical Theory and Computation, 19(18):6151-611) 2023.
[51] Andre Wibisono, Ashia C Wilson, and Michael Jordan. A variational perspective on accelerated methods optimization. proceedings of the National Academy of Sciences, 17755) E7358, 2016.
13

---

[52] Brian DO Anderson. Reverse-timg diffusion equation models. Stochastic Processes and their Applications, 12(3):313-322 1982.
[53] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1611-1674 2011.
[54] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021. URL /ttps:ttp///oooereiviie
net/Eorummid-Stlgiis arCPLP
[55] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Advances in Neural Information Processing Systems, volume 34, pages 21696-21707, 2021.
[56] Giannis Daras, Yuval Dagan, Alex Dimakis, and Constantinos Daskalakis. Consistent diffu- sion models: Mitigating sampling drift by learning to be consistent. In Advances in Neural Information Processing Systems, volume 36, 2024.
[57] Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.0011 00133, 2022.
[58] Chen Wei, Karttikeya Mangalam, Po-Yao Huang, Yanghao Li, Haoqi Fan, Hu Xu, Huiyu Wang, Cihang Xie, Alan Yuille, and Christoph Feichtenhofer. Diffusion models as masked autoencoders. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 16284-16294, 2023.
[59] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pages 5361-5370. PMLR, 2020.
[60] Jason Yim, Brian L Trippe. Valentin De Bortoli, Emile Mathieu, Amnud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. In Proceedings of the 40th International Conference on Machine Learning, pages 40001-40039, 2023.
[61] James JP Stewart. Optimization of parameters for semiempirical methods V: Modification of NDDO approximations and application 70 elements. Journal of Molecular modeling, 13: 1173-1213, 2007.
[62] Wolfgang Kabsch. A solution for the best rotation to relate two sets of vectors. Acta Crystallo- graphica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallography, 32 (5):92299) 129922923,, 1976.
[63] Greg Landrum et al. RDKit: A software suite for cheminformatics, computational chemistry. and predictive modeling. Greg Landrum, 831.11))5281, 2013.
[64] Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie- Liu. Benchmarking graphormer large-scale molecular modeling datasets. arXiv preprint arXiv:2203. 200::240004 2022.
[65] Tianlang Chen, Shengjie Luo, Di He, Shuxin Zheng, Tie-Yan Liu, and Liwei Wang. GeoM- Former: A general architecture for geometric molecular representation learning. In International Conference on Machine Learning, 2024.
[66] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Leurning Representations, San Diega, CA, USA, 2015.
[67] Dávid Bajusz, Anita Rácz, and Károly Héberger. Why Tanimoto index an appropriate choice for fingerprint-based similarity calculations? Journal of cheminformatics, 2015.
14

---

Appendix
A Proof of Proposition [
We first prove the following two preliminary Lemmas.
Lemma A.1. Let e -_. and e, are independent, identically distributed random vectors with e; ~ Is) Q is a random variable uniformly distributed over SO(3). Then the random matrix cQ has the same distribution
Proof. Consider the probability density function Y:
Py(r) -f Pr. Pr.Q(r. Q)QQ pa (a(Q)rri)()))))) Q)dQ.
Since eQ, we have:
imi (2x) exp |
Y:ll
exp 11::1
(y Q) = II
Pe pe(r).
(2 3//2 (2m)
Thus, it follows that:
Py = / pQ (q(Q)priq( QQQQ = / pQ (Q))) -ppp))dddpppp)))...). Pe(
which demonstrates that is identically distributed as e.
Lemma A.2. Assume Q is a random variable uniformly distributed over 50(3). Then EQQ 0.
Proof. For any fixed Qa E SO(3), the distribution of QQo is identical to that of Q due to the uniformity the distribution over the group SO(3). Consequently, we have:
Eq Eq QQo
[[ Qo
Suppose Eq Q [q1, 92: qs), where qı, q2, q3 are the columns of EQQ. Then, follows that:
[q 92, q3 [q1, q2, q3]Qu-
Selecting Q Diag( 11, 11 yields qi qı and qe Q2. which implies that qı q2 O. Similarly, choosing Qo = Diag 1, -1, -1 11} results in q3 Q3, hence qs 0. Therefore, we conclude that Eq Q
We proceed to prove Proposition Let Q be random variable uniformly distributed SO(3). Then according to Lemma A.1 the random matrix y = eQ is identically distributed as e. Then have:
L( E. [|((0) - R Ey [
R R = Ee ( - RR1
E. Q (1 (e)
R 112 Ee
Eq (E) R Q T13
For any e E Ax3 the objective function Eq (c) - R'Q 1| achieves its minimum when (())(e) | R R= T. Since EQQT Eq Q and using Lemma A.2 have EqQ = 0, which implies (:0) (() 0. Therefore, S()) is zero map.
B Implementation Details
B.1 Implementation Details for Consistency Losses
The proposed consistency loss Lopti--oos in Eq. 10 and Lscere-cons in Eq. (II) are designed to update parameters O5 within structure prediction model ((b.0s) To achieve this, we employ the stop gradient operation SG(-) to prevent unnecessary gradient computation for the parameters db in the encoder model gg.7 and 0g in the energy decoder ((())) The detailed implementations of optimality consistency and score consistency are presented in Alg. B.I Alg. B.2 respectively.
15

---

Algorithm B.1 Implementation of Optimality Consistency Loss
Require: Encoder: c(4) &gtt (R., structure decoder ( energy decoder D( Diffusion time step
Ensure: Vo, Lopim-oooonss
1: Sample e and 17.
c(4) 2: Extract the molecular feature h < &G,t=r (e) using the encoder.
3: Apply the stop gradient operation to the molecular feature and obtain h < SG(h) through Pytorch's detach () method.
4: Compute the denoised structure Ry <D$ (h, e) using the structure decoder.
False for the parameters in the energy model Egg (.(0.00) (R)
5: Set requires_gre
) (E
6: Evaluate the loss Loptim--onss max fo, Eg 1
0a) (Ro) _ en) (Ro + n)}
7: Determine the gradient Vo, Lopim cons through automatic differentiation.
Algorithm B.2 Implementation of Score Consistency Loss
Require: Encoder: e()) &g,t" (R), structure decoder ( energyy decoder D$ Diffusion time step
(On)
Ensure: Le.
1: Sample R,
2: Extract the molecular feature hg + $9,1=0 (R,, using the encoder.
<(b)
3: Compute the free energy E < ()(e) (ho) with the energy decoder.
4: Compute the energy gradient F <- V R, E using PyTorch's torch. autograd.
5: Apply the stop gradient operation (.detach() () method in Pytorch) to the energy gradient:
F SG (F).
6: Extract the molecular feature h, + &g,t (R₂) using the encoder.
c(
7: Obtain h SG(h,) using .detach ().
8: Compute the denoised structure Ry + D2 (s)
(SG((,) R,) with structure decoder.
9: Evaluate the loss Lecarr-cons
va. R. 1-a.
7
10: Determine the gradient V Lycore cons through automatic differentiation.
B.2 Model Architecture
The model is composed of an encoder, and two decoders for structure and energy prediction.
The encoder e((p) (R.) a simple modification to the Graphormer model 2 which additionally adopts diffusion time embedding into both node features and pairwise-distan attention bias. The encoder consists of 24 layers of Graphormer, with the dimension of both hidden and feed-forward layers set to 768. utilizes a multi-head attention mechanism with 32 heads and employs 128 Gaussian Basis kernels for enhancing the positional encoding.
For the time embedding, we implement Sinusoida dalPosi Embeddiigg module and TimeStepEncoder module. The former generates dime-dependent sinusoidal embeddings, and the latter refines these embeddings using a forw-forward network with GELU activation function. The resulting time embeddings are then integrated into the node features to inform the model of temporal information.
Additionally, we incorporate time embeddings into the attention mechanism by computing structure- based attention bias. This achieved by calculating the outer product of the time embeddings and using the result as an additive bias in the self-attention layers. This integration allows the model to adapt its attention based on the temporal relationships between nodes the graph.
top of the encoder, the energy decoder ((())) (h) simple MI.P layer concatenated to the invariant node features h.
16

---

Table C.I: Index of main result tables in the paper.
Test
Standard deviation
Paired t-test
Validation
Validation (energy)
Training settings
Coverage
ЕGap
RMSD
(structure)
PM6
Table
Table C.8
Table C.10
Table C.3
Tahle
Table C.6
Table C.5
PM6 w/ force
Table
Tahle C.
Table
C.T0
Table C.3
Table
Table C.6
Table
PM6
finetuning
Table
Table
Table
C.10
Table
N/A
N/A
N/A
PM6 force + finetuning Table
Table
Table
N/A
N/A
N/A
The structure decoder D ( (h, R) adopts the GeoMFormer architecture [65]. which takes the invariant node features h from the output of the encoder, and the atom coordinates R. The output is denoised atom coordinates which are equivariant w.r.t the input coordinates R. These modules are combined to form the energy and structure prediction models following Eq. (
B.3 Training Details
Our pre-training procedure is executed in two discrete stages. Initially, the model subjected to training exclusively utilizing the multi-task loss function, Lomlti-task, for total of 300,000 iterations. Subsequently, in the second stage, we integrate the proposed consistency loss and the force loss, Liores into the training regimen, which then proceeds for additional 200,000 iterations. All the models are trained with the Adam optimizer [66] with batch size 256. The learning rate is set to 2 x 10 with linear warm-up phase in the initial 10,000 steps, which followed by linear decay schedule thereafter.
The weights of the energy loss and the diffusion denoising loss, i.e. the first and the second terms in Eq. 7 are set 1.0 and 0.01, respectively. The weights of the optimality consistency loss Eq. (TO) and the score consistency loss Eq. (11) are set to 0.1 and 1.0, respectively.
We employ a sigmoid schedule across 1,000 diffusion time steps for Bt, with 30 1 x 10-4 and 3T 2 x 10-2 For the optimality consistency loss, the diffusion time step is sampled uniformly from [400, 700). For the score consistency loss, the diffusion time step is sampled uniformly from [5, 300], and kg is set to eV..
The multi-task model is trained 8x Nvidia V100 GPU server for approximately one week. The model with the consistency loss is trained on one 16x Nvidia V100 GPU server.
C Additional Results
In this section, we provide additional results under various combinations settings, and additional metrics and supporting evidence to complement the main results. For clarity, we summarize main result tables in Table C.T for easier indexing.
C.1 Additional Fine- TuninT Results
This section demonstrate more fine-tuning results listed in the main text. Besides the fine-tuned PM6 dataset pre-trained mode, we performed fine-tuning experiments the model pre-trained with the SPICE force and PM6 subset force datasets, as well. The results are presented in Table C.2 Similar to the case in Table we can again observe that while finetuning improves structure prediction accuracy in all settings (compared to Table 2), pre- pre-training the model with consistency loss still enhances the accuracy universally.
C.2 Coverage Evaluation for Structure Generation
Following the common practice in structure-gemeration literature, we also test the coverage of the ground-truth structure over model-generated structures. This metric defined as:
COV (Sgen: Req)
(R
Sgen RRMD((Rgp R.) < 6}
(C.1)
Bgen
17

---

Table C.2: Test RMSD (A; lower is better) after finetuning for structure prediction pre-trained by multi-task learning and consistency learning the PM6 dataset with additional SPICE force data PM6 subset force data.
Test Set
PCQ
QM9
Preennrain Set
Generated by
Denoising
DDIM
Denoising
DDIM
Struct. Stat.
Mcan
Min
Mean
Min
Mean
Min
Mean
Min
PM6 with SPICE force
Multi-Task Consistency
1.161
0.618
0.930
0.219
0.855
0.444
0.505
0.081
1.132
0.581
0.916
0.215
0.832
0.418
0.492
0.073
PM6 with subset force
Multi-Task Consistency 1.099
1.143
0,603
0.927
0.224
0.855
0,441
0.497
0.080
0.542
0.914
0.215
0.822
0.419
0.490
0.076
where Sgen denotes the set of generated structures, Reg denotes the ground-truth equilibrium structure provided from the evaluation dataset, is a threshold parameter, and |-| takes the cardinality of a set. Note that since for the task of equilibrium structure prediction, there is only one ground-truth structure, we only evaluate the so-called precision coverage, since the recall coverage (by switching the roles of gen and Rey in the definition Eq. (C.I)) evaluates to in all cases. Here, the RMSD evaluated after alignment of the two structures by Kabsch algorithm [62] We choose ō as 0.9 and 1. 1.25 for QM9 and PCQ dataset. Same as the RMSD test, we use the same test molecules and sample 200 structure for each molecule. The coverage (COV) of the PM6 dataset model prediction and including force dataset model prediction are shown in Table C.3 The results present consistent results as the RMSD. After adding the consistency loss, all have improvement over the multi-task setting. Besides, incorporating the force data can also improve the accuracy.
Table C.3: Test coverage (higher is better) of structure prediction by multi-task learning and consis- tency learning on the PM6 dataset, and together with additional SPICE force data PM6 subset force data.
Test Set
PCQ
QM9
Training Set
Generated by
Denoising
DDIM
Denoising
DDIM
Struct. Stat.
Mean
Median
Mean
Median
Mean
Median
Mean
Median
Multi-Task
0.597
0,670
0.649
0.675
0.468
0.478
0.717
0.780
PM6
Consistency 0.626
0.695
0.671
0.730
0.604
0.660
0.732
0.825
PM6 with SPICE force
Multi-Task Consistency 0.644
0.614
0.660
0.645
0.685
0.550
0.580
0.715
0.765
0.710
0.675
0.725
0.617
0.705
0.740
0.835
PM6 with subset force
Multi-Task 0.590
0.650
0.662
0.705
0.493
0.520
0.736
0.835
Consistency 0.677
0.775
0.671
0.720
0.643
0.690
0.737
0.845
We also evaluated the coverage the fine-tuned models. The results shown in Table C where we again observe that pre-training with consistency loss improves structure prediction accuracy even after fine-tuning.
C.3 Energy Gap Analysis
Our goal is to predict the equilibrium structure Rey argming ER). It is desirable for the predicted structure to approximate this state of minimal energy. To quantify the proximity of the predicted structure to the equilibrium state, introduce the energy gap metric, which is defined as:
E6 (e)) 0e)
(Rprod) - E
(s).0s) (Rog)
RGap
EE 0.0))) (Rey))
The F.Gap metric serves to evaluate the energy difference between the predicted and equilibrium structures, with smaller energy gap signifying more accurate prediction.
18

---

Table C.4:
Test coverage (higher is better) after finetuming for structure prediction pre-trained by multi-task learning and consistency learning on the PM6 dataset, and together with additional SPICE force data or PM6 subset force data.
Test Set
PCQ
QM9
(Pre-)Trainnl) Generated by
Denoising
DDIM
Denoising
DDIM
Set
Struct. Stat.
Mean
Median
Mean
Median
Mean
Median
Mean
Median
Multi-Task 0.629
0.675
0.721
0.800
0.543
0.595
0.790
0.960
PM6
Consisteney 0.632
0.700
0.724
0.795
0.622
0.690
0.788
0.990
PM6 with SPICE force
Multi-Task Consistency 0.649
0.632
0.705
0.714
0.760
0.597
0.660
0.784
0.980
0.725
0.721
0.795
0.631
0.710
0.791
0.990
PM6 with subset force
Multi-Task Consistency 0.680
0,643
0.705
0.719
0.795
0.600
0.650
0.786
0.990
0.780
0.720
0.794
0.643
0.715
0.789
0.990
To compute this metric, we randomly select 200 molecules from the intersection of PM6 and PCQ dataset, using the structure from the PCQ dataset as Req. The predicted structure Rored is generated using the denoising method. The results are presented in Table We observe that the incorporation of the consistency loss reduces the EGap metric, particularly in cases with additional force labels. These results demonstrate that the consistency loss effectively transfers information from the energy model to the structure model.
Table C.5: Comparison of averaged EGap between structure prediction by multi-task learning and consistency learning. Lower EGap values suggest that the energy of the predicted structure is closer to the theoretical minimum energy.
PM6 with
PM6 with
Training Set
PM6
SPICE force
PM6 subset force
Multi-Task 0.1278
0.0546
0.1163 0.1013
Consistency
0.1172
0.0306
C.4 Validation Results
To make sure that the conclusions drawn from the above results are solid, we further provide validation results for energy and structure.
For the energy. we randomly selected 200 molecules from the intersection of the PM6 dataset (training set) and the PCQ dataset (test set). This choice allows evaluating energy on both PM6 structure and PCQ structure for each molecule, where the former reflects training quality, and the latter reflects the utility for consistency learning of structure prediction. Although the original PCQ dataset [27] does not provide energy labels, we noticed that it is curated based on the PubChemQC Project dataset [26] which provides energy labels on the same PCQ structures under the same DFT settings (B3LYP/6-31G*) as the PM6 energy labels.
The results listed in Table C.6 and boxplots of the distributions of energy prediction MAE (eV) evaluated on the PM6 structures are shown in Fig. C.T. We can observe that energy prediction on PM6 structures are reasonably accurate, and it gets even better when consistency training is activated. This indicates that consistency training does not sacrifice energy prediction. The energy prediction error larger on PCQ structures, which is not surprising since the model does not see data with PCQ structures in training. But as we mentioned in the main paper, better generalization can be expected if foree data are available, which enriches the information on the energy landscape hence improving the generalization on unseen structures. From the table, we can observe that the energy prediction accuracy on PCQ structures is indeed improved. We also notice that including force data in training even improves energy prediction accuracy on PM6 structures, for which a possible explanation is that introducing additional relevant prediction tasks could help the model learn a more comprehensive representation of the input molecular structure, which in turn helps other tasks.
19

---

Table C.6: Validation MAE (cV) of energy prediction trained by multi-task learning and consisteney learning. Validation molecules are randomly selected from the intersection of PM6 and PCQ. and results on both PM6 structures and PCQ structures of the molecules are shown.
Validation structures from
PM6 with
PM6 with
Training Set
PM6
SPICE force
PM6 subset force
Multi-Task Consistency
96.08
83.15
79.00 72.61
PM6 structures
88.58
78.56
Multi-Task 136.66
135.93
117.22
PCQ structures
Consistency
127.06
110.80
108.71
190
190
240
130
350
350
130
220
800
Figure C.I: Box plots for the distributions energy prediction MAE (eV) evaluated on the PM6 structures of randomly selected 200 molecules from the intersection of PM6 and PCQ datasets.
For structure prediction, we validate the trained model by evaluating the RMSD of generated structures on 200 randomly selected PM6 molecules. From the results shown in Table C.7 we can see that the models trained by both multi-task learning and consistency learning achieve reasonable predictions.
Table C.7: Validation RMSD (Å) evaluated the PM6 validation set of structure prediction trained by multi-task learning and consistency learning on the PM6 dataset, and together with additional SPICE force data PM6 subset force data. Predicted structures are generated by the DDIM method.
PM6 with
PM6 with
Training Set
PM6
SPICE force PM6 subset force
Struct. Stat.
Mean
Min
Mean
Min
Mean
Min
Multi-Task 0.970 0.287 0.983 0.314 0.953
0.267 0.290
Consistency 0.931 0.280 0.944 0.286 0.958
C.5 Standard Deviations and Statistical Significance
To evaluate the significance of our comparison results, for each experimental setup, we repeated five independent runs using different random seeds (666666, 666667, 666668, 666669, 666670). The reported results in the main paper are the means of the five repeats, while the standard deviations are listed in Table C.8 and Table C.9 here, which cover the results for the pre- pre-trained models and for fine-tuned models, respectively. We can see that all the standard deviations are around 0.003 Å, which
20

---

is clearly lower than the improvements by consistency training, hence indicating the significance of the effectiveness of the proposed methods.
To more seriously assess the statistical significance, we also conducted a paired t-test for each comparison. The p-values are listed in Vable(C.10) where each number represents the probability of the observed results under the null hypothesis that the means by multi-task learning and consistency learning are the same, hence the lower value the more significant that consistency learning outperforms multi-task learning. Nearly all p-values are well below the 0.05 significance threshold, indicating that the improvement by consistency learning is significant. Exceptional cases almost all correspond to the setting where the result in each repeat is the "Mim'imum RMSD over 200 generated samples, in which case multi-task learning may also has a chance hit the larget structure.
Table C.8: Standard deviations for the test RMSD (A) of structure prediction by multi-task learning and consistency learning on the PM6 dataset (corresponding to Table T), and together with additional SPICE force data or PM6 subset force data (corresponding to Table
Test Set
PCQ
QM9
Training
Generated by
Denoising
DDIM
Denoising
DDIM
Set
Struct. Stat.
Mean
Min
Mean
Min
Mean
Min
Mean
Min
Multi- 0.0013 0.0045 0.0019 0.0059 0.0004 0.0026 0.0017 0.0015
PM6
Consistency 0.0021 0.0021 0,0016 0.0053 0.0018 0.0022 0.0080 0,0016
PM6 with
Multi-Task 0.0021 0.0042
0.0020
0.0067 0.0013
0.0040
0.0084
0.0024
SPICE force
Consistency 0.0026 0.0024
0.0017 0.0030 0.0019 0.0034
0.0025 0.0030
PM6 with subset force
Multi-Task 0.0113 0.0048
0.0044
0.0127 0.0015 0.0134
0.0020 0.0025
Consistency 0.0027 0.0040 0.0023 0.0043 0.0011 0.0030 0.0033 0.0030
Table C.9: Standard deviations for the test RMSD (Å) after finetuning for structure prediction pre-
trained by multi-task learning and consistency learning on the PM6 dataset (corresponding to Table 3
and together with additional SPICE foree data or PM6 subset force data (corresponding to Table C.2
Test Set
PCQ
QM9
(Pre-)
Generated by
Denoising
DDIM
Denoising
DDIM
Set
Struct. Stat.
Mean
Min
Mean
Min
Mean
Min
Mean
Min
Multi-Task 0.0021
0.0059 0.0025
0.0036 0.0008 0.0026 0.0011 0.0044
PM6
Consistency 0.0024
0.0064 0.0020 0.0033
0.0011 0.0028 0.0022 0.0023
PM6 with
Multi-Task 0.0023
0.0053
0.0012
0.0051
0.0016
0.0038
0.0019
0.0033
SPICE force
Consistency 0.0023
0.0047
0.0030
0.0013
0.0018
0.0030
0.0027
0.0015
PM6 with
Multi-Taskk 0.0022
0.0042
0.0016
0.0053
0.0013
0.0036
0.0019
0.0014
subset force
Consistency 0.0019
0.0038
0.0019
0.0072
0.0016
0.0068
0.0016
0.0023
21

---

Table C.10: Paired t-test p-values on structure prediction RMSD means over 5 repeats (standard
deviations are shown in Tables C.8 and C.9) corresponding to the results in Table [ (row 1, pre- training on PM6), Table 2 (rows and 3, pre-training on PM6 together with force labels), and Table 3 (row 4, pre-training on PM6 then finetuning). Values lower than the 0.05 significance threshold are shown in bold.
Test Set
PCQ
QM9
Generated by
Denoising
DDIM
Denoising
DDIM
Struct. Stat.
Mcan
Min
Mean
Min
Mean
Min
Mean
Min
PM6 7.8xa0- 5.6xao
3.8xa0 8.9xa0 4.5x10
8.2xa0
1.1x10
7.4xa0
PM6 w/ SPICE force 2.4x10-5 1.0x10-5
6.0x10-4 3.1x 9.9x10-
4.5x10-1
2.3x10-2 -2 1.5x10-
PM6 subset force 4.5xa0 2.4xau 6 2.9xau
2.4x1u 2 6.5xa0 1.2x10
5. 2 7.5xav
PM6 finetuning 4.2x10-3 2.4x10-1 2.4x10-a
1.7x10-1
1.2x10-7
2.5x10-" 1.0x10-1
8.5x10-a
C.6 Evaluation on Dissimilar Molecules
To further consolidate the effectiveness of our method, we give a closer look into the influence of the similarity of the test dataset to the training dataset. Recall that we have excluded identical molecules appearing in the training dataset from the test dataset, but there remains the possibility that the test dataset may still contain similar molecules as those in the training dataset. For this, we first investigate the Tanimoto similarity [7 between the training and test datasets. This similarity can be computed from the molecular graphs of two molecules, while considering structural similarity between the two molecules. We plot in Fig. [22 where each column visualizes the count of PCQ test molecules that have a certain portion of similar (Tanimoto similarity > 0.7) molecules in the PM6 training dataset. We can observe that most (almost all) of the test molecules only has less than x 10 (2.5 x 10-7) similar molecules in PM6. This indicates that excluding identical molecules appearing in the training dataset from the test dataset already makes the test dataset sufficiently dissimilar to the training dataset, hence the presented results are valid prediction evaluations that not close to "memorizing the training molecules".
Similarity statistics of PCO (test) molecules to PM6 (training) melecules
175
The number molecule
220
20-
portion the maleculas against test PME dataset
Se-6
Figure C.2: Histogram showing the distribution of the portion of similar (Tanimoto similarity > 0.7) molecules in PM6 (the training dataset) over the 200 PCQ test molecules. Note that the x-axis is scaled by 1 x 10
For completely sanitized evaluation on dissimilar molecules, we provide the results evaluated on PCQ test molecules that do not have any similar (Tanimoto similarity > 0.7) molecules in the PM6 (training) dataset. Such molecules make up 49% of the total 200 PCQ test molecules. Table C. shows the results in both RMSD and Coverage using the denoising generation method, corresponding to the settings Tables and and C.3 (row block 1) (pre-training on PM6), and Tables and (row blocks 2 and 3) (pre-training on PM6 together with force labels). The results confirmed the advantage
22

---

of consistency learning on dissimilar molecules that cannot rely on memorizing the training dataset, which further consolidates the conclusion.
Table C.11: Test RMSD (Á; lower is better) and coverage (higher is better) dissimilar molecules for structure prediction pre-trained by multi-task learning and consisteney learning on the PM6 dataset, and together with additional SPICE force data PM6 subset force data.
RMSD (Å)
Cov
Training Set
Method
Mean
Min
Mean
Median
PM6
Multi-Task 1.175 0.642
0.613
0.675
(c.f. Tables or
(row block 1))
Consistency 1.135
0.625 0.644
0.745
PM6 w/ SPICE force (row block 2))
Multi-Task 1.136 0.609
0.639
0.735
(c.f. Tables 2 (row block 1)
C.3
Consistency 1.121
0.579
0.672
0.790
PM6 subset force
Multi-Task 1.174 0.653
0.612
0.660
(cff. Tables 2
(row block 2) or C.3 (row block 3))
Consistency 1.099
0.616
0.697
0.830
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have accurately reflect the paper's contributions and scope in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations of the work in the Sec.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]

---

Justification:
We have provided complete proof for the assumptions in Sec. 3 and Ap-
pendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have disclosed the needed information for reproducibility in Ap- pendix B
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
25

---

reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Releasing data and code requires an internal asset releasing review process in our organization. We have started the process, but cannot guarantee availability during the review period.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
wublic//giide//CoeeSuutii http://ooooebabmssssoooooooooooo Psllcll for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//ipps.cc/pulli//uiiiee htt///CooSSSeiiaaeeeeeee ssisssssicyyyyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:[Yoe
Justification:
We have included all details of the model architecture, data processing, and hyperparameter settings in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We have report the standard deviations and statistical significance of the main results in Appendix C.5
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
26

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The required computational resources are provided in Appendix B.3
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neurips. httpr//spspcpppbpiieccccccccaeee
Answer: [Yes]
Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: This work develops techniques for better solving molecular science tasks using machine learning. The techniques could potentially foster the development related industries, e.g., drug discovery and material design. Possible caveats include inaccurate generated structures or predicted properties which may lead to failure in downstream research.
27

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI...
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper does not release data or models that have high risk for misuse. Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Our paper has properly credited and cited the code and data used in this paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
28

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com/ com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's ' creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
The new asset in this paper is the in-house generated force labels on a subset of PM6 molecules. Releasing the new asset requires internal releasing review process in our organization. We cannot guarantee releasing the asset during the review period.
Justification:
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/codeemmodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing experiments research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
29

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "R5xCZIuRyf", "reviewer": "Reviewer_B2Rx", "summary": "To handle heterogeneity in molecular data and different computational costs, authors propose to exploit molecular tasks that have scientific laws connecting them. Their results show that the more accurate energy data can improve the accuracy of structure prediction.\nAuthors highlight that in contrast to conventional machine learning tasks defined by data, scientific tasks originate in fundamental scientific laws. Scientific laws impose explicit constraints between tasks defining the scientific consistency between model predictions on these tasks. By enforcing such consistency, model predictions for different tasks are connected and can explicitly share the information in the data of one task to the prediction for other tasks, hence bridging data heterogeneity.\nAuthors demonstrate the practical value of the scientific consistency between energy prediction and equilibrium structure prediction. Authors demonstrate the advantages of incorporating the proposed consistency losses into multi-task learning. \nAuthors design consistency losses to enforce scientific laws between inter-atomic potential energy prediction and equilibrium structure prediction.", "questions": "As authors point out the significance of improvement in this work is still limited by the abundancy of involved data. Can authors provide a measure to quantify the abundancy of dada for molecular equilibrium structure prediction?", "limitations": "As authors point out this work is limited to the consistency between energy and molecular structure prediction, while more consistency laws can be considered in molecular science and the significance of improvement in this work is still limited by the abundancy of involved data.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "To handle heterogeneity in molecular data and different computational costs, authors propose to exploit molecular tasks that have scientific laws connecting them. Their results show that the more accurate energy data can improve the accuracy of structure prediction.\nAuthors highlight that in contrast to conventional machine learning tasks defined by data, scientific tasks originate in fundamental scientific laws. Scientific laws impose explicit constraints between tasks defining the scientific consistency between model predictions on these tasks. By enforcing such consistency, model predictions for different tasks are connected and can explicitly share the information in the data of one task to the prediction for other tasks, hence bridging data heterogeneity.\nAuthors demonstrate the practical value of the scientific consistency between energy prediction and equilibrium structure prediction. Authors demonstrate the advantages of incorporating the proposed consistency losses into multi-task learning. \nAuthors design consistency losses to enforce scientific laws between inter-atomic potential energy prediction and equilibrium structure prediction.", "weaknesses": "As authors point out this work is limited to the consistency between energy and molecular structure prediction, while more consistency laws can be considered in molecular science and the significance of improvement in this work is still limited by the abundancy of involved data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your effort in evaluating our paper. Your informative feedback are greatly appreciated.\n\n## Consistency Beyond Energy and Structure\n\nThank you for the opportunity to elaborate on this point. Please refer to the global rebuttal (item 1).\n\n## Abundancy of involved data\n\nIn consistency training, energy landscape matters: the energy model needs to rank different structures of each molecule in optimality consistency, and provide gradient (slope) in score consistency. Although the PM6 dataset [31] is already the largest public DFT-labeled molecular dataset (up to our best knowledge), it provides energy on only one structure per molecule, which may be insufficient for learning the landscape. So we tried leveraging more data (Sec. 3.4, 4.3): we generated force (negative energy gradient) labels on a subset of PM6 molecules, and leveraged the force labels of multiple structures per molecule in the SPICE dataset [11]. They indeed improve energy landscape (Table 7: better energy prediction results) and lead to more accurate structure prediction with consistency (Table 2 vs. Table 1), but data abundancy is still limited: the generated data do not cover multiple structures, and there is a mismatch in the DFT settings of SPICE labels. We did not find public datasets providing energy or force labels under the same DFT setting as the PM6 dataset on multiple structures per molecule, and we did not have sufficient resource to generate such data. We expect more significant benefits of consistency training if such datasets are available.\n\nWe'll include these discussions in our next version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PWQk6dB9DT", "reviewer": "Reviewer_fhSX", "summary": "&nbsp;\n\nThe authors consider the multitask learning setting for molecular structure and energy prediction where the fidelity of the labels differs between tasks [1]. The authors note that they can leverage the relationship between high fidelity labels (energy) and low fidelity labels (structure) to design loss functions, a) the optimality consistency loss and b) the score consistency loss that operate as inductive biases in the multitask learning setting. Given that the proposed method is straightforward and appears to work well empirically, I consider the paper borderline at the moment since the code has not been supplied to reproduce the experimental results. I will revise my score if the code can be provided to ensure the reproducibility of the reported results.\n\n&nbsp;", "questions": "&nbsp;\n\n1. In Section C.1, the authors state that Tables 4 and 5 contain the std of the predictions, however, the caption states that the table contains test coverage values. The caption of Table 4 also indicates that lower is better. Is this the correct direction?\n\n2. Could the authors produce standard errors for the results presented in Table 7 of the appendix?\n\n&nbsp;", "limitations": "&nbsp;\n\n1. One limitation of the work is the realm of applicability of the method since it applies specifically to molecular structure and energy prediction. \n\n2. At the current point the authors have yet to provide code to reproduce the results.\n\n\n&nbsp;", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "&nbsp;\n\nThe method introduced by the authors is straightforward and appears to work well empirically. \n\n&nbsp;", "weaknesses": "&nbsp;\n\n__MAJOR POINTS__\n\n&nbsp;\n\n1. The title is not fully descriptive of the authors' contribution. I would recommend the authors revise the title to be more descriptive of the paper content. Specifically, the contribution does not apply to molecular science as a whole, but rather to structure and energy prediction. From the title alone, it is also unclear what the meaning of scientific consistency is, given that this appears to be a neologism coined by the authors.\n\n&nbsp;\n\n__MINOR POINTS__\n\n&nbsp;\n\n1. It would be great if the references appeared in numbered order.\n\n2. Line 2, typo, \"at scale\".\n\n3. There are some discrepancies in capitalization in the references e.g. 3d and 3D.\n\n4. Line 38, typo, \"hundreds of times more costly\".\n\n5. The original paper on molecule generation with VAEs [2] should probably be cited at some point in the text.\n\n6. Given that the standard deviations are crucial to establishing the improvement afforded by the consistency loss approach, I would recommend providing them in the main text instead of in the appendix. It may also be possible to perform a paired t-test to assess the statistical significance of the results e.g. https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.ttest_rel.html\n\n7. The axis labels of Figure 2 are somewhat confusing. The y-axis label is the predicted energy on the model-generated structure whereas the x-axis label is the energy of the equilibrium structure in the dataset. Would it be more appropriate to label the x-axis ground truth energy? The axis labels are confusing because the \"predictor\" for both energies is different.\n\n8. In Figure 2, why is there a systematic deviation in the predicted energy of R_pred relative to R_eq?\n\n9. 1. ADAM, reference 22, was published at ICLR 2015.\n\n10. The details of the model should be provided in the main paper rather than the appendix.\n\n11. In Section 4.4 the authors refer to the models from Section 4.2. It is not clear what these models are or how they were pre-trained. I would suggest including this in the main text.\n\n12. Line 375, \"by the abundance of the data involved\".\n\n13. It would be worth mentioning [3] as a reference source for multitask learning at some point in the text.\n\n\n&nbsp;\n\n__REFERENCES__\n\n&nbsp;\n\n[1] Peherstorfer, B., Willcox, K. and Gunzburger, M., 2018. Survey of multifidelity methods in uncertainty propagation, inference, and optimization. Siam Review, 60(3), pp.550-591.\n\n[2] Gómez-Bombarelli, R., Wei, J.N., Duvenaud, D., Hernández-Lobato, J.M., Sánchez-Lengeling, B., Sheberla, D., Aguilera-Iparraguirre, J., Hirzel, T.D., Adams, R.P. and Aspuru-Guzik, A., 2018. Automatic chemical design using a data-driven continuous representation of molecules. ACS Central Science, 4(2), pp.268-276.\n\n[3] Caruana, R., 1997. Multitask learning. Machine learning, 28, pp.41-75.\n\n&nbsp;", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your dedicated effort in evaluating our paper! We can feel your careful read, and are grateful for your feedback and suggestions.\n\n## About the title\n\nThank you for your informative feedback. Please check item 2 of the global rebuttal.\n\n## The minor points\n\n1. Thank you for the professional suggestion. We have revised the references in numbered order.\n\n2. Thank you for pointing it out. We have revised it.\n\n3. Thank you for your careful read! We have revised the capitalization.\n\n4. Thank you. We have revised it.\n\n5. Thank you for your suggestion. We have cited the paper in Introduction.\n\n6. Thank you for your suggestion. We have managed to insert standard deviation (in a smaller font) in the main tables.\n\n    Following your instruction, we also conducted a t-test to all results in the main tables (1-3). Please check Table R1 in the pdf file from the global rebuttal. There are only 4 out of 32 cases that have a p-value > 0.05, and those are all under the \"Min\" case where multi-task learning may also have a chance to hit the target structure. The means are indeed close in those 4 cases. We have included these results in the revision.\n\n7. Thank you for your feedback. We'd like to clarify that for the x-axis of Figure 2, the structures are from the PCQ dataset (vs. predicted by the model for the y-axis), but the energies of the structures are predicted by the energy prediction model (this part is the same as the y-axis).\nThe purpose of this figure is to verify that the improved structure prediction accuracy by consistency training is indeed due to that the predicted structures achieve a lower DFT-level energy (instead of e.g. a better fit to the PM6 structures), hence it is the consistency learning that makes the model predict structures closer to DFT-level equilibrium structures.\nWe did not compare the ground-truth energies since we do not have DFT energy labels for the model-predicted structures.\nThe energy prediction model is trained on DFT-level energy labels on PM6 structures, hence can be used as a surrogate to evaluate DFT-level energy.\n\n8. Alignment of predicted energy of R_pred and R_eq (i.e., points are on the diagonal) means the model predicts the same as R_eq, which is not what we expected, since the model does not see any DFT-level equilibrium structures in training (note only the less accurate, PM6-level equilibrium structures are available for training, while R_eq are DFT-level equilibrium structures from the PCQ dataset as the ground truth for evaluation).\nInstead, the point of Figure 2 is that consistency learning (still no DFT-level equilibrium structure data!) drives the model to predict structures closer to R_eq, which is verified in Figure 2 as the orange points (consistency) lie closer to the diagonal than the blue points (multi-task).\n\n9. Thank you. We have revised the citation.\n\n10. Thank you for your suggestion. We have added model details in Sec. 4.1.\n\n11. Thank you for your feedback. We meant that the pre-trained models for the finetuning experiments in Sec. 4.4 are those trained under the settings described in Sec. 4.2 (the models corresponding to the results in Table 1).\nThe detailed training settings for Sec. 4.2 are provided in Appendix B.3, which we have moved to Sec. 4.1 in the revision.\n\n12. Thank you. We have revised it.\n\n13. Thank you for your suggestion. We have added the reference in the introduction of multitask learning.\n\n## About the questions\n\n1. Thank you for your careful read! We apologize for the confusion. We have revised Sec. C.1 such that Tables 4 and 5 contain the coverage results, and revised the caption such that it is the higher the better. The corresponding std results are presented in Tables 9 and 10.\n\n2. As the scale of energy depends on the size of molecule (energy is an extensive quantity), which also affects the standard deviation of the error, we provide a box plot (Figure R1 in the pdf file from the global rebuttal) for the energy MAE in each case corresponding to Table 7, which provides more details about the error distributions.\nComparing results in each column, we see that training with consistency losses does not hurt energy prediction.\nComparing results in each row, we find including force data in training leads to a more accurate energy model, which explains that consistency learning performs better in this case (comparing Table 2 to Table 1).\nThis aligns with the observations from Table 7.\n\n## About the limitations\n\n1. Please refer to the global rebuttal (item 1) for broader applicability.\n\n2. As we mentioned in the Paper Checklist (item 5), releasing code requires an internal asset release review process within our organization. We have started the process, but cannot guarantee availability during the review period. We have been pushing the process, but unfortunately it is still not completed. We have provided implementation details in Appendix B, and will provide more to guarantee reproducibility. We will push further to guarantee the release before the time of publication.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tqJg7ZApyZ", "reviewer": "Reviewer_p18M", "summary": "The paper proposed a scientific consistency based improvement of molecule structure and energy prediction task. Upon the overall diffusion process for structure prediction, the authors incorporated energy-guided losses, which enables direct information exchange between the two tasks. Based on the two benchmark datasets, the proposed consistency losses improved the naive multi-task learning framework.", "questions": "- Do the authors think the proposed method would be useful for other tasks besides molecular structure and energy? If not, how about limiting the scope of the paper to molecular structure?\n- How was the number 200 determined for the test molecules?\n- What do the authors think the results would be if molecules from PCQ and QM9 with a similarity above a certain threshold (e.g., Tanimoto similarity 0.7) were additionally removed and the experiments were conducted?", "limitations": "- The effectiveness of the proposed method is uncertain as there are no experiments applying it to the latest diffusion-based methods.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- Well-written in Methods\n- It is notable to improve prediction performance based on the scientific correlation between molecular structure and energy without additional data", "weaknesses": "- The title and abstract of the paper are too general in relation to the specific task performed.\n- The proposed method's utility is limited as it can only be applied to tasks with a strong correlation, such as molecular structure and energy.\n- The paper specifies 200 test molecules, but it does not explain the criteria for their selection or why the number 200 was chosen.\n- In the experimental process, PM6 was used for pretraining, and identical molecules from PCQ and QM9 were removed; however, there is no mention of removing structurally similar molecules.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your devoted effort in evaluating our paper! We appreciate your informative feedback and solid suggestions.\n\n## About the title and abstract\n\nThank you for your feedback. Please refer to the global rebuttal (item 2).\n\n## Utility of the proposed method\n\nThank you for the opportunity to elaborate on this point.\nPlease refer to the global rebuttal (item 1).\n\n## About test molecules\n\nThe 200 test molecules are uniformly randomly selected from the PCQ (or QM9) dataset to guarantee that they are in-distribution with the whole PCQ (or QM9) dataset, while excluding molecules that also appear in the PM6 (training) dataset.\nThis setting, including the number 200, follows previous mainstream structure generation works, including ConfGF [39], GeoDiff [53], and DMCG [62].\n\n## Structurally dissimilar test molecules\nThank you for the insightful suggestion! It makes sense to exclude structurally similar molecules from the test set. Nevertheless, we found our selected 200 test molecules are already sufficiently dissimilar from the PM6 (training) dataset.\nFor each test molecule, we take the portion of PM6 molecules that have a Tanimoto similarity larger than 0.7 with the test molecule as a measure of similarity to the PM6 dataset.\nFigure R2 in the pdf file from the global rebuttal shows the distribution of this measure. We can see that most (almost all) of the 200 molecules only has less than 1e-7 (2.5e-7) similar molecules in PM6. This indicates that the presented results already show the performance on structurally dissimilar test molecules.\n\n## Applicability to latest diffusion-based methods\n\nWe'd like to mention that our consistency learning methods are model-agnostic and can be applied to any diffusion-based model.\nThis is because they all predict the score at each diffusion time step (predicting the noise or clean sample (denoising) are equivalent to predicting the score), so at small time steps, the score should align with the energy gradient (score consistency), and at larger time steps, the score can be used to predict the equilibrium structure through the denoising formulation (Eq. 4) (optimality consistency).\n\nDue to limited time and computational resources during the rebuttal period, we are unable to provide results on more diffusion-based methods, but we'd be happy to try if you could specify the method in your mind.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "277H7zg9sH", "reviewer": "Reviewer_3Qs1", "summary": "This paper presents a method to improve different tasks simultaneously by combining datasets of different fidelities and focuses, utilizing scientific laws that connect the tasks. Predicting the energy and equilibrium structure of molecules is used as an example. Two forms of consistency losses are developed based on (1) the rule that equilibrium structure minimizes energy, and (2) the relation between the probability distribution of structure and energy. The benefits of consistency losses in multi-task learning are demonstrated on quantum chemistry datasets.", "questions": "- This work essentially incorporates scientific laws by modifying loss functions. There are other works following a similar approach, e.g., PINNs, should they be discussed as related works?\n- In the experiments, the demonstrated benefits of consistency learning are mainly about accuracy. While for structure prediction via diffusion, efficiency (e.g., the diffusion steps required to attain a near-equilibrium structure) is also important. Could you comment on the efficiency?\n- Training on 8× Nvidia V100 GPU for a week is a considerable computational cost. How does consistency loss affect computational cost and model convergence?", "limitations": "Discussed in Sec. 5.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Developing consistency loss functions is an elegant approach to incorporating scientific laws into machine learning.\n- The proposed method does not require additional data to improve performance.", "weaknesses": "- The demonstration of the method is limited to one scenario (energy + equilibrium structure), where the formulation of consistency losses is ad hoc. The broader applicability is thus questionable.\n- The review of related works is not comprehensive (see Questions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your dedicated effort in reviewing our paper! We deeply appreciate your acknowledgement of our contributions, as well as informative feedback and suggestions.\n \n## Broader applicability\n\nThank you for the opportunity to elaborate on this point.\nPlease refer to the global rebuttal (item 1).\n\n## Related work about PINN\nThank you for the suggestion! We have included a discussion on related works on PINN as follows:\n\nPhysics-informed neural networks (PINNs) [R1] are another example of incorporating physical laws in neural network training.\nThe principle of PINNs involves representing the unknown target function with a neural network and optimizing it using a loss function derived from a system of partial differential equations (PDEs), such as the variational form of the PDE [R2]. This approach has shown promise in solving PDEs across various applications, including higher-dimensional problems [R3, R4]. PINNs are also applied to solve inverse problems by optimizing the parameters of PDEs [R5, R6]. More relevant to our case, PINNs can also be used to tackle data heterogeneity. For instance, HFM [R7] uses the Navier-Stokes equation to establish a connection between the concentration of contrast agents in the bloodstream, and the dynamic quantities of blood flow such as velocity and pressure. It is used to derive the latter from the inferred results of the former from concentration variations in medical imaging. Similarly, PhySR [R8] utilizes the underlying physical laws in the system, such as those governing the 2D Rayleigh-Bénard convection system, to reconstruct high-resolution results from low-resolution data.\n\nWhile sharing the same spirit, our work has some technical differences. Physical laws in molecular properties often do not come in the form of PDEs (unless solving the Schrödinger equation in the original form) but as algebraic or statistical equations. There is no need of grids, but involves multiple quantities, e.g., energy and structure in our case, so the laws are used to bridge different prediction models rather than learning one model.\nMoreover, neural networks can be treated as black-box models in PINNs, while in our case, in-depth analyses are needed to connect model output to the desired quantities: Sec. 3.2 analyzed how to produce a rough structure prediction from the output of the denoising model for optimality consistency, and Sec. 3.3 analyzed how to compute the score using the denoising model for score consistency.\n\n[R1] Raissi M, Perdikaris P, Karniadakis G E. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations[J]. Journal of Computational Physics, 2019, 378: 686-707.\n\n[R2] Yu B. The deep Ritz method: a deep learning-based numerical algorithm for solving variational problems[J]. Communications in Mathematics and Statistics, 2018, 6(1): 1-12.\n\n[R3] Lu L, Meng X, Mao Z, et al. DeepXDE: A deep learning library for solving differential equations[J]. SIAM review, 2021, 63(1): 208-228.\n\n[R4] Han J, Jentzen A, E W. Solving high-dimensional partial differential equations using deep learning[J]. Proceedings of the National Academy of Sciences, 2018, 115(34): 8505-8510.\n\n[R5] Lu L, Pestourie R, Yao W, et al. Physics-informed neural networks with hard constraints for inverse design[J]. SIAM Journal on Scientific Computing, 2021, 43(6): B1105-B1132.\n\n[R6] Yu J, Lu L, Meng X, et al. Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems[J]. Computer Methods in Applied Mechanics and Engineering, 2022, 393: 114823.\n\n[R7] Raissi M, Yazdani A, Karniadakis G E. Hidden fluid mechanics: Learning velocity and pressure fields from flow visualizations[J]. Science, 2020, 367(6481): 1026-1030.\n\n[R8] Ren P, Rao C, Liu Y, et al. PhySR: Physics-informed deep super-resolution for spatiotemporal data[J]. Journal of Computational Physics, 2023, 492: 112438.\n\n## Efficiency for structure prediction via diffusion\n\nThank you for bringing out the efficiency consideration.\nOur methods would introduce some cost in the training stage for going beyond the level of accuracy of training data, but they do not alter the way to generate structures in the inference stage, hence not affecting the efficiency of structure prediction.\nTo improve the efficiency, we can directly leverage prevailing techniques to reduce the number of diffusion steps, e.g., DDIM [43], Heun's method [R9], and DPM-Solver [R10].\n\n[R9] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas. Gotta go fast when generating data with score-based models. CoRR, abs/2105.14080, 2021.\n\n[R10] Lu C, Zhou Y, Bao F, Chen J, Li C, Zhu J. DPM-Solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems. 2022 Dec 6;35:5775-87.\n\n## Computational cost\n\nThank you for your careful read! Consistency training indeed introduces more training effort, but we made an implementation design as described in Appendix B.3, so that the additional cost is still manageable. The main reason that we need such a computational cost is that we used a relatively large model (around 130M parameters) to sufficiently capture the information on PM6, which is the largest public molecular dataset with DFT-level energy label.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their careful read, informative feedback, and sincere suggestions. We provide responses to two common questions in this global rebuttal.

1. Applicability beyond energy and structure (for Reviewers B2Rx, fhSX, p18M, 3Qs1)

    Within the presented content, beyond energy and equilibrium structure, we'd like to mention that we also showed in Sec. 3.4, 4.3 that leveraging _force_ labels on _off-equilibrium structures_ in the proposed consistency learning methods can further improve the accuracy of equilibrium structure prediction.

    As we mentioned in Sec. 5, the proposed methods can be used for connecting energy and thermodynamic distribution (beyond predicting a single equilibrium structure but generating a thermodynamic ensemble of structures), since the score consistency still holds, and the optimality consistency can be adjusted to match model-derived structure statistics to macroscopic observations. This consistency training can potentially improve the accuracy of distribution beyond data-based training, as data samples are often only available from unconverged simulations (hence biased). We treated it as future work due to more tricky and elaborate training settings, evaluation protocols and benchmarks, and chose to develop and demonstrate the methods in the structure prediction scenario, which has already taken the capacity of a conference paper.

    We also mentioned broader possibilities following the same idea. Molecular properties (e.g., energy) are derived from electronic structure of the molecule following physical laws, and coarse-grained structure distribution is a partial integration of fine-grained structure distribution. We can design consistency training losses according to such laws to connect these tasks, tackling data heterogeneity in these cases.

2. About the title (for Reviewers fhSX, p18M)

    To better describe the content, we plan to revise the title to "Tackling Data Heterogeneity in Molecular Energy and Structure by Enforcing Physical Consistency". We have also revised the abstract to focus on molecular energy and structure and their consistency.

    * The new title highlights the major scenario ("Energy and Structure") for the proposed consistency techniques. We'd like to mention that we also involved force and off-equilibrium structure data (Sec. 3.4, 4.3), extending the relevance beyond energy and (equilibrium) structure.

    * We hope that "Enforcing Physical Consistency" could convey the sense that there is a _physical_ law connecting the two quantities, which defines a _consistency_ between the two quantities and can be _enforced_ by a loss. It is challenging to convey the precise meaning in the title, and we'd be more than glad if you have any suggestions.

Meta-review: The paper introduces the novel approach of leveraging scientific laws to enforce scientific consistency between tasks in multi-task learning. The current work applies this approach to energy prediction and structure prediction - exploiting the the relationship between structure and energy - and demonstrates improvement over vanilla multi-task learning. While the presented application is limited to energy prediction and structure prediction, the idea has broad applicability to other fields especially for machine learning in scientific domains. Concerns regarding regarding presentation and generalization to new chemical space are also sufficiently addressed.
Decision: Accept (poster)