On Accelerating Diffusion-based Molecular Conformation Generation in SE(3)-invariant Space

1. The paper is not well-motivated. The proposed sampling method is specifically designed for the task of molecular conformation generation in SE(3)-invariant space. However, the SOTA method of conformation generation like Torsional Diffusion which operates on the hypertorus achieves much better performance than GeoDIFF which operates on Euclidean space. Moreover, the sampling process of Torsional Diffusion only needs 20 steps. Therefore, focusing on accelerating the sampling process of GeoDIFF/SDDiff seems meaningless.

We do not agree with that accelerating the sampling process of GeoDIFF/SDDiff seems meaningless. Different from Torsional Diffusion, GeoDIFF and SDDiff are purely data-driven without any prior knowledge and preprocess(like RDkit in Torsinal Diffusion) and their distance-based modeling is effortless, indicating that they are not limited in the conformation generation tasks. As for our acceleration method, from the derivation of our method is also applicable to any coordinates generation tasks that take SE(3)-invariance into consideration. For example, our method is also applicable to sparse cloud point generation tasks. However, torsional diffusion is inapplicable in this case since it cannot first generate a local structure. In a nutshell, although focusing on molecular conformation generation, our contribution in deed lies in that we rethink the modeling occurring in the general SE(3)-based methods and propose a scheme for accelerating general diffusion-based generation in SE(3)-invariant space, and we plan to generalize it in future works.

2. The contribution of the introduced two minor modifications is limited. The operations just introduced two coefficients into existing sampling methods[1] to compensate for the poor performance on the conformation generation task by using [1] directly. From the experiments, the proposed methods don't get surprising results.

Directly using the DPM solver is inapplicable rather than just poor performance, we have demonstrated in Sec 4 around Eq.4 "...cannot produce substantial conformations. ". We found that the results are chaos. Ignoring any one of our proposed two modifications would lead to zero performance w.r.t COV and nan w.r.t. MAT. Our proposed modifications are minor but non-trivial. Such a modeling issue has existed since ConfGF (2021) and remains unsolved. The mathematical proofs behind such two simple modifications are intricate, and our experiments strongly proved that our method gains a roughly 50-150x speedup on comparable performance.

3. The paper is over-clams. The paper doesn’t provide the theoretical results about the statement “We analyze current modeling methods in SE(3) (Shi et al., 2021; Xu et al., 2022; Zhou et al., 2023) and theoretically pose crucial mistakes shared in these methods, which inevitably bring about the failure of acceleration” and don’t give the reasons why the modification solve the theoretical problems. If the proposed modifications can solve the problems, at least it can be directly used to improve the existing sampling methods applied in (Shi et al., 2021; Xu et al., 2022; Zhou et al., 2023).

The current modeling of the transformation from distance score to conformation score is heuristic. In Sec 4.1 we proposed revised modeling via differential geometry and proposed modification 1. Then we consider remedying the approximation error of the conformation score by an additional scaling factor in Sec 4.2.

We do not understand what the reviewer means by ``at least it can be directly used to improve the existing sampling methods applied in (Shi et al., 2021; Xu et al., 2022; Zhou et al., 2023)''. We have directly used our method to accelerate the existing sampling methods applied in GeoDiff and SDDiff (Xu et al., 2022; Zhou et al., 2023) by two slight modifications, and the results are shown in experiments. In the experimental section, we showed that our method indeed improves the existing sampling methods since originally, using the traditional method to accelerate these models would all lead to 0 performance.

4. The paper is not well-written. The paper hasn’t explained clearly why the introduced two modifications can accelerate the sampling process significantly. For modification one, the introduced $degree$ can make the sparse conformation into a fully connected conformation. Why connected conformation can outperform sparse one?

Sec 4.1 and Sec 4.2 have introduced the theory behind our modifications. The first modification is not designed to transform the sparse conformation into a fully connected conformation. Instead, we argued that the previous computation of the conformation score is wrong, i.e., the conformation score cannot be directly computed by the chain rule. We proposed that when computing the conformation score, one of the sum operators should be replaced by a mean operator.

