Structure-Aware E(3)-Invariant Molecular Conformer Aggregation Networks

Message-passing neural networks (MPNN) learn $d$-dimensional real-valued vector representations for each vertex in a graph by exchanging and aggregating information from neighboring nodes. Each vertex $v$ is annotated with a feature $\mathbf{h}_{v}^{(0)}$ in $\mathbb{R}^{d}$ representing characteristics such as atom positions and numbers in the case of chemical molecules. In addition, each edge $(u,v)$ is associated with a feature vector $\bm{e}(u,v)$. An MPNN architecture consists of a composition of permutation-equivariant parameterized functions.

Following Gilmer et al. (2017a) and Scarselli et al. (2009), in each layer, $\ell>0$, we compute vertex features

where $\mathsf{UPD}^{(\ell)}$, $\mathbf{M}^{(\ell)}$, and $\mathsf{AGG}^{(\ell)}$ are differentiable parameterized functions. In the case of graph-level regression problems, one uses

to compute a single vectorial representation based on learned vertex features after iteration $L$ where $\mathsf{READOUT}$ can be a differentiable parameterized function.

Molecules are $3$-dimensional structures that can be represented by geometric graphs, capturing each atom’s 3D position. To obtain more expressive representations, we also consider geometric input attributes and focus on vectorial features $\vv{\mathbf{v}}_{v},\vv{\mathbf{v}}_{u}$ of nodes. Since we address the problem of molecular property prediction, where we assume the properties to be invariant to actions of the group $E(3)$, we focus on $E(3)$-invariant MPNNs for geometric graphs.

Fused Gromov-Wasserstein. An undirected attributed graph $G$ of order $n$ in the optimal transport context is defined as a tuple $G:=({\bm{H}},{\bm{A}},\bm{\omega})$, where ${\bm{H}}\in\mathbb{R}^{n\times d}$ is a node feature matrix and ${\bm{A}}$ is a matrix encoding relationships between nodes, and $\bm{\omega}\in\Delta_{n}$ denotes the probability measure of nodes within the graph, which can be modeled as the relative importance weights of graph nodes. Without any prior knowledge, uniform weights can be chosen $(\bm{\omega}=\mathbf{1}_{n}/n)$ (Vincent-Cuaz et al., 2022). The matrix ${\bm{A}}$ can be the graph adjacency matrix, the shortest-path matrix or other distance metrics based on the graph topologies (Peyré et al., 2016; Titouan et al., 2019, 2020). Given two graphs $G_{1},G_{2}$ of order $n_{1},n_{2}$, respectively, Fused Gromov-Wasserstein (FGW) distance can be defined as follows:

Here ${\bm{M}}:=\left(d_{f}({\bm{H}}_{1}[i],{\bm{H}}_{2}[j])^{p}\right)_{n_{1}\times n_{2}}\in{\mathbb{R}}^{n_{1}\times n_{2}}$ is the pairwise node distance matrix, ${\bm{\mathsfit{L}}}\left({\bm{A}}_{1},{\bm{A}}_{2}\right)=(|{\bm{A}}_{1}[i,j]-{\bm{A}}_{2}[l,m]|^{p})_{ijlm}$ the 4-tensor representing the alignment cost matrix, and $\bm{\Pi}(\bm{\omega}_{1},\bm{\omega}_{2}):=\{\bm{\pi}\in\mathbb{R}_{+}^{n_{1}\times n_{2}}|\bm{\pi}\mathbf{1}_{n_{2}}=\bm{\omega}_{1},~{}\bm{\pi}\mathbf{1}_{n_{1}}=\bm{\omega}_{2}\}$ is the set of all valid couplings between node distributions $\bm{\omega}_{1}$ and $\bm{\omega}_{2}$. Moreover, $d_{f}(\cdot,\cdot)$ is the distance metric in the feature space, and $\alpha\in[0,1]$ is the weight that trades off between the Gromov-Wasserstein cost on the graph structure and Wasserstein cost on the feature signal. In practice, we usually choose $p=2$, Euclidean distance for $d_{f}(\cdot,\cdot)$, and $\alpha=0.5$ to calculate FGW distance.

