Molecular representation learning (MRL) has gained tremendous attention due to its critical role in learning from limited supervised data for applications like drug design. In most MRL methods, molecules are treated as 1D sequential tokens or 2D topology graphs, limiting their ability to incorporate 3D information for downstream tasks and, in particular, making it almost impossible for 3D geometry prediction/generation. In this paper, we propose a universal 3D MRL framework, called Uni-Mol, that significantly enlarges the representation ability and application scope of MRL schemes. Uni-Mol contains two pretrained models with the same SE(3) Transformer architecture: a molecular model pretrained by 209M molecular conformations; a pocket model pretrained by 3M candidate protein pocket data. Besides, Uni-Mol contains several finetuning strategies to apply the pretrained models to various downstream tasks. By properly incorporating 3D information, Uni-Mol outperforms SOTA in 14/15 molecular property prediction tasks. Moreover, Uni-Mol achieves superior performance in 3D spatial tasks, including protein-ligand binding pose prediction, molecular conformation generation, etc. The code, model, and data are made publicly available at https://github.com/dptech-corp/Uni-Mol.

Graph representation is a challenging and significant problem for many real-world applications. In this work, we propose a novel paradigm called “Gromov-Wasserstein Factorization” (GWF) to learn graph representations in a flexible and interpretable way. Given a set of graphs, whose correspondence between nodes is unknown and whose sizes can be different, our GWF model reconstructs each graph by a weighted combination of some “graph factors” under a pseudo-metric called Gromov-Wasserstein (GW) discrepancy. This model leads to a new nonlinear factorization mechanism of the graphs. The graph factors are shared by all the graphs, which represent the typical patterns of the graphs’ structures. The weights associated with each graph indicate the graph factors’ contributions to the graph's reconstruction, which lead to a permutation-invariant graph representation. We learn the graph factors of the GWF model and the weights of the graphs jointly by minimizing the overall reconstruction error. When learning the model, we reparametrize the graph factors and the weights to unconstrained model parameters and simplify the backpropagation of gradient with the help of the envelope theorem. For the GW discrepancy (the critical training step), we consider two algorithms to compute it, which correspond to the proximal point algorithm (PPA) and Bregman alternating direction method of multipliers (BADMM), respectively. Furthermore, we propose some extensions of the GWF model, including (i) combining with a graph neural network and learning graph representations in an auto-encoding manner, (ii) representing the graphs with node attributes, and (iii) working as a regularizer for semi-supervised graph classification. Experiments on various datasets demonstrate that our GWF model is comparable to the state-of-the-art methods. The graph representations derived by it perform well in graph clustering and classification tasks.

As a valid metric of metric-measure spaces, Gromov-Wasserstein (GW) distance has shown the potential for matching problems of structured data like point clouds and graphs. However, its application in practice is limited due to the high computational complexity. To overcome this challenge, we propose a novel importance sparsification method, called Spar-GW, to approximate GW distance efficiently. In particular, instead of considering a dense coupling matrix, our method leverages a simple but effective sampling strategy to construct a sparse coupling matrix and update it with few computations. The proposed Spar-GW method is applicable to the GW distance with arbitrary ground cost, and it reduces the complexity from $O(n^4)$ to $O(n^{2+\delta })$ for an arbitrary small $\delta >0$. Theoretically, the convergence and consistency of the proposed estimation for GW distance are established under mild regularity conditions. In addition, this method can be extended to approximate the variants of GW distance, including the entropic GW distance, the fused GW distance, and the unbalanced GW distance. Experiments show the superiority of our Spar-GW to state-of-the-art methods in both synthetic and real-world tasks. Supplementary materials for this article are available online.