Thanks for your reply and explanation. I'm maintaining my score because I don't feel that the paper's contribution is currently sufficient to be accepted by ICLR.

Questions

1. On FrameDiff (Yim et al., 2023) in Section 2 and MAT-R in Appendix G.

Thanks for your pointing out it. We have revised it.

2. On the top of page 5 , symbols α_ij and e_ij are introduced without explanation. Also, it would be nice to provide a brief explanation or reference to the generalized multidimensional scaling.

These notations are commonly used in linear algebra. $e_{ij}$ is a square matrix such that its entry in the i-th row and j-th column equals 1 and the other entries equal to 0. $\alpha_{ij}$ denotes the magnitude of change in edge that connects node $i$ and node $j$. The derivation for the Eq 8-9 can be found in the Appendix D.2. Generalized multidimensional scaling extends metric multidimensional scaling by accommodating a non-Euclidean target space. GMDS facilitates the discovery of the most accurate embedding of one surface into another when dissimilarities represent distances on a surface and the target space is an alternative surface.

3. In the sentence just before eq 11 , the meaning of the phrase "the pairwise distance matrix is sparse" is not clear.

This means that in practice, we usually consider a partially connected graph, hence, we only consider a proportion of entries in the full adjacent matrix.

4. It might be better to reformulate the optimization problem in eq 8 by replacing the summation over all $i<j$ pairs with connected $(i, j)$ pairs.

Thanks for your suggestions, we will consider modifying this.

Weakness

1. About manifold construct. In this paper, we do not need to explicitly construct the manifold of the adjacent matrix and the SE(3)-invariant coordinates manifold. If one is interested in such a construction, one can easily construct the manifold of the adjacent matrix manifold by considering the spectral theorem. More specifically, for any feasible $D$, there corresponds to a Gram matrix $M \in \mathbb{R}^{n \times n}$ by $M_{ij} = \frac{1}{2} (D_{1j} + D_{i1} - D_{ij})$ and vice versa $D_{ij} = M_{ii} + M_{jj} - 2 M_{ij}$ and the matrix $M$ has a specific structure

$

M=\left(\begin{array}{ll}
0 & \mathbf{0}^{\top} \\
\mathbf{0} & L
\end{array}\right)

$

with $L \in \mathbb{R}^{(n-1) \times (n-1)}$ being symmetric and positive semi-definite and $\operatorname{rank}(L)=3$. Then, we define the mapping from $L$ to $D$ by $h$. One can define the manifold of the adjacent matrix to be $M = \{h(XX^T) \mid X \in \mathbb{R}^{(n-1) \times 3}\}$. To construt a SE(3)-invariant coordinate manifold, we pick any $\mathcal{C} \in \mathbb{R}^{n \times 3}$, then define $N_{\mathcal{C}} = \{ \hat{\mathcal{C}} = [\hat{\mathbf{x}}_1, \ldots, \hat{\mathbf{x}}_n] \mid \hat{\mathbf{x}}_i = \mathbf{x}_i + \sum_{j=1}^n \alpha_{ij} \frac{\mathbf{x}_i - \mathbf{x}_j}{\| \mathbf{x}_i - \mathbf{x}_j \|}, \alpha_{ij} = \alpha_{ji}, \alpha_{ii} = 0, \alpha_{ij} \in \mathbb{R} \}$. In this paper, we do not need to explicitly define the above two manifolds, hence, we do not consider such constructions.

2. About manifold assumption. We assume that the tangent space of the adjacent matrix manifold is $\mathbb{R}^{n \times n}$. Such an assumption is certainly wrong. However, since we did not apply any constraints on the model's output. When our model predicts the tangent vector of the tangent space of the distance manifold, the output may not necessarily lay in the tangent space. This pushes us to consider the tangent space as the $\mathbb{R}^{n \times n}$. The above argument for why we assume the tangent space is $\mathbb{R}^{n \times n}$ can also be found in Appendix D.1.

3. About factor $\frac{1}{2(n-1)}$. This factor is composed of two parts. One is $2$, which comes from the bipartition of the distance change, as illustrated in Fig. 6 by $-\frac{\delta}{2} \boldsymbol{\lambda}_{uv}$ and $\frac{\delta}{2} \boldsymbol{\lambda}_{uv}$. The other is $n-1$, which is the node degree in a fully connected graph.