Entropic Fused Gromov-Wasserstein. The entropic FGW distance adds an entropic term (Cuturi, 2013) as

where the entropic scalar $\epsilon$ facilitates the tunable trade-off between solution accuracy and computational performance (w.r.t. lower and higher $\epsilon$, respectively). Solving this entropic FGW involves iterations of solving the linear entropic OT problem Equation 37 with (stabilized) Sinkhorn projections (Proposition 2 (Peyré et al., 2016)), described in Appendix C and Algorithm 2.

To efficiently generate conformers, we employ distance geometry-based algorithms, which convert distance constraints into Cartesian coordinates. For atomistic systems, constraints typically define lower and upper bounds on interatomic squared distances. In a 2D input graph, covalent bond distances adhere to known ranges, while bond angles are determined by corresponding geminal distances. Adjacent atoms or functional groups adhere to cis/trans limits for rotatable bonds or set values for rigid groups. Other distances have hard sphere lower bounds, usually chosen approximately 10$\%$ below van der Waals radii (Hawkins, 2017). Chirality constraints are applied to every rigid quadruple of atoms.

A distance geometry algorithm now randomly generates a $3$-dimensional conformation satisfying the constraints. To bias the generation towards low-energy conformations, a simple and efficient force field is typically applied. We use efficient implementations from the RDKit package (Landrum, 2016).

We propose a new MPNN-based neural network that consists of three parts as depicted in Figure 1. First, a 2D MPNN model is used to capture the general molecular features such as covalent bond structure and atom features. Second, a novel FGW barycenter-based implicit $E(3)$-invariant aggregation function that integrates the representations of molecular 3D conformations computed by geometric message-passing neural networks. Finally, a permutation and $E(3)$-invariant aggregation function will be used to combine the 2D graph and 3D conformer representations of the molecules.

2D Molecular Graph Message-Passing Network. Each molecule is represented by a 2D graph $G=(V,E)$ with nodes $V$ representing its atoms and edges $E$ representing its covalent bonds, annotated with molecular features ${\bm{h}}_{v}^{(0)}$ and ${\bm{e}}_{v,u}$, respectively (see Section 6 for details). To propagate features across a molecule and get 2D molecular representations, we use GAT layers, which utilize a self-attention mechanism in message-passing with the following operations:

and where $\alpha_{v,u}$ are the GAT attention coefficients and $\mathbf{W}$ a learnable parameter matrix. Following Veličković et al. (2018), the attention mechanism is implemented with a single-layer feedforward neural network. To obtain a per-molecule embedding, we compute $\mathbf{h}^{\mathtt{2D}}_{G}=\sum_{v\in V}\mathbf{h}_{v}^{(L)}$, where $L$ is the number of message-passing layers.

3D Conformer Message-Passing Network. A conformer (atomic structure) of a molecule is defined as $S=\{\mathbf{r}_{i},Z_{i}\}_{i=1}^{N}$ where $N$ is the number of atoms, $\mathbf{r}_{i}\in\mathbb{R}^{3}$ are the Cartesian coordinates of atom $i$, and $Z_{i}\in\mathbb{N}$ is the atomic number of atom $i$. We use weighted adjacency matrices $\mathbf{A}\in\mathbb{R}^{n\times n}$ to represent pairwise atom distances. In some cases we will apply a cutoff radius to these distances. We employ the geometric MPNN SchNet (Schütt et al., 2017), although it is worth noting that alternative $E(3)$-invariant neural networks could be seamlessly integrated. The selection of SchNet is motivated not only by its proficient balance between model complexity and efficacy but also by its proven utility in previous works (Axelrod & Gómez-Bombarelli, 2023). SchNet performs $E(3)$-invariant message-passing by using radial basis functions to incorporate the distances of the geometric node features $\vv{\mathbf{v}}_{v},\vv{\mathbf{v}}_{u}$. We refer the reader to Section D.1 for more details. We denote the matrix whose columns are the atom-wise features of SchNet from the last message-passing layer $L$ with $\mathbf{H}$, that is, $\mathbf{H}[v]=\mathbf{h}^{(L)}_{v}$.

