Abstract
Accurately predicting molecular properties requires effective integration of structural information from both 2D molecular graphs and their corresponding equilibrium conformer ensembles. In this work, we propose a scalable Structure-Aware Graph Transformer that efficiently aggregates features from multiple 3D conformers while incorporating fragment-level information from 2D graphs. Unlike prior methods that depend on static geometric solvers or rigid fusion strategies, our approach employs a trainable attention-based mechanism within a graph transformer to dynamically fuse 2D and 3D representations. We further enhance this mechanism by injecting fragment-specific structural biases into the attention layers, enabling the model to capture fine-grained molecular details. Our method scales to large datasets, handling up to 75,000 molecules and hundreds of thousands of conformers, and achieves state-of-the-art results in molecular property prediction and reaction-level modeling. It is particularly effective on chemically diverse compounds, including organocatalysts and transition-metal complexes.
1 Introduction
Machine learning has become a powerful tool for predicting molecular properties, with wide-ranging applications in drug discovery and materials science (Choudhary et al., 2022; Fedik et al., 2022; Batatia et al., 2023). Most existing models rely either on 2D molecular graphs, which efficiently capture topological connectivity (Xu et al., 2018; Veličković et al., 2018), or on 3D representations derived from a single conformer (Schütt et al., 2017; Batzner et al., 2022; Batatia et al., 2022). While 2D graphs are computationally efficient, they lack geometric information that is often critical for accurate property prediction. Incorporating 3D conformers helps address this by introducing spatial features such as bond lengths, and torsion angles. However, relying on a single conformer still fails to capture the intrinsic flexibility of molecular structures.
In reality, molecules dynamically sample a range of thermodynamically accessible conformations due to bond rotations, vibrations, and environmental interactions (Ramsundar et al., 2019). As a result, many experimentally observable properties such as solubility and binding affinity depend on the full ensemble of conformers a molecule can adopt (Perola & Charifson, 2004). Yet, fully modeling this distribution is computationally prohibitive, as generating and evaluating large numbers of conformers using quantum methods is costly (Medrano Sandonas et al., 2024). This has motivated hybrid models that combine the structural efficiency of 2D graphs with the geometric richness of a small and representative subset of 3D conformers. By jointly capturing topological and spatial variation, hybrid models offer scalable and expressive frameworks for molecular representation, enabling more accurate prediction of conformation-sensitive properties across a range of chemical and biological tasks.
Building on this hybrid paradigm, recent methods have introduced hybrid models that integrate 2D molecular graphs with 3D conformer information to capture both topological and spatial features (Zhu et al., 2024b; Axelrod & Gomez-Bombarelli, 2023). Despite the successes, these methods often assume conformers contribute equally or can be reweighted without considering deeper geometric context. In practice, only a subset of conformers may be thermodynamically or functionally relevant, and naive aggregation overlooks their spatial relationships, such as alignment or structural similarity. Moreover, current strategies rarely leverage interactions between 2D structural priors and 3D conformational variability, hindering the formation of truly expressive representations.
To address this, structure-aware ensemble methods based on optimal transport, especially those using fused Gromov-Wasserstein (FGW) alignment, have shown promise Ma et al. (2023); Nguyen et al. (2024a). By aligning both feature and geometric spaces, these models better preserve spatial correspondences across conformers and enable expressive ensemble aggregation. However, such methods are computationally expensive and struggle to scale to large molecular datasets such as Drugs-75k Zhu et al. (2023); Axelrod & Gomez-Bombarelli (2022), limiting their utility for high-throughput applications in generative biology.
To address scalability challenges in geometry-aware molecular modeling, we introduce a novel approach that replaces expensive FGW alignment with efficient attention-based conformer aggregation. By supervising the model with FGW distances during training, we learn a latent embedding space where conformer similarities reflect both topological and geometric structure. This enables fast, permutation-invariant conformer integration suitable for large-scale generative pipelines. Beyond efficiency, we further enrich our model with fragment-level structural priors from 2D molecular graphs, injecting chemically meaningful hierarchies into both message passing and 3D attention layers. This unified 2D–3D framework captures fine-grained spatial and topological interactions essential for applications such as molecular property prediction, virtual screening, and functional optimization. In summary, our key contributions are:
-
•
We propose a scalable, geometry-aware conformer aggregation framework, denoted as FACET, that replaces costly FGW alignment with a trainable Graph Transformer, enabling efficient, deterministic attention-based inference. We further provide theoretical bounds on the approximation error relative to FGW distances.
-
•
We introduce a unified 2D–3D representation learning approach that embeds fragment-level structural priors into both 2D message passing and 3D spatial self-attention, capturing multi-scale interactions between molecular topology and geometry.
-
•
Our method delivers over faster aggregation than prior geometry-aware baselines and achieves state-of-the-art performance across six benchmarks, including molecular property prediction and Boltzmann-weighted ensemble tasks, demonstrating robustness across diverse molecular scenarios and dataset scales.
2 Related Work
2.1 Conformer Ensemble Learning in Molecular Representations
Traditional molecular representations span connectivity fingerprints (Morgan, 1965), 1D string encodings (Ahmad et al., 2022; Wang et al., 2019), 2D topological graphs (Yang et al., 2019a; Rong et al., 2020), and 3D geometric graphs (Fang et al., 2021; Zhou et al., 2023). 3D models typically rely on a single conformer, overlooking the fact that molecules often adopt multiple low-energy conformations, which can serve as informative features, particularly in capturing thermodynamic properties. Hybrid approaches now combine 2D graphs with ensembles of 3D conformers (Zhu et al., 2024b; Wang et al., 2024), aggregated via mean pooling, DeepSets (Zaheer et al., 2017), or self-attention (Vaswani et al., 2017). More advanced geometry-aware methods based on Fused Gromov-Wasserstein (FGW) alignment (Ma et al., 2023; Nguyen et al., 2024a) capture both feature and structural similarity across conformers, but remain computationally costly and scale poorly to large datasets (e.g., Drugs-75k) or foundation models (Zhou et al., 2023; Chithrananda et al., 2020). To address this, we propose a scalable framework that learns latent embeddings of 3D conformers with graph transformers, integrating geometry-aware signals inspired by FGW and hierarchical fragment-level features. This yields a permutation-invariant, expressive, and efficient method.
2.2 Scalable Optimal Transport for Graph Learning
Learning-based approximations of Optimal Transport (OT) have emerged as efficient alternatives to traditional solvers. Early works introduced differentiable Sinkhorn distances with entropic regularization for stability and scalability (Cuturi, 2013; Feydy et al., 2019; Genevay et al., 2018). Later methods improved efficiency via structural assumptions - e.g., low-rank factorization (Scetbon et al., 2021; Cuturi et al., 2020) and spatial geometry (Bachmann et al., 2022; Solomon et al., 2015). Meta-learning approaches further accelerated convergence by learning initialization schemes (Amos et al., 2023). More recently, neural OT surrogates trained directly on data have bypassed iterative solvers entirely (Courty et al., 2017; Tong et al., 2021; Haviv et al., 2024).
However, prior works focus on standard OT and fail to extend to structure-aware variants like FGW, which jointly capture node attributes and graph topology. To address this, we introduce the first learned approximation of FGW with a graph transformer, enabling scalable, geometry-aware conformer aggregation. By embedding fragment-level priors into both 2D and 3D encoders, our approach supports multi-scale reasoning across topological and spatial hierarchies, effectively bridging molecular graphs with 3D conformational diversity.
2.3 Fragment-biases in Molecular GNN
Fragment-level substructures - such as rings, functional groups, and pharmacophores - are key to molecular property prediction and drug design (Merlot et al., 2003; Varnek et al., 2005). Recent works have leveraged these motifs for scaffold-aware drug discovery (Lee et al., 2024; Chan et al., 2024), self-supervised learning via fragment-based masking or contrastive tasks (Rong et al., 2020; Zhang et al., 2021; Wen et al., 2024), and GNN architectures that encode fragment-level inductive biases (Wang et al., 2025; Wollschläger et al., 2024). These methods show that fragments enhance generalization, interpretability, and data efficiency. Building on these insights, we explore a complementary direction: integrating fragment-level priors into hybrid 2D–3D ensemble models. In our approach, fragment hierarchies are embedded into both 2D message-passing and 3D spatial attention layers, enabling multi-scale processing across molecular topology and geometry. This design improves conformer aggregation and yields more expressive, geometry-aware representations suited for conformation-sensitive tasks.
3 Fragment-Aware Conformer Ensemble Transformer
Notation. Let denote the probability simplex, where is the all-ones vector in . For , is the Dirac measure at . We write for , and use to denote the Frobenius inner product. For a tensor and matrix , define the contraction . A graph has nodes and edges . An attributed graph is given by , where is the node feature matrix (with row for node ), encodes structure (e.g., adjacency or shortest-path distance matrix), and is a node weight distribution.
Given two graphs and with and nodes, the Fused Gromov-Wasserstein (FGW) distance (Peyré et al., 2016; Titouan et al., 2019, 2020) is: where is the set of valid couplings, is the distance between feature of node in and of node in , captures structural mismatch, and balances feature and structure alignment.
3.1 Conformer Generation
Following prior work, we generate molecular conformers using distance geometry methods that convert interatomic constraints - such as bond lengths, angles, stereochemistry, and steric limits - into 3D coordinates (Hawkins, 2017). A lightweight force field refines the structures toward low-energy conformations. Compared to quantum methods like DFT, this approach is highly scalable and efficient for large datasets. As in prior studies (Raza et al., 2022; Nguyen et al., 2024b), we use RDKit (Landrum, 2016) for fast and reliable conformer generation.
3.2 Framework Overview
We propose a neural architecture with three main components (Fig. 1). First, a 2D message passing neural network (MPNN) captures molecular topology, while another 2D-MPNN operates on a fragment-level graph, which consists of pairwise edges between fragment nodes, to encode higher-order structural priors (Sec. 3.3). Their outputs are fused and refined through a lightweight adaptor module before entering a pre-trained FGW-guided graph transformer (Sec. 3.4). For 3D information, a set of conformers is sampled from the input molecule, and a 3D-MPNN extracts conformer embeddings (Sec. 3.4.1), which are also calibrated by an adaptor layer to handle variability between 2D and 3D features. Then, conformer embeddings are fed into the graph transformer, where each node attends to all other nodes, taking into account the conformers graph structure and fragment-level information (Sec. 3.4.2). In essence, the graph transformer encodes conformer embeddings into another space where their pairwise Euclidean distance is equal FGW distance (Sec. 3.4.3). Finally, a permutation- and E(3)-invariant fusion module unifies the 2D and 3D features into a single embedding for downstream tasks (Sec. 3.4.4).
3.3 Fragment-Enhanced 2D Molecular Graph
Each molecule is represented as a 2D graph , where nodes correspond to atoms and edges to covalent bonds. Atom features encode properties like atom type and valence, while bonds are annotated with features (Scarselli et al., 2008; Gilmer et al., 2017). We adopt a 2D message-passing neural network (MPNN) that updates node embeddings layer-wise:
| (1) |
where is a message function, is sum aggregation, and is identity or multilayer perception layers. We use Graph Attention Networks (GATs) (Veličković et al., 2017), where messages are computed as:
| (2) |
After layers, we obtain final atom-level features for each atom used for downstream tasks.
Fragment-Based Structural Augmentation.
To enhance atomic representations with higher-order structural context, we construct a fragment-level graph from the input molecular graph using ring-path decomposition (Kong et al., 2022; Geng et al., 2023; Wollschläger et al., 2024) to identify key substructures such as aromatic rings and functional groups (Fig. 4). Each fragment is treated as a node in a new graph , where nodes correspond to fragments and edges are induced from the connectivity in , two fragments are connected if they share an atom or are directly bonded. In this work, we specifically follow the approach proposed in (Wollschläger et al., 2024), as it offers a good balance of simplicity and effectiveness for our use case.
We apply the same GAT formulation in Eq. (1) to the fragment graph to obtain fragment embeddings . Then for each atom that belongs to its fragment , we fuse their atom-level representations with by:
| (3) |
where is a learnable feedforward network that projects fragment-level context into the same space as atom features. Finally, we define a fragment-enhanced graph level representation that is computed by applying a readout function Intuitively, the dual-level encoding combining local atomic features and global fragment-level context as Eq.(3) allows the model to reason over both fine-grained and coarse-grained structures, enhancing the expressivity of the molecular representation.
3.4 Learning Graph Transformer for 3D Molecule Aggregations
A molecular conformer is represented as a set , where denotes the number of atoms, corresponds to the 3D Cartesian coordinates of atom , and indicates its atomic number.
3.4.1 3D conformer feature representation.
For each conformer , we can define its graph and compute its 3D feature embedding by using the geometric MPNN SchNet (Schütt et al., 2017), though other -invariant neural architectures can be readily substituted without modification (Table 3). We represent the matrix of atom-level features from the final message-passing layer of SchNet as , where each column corresponds to the feature vector of atom . We then compute the vector representation of a conformer as with and are learnable vectors. Given a set of conformers , we define as the feature embedding of the -th conformer. The matrix thus summarizes the feature representations of all conformers in the set.
3.4.2 Fragment-aware Graph Transformer.
Given the atom-wise feature matrix for each conformer , we aim to learn structure-encoded latent representations using Graph Transformer architectures (Ying et al., 2021; Kreuzer et al., 2021; Luo et al., 2024). We adopt the architecture from Ying et al. (2021) due to its strong expressiveness on small molecular graphs, and further extend its attention mechanism with fragment sub-structures (Fig .4). It is important to note that our framework is flexible and can incorporate alternative transformer-based models.
In particular, we compose transformer layers (Vaswani et al., 2017), each consisting of a self-attention mechanism followed by a position-wise feed-forward network. Given computed in Section 3.4.1 by a 3D-MPNN, where is the vector embedding of an atom with dimensions. We compute self-attention, by linearly projecting into query (), key (), and value () matrices using learned weights :
| (4) |
Here, denotes the attention score matrix representing pairwise similarities between tokens. For clarity, we present the single-head version; extending to multi-head attention is straightforward. Bias terms are omitted for brevity.
While the attention in Eq. (4) operates only on feature nodes, leveraging the structural information of the 3D conformer graph is essential. Following Ying et al. (2021), we incorporate (i) centrality encoding, which measures the importance of a node in the graph via its degree, and (ii) spatial encoding, which captures the spatial relation between two nodes and in using the shortest path distance (SPD) (Cormen et al., 2022; Balaban, 1985), augmented with a learnable weight assigned to each edge along the SPD. Specifically, we incorporate (i) by:
| (5) |
where are learnable embedding vectors specified by the indegree and outdegree of atom respectively. Here is implemented as a lookup table of learnable embeddings indexed by shortest-path distance (i.e., degree) from node to others in the graph. The shortest-path distance (SPD) matrix is first computed, and these distances are used to retrieve the corresponding embeddings, which are then integrated into the attention mechanism to inject topology-aware structural bias. Assume as the -element of the Query-Key product matrix , the condition (ii) extends as:
| (6) |
where is a learnable scalar indexed by the SPD distance and shared across layers; , with the feature of edge in , its weight embedding, and the dimensionality of edge features, computed as the difference between the feature embeddings of its incident nodes.
While the spatial encoding in Eq.(6) is implicated by the SPD, we argue that this might inadequately capture chemically meaningful substructures (ablation in Tab. 3). This motivates us to extend attention scores in Eq. (6) using values derived from (iii) fragment-level node features computed on 2D topology graph in Eq. (3), directly guiding attention toward structurally and functionally relevant regions such as rings, functional groups, or scaffolds. To this end, we compute an adjacency-like matrix using cosine distance over the final node embeddings . Specifically, for each pair of atoms in the 2D molecular graph, we define
| (7) |
which quantifies their directional dissimilarity in the embedding space. Finally, we compute the attention score as:
| (8) |
3.4.3 Learning to Approximate FGW distance.
We denote as the graph transformer model whose attention operation is Eq.(8), our goal is to train to map the feature representation of each conformer into a latent space where the distance between any pair , approximates their FGW distance - an effective, yet computationally expensive, geometry-aware metric (Ma et al., 2023; Nguyen et al., 2024a). To this end, given a set of of generated conformers, we sample conformers from , then compute their encoding features by for each . These outputs are compared with their pair-wise FGW distance to optimize the loss:
| (9) |
By minimizing the loss , we update the parameters of the transformation module using gradient descent: Once trained, we freeze and incorporate it back into the framework to compute a geometry-aware representation across conformers as follows: where denotes the aggregated structural embedding. However, the 3D conformer feature distribution, extracted by 3D-MPNN, used to train (Eq. 9) may experience a domain shift when co-trained with other components in the full framework (Sec. 3.4) due to the continuous updating of 3D-MPNN. To address this, we design adapter layers as simple FFN layers to transform the input features in Eq. (9), aligning them to the seen distribution during training .
3.4.4 Invariant Aggregation of 2D and 3D Representation.
We integrate representations from the 2D molecular graph and multiple 3D conformers using both average pooling and a GraphTransformer-based aggregation. The transformer captures rich spatial interactions while ensuring permutation invariance across conformers and E(3) equivariance, preserving robustness to 3D transformations. Given conformers, using as the GraphTransformer (GT)-aggregated atom features. We compute the global GT representation as: where and are learnable parameters. We then define and be the matrices whose columns are, respectively, copies of the 2D feature (Sec.3.3) and representations from previous section. We fuse those representations with the 3D conformer features to produce the final atom-wise embedding: where each are trainable projection matrix. The combined embedding is fed into a final FFN layer to predict the target property (Sec.J Appendix).
4 Theoretical Bounds for Embedding Non-Euclidean FGW
Learning a Transformer to predict the FGW problem is closely related to multidimensional scaling (MDS) (Torgerson, 1952). Building on recent advances (Haviv et al., 2024; Sonthalia et al., 2021), we extend MDS theory to derive bounds on the error of embedding non-Euclidean distances, specifically Wasserstein and FGW, into a Euclidean space suitable for graph transformer integration. While computing FGW barycenters is costly, our embedding enables efficient approximation via averaging and decoding in latent space. Prior work (Haviv et al., 2024) validated this approach for Wasserstein distances; we generalize it to FGW and provide theoretical justification, offering a scalable path for structure-aware graph alignment.
Cumulative Stress Optimization Problem via Pairwise FGW Distance Matrix. We define the pairwise FGW distance matrix for a set of distributions as for all , following Section 3.4. The empirical FGW barycenter is given by , where denotes the space of attributed graphs with finite -th order FGW distance.
To approximate this barycenter in embedding space, we require for all , where is the mean embedding and is the learned representation. To assess how well the embeddings preserve both pairwise FGW distances and barycenter structure, we define the cumulative stress: . This objective encourages faithful reconstruction of both the distance structure and the barycenter alignment in the learned embedding space, as formalized in Theorem 1, which is proved in Appendix I.
Theorem 1.
Let denote the pairwise distance matrix, and let represent the eigendecomposition of the associated criterion matrix , where is the centering matrix. The optimal stress value, denoted by , is bounded as follows: , where , , . Here, denotes the -th component of the -th eigenvector of , and quantifies the approximation error between the empirical barycenter in the Euclidean embedding space and the one in the original space of undirected attributed graphs.
5 Experiments
5.1 Implementation Details
General pipeline. Our training consists of three stages. Stage 1: We train the 2D and 3D MPNNs independently for 150 epochs and the learning rate of to extract features from 2D molecular graphs and 3D conformers, used to predict molecular properties by regression loss. These extracted features also serve as a dataset to supervise the training of Graph Transformer for approximating the FGW distance. Stage 2: The Graph Transformer is trained separately to approximate the computationally expensive FGW distance between pairs of conformers, using the learned representations from Stage 1. We use the architecture of Graphormer (Ying et al., 2021), with 12 attention layers, 8 heads, and a hidden size of 64 (372k parameters). It is trained for 1000 epochs with a learning rate of . Stage 3: We train the full model end-to-end with 2D/3D MPNNs and the Graph Transformer (300 epochs, learning rate ). We further discuss this training scheme in Section G of the Appendix. To mitigate feature shift during finetuning, MLP-based adaptors map 3D conformers into 64-d refined embeddings, applied to both 2D and 3D features before the Graph Transformer.
5.2 Approximation of FGW Distance via Graph Transformer
Beyond theoretical estimation, we empirically evaluate how well the Graph Transformer approximates FGW distances between conformers in Euclidean space. As shown in Figure 2, results on the MoleculeNet benchmarks reveal a strong correlation between learned embeddings and true FGW distances, validating the transformer’s effectiveness in simulating costly FGW computations. While correlation varies slightly across datasets, the results consistently highlight the model’s reliability as a fast FGW surrogate, especially as the number of conformers in the aggregation increases.
5.3 Scaling Fragment Geometry-Aware Aggregation
To validate the scalability of FACET model, based on a Graph Transformer for structure-aware aggregation, we compare it against Conan-FGW (Nguyen et al., 2024a), a method computing FGW distances on-the-fly during training and inference. We evaluate two key aspects: (i) inference-time efficiency with varying numbers of conformers, and (ii) average training time per epoch at different dataset scales. For inference, we measure the time required to generate output embeddings from conformers () using a single GPU. Experiments are conducted on FreeSolv and BACE, which differ in node/edge distributions, to assess performance across molecular graph complexities.
In the second setting, we compare the average per-epoch training time of FACET and Conan-FGW on two datasets of different scales: Kraken (1,086 molecules) and Drugs-75k (52,569 molecules). As summarized in Figures 7 (appendix) and 3, FACET exhibits linear scaling with the number of conformers and achieves 5–6× faster runtime on average than Conan-FGW. This efficiency is critical for scaling to large datasets and longer training schedules - for e.g, training Conan-FGW on Drugs-75k for 300 epochs requires 1,107.58 GPU hours, while FACET only takes 214 hours. This can be further reduced to 26.75 hours with 8 GPUs, compared to 138 hours for Conan-FGW under the same hardware setup. We present in Section D more analysis on this scaling factor.
| Model | Lipo | ESOL | FreeSolv | BACE | Drugs-75k | Kraken |
|---|---|---|---|---|---|---|
| Train | 2940 | 789 | 449 | 1059 | 52569 | 1086 |
| Valid. | 420 | 112 | 64 | 151 | 7509 | 155 |
| Test | 2940 | 227 | 129 | 303 | 15021 | 311 |
| Total | 4200 | 1128 | 642 | 1513 | 75099 | 1552 |
| Settings | FACET | w/o Frag. | w/o Frag. in Trans. | w/o Adap. |
|---|---|---|---|---|
| ESOL | 0.516 | 0.531 | 0.525 | 0.546 |
| FreeSolv | 0.967 | 1.072 | 0.973 | 1.085 |
| Model | Lipo | ESOL | FreeSolv | BACE |
|---|---|---|---|---|
| ConAN (VisNet) | ||||
| ConAN-FGW | ||||
| FACET | 0.48 0.01 | 0.53 0.05 | 0.61 0.02 | 0.47 0.01 |
| ConAN (SchNet) | ||||
| ConAN-FGW | ||||
| FACET | 0.42 0.01 | 0.52 0.04 | 0.97 | 0.50 0.03 |
5.4 State-of-the-Art Performance Comparison on Molecular Tasks
| Model | Lipo | ESOL | FreeSolv | BACE |
|---|---|---|---|---|
| 2D-GAT | 1.387 0.206 | 2.288 0.017 | 8.564 1.345 | 1.844 0.33 |
| D-MPNN | 0.534 0.022 | 0.923 0.045 | 4.213 0.068 | 0.723 0.021 |
| Attentive FP | 0.520 0.001 | 0.771 0.026 | 4.197 0.193 | - |
| PretrainGNN | 0.545 0.003 | 1.210 0.005 | 6.392 0.003 | - |
| GROVER_large | 0.676 0.012 | 0.798 0.018 | 5.162 0.047 | - |
| ChemBERTa-2* | 0.639 0.006 | 0.795 0.033 | - | 1.858 0.029 |
| ChemRL-GEM | 0.486 0.008 | 0.706 0.061 | 3.924 0.436 | - |
| MolFormer | 0.492 0.012 | 0.766 0.026 | 5.485 0.045 | 1.091 0.021 |
| ConfNet | 1.360 0.038 | 2.115 0.484 | - | 1.329 0.042 |
| UniMol | 0.374 0.012 | 0.741 0.014 | 2.867 0.186 | - |
| SchNet-scalar | 0.704 0.032 | 0.672 0.027 | 1.608 0.158 | 0.723 0.100 |
| SchNet-emb | 0.589 0.022 | 0.635 0.057 | 1.587 0.136 | 0.692 0.028 |
| ChemProp3D | 0.602 0.035 | 0.681 0.023 | 2.014 0.182 | 0.815 0.170 |
| ConAN | 0.556 0.013 | 0.571 0.019 | 1.496 0.158 | 0.635 0.051 |
| ConAN-FGW | 0.422 0.016 | 0.529 0.022 | 1.068 0.083 | 0.549 0.016 |
| FACET | 0.424 0.009 | 0.516 0.044 | 0.967 0.082 | 0.495 0.115 |
Datasets. We evaluate molecular property regression on the MoleculeNet (Wu et al., 2018) and MARCEL (Zhu et al., 2024a) benchmarks. MoleculeNet includes four datasets, ESOL, BACE, Lipo, and FreeSolv, with targets covering solubility, inhibitory concentration (pIC50), lipophilicity, and hydration free energy. MARCEL consists of Drugs-75K and Kraken, where the goal is to predict the Boltzmann-averaged property from sampled conformers. Drugs-75K uses quantum descriptors (IP, EA, ), while Kraken focuses on Sterimol features (B5, L, and their buried forms). The Boltzmann average is computed as a weighted sum over conformer-specific values with probabilities . All datasets follow the original random split settings, using the provided sampled conformers.
Baselines. For the MoleculeNet benchmark (Wu et al., 2018), we compare FACET with a wide range of baselines, including (i) 2D supervised methods (e.g., GAT (Veličković et al., 2018), D-MPNN (Yang et al., 2019a), AttentiveFP (Xiong et al., 2019)), (ii) pre-training approaches (e.g., PretrainGNN (Hu et al., 2020b), GROVER (Rong et al., 2020), ChemBERTa-2* (Ahmad et al., 2022), ChemRL-GEM (Fang et al., 2022), MolFormer (Ross et al., 2022)), (iii) 3D-conformers based models (ConfNet (Liu et al., 2021), UniMol (Zhou et al., 2023), SchNet (Schütt et al., 2017), ChemProp3D (Axelrod & Gómez-Bombarelli, 2023),ConAN-FGW (Nguyen et al., 2024b)). Training follows the setup in ConAN-FGW (Nguyen et al., 2024b).
For the MARCEL benchmark (Zhu et al., 2024a), we compare FACET against 2D models (e.g., GIN (Xu et al., 2019), GIN+VN (Hu et al., 2020a), ChemProp (Yang et al., 2019b), GraphGPS (Rampášek et al., 2022)), 3D models (e.g., SchNet (Schütt et al., 2017), DimeNet++ (Klicpera et al., 2020), GemNet (Gasteiger et al., 2021), PaiNN (Schütt et al., 2021), ClofNet (Du et al., 2022), LEFTNet (Du et al., 2023)), and ensemble strategies such as DeepSets-based ensemble (Zaheer et al., 2017), self-attention (Vaswani et al., 2017), etc. All methods are evaluated under the same settings as described in the MARCEL benchmark.
5.4.1 Results
| Category | Model | Drugs-75K | Kraken | |||||
| IP | EA | B5 | L | BurB5 | BurL | |||
| 2D models | GIN | 0.4354 | 0.4169 | 0.2260 | 0.3128 | 0.4003 | 0.1719 | 0.1200 |
| GIN+VN | 0.4361 | 0.4169 | 0.2267 | 0.3567 | 0.4344 | 0.2422 | 0.1741 | |
| ChemProp | 0.4595 | 0.4417 | 0.2441 | 0.4850 | 0.5452 | 0.3002 | 0.1948 | |
| GraphGPS | 0.4351 | 0.4085 | 0.2212 | 0.3450 | 0.4363 | 0.2066 | 0.1500 | |
| 3D models | SchNet | 0.4394 | 0.4207 | 0.2243 | 0.3293 | 0.5458 | 0.2295 | 0.1861 |
| DimeNet++ | 0.4441 | 0.4233 | 0.2436 | 0.3510 | 0.4174 | 0.2097 | 0.1526 | |
| GemNet | 0.4069 | 0.3922 | 0.1970 | 0.2789 | 0.3754 | 0.1782 | 0.1635 | |
| PaiNN | 0.4505 | 0.4495 | 0.2324 | 0.3443 | 0.4471 | 0.2395 | 0.1673 | |
| ClofNet | 0.4393 | 0.4251 | 0.2378 | 0.4873 | 0.6417 | 0.2884 | 0.2529 | |
| LEFTNet | 0.4174 | 0.3964 | 0.2083 | 0.3072 | 0.4493 | 0.2176 | 0.1486 | |
| Ensemble Strategy with DeepSets | SchNet | 0.4452 | 0.4232 | 0.2243 | 0.2704 | 0.4322 | 0.2024 | 0.1443 |
| DimeNet++ | 0.4126 | 0.3944 | 0.2267 | 0.2630 | 0.3468 | 0.1783 | 0.1185 | |
| GemNet | 0.4066 | 0.3910 | 0.2027 | 0.2313 | 0.3386 | 0.1589 | 0.0947 | |
| PaiNN | 0.4466 | 0.4269 | 0.2294 | 0.2225 | 0.3619 | 0.1693 | 0.1324 | |
| ClofNet | 0.4280 | 0.4033 | 0.2199 | 0.3228 | 0.4485 | 0.2178 | 0.1548 | |
| LEFTNet | 0.4149 | 0.3953 | 0.2069 | 0.2644 | 0.3643 | 0.2017 | 0.1386 | |
| FACET | SchNet | 0.4235 | 0.3971 | 0.2155 | 0.2508 | 0.3982 | 0.1803 | 0.1245 |
| GemNet | 0.3891 | 0.3852 | 0.1970 | 0.2225 | 0.3402 | 0.1503 | 0.0952 | |
MoleculeNet. As shown in Table 4, FACET achieves state-of-the-art performance on three molecular property regression tasks (ESOL, FreeSolv, BACE), with the lowest MSEs: , , and , respectively. Its consistent gains over CONAN-FGW indicate that, beyond geometry-aware aggregation, FACET’s use of fragment substructures (Figure 4) enhances attention to localized chemical contexts. This demonstrates the advantage of combining 3D spatial information with chemically meaningful substructures for molecular property prediction.
MARCEL. In Table 5, we evaluate FACET on two backbones, SchNet and GemNet. FACET consistently boosts both, confirming the benefits of structure-aware aggregation and fragment-level hierarchy. Unlike CONAN-FGW, which struggles to scale on the large MARCEL benchmark, FACET remains efficient and achieves near-SOTA performance across all targets, demonstrating robust effectiveness in diverse molecular property prediction tasks.
5.5 Ablation study
In this section, we analyze the key components of FACET through ablation studies. Specifically, we evaluate the impact of: (i) removing fragment structures from both the 2D MPNN and the self-attention mechanism in the graph transformer (w/o Frag); (ii) using fragments only in the 2D MPNN but not in the graph transformer (w/o Frag in Trans.); and (iii) omitting the trainable adaptor (w/o Adap.) that aligns 3D conformer features with the graph transformer, which can lead to performance degradation due to domain shift during training. As shown in Table 3, the absence of (i) significantly reduces performance, making FACET comparable to CONAN-FGW but with better scalability. Incorporating fragments into both components (ii) provides further gains, while (iii) proves essential for mitigating the domain shift introduced by changes in the 3D MPNN during training.
6 Conclusion
We present FACET, a scalable method that integrates 3D conformer features with fragment-level 2D graph information. Using an FGW-guided trainable attention mechanism, FACET dynamically fuses 2D and 3D representations, outperforming FGW-based baselines across all MoleculeNet tasks. It also scales to 75,000 molecules and large conformer ensembles in the MARCEL benchmark, achieving state-of-the-art results in property and reaction prediction with efficient runtimes. A discussion of the current limitations of FACET is presented in Section H Appendix.
References
- Ahmad et al. (2022) Walid Ahmad, Elana Simon, Seyone Chithrananda, Gabriel Grand, and Bharath Ramsundar. Chemberta-2: Towards chemical foundation models, 2022. ↩ a b
- Amos et al. (2023) Brandon Amos, Samuel Cohen, Giulia Luise, and Ievgen Redko. Meta optimal transport. International Conference on Machine Learning, 2023. ↩ a
- Axelrod & Gomez-Bombarelli (2022) Simon Axelrod and Rafael Gomez-Bombarelli. Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022. ↩ a
- Axelrod & Gomez-Bombarelli (2023) Simon Axelrod and Rafael Gomez-Bombarelli. Molecular machine learning with conformer ensembles. Machine Learning: Science and Technology, 4(3):035025, 2023. ↩ a
- Axelrod & Gómez-Bombarelli (2023) Simon Axelrod and Rafael Gómez-Bombarelli. Molecular machine learning with conformer ensembles. Mach. Learn.: Sci. Technol., 4(3):035025, September 2023. ISSN 2632-2153. doi: 10.1088/2632-2153/acefa7. ↩ a
- Bachmann et al. (2022) Fynn Bachmann, Philipp Hennig, and Dmitry Kobak. Wasserstein t-sne. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 104–120. Springer, 2022. ↩ a
- Balaban (1985) Alexandru T Balaban. Applications of graph theory in chemistry. Journal of chemical information and computer sciences, 25(3):334–343, 1985. ↩ a
- Batatia et al. (2022) Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gabor Csanyi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 11423–11436. Curran Associates, Inc., 2022. ↩ a
- Batatia et al. (2023) Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena, Dávid P. Kovács, Janosh Riebesell, Xavier R. Advincula, Mark Asta, William J. Baldwin, Noam Bernstein, Arghya Bhowmik, Samuel M. Blau, Vlad Cărare, James P. Darby, Sandip De, Flaviano Della Pia, Volker L. Deringer, Rokas Elijošius, Zakariya El-Machachi, Edvin Fako, Andrea C. Ferrari, Annalena Genreith-Schriever, Janine George, Rhys E. A. Goodall, Clare P. Grey, Shuang Han, Will Handley, Hendrik H. Heenen, Kersti Hermansson, Christian Holm, Jad Jaafar, Stephan Hofmann, Konstantin S. Jakob, Hyunwook Jung, Venkat Kapil, Aaron D. Kaplan, Nima Karimitari, Namu Kroupa, Jolla Kullgren, Matthew C. Kuner, Domantas Kuryla, Guoda Liepuoniute, Johannes T. Margraf, Ioan-Bogdan Magdău, Angelos Michaelides, J. Harry Moore, Aakash A. Naik, Samuel P. Niblett, Sam Walton Norwood, Niamh O’Neill, Christoph Ortner, Kristin A. Persson, Karsten Reuter, Andrew S. Rosen, Lars L. Schaaf, Christoph Schran, Eric Sivonxay, Tamás K. Stenczel, Viktor Svahn, Christopher Sutton, Cas van der Oord, Eszter Varga-Umbrich, Tejs Vegge, Martin Vondrák, Yangshuai Wang, William C. Witt, Fabian Zills, and Gábor Csányi. A foundation model for atomistic materials chemistry, 2023. ↩ a
- Batzner et al. (2022) Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nat. Commun., 13(1):2453, May 2022. ISSN 2041-1723. ↩ a
- Chan et al. (2024) Bill WGL Chan, Nicholas B Lynch, Wendy Tran, Jack M Joyce, G Paul Savage, Wim Meutermans, Andrew P Montgomery, and Michael Kassiou. Fragment-based drug discovery for disorders of the central nervous system: designing better drugs piece by piece. Frontiers in Chemistry, 12:1379518, 2024. ↩ a
- Chithrananda et al. (2020) Seyone Chithrananda, Gabriel Grand, and Bharath Ramsundar. Chemberta: Large-scale self-supervised pretraining for molecular property prediction, 2020. ↩ a
- Choudhary et al. (2022) Kamal Choudhary, Brian DeCost, Chi Chen, Anubhav Jain, Francesca Tavazza, Ryan Cohn, Cheol Woo Park, Alok Choudhary, Ankit Agrawal, Simon J. L. Billinge, Elizabeth Holm, Shyue Ping Ong, and Chris Wolverton. Recent advances and applications of deep learning methods in materials science. npj Comput. Mater., 8(1):59, Apr 2022. ISSN 2057-3960. ↩ a
- Chung et al. (2019) Yongchul G Chung, Emmanuel Haldoupis, Benjamin J Bucior, Maciej Haranczyk, Seulchan Lee, Hongda Zhang, Konstantinos D Vogiatzis, Marija Milisavljevic, Sanliang Ling, Jeffrey S Camp, et al. Advances, updates, and analytics for the computation-ready, experimental metal–organic framework database: Core mof 2019. Journal of Chemical & Engineering Data, 64(12):5985–5998, 2019. ↩ a
- Cormen et al. (2022) Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022. ↩ a
- Courty et al. (2017) Nicolas Courty, Rémi Flamary, and Mélanie Ducoffe. Learning wasserstein embeddings. International Conference on Learning Representations (ICLR), 2017. ↩ a
- Cuturi (2013) Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013. ↩ a
- Cuturi et al. (2020) Marco Cuturi, Olivier Teboul, Jonathan Niles-Weed, and Jean-Philippe Vert. Supervised quantile normalization for low rank matrix factorization. In International Conference on Machine Learning, pp. 2269–2279. PMLR, 2020. ↩ a
- Degen et al. (2008) Jorg Degen, Christof Wegscheid-Gerlach, Andrea Zaliani, and Matthias Rarey. On the art of compiling and using’drug-like’chemical fragment spaces. ChemMedChem, 3(10):1503, 2008. ↩ a
- Du et al. (2022) Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie-Yan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning, pp. 5583–5608. PMLR, 2022. ↩ a
- Du et al. (2023) Yuanqi Du, Limei Wang, Dieqiao Feng, Guifeng Wang, Shuiwang Ji, Carla P Gomes, Zhi-Ming Ma, et al. A new perspective on building efficient and expressive 3d equivariant graph neural networks. Advances in neural information processing systems, 36:66647–66674, 2023. ↩ a
- Fang et al. (2021) Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, and Haifeng Wang. Chemrl-gem: Geometry enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 2021. doi: 10.48550/ARXIV.2106.06130. ↩ a
- Fang et al. (2022) Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, and Haifeng Wang. Geometry-enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 4(2):127–134, 2022. ↩ a
- Fedik et al. (2022) Nikita Fedik, Roman Zubatyuk, Maksim Kulichenko, Nicholas Lubbers, Justin S. Smith, Benjamin Nebgen, Richard Messerly, Ying Wai Li, Alexander I. Boldyrev, Kipton Barros, Olexandr Isayev, and Sergei Tretiak. Extending machine learning beyond interatomic potentials for predicting molecular properties. Nat. Rev. Chem., 6(9):653–672, Sep 2022. ISSN 2397-3358. doi: 10.1038/s41570-022-00416-3. ↩ a
- Feydy et al. (2019) Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference on aRtIfIcIaL InTeLlIgEnCe and Statistics, pp. 2681–2690. PMLR, 2019. ↩ a
- Gasteiger et al. (2021) Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790–6802, 2021. ↩ a
- Genevay et al. (2018) Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608–1617. PMLR, 2018. ↩ a
- Geng et al. (2023) Zijie Geng, Shufang Xie, Yingce Xia, Lijun Wu, Tao Qin, Jie Wang, Yongdong Zhang, Feng Wu, and Tie-Yan Liu. De novo molecular generation via connection-aware motif mining. arXiv preprint arXiv:2302.01129, 2023. ↩ a b
- Gilmer et al. (2017) Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp. 1263–1272, 2017. ↩ a
- Gower (1985) J.C. Gower. Properties of Euclidean and non-Euclidean distance matrices. Linear Algebra and its Applications, 67:81–97, June 1985. ISSN 0024-3795. doi: 10.1016/0024-3795(85)90187-9. ↩ a
- Haviv et al. (2024) Doron Haviv, Russell Zhang Kunes, Thomas Dougherty, Cassandra Burdziak, Tal Nawy, Anna Gilbert, and Dana Pe’Er. Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformer. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 17697–17718. PMLR, July 2024. ↩ a b c
- Hawkins (2017) Paul CD Hawkins. Conformation generation: the state of the art. Journal of chemical information and modeling, 57(8):1747–1756, 2017. ↩ a
- Hetzel et al. (2023) Leon Hetzel, Johanna Sommer, Bastian Rieck, Fabian Theis, and Stephan Günnemann. MAGNet: Motif-agnostic generation of molecules from shapes. arXiv preprint arXiv:2305.19303, 2023. ↩ a
- Higham (1988) Nicholas J. Higham. Computing a nearest symmetric positive semidefinite matrix. Linear Algebra and its Applications, 103:103–118, May 1988. ISSN 0024-3795. doi: 10.1016/0024-3795(88)90223-6. ↩ a
- Hu et al. (2020a) Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118–22133, 2020a. ↩ a
- Hu et al. (2020b) Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. In International Conference on Learning Representations, 2020b. ↩ a
- Klicpera et al. (2020) Johannes Klicpera, Janek Groß, Stephan Günnemann, et al. Directional message passing for molecular graphs. In International Conference on Learning Representations, pp. 1–13, 2020. ↩ a
- Kong et al. (2022) Xiangzhe Kong, Wenbing Huang, Zhixing Tan, and Yang Liu. Molecule generation by principal subgraph mining and assembling. Advances in Neural Information Processing Systems, 35:2550–2563, 2022. ↩ a b
- Kreuzer et al. (2021) Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Létourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34:21618–21629, 2021. ↩ a
- Landrum (2016) G Landrum. Rdkit: open-source cheminformatics http://www.rdkit.org, 2016. ↩ a
- Lee et al. (2024) Seul Lee, Seanie Lee, Kenji Kawaguchi, and Sung Ju Hwang. Drug discovery with dynamic goal-aware fragments. International Conference on Machine Learning, 2024. ↩ a
- Liu et al. (2021) Meng Liu, Cong Fu, Xuan Zhang, Limei Wang, Yaochen Xie, Hao Yuan, Youzhi Luo, Zhao Xu, Shenglong Xu, and Shuiwang Ji. Fast quantum property prediction via deeper 2d and 3d graph networks. arXiv preprint arXiv:2106.08551, 2021. ↩ a
- Liu et al. (2015) Zhihai Liu, Yan Li, Li Han, Jie Li, Jie Liu, Zhixiong Zhao, Wei Nie, Yuchen Liu, and Renxiao Wang. Pdb-wide collection of binding data: current status of the pdbbind database. Bioinformatics, 31(3):405–412, 2015. ↩ a
- Luo et al. (2024) Yuankai Luo, Hongkang Li, Lei Shi, and Xiao-Ming Wu. Enhancing graph transformers with hierarchical distance structural encoding. Advances in Neural Information Processing Systems, 37:57150–57182, 2024. ↩ a
- Ma et al. (2023) Xinyu Ma, Xu Chu, Yasha Wang, Yang Lin, Junfeng Zhao, Liantao Ma, and Wenwu Zhu. Fused gromov-wasserstein graph mixup for graph-level classifications. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. ↩ a b c
- Medrano Sandonas et al. (2024) Leonardo Medrano Sandonas, Dries Van Rompaey, Alessio Fallani, Mathias Hilfiker, David Hahn, Laura Perez-Benito, Jonas Verhoeven, Gary Tresadern, Joerg Kurt Wegner, Hugo Ceulemans, et al. Dataset for quantum-mechanical exploration of conformers and solvent effects in large drug-like molecules. Scientific Data, 11(1):742, 2024. ↩ a
- Merlot et al. (2003) Cédric Merlot, Daniel Domine, Christophe Cleva, and Dennis J Church. Chemical substructures in drug discovery. Drug Discovery Today, 8(13):594–602, 2003. ↩ a
- Morgan (1965) H. L. Morgan. The generation of a unique machine description for chemical structures-a technique developed at chemical abstracts service. Journal of Chemical Documentation, 5(2):107–113, May 1965. ISSN 1541-5732. doi: 10.1021/c160017a018. ↩ a
- Nguyen et al. (2024a) Duy MH Nguyen, Nina Lukashina, Tai Nguyen, An T Le, TrungTin Nguyen, Nhat Ho, Jan Peters, Daniel Sonntag, Viktor Zaverkin, and Mathias Niepert. Structure-aware e (3)-invariant molecular conformer aggregation networks. International Conference on Machine Learning, 2024a. ↩ a b c d
- Nguyen et al. (2024b) Duy Minh Ho Nguyen, Nina Lukashina, Tai Nguyen, An Thai Le, Trungtin Nguyen, Nhat Ho, Jan Peters, Daniel Sonntag, Viktor Zaverkin, and Mathias Niepert. Structure-aware E(3)-invariant molecular conformer aggregation networks. In International Conference on Machine Learning, pp. 37736–37760. PMLR, 2024b. ↩ a b c d e f g
- Otsuka et al. (2011) Shingo Otsuka, Isao Kuwajima, Junko Hosoya, Yibin Xu, and Masayoshi Yamazaki. Polyinfo: Polymer database for polymeric materials design. In 2011 International Conference on Emerging Intelligent Data and Web Technologies, pp. 22–29. IEEE, 2011. ↩ a
- Perola & Charifson (2004) Emanuele Perola and Paul S Charifson. Conformational analysis of drug-like molecules bound to proteins: an extensive study of ligand reorganization upon binding. Journal of medicinal chemistry, 47(10):2499–2510, 2004. ↩ a
- Peyré et al. (2016) Gabriel Peyré, Marco Cuturi, and Justin Solomon. Gromov-Wasserstein Averaging of Kernel and Distance Matrices. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 2664–2672, New York, New York, USA, June 2016. PMLR. ↩ a
- Peyré et al. (2019) Gabriel Peyré, Marco Cuturi, and others. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 11(5-6):355–607, 2019. Publisher: Now Publishers, Inc. ↩ a
- Rampášek et al. (2022) Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35:14501–14515, 2022. ↩ a
- Ramsundar et al. (2019) Bharath Ramsundar, Peter Eastman, Pat Walters, and Vijay Pande. Deep learning for the life sciences: applying deep learning to genomics, microscopy, drug discovery, and more. O’Reilly Media, 2019. ↩ a
- Raza et al. (2022) Ali Raza, E Adrian Henle, and Xiaoli Fern. Non-equilibrium molecular geometries in graph neural networks. arXiv preprint arXiv:2203.04697, 2022. ↩ a
- Rong et al. (2020) Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems, 33:12559–12571, 2020. ↩ a b c
- Ross et al. (2022) Jerret Ross, Brian Belgodere, Vijil Chenthamarakshan, Inkit Padhi, Youssef Mroueh, and Payel Das. Large-scale chemical language representations capture molecular structure and properties. Nature Machine Intelligence, 4(12):1256–1264, 2022. ↩ a
- Scarselli et al. (2008) Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE transactions on neural networks, 20(1):61–80, 2008. ↩ a
- Scetbon et al. (2021) Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pp. 9344–9354. PMLR, 2021. ↩ a
- Schütt et al. (2017) Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 991–1001, 2017. ↩ a
- Schütt et al. (2017) Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. ↩ a b c
- Schütt et al. (2021) 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, pp. 9377–9388. PMLR, 2021. ↩ a
- Solomon et al. (2015) Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1–11, 2015. ↩ a
- Sommer et al. (2023) Johanna Sommer, Leon Hetzel, David Lüdke, Fabian Theis, and Stephan Günnemann. The power of motifs as inductive bias for learning molecular distributions. arXiv preprint arXiv:2306.17246, 2023. ↩ a
- Sonthalia et al. (2021) Rishi Sonthalia, Greg Van Buskirk, Benjamin Raichel, and Anna Gilbert. How can classical multidimensional scaling go wrong? In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 12304–12315. Curran Associates, Inc., 2021. ↩ a
- Titouan et al. (2019) Vayer Titouan, Nicolas Courty, Romain Tavenard, Chapel Laetitia, and Rémi Flamary. Optimal Transport for structured data with application on graphs. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 6275–6284. PMLR, June 2019. ↩ a
- Titouan et al. (2020) Vayer Titouan, Laetitia Chapel, Remi Flamary, Romain Tavenard, and Nicolas Courty. Fused Gromov-Wasserstein Distance for Structured Objects. Algorithms, 13(9):212, August 2020. ISSN 1999-4893. doi: 10.3390/a13090212. ↩ a
- Tong et al. (2021) Alexander Y Tong, Guillaume Huguet, Amine Natik, Kincaid MacDonald, Manik Kuchroo, Ronald Coifman, Guy Wolf, and Smita Krishnaswamy. Diffusion earth mover’s distance and distribution embeddings. In International Conference on Machine Learning, pp. 10336–10346. PMLR, 2021. ↩ a
- Torgerson (1952) Warren S. Torgerson. Multidimensional Scaling: I. Theory and Method. Psychometrika, 17(4):401–419, 1952. ISSN 0033-3123. doi: 10.1007/BF02288916. Edition: 2025/01/01 Publisher: Cambridge University Press & Assessment. ↩ a
- Varnek et al. (2005) Alexandre Varnek, Denis Fourches, Frank Hoonakker, and Vitaly P Solov’ev. Substructural fragments: an universal language to encode reactions, molecular and supramolecular structures. Journal of computer-aided molecular design, 19:693–703, 2005. ↩ a
- Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. ↩ a b c
- Veličković et al. (2017) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.10903, 2017. ↩ a
- Veličković et al. (2018) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations, 2018. ↩ a
- Veličković et al. (2018) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In ICLR, 2018. ↩ a
- Wang et al. (2025) Jiaxi Wang, Yaosen Min, Miao Li, and Ji Wu. Fragformer: A fragment-based representation learning framework for molecular property prediction. Transactions on Machine Learning Research, 2025. ↩ a
- Wang et al. (2019) Sheng Wang, Yuzhi Guo, Yuhong Wang, Hongmao Sun, and Junzhou Huang. Smiles-bert: Large scale unsupervised pre-training for molecular property prediction. In Proceedings of the 10th ACM International Conference on Bioinformatics, Computational Biology and Health Informatics, BCB ’19, pp. 429–436, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450366663. doi: 10.1145/3307339.3342186. ↩ a
- Wang et al. (2024) Zeyu Wang, Tianyi Jiang, Jinhuan Wang, and Qi Xuan. Multi-modal representation learning for molecular property prediction: Sequence, graph, geometry. arXiv preprint arXiv:2401.03369, 2024. ↩ a
- Wen et al. (2024) Qianlong Wen, Mingxuan Ju, Zhongyu Ouyang, Chuxu Zhang, and Yanfang Ye. From coarse to fine: enable comprehensive graph self-supervised learning with multi-granular semantic ensemble. In Forty-first International Conference on Machine Learning, 2024. ↩ a
- Wollschläger et al. (2024) Tom Wollschläger, Niklas Kemper, Leon Hetzel, Johanna Sommer, and Stephan Günnemann. Expressivity and generalization: Fragment-biases for molecular gnns. International Conference on Machine Learning, 2024. ↩ a b c d
- Wu et al. (2018) Z. Wu, B. Ramsundar, E. N. Feinberg, J. Gomes, C. Geniesse, A. S. Pappu, K. Leswing, and V. Pande. MoleculeNet: A benchmark for molecular machine learning. Chemical Science, pp. 513–530, 2018. ↩ a b
- Xiong et al. (2019) Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li, Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism. Journal of medicinal chemistry, 63(16):8749–8760, 2019. ↩ a
- Xu et al. (2018) Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning on graphs with jumping knowledge networks. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 5453–5462. PMLR, 10–15 Jul 2018. ↩ a
- Xu et al. (2019) Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations, 2019. ↩ a
- Yang et al. (2019a) Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman-Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, Andrew Palmer, Volker Settels, Tommi Jaakkola, Klavs Jensen, and Regina Barzilay. Analyzing learned molecular representations for property prediction. Journal of Chemical Information and Modeling, 59(8):3370–3388, July 2019a. ISSN 1549-960X. doi: 10.1021/acs.jcim.9b00237. ↩ a b
- Yang et al. (2019b) Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman-Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8):3370–3388, 2019b. ↩ a
- Ying et al. (2021) Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877–28888, 2021. ↩ a b c d e f
- Zaheer et al. (2017) Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017. ↩ a b
- Zhang et al. (2021) Zaixi Zhang, Qi Liu, Hao Wang, Chengqiang Lu, and Chee-Kong Lee. Motif-based graph self-supervised learning for molecular property prediction. Advances in Neural Information Processing Systems, 34:15870–15882, 2021. ↩ a
- Zhou et al. (2023) Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023. ↩ a b c
- Zhu et al. (2024a) Y Zhu, J Hwang, K Adams, Z Liu, B Nan, B Stenfors, Y Du, J Chauhan, O Wiest, O Isayev, et al. Learning over molecular conformer ensembles: Datasets and benchmarks. In International Conference on Learning Representations, 2024a. ↩ a b c
- Zhu et al. (2023) Yanqiao Zhu, Jeehyun Hwang, Keir Adams, Zhen Liu, Bozhao Nan, Brock Stenfors, Yuanqi Du, Jatin Chauhan, Olaf Wiest, Olexandr Isayev, Connor W. Coley, Yizhou Sun, and Wei Wang. Learning over molecular conformer ensembles: Datasets and benchmarks, 2023. ↩ a
- Zhu et al. (2024b) Yanqiao Zhu, Jeehyun Hwang, Keir Adams, Zhen Liu, Bozhao Nan, Brock Stenfors, Yuanqi Du, Jatin Chauhan, Olaf Wiest, Olexandr Isayev, et al. Learning over molecular conformer ensembles: Datasets and benchmarks. International Conference on Learning Representations (ICLR), 2024b. ↩ a b
Appendix A Implementation Details
Our training pipeline includes three stages: In the first stage, we train only the 2D and 3D MPNNs to learn corresponding features from 2D molecular graph and 3D conformers. The features in this stage also serve as a dataset for approximating Graph Transformer to the FGW distance. In the next stage, the Graph Transformer is trained separately to simulate the costly computation of FGW distance between two conformers using learned features from stage 1. In the last stage, Graph Transformer is integrated in a single end-to-end training with 2D and 3D MPNNs. At this stage, only 2D and 3D MPNNs are trained. As a result of changing MPNNs during the last stage, a shift in the distribution of the Graph Transformer input might occur. We solve this problem by adding an adaptor layer using an MLP on both 3D and 2D features before feeding them to the GraphTransformer. For all experiments on the MoleculeNet and MARCEL benchmarks, we use the same number of conformers as specified in their original settings.
In all stages, we use Adam as our optimizer. We train our model on an 8 V100-GPUs cluster.
Stage 1. Learning 2D and 3D features.
For each molecule, we define by , we then train for 150 epochs and set the learning rate to . to optimize target property tasks where be the ground-truth value and be our predicted value defined by:
| (10) |
with and are learnable parameters and is number of conformers.
Stage 2. Training Graph Transformer to approximate FGW distance.
The Graph Transformer is trained separately in the second stage to approximate the FGW distance by Euclidean embedding space. For the Graph Transformer architecture, we employ the same setting as Graphormer from Ying et al. (2021). Specifically, a number of attention layers, a number of attention heads, and the hidden dimension of the transformer are set to 12, 8, and 64, respectively, which makes the total number of parameters of the Graph Transformer 372k. In our attention, we use the shortest-path distance (SPD) between a pair of nodes. Following practical implementation in Ying et al. (2021), we pre-compute SPD distance for each 3D molecule graph and load these values during training and inference. We set a learning rate of and train for 1000 epochs with the following loss function:
| (11) |
Stage 3. Training Fragment-aware Graph Transformer. In the final stage, we freeze the trained GraphTransformer and use it to compute aggregated features from 3D conformer embeddings generated by the 3D-MPNN. To accommodate potential distribution shifts, we add lightweight FFN adaptor layers on top of both the 2D- and 3D-MPNNs used in , while continuing to update the MPNNs during training. The full model is trained for 300 epochs with a reduced learning rate to optimize the training loss where
| (12) |
is final atom-wise embedding.
Appendix B Further Visualization Fragment Outputs
Fragment Generation Algorithms.
We use a structural fragmentation method based on RingPath algorithms (Kong et al., 2022; Geng et al., 2023; Wollschläger et al., 2024) that decompose a molecular graph , where denotes atoms and denotes covalent bonds, into a set of chemically interpretable fragments. The fragmentation process identifies a set of ring fragments using RDKit’s cycle basis algorithm (SSSR), where each ring is encoded by its atom indices and size class.
Next, all bonds not part of any ring are grouped into acyclic path fragments , where each is a linear chain of nodes, extracted via depth-first search under a degree constraint. Each fragment is assigned a type (representing ring, path, or junction) and a type index within a fixed vocabulary of size . Fragments whose sizes exceed a predefined threshold are mapped to the final index of their category to preserve bounded dimensionality.
We define a fragment-atom incidence matrix , where if atom belongs to fragment . From this, we derive a fragment-level graph , where each node represents a molecular fragment and an edge is added if two fragments share at least one atom or are directly bonded.
Compared to traditional fragmentation algorithms like BRICS (Degen et al., 2008), BBB (Sommer et al., 2023), or MagNet (Hetzel et al., 2023), the RingPath algorithm offers a more topology-aware decomposition by explicitly capturing key structural motifs such as rings, paths, and linkers. While BRICS and BBB often generate chemically meaningful fragments based on retrosynthetic rules, they may overlook contextual connectivity critical for graph-based learning. In contrast, RingPath preserves the relational structure between fragments, aligning closely with how molecules are built and understood in topological space—making it particularly beneficial for tasks requiring structural interpretability and generalization in graph neural networks. The advantages of RingPath have also been empirically validated in recent studies, demonstrating improved performance across various molecular property prediction benchmarks.
Visualization of Typical Extracted Fragment Graphs.
Figures 5 and 6 illustrate representative examples of fragment extraction using the RingPath algorithm on the Kraken and Drug-75k datasets. The top row displays the original 2D molecular structures, while the bottom row shows the corresponding RingPath decompositions. Each colored region highlights a distinct structural fragment, such as a ring or path, demonstrating the algorithm’s ability to segment complex molecules into chemically meaningful and interpretable components.
Appendix C Additional Analysis of FACET’s Scalability and Performance with More 3D Conformers
In this section, we further analyze FACET’s scalability on the following two factors:
C.1 Inference Time when Increasing the Number of 3D Conformers for Each Molecule.
We compare FACET against two versions of ConAN-FGW in running time to extract structure-aware embedding aggregation with different input of 3D conformers. We use two variations of ConAN-FGW, including a single GPU version and another relaxed solver that permits running Sinkhorn iterations on GPUs by matrix multiplication, thus supporting distributed multi-GPUs acceleration. The experiments are conducted on a single GPU using a batch size of molecules, each with different conformers ranging from , , , , and , and another experiment with four GPUs on the same batch size, i.e., 8 molecules per GPU.
Figure 7 indicates our observations across four datasets of MoleculeNet benchmark, where we report the required time to extract embedding aggregations for all molecules in the test set. We see that (i) FACET demonstrates excellent scalability where its runtime remains nearly constant regardless of the number of conformers, both in single-GPU and multi-GPU settings. In contrast, ConAN-FGW shows poor scalability where runtime increases steeply with the number of conformers. While the multi-GPU usage improves runtime over single-GPU, the growth trend remains significant, with runtimes still exceeding 30 seconds at 20 conformers (e.g., with ESOL dataset).
Secondly, the nearly identical runtime of FACET across single- and multi-GPU settings, as shown in the plot, can be attributed to its computational efficiency and the relatively small workload in this experiment. In such cases, the overhead introduced by multi-GPU parallelization - such as inter-GPU communication and data synchronization - can outweigh its potential speedup benefits. Therefore, we argue that multi-GPU acceleration for FACET becomes advantageous only under substantially larger workloads, such as batch processing of thousands to millions of molecules or handling complex input representations that exceed the memory capacity of a single GPU.
C.2 Average Training Time per Epoch as a Function of Dataset Size.
We analyze the scalability of FACET with respect to the number of training molecules. To this end, we report the average training time per epoch across four datasets from the MoleculeNet benchmark. Figure 8 compares the training time of FACET and ConAN-FGW on a single GPU, using a batch size of 256 and 5 conformers per molecule. As shown in the figure, FACET achieves a 2.28× to 3.17× speedup over ConAN-FGW. Notably, this speedup is roughly proportional to the number of training molecules in each dataset, as reported in Table 3.
C.3 Ablation Study on the Impact of Increasing the Number of 3D Conformers in FACET
We provide below a comprehensive ablation study on the impact of using an increasing number of RDKit-generated conformers across four datasets (ESOL, FreeSolv, BACE, and Lipo).
As shown in Table 6, we observe a consistent trend across datasets: increasing the number of conformers from 3 to 5 leads to improved regression performance (lower values indicate better results). However, beyond 5 conformers, the performance tends to converge or slightly fluctuate, confirming that our geometry-aware embedding approach using the FGW distance provides stable and reliable approximations. This aligns with the theoretical expectation that the approximation error scales with , where is the number of conformers used.
| Settings | 3 conf. | 5 conf. (default) | 10 conf. | 15 conf. | 20 conf. |
|---|---|---|---|---|---|
| ESOL | 0.539 ± 0.06 | 0.516 ± 0.04 | 0.501 ± 0.02 | 0.511 ± 0.03 | 0.546 ± 0.02 |
| FreeSolv | 0.977 ± 0.25 | 0.967 ± 0.08 | 0.933 ± 0.23 | 0.946 ± 0.24 | 0.949 ± 0.21 |
| BACE | 0.542 ± 0.05 | 0.495 ± 0.03 | 0.513 ± 0.02 | 0.519 ± 0.01 | 0.517 ± 0.03 |
| Lipo | 0.445 ± 0.02 | 0.424 ± 0.01 | 0.444 ± 0.02 | 0.447 ± 0.08 | 0.445 ± 0.01 |
When 3D conformers generated by RDKit are not used, our FACET model simplifies significantly. In this configuration, the model only receives 2D molecular graphs along with fragment-level information, and key components such as the Graph Transformer are removed. Table 7 presents the performance comparison between the full FACET model and its 2D-only variant across four benchmark datasets:
| Method | ESOL(↓) | FreeSolv(↓) | BACE(↓) | Lipo(↓) |
|---|---|---|---|---|
| FACET | 0.516 ± 0.04 | 0.967 ± 0.08 | 0.495 ± 0.03 | 0.424 ± 0.01 |
| w/o 3D conformers | 0.546 ± 0.03 | 1.197 ± 0.09 | 0.584 ± 0.03 | 0.543 ± 0.02 |
These results clearly demonstrate that incorporating 3D conformers, even those generated by RDKit, is critical to the expressiveness and performance of FACET. The full model consistently outperforms its 2D-only counterpart, highlighting the importance of 3D geometry in learning accurate molecular representations.
Appendix D Comparison of Training Time between FACET and CoNan-FGW
To provide a comprehensive comparison, we conducted additional experiments to compare the training time of FACET and CoNAN-FGW, with the addition of GemNet, a strong state-of-the-art 3D molecular model, on two benchmark datasets: BACE (1,059 molecules) and LIPO (2,940 molecules). All models were trained for 200 epochs under the same settings. Since both FACET and CoNAN-FGW are originally built on the SchNet architecture, which is generally less expressive than GemNet, we also report the performance of FACET when upgraded to use GemNet as its backbone. From the results listed in Table 8, we have the following key observations:
-
•
FACET vs. CoNAN-FGW: FACET consistently shows reduced training time compared to CoNAN-FGW, though the degree of reduction varies by dataset size.
-
–
On BACE: the time savings are marginal due to the additional cost introduced by the Graph Transformer component in FACET, which is trained using the pre-computed FGW distances from the optimal transport solver.
-
–
On LIPO: the training time reduction is more substantial. This is because CoNAN-FGW incurs a high computational cost from directly computing FGW distances between sets of 3D conformers in every forward pass. In contrast, FACET leverages pre-learned geometry-aware embeddings, where the corresponding operation reduces to a lightweight matrix multiplication in the Graph Transformer.
-
–
-
•
FACET vs. GemNet: FACET represents a balanced trade-off between CoNAN-FGW and GemNet in terms of training time. Despite using the simpler SchNet backbone, FACET achieves competitive, sometimes better, performance compared to GemNet, thanks to its geometry-aware aggregation via FGW-based embeddings. This efficiency stems from replacing costly pairwise conformer comparisons with a latent-space transformer that captures 3D geometric information in a more scalable manner.
-
•
FACET (GemNet) vs. GemNet: When both models share the same GemNet architecture, FACET outperforms GemNet in terms of predictive accuracy on both datasets. We observe that (i) the additional training time incurred by FACET is relatively modest: approximately +21% on BACE and +33% on LIPO, and (ii) given the performance gains, this extra time remains acceptable in practical scenarios and demonstrates FACET’s scalability and effectiveness.
| Model | Metric | BACE | LIPO |
|---|---|---|---|
| GemNet | MSE | 0.51 ± 0.07 | 0.45 ± 0.01 |
| Time | 2.04 hours | 4.8 hours | |
| Conan-FGW (SchNet) | MSE | 0.55 ± 0.02 | 0.42 ± 0.02 |
| Time | 2.5 hours | 6.3 hours | |
| FACET (SchNet) | MSE | 0.50 ± 0.03 | 0.42 ± 0.01 |
| Time | 2.3 hours | 5.05 hours | |
| FACET (GemNet) | MSE | 0.46 ± 0.03 | 0.39 ± 0.02 |
| Time | 2.47 hours | 6.4 hours |
Appendix E Performance of FACET and CoNan-FGW on MARCEL benchmark
To provide a meaningful comparison, we benchmarked FACET against CoNAN-FGW on 10% of the Drug-75k dataset and on the Kraken dataset, which serve as representative subsets. The results (provided below) show that FACET performs competitively or outperforms CoNAN-FGW, even under these reduced-scale settings, reinforcing the efficiency and effectiveness of our approach. The results are shown in Tables 10 and 11.
Appendix F Comparisons with SOTA methods in 2D (or 3D)
FACET is designed as a modular framework for enhancing molecular property prediction by integrating structure-aware aggregation over multiple conformers. A central strength of this design is that it can be plugged into a variety of existing backbone architectures, whether 2D or 3D, thus offering a complementary mechanism rather than an alternative to these models.
FACET improves standalone 3D architectures
We integrated FACET with established 3D models such as SchNet, GemNet, and VisNet, and consistently observed performance improvements across datasets. Table 9 demonstrates that FACET’s geometry-aware aggregation over multiple conformers complements even strong 3D baselines, validating its utility beyond what these models achieve on their own.
| Model | BACE(↓) | LIPO(↓) |
|---|---|---|
| SchNet | 0.64 ± 0.05 | 0.56 ± 0.01 |
| FACET (SchNet) | 0.50 ± 0.03 | 0.42 ± 0.01 |
| GemNet | 0.51 ± 0.07 | 0.45 ± 0.01 |
| FACET (GemNet) | 0.46 ± 0.03 | 0.39 ± 0.02 |
| VisNet | 0.61 ± 0.15 | 0.55 ± 0.45 |
| FACET (VisNet) | 0.47 ± 0.01 | 0.48 ± 0.01 |
FACET enhances simple 2D MPNNs
We also applied FACET to a lightweight 2D message-passing neural network and found that incorporating FACET’s fragment-level structure-aware aggregation significantly improved performance. This result underscores the compatibility of FACET with 2D backbones and its ability to enhance models that do not explicitly process 3D information.
Appendix G Unified training pipeline
We investigated the performance of the proposed method when combining all training steps into an end-to-end pipeline. Below, we summarize our findings step by step:
-
•
Step 1 – Pretraining 2D and 3D MPNNs: As suggested in prior work like CoNAN-FGW, we begin by pretraining the 2D and 3D MPNNs independently. This initial phase is critical to ensure that the encoders, especially the 3D MPNN, converge to a stable and meaningful representation before introducing structure-aware aggregation. To test the necessity of this stage, we experimented with a variant where all three stages were co-trained from scratch. The results showed substantially lower performance, confirming that Stage 1 is crucial for learning rich, aligned, and stable representations.
-
•
Steps 2 and 3 – Co-training Graph Transformer and Downstream Fine-tuning: While our default setup trains Step 2 (Graph Transformer with FGW supervision) and Step 3 (fine-tuning on molecular properties) sequentially, we explored an alternative setup where both steps are co-trained. To manage the computational cost of FGW supervision, we adopted an alternating strategy: after every five steps of property prediction optimization, we update the Graph Transformer to approximate FGW distances. This reduces the training burden compared to full FGW supervision at every iteration (as in CoNAN-FGW).
| L | BurL | B5 | BurB5 | |
|---|---|---|---|---|
| Conan-FGW (SchNet) | 0.397 | 0.117 | 0.272 | 0.195 |
| FACET (SchNet) | 0.398 | 0.125 | 0.251 | 0.180 |
| IP | EA | ||
|---|---|---|---|
| Conan-FGW (SchNet) | 0.374 | 0.541 | 0.587 |
| FACET (SchNet) | 0.365 | 0.535 | 0.552 |
As shown in Table 12, without separately training Step 1, the model got low performance, confirming that this stage helps the model ensure rich, aligned, and stable molecular representations before incorporating more advanced structure awareness. Secondly, on four MoleculeNet datasets, co-training Steps 2 and 3 produced slightly improved performance over the default FACET setup. For example, on ESOL, performance improved from 0.505 to 0.516, and on FreeSolv, from 0.867 to 0.967. This improvement can be attributed to the model’s ability to jointly adapt the 2D/3D encoders and the Graph Transformer, leading to more aligned, task-relevant representations. However, there is a trade-off. This co-training strategy comes with an increased training cost, as FGW distances must still be computed periodically. As a result, while training is slower than the default FACET setup, it remains significantly faster than CoNAN-FGW, and achieves a strong balance between efficiency and predictive performance, especially on large-scale datasets like Drug-75k.
| ESOL(↓) | FreeSolv(↓) | BACE(↓) | Lipo(↓) | |
|---|---|---|---|---|
| ConAN-FGW | 0.529 ± 0.022 | 1.068 ± 0.083 | 0.549 ± 0.016 | 0.422 ± 0.016 |
| FACET | 0.516 ± 0.044 | 0.967 ± 0.082 | 0.495 ± 0.115 | 0.424 ± 0.009 |
| FACET (Merge all steps) | 0.567 ± 0.023 | 1.264 ± 0.094 | 0.591 ± 0.062 | 0.530 ± 0.013 |
| FACET (Merge steps 2-3) | 0.505 ± 0.014 | 0.867 ± 0.102 | 0.497 ± 0.035 | 0.44 0± 0.014 |
Appendix H Limitations of FACET
H.1 FACET Operates on a Predefined Set of 3D Conformers.
Our method enables efficient geometry-aware aggregation without requiring expensive alignment procedures at inference time. While FACET demonstrates improved performance even with a small subset of conformers, the quality and representativeness of this subset can still influence downstream predictions. In particular, if the selected conformers are heavily biased or fail to capture key structural variations, some aspects of molecular flexibility may be underrepresented. Addressing this challenge through better conformer sampling strategies or task-aware selection mechanisms could further enhance model robustness, especially for highly flexible molecules.
Future direction: A promising extension would be to develop end-to-end models that can learn to generate conformers dynamically during training, using gradient feedback from downstream prediction losses. Such a differentiable conformer generation module could enable task-aware structural modeling, ensuring that the generated conformers are optimized not just for physical plausibility, but also for relevance to the predictive task at hand.
H.2 Limitations in Scope: Focus on Small Molecules
FACET has primarily been evaluated on standard molecular property prediction benchmarks such as those in MoleculeNet, which consist mostly of small, drug-like molecules. While this setup is well-suited for many pharmacological applications, it limits the assessment of FACET’s generalizability to more complex molecular systems. For example, biomacromolecules (e.g., peptides, proteins, nucleic acids) exhibit high flexibility, long-range dependencies, and hierarchical organization that are not present in small molecules. Polymers and materials often involve much larger structures without well-defined conformers, challenging FACET’s reliance on discrete 3D inputs. Additionally, FACET currently models only single-molecule properties and has not been extended to multi-molecular interactions, such as protein-ligand binding.
Future direction: To broaden FACET’s applicability, several promising future directions can be explored. First, incorporating efficient attention to capture both local fragment-level information and long-range structural dependencies is essential for handling large biomolecules. Second, adapting FACET to support flexible input formats, such as voxel grids or material-specific graphs, would allow it to process polymers and crystalline materials that lack stable conformers. Third, extending FACET to jointly model molecular interactions through cross-graph attention or co-embedding mechanisms could open applications in drug docking and molecular complex prediction. Finally, applying and evaluating FACET on broader datasets, such as PDBbind (Liu et al., 2015), PolyInfo (Otsuka et al., 2011), or CoRE MOF 2019 (Chung et al., 2019), would provide a more comprehensive understanding of its strengths and limitations across molecular domains.
Appendix I Proof of Theorem 1
Recall that we aim to establish the following novel theoretical bounds: Let denote the pairwise distance matrix, and let represent the eigendecomposition of the associated criterion matrix , where is the centering matrix. The optimal stress value, denoted by , is bounded as follows: , where
Here, denotes the -th component of the -th eigenvector of , and quantifies the approximation error between the empirical barycenter in the Euclidean embedding space and its counterpart in the original space of undirected attributed graphs. This is equivalent to that given , our objective is to derive lower and upper bounds for the following cumulative stress:
| (13) | ||||
| (14) | ||||
| (15) |
To this end, we begin by specifying and formally defining the following important concepts in Section I.1.
I.1 Non-Euclidean Nature of Pairwise FGW Distance Matrix
Definition 1 (Euclidean Distance Matrix).
A distance matrix is said to be Euclidean if there exists a set of points in some Euclidean space such that
The space of all Euclidean distance matrices (EDM) is denoted by .
Fact 1 (Conditions for Euclidean Distance Matrix, see, e.g., Gower (1985)).
A matrix is an EDM if and only if it satisfies the following three conditions:
-
(i)
Non-negativity: for all ,
-
(ii)
Hollow diagonal: for all ,
-
(iii)
Positive semidefiniteness: the associated double-centered matrix is positive semidefinite (PSD), where is the centering matrix, and denotes the -dimensional vector of ones.
Recall that the pairwise FGW distance matrix for a collection of distributions is defined entry-wise by for all , as introduced in Section 3. The following result establishes that this matrix does not correspond to a Euclidean distance matrix:
Lemma 1 (Non-Euclidean Nature of Pairwise FGW Distance Matrix).
Consider the case where . Then the FGW distance matrix , whose entries are given by
with , does not define a Euclidean distance matrix.
As established in Lemma 1, which is proved in Section I.4, the distance is not a Euclidean distance. Therefore, we are interested in quantifying how accurately non-Euclidean distance matrices can be approximated by pairwise distances between learned embeddings. To this end, we analyze the lower and upper bound of the set in Sections I.2 and I.3, respectively.
I.2 Lower Bounds on Embedding non-Euclidean FGW Distances
We would like to find the lower bound of . We note that the original formulation is non-convex, making it analytically intractable. Nonetheless, by reparameterizing the objective as a function of the pairwise squared distances and induced by the embedding, and by incorporating the necessary conditions to ensure that corresponds to a valid Euclidean distance matrix, the reformulated problem becomes convex for . Note that we can prove that has a lower bound at , where is a minimizer of , that is,
| (16) | ||||
| (17) |
Indeed, given the previous reformulation of , we can establish the following lower bound via Proposition 1. Notably, to simplify the problem, in Proposition 1, we relax the EDM constraint by considering , containing by keeping only the PSD property from the EDM definition in Fact 1. We will reintroduce the missing constraints in and use the solution for the simplified problem to construct an upper bound in Section I.3.
Proposition 1 (Error Lower Bound of ).
The lower bound of is provided as follows:
| (18) | ||||
| (19) | ||||
| (20) |
Here contains by keeping only the PSD property from the EDM definition in Fact 1.
Proof of Proposition 1.
Note that if is minimized at , that is,
| (21) |
We then can find the lower bound of via the minimizer .
Using the definition of Frobenius norm and , we can obtain:
We then obtain the following decomposition:
where is the centering matrix and is the all-ones matrix. Indeed, using the definition of the centering matrix , we have .
Here we used the fact that the matrix product is invariant under cyclic permutation:
and
Under only the PSD constraint, the optimal solution that minimizes can be decomposed as:
where and respectively minimize the terms and independently.
In particular, using the definition of the centering matrix , the entries of are given by:
where denotes the mean of the -th row (or column) of , and is the global mean of all elements in . Therefore, the rows/columns mean of equal those of itself, and hence
Therefore,
Here we used the fact that the matrix is given by for all and the empirical FGW barycenter is given by
where denotes the space of attributed graphs with finite -th order FGW distance. To approximate this barycenter in embedding space, we require
where is the mean embedding and is the learned representation.
Now we would like to find a local analytic solution minimizing such that the global solution minimizes both terms and simultaneously. That is,
Here we used the fact that by definition of , it holds that . Hence, the optimization becomes:
This is in fact the problem of computing the nearest PSD approximation to a symmetric matrix . Using the result from Higham (1988), we find the analytic solution as follows:
| (22) |
Here are the eigenvalues and eigenvectors of . Because has rows/columns means , the ones vector is an eigenvector of with eigenvalue . This leads to is also in the null space and:
Therefore,
Combining all together, Proposition 1 is derived as follows:
∎
I.3 Upper Bounds on Embedding of Pairwise Empirical FGW Barycenter Distances
As discussed in Section I.2, the lower bound stated in Proposition 1 is derived by simplifying the problem and relaxing the EDM constraint. Specifically, this relaxation involves considering the set , which contains but retains only the PSD requirement from the EDM characterization given in Fact 1. In Proposition 2, we reintroduce the missing constraints excluded in and leverage the closed-form solution obtained from the relaxed problem to construct an upper bound under the original EDM constraint set .
Proposition 2 (Error Upper Bound of ).
There exists a matrix such that the following upper bounds hold:
| (23) | ||||
| (24) | ||||
| (25) |
where the aggregated error term is defined as:
We aim to exploit the information derived from the truncation of the negative eigenspace of the matrix , specifically the matrix introduced in Eq.(22), defined as:
where denote the eigenvalues and corresponding eigenvectors of the matrix .
Recall that the entries of are given by:
As a consequence, the sum may not be strictly hollow or PSD when is not an EDM. To address this, we seek to construct a symmetric matrix to be added to , resulting in the matrix , which is both hollow and PSD. This adjustment is designed to avoid any additional penalty on the term , though it may introduce some approximation errors in and in the quantity . These induced errors contribute to the upper bound for the optimal score .
We begin with the requirement that the matrix does not contribute any additional penalty to the term . This can be ensured by imposing the constraint . Under this condition, the matrix remains a minimizer of , as demonstrated below:
where the final equality holds due to the constraint .
This leads to the condition , implying that lies in the left null space of , and lies in its right null space. As a result, all rows of must be constant, and this expression can be written as:
where is a column vector to be defined subsequently. Here, we have used the fact that is the centering matrix defined by .
Multiplying both sides on the left by yields:
This leads to
Indeed, via the definition of , we can verify this as follows:
Hence,
Since is a column vector, to satisfy this constraint, must be of the form:
where is a vector of free parameters, and denotes its average. This construction implies that has only degrees of freedom. However, to ensure that has zero diagonal (i.e., the resulting matrix is hollow), the diagonal entries of must satisfy the following linear constraints:
Solving this linear system yields:
where we have used the fact that , and hence its diagonal entries are given by .
Consequently, the resulting matrix can be expressed element-wise as:
where the inequality follows from the fact that is negative semi-definite.
In summary, the matrix satisfies all three constraints specified in Definition 1.
Although preserves the value of , it differs from and introduces approximation errors in the term and the term. Note that the sum of the untruncated version of and the matrix
is equal to and remains hollow. Recall the decomposition:
where is the centering matrix and is the all-ones matrix.
The matrix
can be written similarly to by including the contributions from the negative eigenvalues, resulting in the matrix , parameterized by:
Define the correction due to negative eigenvalues as:
The resulting error in the term is given by:
Furthermore, the contribution to is bounded as:
where the aggregated error term is defined as:
I.4 Proof of Lemma 1
The proof is proved via leveraging Proposition 8.2 from Peyré et al. (2019), applied to the specific case , and relies on the relationships among FGW, Wasserstein (W), and Gromov-Wasserstein (GW) distances.
The FGW cost is defined via two components: the node feature cost matrix , and the structural discrepancy tensor .
Let and be two attributed graphs with and nodes, respectively. Their associated probability measures are
We define the marginals and (and analogously for and ) as projections of and onto the feature and structural spaces, respectively.
Using these definitions, we introduce the following notation:
| (26) | ||||
| (27) | ||||
| (28) | ||||
| (29) |
Let be any admissible coupling. If both and are defined over a common metric space , then the FGW distance is given by:
| (30) |
We now derive the following fundamental identity:
| (31) |
Moreover, let denote the optimal coupling that minimizes the FGW objective . Then the FGW distance admits the following decomposition:
| (32) |
This inequality follows from the optimality of the W and GW distances with respect to the cost functions and , respectively, and highlights the interpolation nature of the FGW distance between these two metrics as governed by the parameter .
The generalized FGW cost admits the following explicit formulation:
Based on the formulation above, we obtain the following upper bound on the FGW distance:
| (33) |
where the second inequality follows from the convexity of the function for and an application of Minkowski-type bounds on the structural term. Importantly, inequality in Eq.(33) holds for any admissible coupling , and in particular, it remains valid when , the optimal coupling associated with the Wasserstein distance over the product metric space . Here, the effective distance between structured nodes and is defined as:
Combining this with the Wasserstein formulation in Eq.(29), we observe the following inequality:
| (34) |
Appendix J E(3) Invariant Property
We utilize a 2D-MPNN, where node embeddings are iteratively refined across layers as follows:
| (35) |
with denoting the message function, representing aggregation by summation, and implemented as either the identity function or a multilayer perceptron. The final atom-level representation is obtained by integrating three modalities: the 2D molecular graph embeddings , the 3D conformational features , and the geometry-based structural descriptors . This fusion is performed using trainable linear projections:
| (36) |
where and are trainable parameter matrices. Assuming that and are composed of repeated copies of their respective feature vectors, we compute the fused representation as:
| (37) |
where is a hyperparameter controlling the influence of the barycenter features. This fusion scheme allows balanced contributions from all modalities, which is empirically beneficial.
To predict the molecular property, we perform a mean pooling over the conformations and apply a linear transformation:
| (38) |
where and are the weight matrix and bias vector used for the final prediction.
We demonstrate that the function specified in Eq.(35) through Eq.(38)remains invariant under both the action of the and permutations of the input conformers.
Theorem 2 (E(3) Invariant Property).
Proof of Theorem 2.
We establish the result in several steps. First, we consider the invariance properties of the geometric representation . By construction, the geometry-aware embedding aggregation used to obtain is invariant under permutation of conformers. Additionally, because transformations preserve Euclidean distances and given that the upstream 3D MPNN is assumed to be -invariant, the generated features are likewise invariant under such transformations.
Next, consider the aggregated representation defined in Eq.(37):
From the prior step, we know that is invariant under both actions and conformer permutations. Additionally, inherits invariance from the 3D MPNN and is permutation equivariant, i.e., permuting the conformer inputs permutes the columns of accordingly. However, because the final prediction in Eq.(38) is based on an average over the conformer-wise features:
which is invariant to column permutations of the matrix , leading to the final is invariant to group and permutation of 3D conformers. ∎