4. About error bound analysis. The error-bound analysis aims to show that when the node number increases, the approximation error is still bounded. However, using GeoDiff's modeling of the conformation score, the error bound explodes when $n$ becomes sufficiently large. For the naive choice of the original coordinates, we can show that by experimental results on the toy dataset, the proposed approximation has a smaller average error compared with using the original coordinates.

5. About prediction error. Since we have no enough information about the structure of the prediction error, we can only apply a common assumption that the prediction error follows a Gaussian distribution.

6. About abalation study. This is a fantastic point. We reduced the $\sigma$ sequence from 5000 to 100 equally spaced points utilizing linear interpolation, the initial and final values from the original set are included. We meticulously select the optimal step size $h$, for each model, consistently converging around a value around 1e-4.

We thank reviewer vTuX for giving feedback timely. Reviewer vTuX mainly has concerns about notation definitions, while we have explicitly defined them in the Appendix. But possibly due to our writing, one may find these notations confusing. We will re-organize our formula in the final version so that notations become more straightforward. We present our response as follows.

1. What are the meanings of $\mathrm{TM}, \mathrm{TN}, \pi_M, \pi_N$ in Figure 1?

$TM$ and $TN$ are tangent space of the manifold of $M$ and $N$, respectively. $\pi_M$ and $\pi_N$ are the projection function that receives a vector in the tangent space and project it into the manifold space. A detailed introduction of the differential geometry definitions and notations can be found in Appendix C.

2. Why $f=\pi_N \circ d \phi=\phi \circ \pi_M$? What processes are described by $\pi_N \circ d \phi$ and $\phi \circ \pi_M$?

$f$ is defined to be the mapping from $TM$ to $N$. There are two ways to transform from $TM$ to $N$. One is by $TM \rightarrow TN \rightarrow N$ and the other one is by $TM \rightarrow M \rightarrow N$. The former transformation is defined by the composition $\pi_N \circ d \phi$ and the latter one is defined by $\phi \circ \pi_M$. Hence, $f=\pi_N \circ d \phi=\phi \circ \pi_M$. Such a $f$ function is only defined to compute $d \phi$ when the rest three functions are known. Such a method is commonly used in linear algebra such as when determining the properties of the quotient space.

3. Why $\pi_M$ is chosen to be an identical mapping, $\pi_N$ is chosen to be a GMD?

$C_t$ should always stay in the manifold $N$. Mathematically speaking, our reversed ODE should be replaced with a projected dynamical system in the form of $\frac{\mathrm{d} C_t}{\mathrm{d} t}=\Pi_N(C_t, -\frac{1}{2} \frac{\mathrm{d} \sigma_t^2}{\mathrm{d} t} \nabla_{C_t} \log q_{0 t}\left(C_t \right))$, where $\Pi_N(x, v)=\lim_{\delta \rightarrow 0^{+}} \frac{\pi_N(x+\delta v)-x}{\delta}$. Comparing with the exact update form $\frac{\partial C_t}{\partial t}=-\frac{1}{2} \frac{\mathrm{d} \sigma_t^2}{\mathrm{~d} t} \nabla_{C_t} \log q_{0 t}\left(C_t \mid C_0\right)$, we can take $\pi_N$ to be the identical mapping.

$\pi_N$ is the projection from the tangent space to the manifold. By the definition of the projection, we must project a matrix to a feasible adjacent matrix while minimizing the distance under a certain metric, which is essentially the same as the problem of the GMD.

4. What does $e_{i j}$ mean in the formula of $\hat{d}$, and how the formula of $\pi_{M, \tilde{d}}(\hat{d})_{i j}$ and subsequently Equation (8) is derived?

$e_{ij}$ is a square matrix such that its entry in the i-th row and j-th column equals 1 and the other entries equal to 0. The derivation for the Eq 8-9 can be found in the Appendix D.2.