To compute the vector representation for a conformer $S$, we aggregate the atom-wise embeddings obtained from the last message-passing layer $L$ of SchNet into a single vector representation as $\mathbf{h}_{S}^{\mathtt{3D}}=\sum_{v\in V}\left(\mathbf{A}\mathbf{h}_{v}^{(L)}+{\bf a}\right)$, where $V$ is the set of atoms and $\mathbf{A}$ and ${\bf a}$ learned during training. For a set of $K$ conformers, the output of our 3D MPNN models is a matrix whose columns are the embeddings $\mathbf{h}_{S_{k}}$ for conformer $k$, that is, $\mathbf{H}^{\mathtt{3D}}[k]=\mathbf{h}_{S_{k}}^{\mathtt{3D}}$.

FGW Barycenter Aggregation. We now introduce an implicit and differentiable neural aggregation function whose output is determined by solving an FGW barycenter optimization problem. Its input is $K$ graphs $G_{k}=(\mathbf{H}_{k},{\bm{A}}_{k},\bm{\omega}_{k})$ for each conformer $S_{k}=\{\mathbf{r}_{k,i},Z_{k,i}\}_{i=1}^{N}$, with features $\mathbf{H}_{k}$ computed by an $E(3)$-invariant MPNN, with weighted adjacency matrix ${\bm{A}}_{k}$ of pairwise atomic distances, and the probability mass of each atom $\bm{\omega}_{k}$, typically set to $1/N$. The output of the barycenter conformer, denoted as $\overline{G}=(\overline{\mathbf{H}},\overline{{\bm{A}}},\overline{\bm{\omega}})$, represents the geometric mean of the input conformers, incorporating both their structural characteristics and features (Figure 1). The barycenter $\overline{G}$ is the conformer graph that minimizes the sum of weighted FGW distances among the conformer graphs $(G_{k})_{k\in[K]}$ with feature matrices $(\mathbf{H}_{k})_{k\in[K]}$, structure matrices $({\bm{A}}_{k})_{k\in[K]}$, and base histograms $(\bm{\omega}_{k})_{k\in[K]}\in\Delta_{n}^{K}$. That is, given any fixed $K\in\mathbb{N}$ and any $\bm{\lambda}\in\Delta_{K}$, the FGW barycenter is defined as

where $\text{FGW}_{p,\alpha}(G,G_{k})$ is the fused Gromov-Wasserstein distance defined in Equation 3, and where we set, for each pair of conformer graphs $G=(\mathbf{H},{\bm{A}},\bm{\omega})$ and $G_{k}=(\mathbf{H}_{k},{\bm{A}}_{k},\bm{\omega}_{k})$, ${\bm{M}}:=\left(\left(\mathbf{H}[i]-\mathbf{H}_{k}[j]\right)^{2}\right)_{n\times n}\in{\mathbb{R}}^{n\times n}$ as the feature distance matrix, and ${\bm{\mathsfit{L}}}\left({\bm{A}},{\bm{A}}_{k}\right)=\left({\bm{A}}[i,j]-{\bm{A}}_{k}[l,m]|\right)_{ijlm}$ as the 4-tensor representing the structural distance when aligning atoms $i$ to $l$ and $j$ to $m$ (Figure 2). Solving Equation 6, we obtain a unique FGW barycenter graph $\overline{G}=(\overline{\mathbf{H}},\overline{{\bm{A}}},\overline{\bm{\omega}})$ with representation $\overline{\mathbf{h}}_{v}=\overline{\mathbf{H}}[v]$ for each atom $v$. We aggregate the atom-wise embeddings obtained from the FGW barycenter $\overline{G}$ into a single vector representation using $\mathbf{h}^{\mathtt{BC}}_{\overline{G}}=\sum_{v\in V}\left(\overline{\mathbf{A}}~{}\overline{\mathbf{h}}_{v}+\overline{{\bf a}}\right)$.