Traditional multi-view learning methods often rely on two assumptions: ( $i$ ) the samples in different views are well-aligned, and ( $ii$ ) their representations obey the same distribution in a latent space. Unfortunately, these two assumptions may be questionable in practice, which limits the application of multi-view learning. In this work, we propose a differentiable hierarchical optimal transport (DHOT) method to mitigate the dependency of multi-view learning on these two assumptions. Given arbitrary two views of unaligned multi-view data, the DHOT method calculates the sliced Wasserstein distance between their latent distributions. Based on these sliced Wasserstein distances, the DHOT method further calculates the entropic optimal transport across different views and explicitly indicates the clustering structure of the views. Accordingly, the entropic optimal transport, together with the underlying sliced Wasserstein distances, leads to a hierarchical optimal transport distance defined for unaligned multi-view data, which works as the objective function of multi-view learning and leads to a bi-level optimization task. Moreover, our DHOT method treats the entropic optimal transport as a differentiable operator of model parameters. It considers the gradient of the entropic optimal transport in the backpropagation step and thus helps improve the descent direction for the model in the training phase. We demonstrate the superiority of our bi-level optimization strategy by comparing it to the traditional alternating optimization strategy. The DHOT method is applicable for both unsupervised and semi-supervised learning. Experimental results show that our DHOT method is at least comparable to state-of-the-art multi-view learning methods on both synthetic and real-world tasks, especially for challenging scenarios with unaligned multi-view data.

Real-world 3D structured data like point clouds and skeletons often can be represented as data in a 3D rotation group (denoted as $\mathbb {SO}(3)$ ). However, most existing neural networks are tailored for the data in the Euclidean space, which makes the 3D rotation data not closed under their algebraic operations and leads to sub-optimal performance in 3D-related learning tasks. To resolve the issues caused by the above mismatching between data and model, we propose a novel non-real neuron model called quaternion product unit (QPU) to represent data on 3D rotation groups. The proposed QPU leverages quaternion algebra and the law of the 3D rotation group, representing 3D rotation data as quaternions and merging them via a weighted chain of Hamilton products. We demonstrate that the QPU mathematically maintains the $\mathbb {SO}(3)$ structure of the 3D rotation data during the inference process and disentangles the 3D representations into “rotation-invariant” features and “rotation-equivariant” features, respectively. Moreover, we design a fast QPU to accelerate the computation of QPU. The fast QPU applies a tree-structured data indexing process, and accordingly, leverages the power of parallel computing, which reduces the computational complexity of QPU in a single thread from $\mathcal {O}(N)$ to $\mathcal {O}(\log N)$ . Taking the fast QPU as a basic module, we develop a series of quaternion neural networks (QNNs), including quaternion multi-layer perceptron (QMLP), quaternion message passing (QMP), and so on. In addition, we make the QNNs compatible with conventional real-valued neural networks and applicable for both skeletons and point clouds. Experiments on synthetic and real-world 3D tasks show that the QNNs based on our fast QPUs are superior to state-of-the-art real-valued models, especially in the scenarios requiring the robustness to random rotations. The code of this work is available at https://github.com/SuferQin/Fast-QPU .

Distribution comparison plays a central role in many machine learning tasks like data classification and generative modeling. In this study, we propose a novel metric, called Hilbert curve projection (HCP) distance , to measure the distance between two probability distributions with low complexity. In particular, we first project two high-dimensional probability distributions using Hilbert curve to obtain a coupling between them, and then calculate the transport distance between these two distributions in the original space, according to the coupling. We show that HCP distance is a proper metric and is well-defined for probability measures with bounded supports. Furthermore, we demonstrate that the modified empirical HCP distance with the $L_{p}$ cost in the $d$ -dimensional space converges to its population counterpart at a rate of no more than $O(n^{-1/2\max \lbrace d,p\rbrace })$ . To suppress the curse-of-dimensionality, we also develop two variants of the HCP distance using (learnable) subspace projections. Experiments on both synthetic and real-world data show that our HCP distance works as an effective surrogate of the Wasserstein distance with low complexity and overcomes the drawbacks of the sliced Wasserstein distance.