5. In Appendix D.2 (proof of theorem 1), what is $\lambda_{u v}$ ? Why does the second $\leq$ (from line $30 \mathrm{~b}-30 \mathrm{c}$ to line $30 d - 30 g$ ) hold?

$\boldsymbol{\lambda}_{u v}=\frac{\tilde{\mathbf{x}}_u-\tilde{\mathbf{x}}_v}{\left\|\tilde{\mathbf{x}}_u-\tilde{\mathbf{x}}_v\right\|}$, which can be found in the Appendix D.1. The second $\leq$ should be $=$. We apologize for this typo and have corrected this in the manuscript.

6. In Section 4.2, authors are recommended to clarify what is $\tilde{C}_{t+\lambda(s-t)}$ ?

After one update, we are expecting that the sample follows the distribution w.r.t $p_t$. However, since there is some inevitable error in the denoising process, we consider such noise acts like applying a forward diffusion process on $C_t$ so that the sample follows the distribution for some $p_{t+\lambda (s-t)}$, i.e., a time in the interval of $(t, s)$.

7. Also, it would be better to discuss how much error will be introduced by the approximation $k_{s_\theta}\left(d_s, s, t\right) \approx k_{s_\theta}\left(p_{\text {data }}\right)$ ?

This is a good question. Unfortunately, we do not know how to do this so far and we leave it as a future work.

8. It will make the experimental results stronger if authors do ablation study to verify that every single strategy can improve the performance.

Thanks for your suggestion, but it is not proper to conduct an ablation study. The main contribution of our work is to propose two necessary modifications for acceleration, which means that the absence of either one disables acceleration. Modification 1 (change ''add'' to ''mean'' on the transformation from distances to coordinates in Eq. 5) is a revision of the modeling to enable acceleration. In our experiment, the results are all from the models under Modification 1. Modification 2 (adds a multiplier ''scale'', $k_\mathbf{s}$ to the sampling process) is necessary for fast sampling. We found that without $k_{\mathbf{s}_{\boldsymbol{\theta}}}$, the fast sampling, i.e. sampling via DPM-Solver's proposed iteration in Eq. 4 will fail, and we conduct sampling via Langevin Dynamics (LD).

9. Also, authors are recommended to make an in-depth discussion about why the proposed method achieves poor performance in precision metric.

We have put the analysis on Precision metrics on Sec.5.4. ''This is due to the deteriorating output of the network.'' We observe that part of the output occasionally starts exploding from a certain sampling iteration, which leads to some unstructured conformers corresponding to one molecule. The recall metrics are not affected much because they only consider the coverage status of ground truth, while precision metrics calculate all generated conformers. This phenomenon is observed for both LD sampling and our fast sampling method, especially when the step size or scale $k_{\mathbf{s}_{\boldsymbol{\theta}}}$ is higher. More low-quality conformers are generated while sampling via the acceleration method, it is reasonable because fast sampling requires a larger step in each iteration, thus the precision metrics are worse. We think a stronger network or special design can solve the deteriorating output, but this is beyond the scope of this study. We'd like to conduct more stronger networks in further work.

10. The proposed strategies are mainly used for diffusion models that calculate coordinate scores from distance scores. However, current state-of-the-art molecular conformation generation method, Torsional Diffusion [1], focusing on generating torsion angles. Can GeoDiff or SDDiff compete with Torsional Diffusion in generation speed when the proposed strategies are used? Can the proposed strategies or similar strategies be also applied to Torsional Diffusion?

Current SOTA Torsional diffusion uses the RDkit to first generate the local structures and then apply the diffusion process on the torsional angles. Since they use the RDkit to generate the local structures, the main difficulties: considering the SE(3)-invariance is avoided. However, our method is a purely data-driven method, which cannot easily avoid the SE(3)-invariance. Hence, the computational speed is relatively slower than Torisonal diffusion's. Since the modeling manifold of the torsional diffusion is the torsion angle, which is greatly different from the distance manifold, hence, our method is inapplicable. However, for those generation tasks like cloud point generation, which takes SE(3)-invariance into consideration, our method applies to accelerate the reversed diffusion process.

11. About typos.

Thank you for bringing those typos to our attention. We are committed to making corrections.