Intuitively, barycenter-based aggregation in Eq.(6) can be seen as a more distance (structure) preserving pooling operation rather than standard mean aggregation. For instance, consider two conformers, where one is a 180-degree rotation of the other. Averaging their coordinates collapses the hydrogen atoms into the same position, creating an unphysical structure. On the contrary, employing the FGW Barycenter might prevent such issues.

Invariant Aggregation of 2D and 3D Representations. We integrate the representations of the 2D graph and the 3D conformer graphs using an average aggregation as well as the barycenter-based aggregation. The requirement for this aggregation is that it is invariant to the order of the input conformers; that is, it treats the conformers as a set as well as invariant to actions of the group $E(3)$.

Let $\mathbf{H}^{\mathtt{2D}}$ and $\mathbf{H}^{\mathtt{BC}}$ be the matrices whose columns are, respectively, $K$ copies of the 2D and barycenter representations from previous sections. Using learnable weight matrices $\mathbf{W}^{\mathtt{2D}}$, $\mathbf{W}^{\mathtt{3D}}$, and $\mathbf{W}^{\mathtt{BC}}$, we obtain the final atom-wise feature matrices as

where $\gamma$ is a hyper-parameter controlling the contribution of the barycenter-based feature. Intuitively, this aggregation function, where we use multiple copies of the 2D graph and barycenter representations, provides a balanced contribution of the three types of representations and is empirically highly beneficial. Finally, to predict a molecular property, we apply a linear regression layer on a mean-aggregation of the per-conformations embedding as:

We can show that the function defined by Section 3.2 to Equation 8 is invariant to actions of the group $E(3)$ and permutations acting on the sequence of input conformers.

This work establishes a novel fast rate convergence for empirical barycenters in the FGW space via Theorem 4.1, which is proved in Appendix B. To the best of our knowledge, this is new in the literature, where only the result for Wasserstein space exists in Le Gouic et al. (2022).

The upper bound in Equation 12 implies that the empirical barycenter converges to the target distribution at a rate of $\mathcal{O}(1/K)$, where $K$ is the number of 3D conformers. This suggests utilizing small values of $K$, such as $K\in\{5,10\}$, would yield a satisfactory approximation for $\overline{\mu}_{0}$. We confirm this empirically in experiments in Section 6.5.

To train on large-scale problems, we propose to solve the entropic relaxation of Equation 6 to take advantage of GPU computing power (Peyré et al., 2019). Given a set of conformer graphs $\{G_{s}:=({\bm{H}}_{s},{\bm{A}}_{s},\bm{\omega}_{s})\}_{s=1}^{K}$, we want to optimize the entropic barycenter $\overline{G}$, where we fixed the prior on nodes $\overline{\bm{\omega}}$

with $\lambda_{s}=1/K,\,\forall s\in[1,K]$. Titouan et al. (2019) solve Equation 13 using Block Coordinate Descent, which iteratively minimizes the original FGW distance between the current barycenter and the graphs $G_{s}$. In our case, we solve for $K$ couplings of entropic FGW distances to the graphs at each iteration, then following the update rule for structure matrix (Proposition 4, (Peyré et al., 2016))

and for the feature matrix (Titouan et al., 2019; Cuturi & Doucet, 2014)

leading to Algorithm 1. More details on practical implementations and algorithm complexity are in Appendix C.

