On Diffusion Process in SE(3)-invariant Space

TL;DR

This study mathematically delineates the diffusion mechanism under SE(3)-invariance, via zooming into the interaction behavior between coordinates and the inter-point distance manifold through the lens of differential geometry, and proposes accurate and projection-free diffusion SDE and ODE accordingly.

Abstract

Sampling viable 3D structures (e.g., molecules and point clouds) with SE(3)-invariance using diffusion-based models proved promising in a variety of real-world applications, wherein SE(3)-invariant properties can be naturally characterized by the inter-point distance manifold. However, due to the non-trivial geometry, we still lack a comprehensive understanding of the diffusion mechanism within such SE(3)-invariant space. This study addresses this gap by mathematically delineating the diffusion mechanism under SE(3)-invariance, via zooming into the interaction behavior between coordinates and the inter-point distance manifold through the lens of differential geometry. Upon this analysis, we propose accurate and projection-free diffusion SDE and ODE accordingly. Such formulations enable enhancing the performance and the speed of generation pathways; meanwhile offering valuable insights into other systems incorporating SE(3)-invariance.

Figures (15)

• Figure 1: We utilized a sampling process with the off-the-shelf DPM-solver lu2022dpm and our method in reduced sampling iterations (1000 steps) on a GeoDiff xu2022geodiff model that was trained for 5000 steps. An example of a molecular conformer sampled via DPM-solver appears to be unreasonable (left), while our method produces a high-quality conformer (right).
• Figure 2: Mathematical framework for the reverse diffusion in the SE(3)-invariant space $\mathbb{R}^{n \times 3} / \operatorname{SE}(3)$. The grey rectangle denotes the coordinate embedding matrix manifold $L = \mathbb{R}^{n \times 3}$. The green ball denotes an SE(3)-invariant manifold $N$ of spectral coordinates (see Theorem \ref{['thm: mapping from adjacency matrix to coordinates']}) that is a submanifold of $L$. The blue ball denotes the adjacency matrix manifold $M$. Detailed definitions for the above three manifolds can be found in Appendix \ref{['sec: manifold definitions']}. Each coordinate in the spectral coordinate manifold $N$ has a unique adjacency matrix in the adjacency matrix manifold $M$, and vice versa. During each step of the denoising process, given a noised sample $\mathcal{C}_s \in L$, our model predicts the change of the adjacency matrix $d_s = \operatorname{adj} \left( \mathcal{C}_s \right)$, illustrated by the orange tangent vector $-F^{\mathcal{N}}\left( d_s, s \right)$ in the tangent space $T_{d_s}M$. Then we transform such a change of the adjacency matrix to the spectral coordinate shift by $\frac{-F^{\mathcal{N}}}{2(n-1)} \circledast \frac{\partial d_s}{\partial \mathcal{C}_s}$ (denoted by the red tangent vector in the tangent space of the spectral coordinate manifold). At the implementation level, we can further transform the tangent vector of spectral coordinates into the change of coordinates in the coordinate embedding space $L$ since this operation does not affect the values of the adjacency matrix.
• Figure 3: Visualization of sampling results obtained through LD and our proposed SDE/ODE with different sampling steps is presented. The human poses are labeled as "eating". Samples are obtained from the exactly same trained model but drawn from different sampling algorithms. We grid search the hyperparameters in each sampling algorithm for a fair comparison.
• Figure 4: Mappings between two manifolds.
• Figure 5: The illustration of how the choice of $s_u$ and $s_v$ affects the optimal values of $\varphi \circ \pi_{M, \tilde{d}}$. Given fixed $\delta, d, \hat{d}$, we set the optimal values as $\underset{\hat{\mathcal{C}} = [\mathbf{x}_1, \ldots, \mathbf{x}_n]}{\min} \sum_{i<j} \left( \left\| \mathbf{x}_i - \mathbf{x}_j \right\| - \hat{d}_{ij} \right)^2 / \delta^2$. In each figure, the matrix grids $(i, j)$ denote the entries of $( \operatorname{adj} _{ij}(\mathcal{C}) - \hat{d}_{ij} )^2 / \delta^2$. We want to minimize the sum of grid values. The left figure denotes the case where we take $\hat{\mathcal{C}} = \tilde{\mathcal{C}}$. In this case, only the entry $(4, 5)$ introduces the loss to the optimal value but the value of the entry $(4, 5)$ is comparably large. The middle figure denotes the case where we take $s_{u} = s_{v} = \delta / 2$. In this case, the value of the entry $(4, 5)$ is reduced to zero but other entries become non-zero, and the total loss increases. In the right figure, we take $s_{u} = \delta / (2 \operatorname{degree}_u)$ and $s_{v} = \delta / (2 \operatorname{degree}_v)$. At this time, although the loss introduced by the entry $(u, v)$ is not reduced to $0$, the value of the entry sum is reduced.

Theorems & Definitions (8)

• Theorem 1
• Remark 2: Necessity of alternating between coordinates and distances
• Definition 3
• Definition 4
• Definition 5
• Definition 6
• Definition 7
• Theorem 8: Adjacency matrix spectrum decomposition schoenberg1935remarksgower1982euclideanhoffmann2019generating