We encode each molecule in the SMILES format and employ the RDKit package to generate 3D conformers. We set the size of the latent dimensions of GAT (Veličković et al., 2018) to $128/256$. Node features are initialized based on atomic properties such as atomic number, chirality, degree, charge, number of hydrogens, radical electrons, hybridization, aromaticity, and ring membership, while edges are represented as one-hot vectors denoting bond type, stereo configuration, and conjugation status. Each 3D conformer generated by RDKit comprises $n$ atoms with the corresponding 3D coordinates representing their spatial positions. Subsequently, we establish the graph structure and compute atomic embeddings utilizing the force-field energy-based SchNet model (Schütt et al., 2017), extracting features prior to the $\mathsf{READOUT}$ layer. Our SchNet configuration incorporates three interaction blocks with feature maps of size $F=128$, employing a radial function defined on Gaussians spaced at intervals of $0.1\mathring{{A}}$ with a cutoff distance of 10 $\mathring{A}$. The output of each conformer $k\in[K]$ forms a graph $G_{k}$, utilized in solving the FGW barycenter $\overline{G}$ as defined in Eq. (6). Subsequently, we aggregate features from 2D, 3D, and barycenter molecule graphs using Eqs. (7-8), followed by MLP layers. Leveraging Sinkhorn iterations in our barycenter solver (Algorithm 1), we speed up the training process across multiple GPUs using PyTorch’s distributed data-parallel technique. Training the entire model employs the Adam optimizer with initial learning rates selected from ${1e^{-3},1e^{-3}/2,1e^{-4}}$, halved using ReduceLROnPlateau after $10$ epochs without validation set improvement. Further experimental details are provided in the Appendix.

To accelerate the training process, especially in large-scale settings (e.g., BDE dataset), we first train the model with 2D and 3D features for some epochs, and then load the saved model and continue to train with full configurations as in Eq.(7) till converge. We set empirically $\gamma$ in Eq.(7) is $0.2$.

Dataset. We use four datasets Lipo, ESOL, FreeSolv, and BACE in MoleculeNet benchmark (Table 1), spanning on various molecular characteristics such as physical chemistry and biophysics. We split data using random scaffold settings as baselines and reported the mean and standard deviation of mean square error (mse) by running on five trial times. More information for datasets is in Section D.2 Appendix.

Baselines. We compare against various benchmarks, including both supervised, pre-training, and multi-modal approaches. (i) The supervised methods are 2D graph neural network models including 2D-GAT (Veličković et al., 2018), D-MPNN (Yang et al., 2019), and AttentiveFP (Xiong et al., 2019); (ii) 2D molecule pretraining methods are PretrainGNN (Hu et al., 2020a), GROVER (Rong et al., 2020), MolCLR (Wang et al., 2022), ChemRL-Gem (Fang et al., 2022), ChemBERTa-2 (Ahmad et al., 2022), and MolFormer (Ross et al., 2022). It’s important to note that these models are pre-trained on a vast amount of data; for example, MolFormer is learned on $1.1$ billion molecules from PubChem and ZINC datasets. We also compare with the (iii) 2D-3D aggregation ConfNet model (Liu et al., 2021), which is one of the winners of KDD Cup on OGB Large-Scale Challenge (Hu et al., 2021). Finally, we benchmark with 3D conformers-based models such as UniMol (Zhou et al., 2023), SchNet, and ChemProp3D (Axelrod & Gómez-Bombarelli, 2023). Among this, UniMol is pre-trained on $209$ M molecular conformation and requires 11 conformers on each downstream task. We train SchNet with $5$ conformers ($10$ for FreeSolv) and test with two versions: (a) taking output at the final layer and averaging different conformers (SchNet-scalar), (b) using feature node embeddings before $\mathsf{READOUT}$ layers and aggregating conformers by an MLP layer (SchNet-em). In ChemProp3D, we replace the classification header with an MLP layer for regression tasks, training with a 2D molecular graph and $10$ conformers. With the ConfNet, we use $20$ conformers in the training step and provide results for $20\textrm{ and }40$ conformers for the evaluations step, followed by configurations in (Liu et al., 2021).

Results. Table 2 presents the experimental findings of ConAN, alongside competitive methods, with the best results highlighted in bold. Baseline outcomes from prior studies (Zhou et al., 2023; Fang et al., 2022; Chang & Ye, 2023) are included, while performance for other models is provided through public codes. ConAN  version denotes the aggregation of 2D and 3D features as per Eq. (7) without employing the barycenter, whereas ConAN-FGW signifies full configurations. We employ a number of conformers $\{5,5,10,5\}$ and $\{5,5,5,5\}$ for ConAN, and ConAN-FGW, respectively, based on validation results for Lipo, ESOL, FreeSolv, and BACE. From the experiments, several observations emerge: (i) ConAN  proves more effective than relying solely on 2D or 3D, as shown by Conan’s performance, achieving second-best rankings on three datasets compared to models using only 2D (ChemRL-GEM) or 3D representations (UniMol). (ii) ConAN-FGW consistently outperforms baselines across all datasets, despite employing significantly fewer 3D conformers than ConAN. This underscores the importance of leveraging the barycenter to capture invariant 3D geometric characteristics.

Dataset.  We evaluate ConAN on two datasets Cov-2 3CL and Cov-2 (Table 1), focusing on molecular classification tasks. The same splitting for training and testing is followed (Axelrod & Gómez-Bombarelli, 2023). We also apply the CREST (Grimme, 2019) to filter generated conformers by RDKit as (Axelrod & Gómez-Bombarelli, 2023) for fair comparisons. Model performance is reported with the receiver operating characteristic area under the curve (ROC) and precision-recall area under the curve (PRC) over three trial times.

Baselines. We compare with three models, namely, SchNetFeatures, ChemProp3D, CP3D-NDU, each with two different attention mechanisms to ensemble 3D conformers and 2D molecular graph feature embedding as proposed by Axelrod & Gómez-Bombarelli (2023). These baselines generate $200$ conformers for their input algorithms. Additionally, the ConfNet (Liu et al., 2021) is also evaluated using $20$ or $40$ conformers in testing.

Results. Table 3 presents performance of ConAN  and ConAN-FGW with the number of conformers $10\text{ or }5$. It can be seen that ConAN-FGW delivers the best performance on ROC metric on two datasets and holds the second-best rank with PRC on CoV-2-3CL while requiring only $10\text{ or }5$ conformers compared with $200$ conformers as CP3D-NDU. These results underscore the efficacy of incorporating barycenter components over merely aggregating 2D and 3D conformer embeddings, as observed in ConAN.

Dataset. We run ConANon the BDE dataset (Table 1), which is the second-largest setting in (Zhu et al., 2023) aim to predict reaction-level molecule properties.
Baselines. ConAN is compared with state-of-the-art conformer ensemble strategies presented in Zhu et al. (2023), including SchNet (Schütt et al., 2017), DimeNet++ (Gasteiger et al., 2020a), GemNet (Gasteiger et al., 2021), PaiNN (Schütt et al., 2021), ClofNet (Du et al., 2022), and LEFTNet (Du et al., 2024). All these approaches employ $20$ conformers in training and testing. We provide two results of ConAN  using only $10$ conformers and based on two architectures, SchNet and LEFTNet.
Results. Table 4 summarizes our achieved scores where the ConAN-FGW using LEFTNet backbone holds the second rank overall while using half the number of conformers. Additionally, it can be seen that ConAN-FGW improves with significant margins over both base models like SchNet ($1.9737\rightarrow 1.6047$) and LEFTNet ($1.5276\rightarrow 1.4829$), demonstrating the generalization of the proposed aggregation.

Contribution of 3D Conformer Number. One of the building blocks of our model is the use of multiple 3D conformations of a molecule. Each molecule is represented by $K$ conformations, so the choice of $K$ affects the model’s behavior. We treat $K$ as a hyperparameter and conduct experiments to validate the impact on model performance. To this end, we test on the ConAN  version with different $K$ ($K=0$ is equivalent to the 2D-GAT baseline) and report performance in Table 7 Appendix. We can observe that using 3D conformers with $K\geq 1$ clearly improves performance compared to using only 2D molecular graphs as 2D-GAT. Furthermore, there is no straightforward dependency between the number of conformations in use and the accuracy of the model. For e.g., the performance tends to increase when using $K=10$ (Lipo and FreeSolv), but overall, the best trade-off value is $K=5$.

Contribution of FGW Barycenter Aggregation.  We examine the effect of barycenter aggregation when varying the number of conformers $K$. Figure 3 summarizes results for those settings where we report average RMSE over four datasets in the MoleculeNet benchmark. We draw the following observations. First, ConAN-FGW shows notable enhancements as the number of conformers increases, with $K$ values ranging within the set ${3,5,10}$; however, when as $K=20$, discernible disparities compared to the results obtained at $K=10$ diminish. We argue that this phenomenon aligns consistently with theoretical results in Theorem 4.1 suggesting that employing a sufficiently large $K$ facilitates a precise approximation of the target barycenter.

Secondly, upon examining various datasets, it becomes evident that ConAN-FGW consistently demonstrates enhanced performance with the utilization of larger conformers, a phenomenon not uniformly observed in the case of ConAN. This observation validates the robustness and resilience inherent in ConAN-FGW. We attribute this advantage to the efficacy of its geometry-informed aggregation strategy in ensemble learning with diverse 3D conformers.

Generalization to other Backbone Model. We investigate ConAN  and ConAN-FGW performance using the VisNet backbone (Wang et al., 2024a), an equivariant geometry-enhanced graph neural network for 3D conformer embedding extraction. Results in Table 5 confirm that ConAN-FGW still advances ConAN  performance. Between VisNet and SchNet, there is no universal best choice over datasets.

We check diversity conformers randomly selected from a set of conformers generated by RDKit. For each pair of 3D conformers, we compute the optimal root mean square distance, which first aligns two molecules before measuring distance. Two settings are conducted: (i) estimating the mean, variance, max, and min distances distribution for conformers sampled by ConAN  over $200$ conformers generated by RDKit. and (ii) estimate distribution for those values in case they are the top closest conformers. Figure 4 below shows our observation with a box plot on the validation set of Fressolv.

We observe that the distribution on the left ranges from 0.1 to 1.5, while in the worst case, the distance is between 0.01 and 0.08. Additionally, there’s a large gap between $d_{\mathrm{max}}$ and $d_{\mathrm{min}}$ on the left, whereas on the right, their means are close. It, therefore, can be seen that ConAN  sampling strategy, given 200 RDKit-generated conformers, remains consistent and diverse.

We contrast our entropic solver (Algorithm 1) with FGW-Mixup (Ma et al., 2023) for the $K$ barycenter problem. FGW-Mixup accelerates FGW problem-solving by relaxing coupling feasibility constraints. However, as the number of conformers $K$ increases, FGW-Mixup requires more outer iterations due to compounding marginal errors in solving $K$ FGW distances. In contrast, our approach employs an entropic-relaxation FGW formulation ensuring that marginal constraints are respected, resulting in a less noisy FGW subgradient. Furthermore, we implement our algorithm with distributed computation on multi-GPUs, as highlighted in Fig. 5. This figure illustrates epoch durations during both forward and backward steps of training, showcasing the performance across various conformer setups on FreeSolv and CoV-2 3CL datasets. Utilizing a batch size of 32 conformers, all three algorithms employ early termination upon reaching error tolerance. Notably, our solver exhibits linear scalability with $K$, while FGW-Mixup shows exponential growth, presenting challenges for large-scale learning tasks. More details are in Section D.5.