From Fragments to Geometry: A Unified Graph Transformer for Molecular Representation from Conformer Ensembles

Duy Minh Ho Nguyen; Trung Quoc Nguyen; Ha Thi Hong Le; Mai Thanh Nhat Truong; TrungTin Nguyen; Nhat Ho; Khoa D. Doan; Duy Duong-Tran; Li Shen; Daniel Sonntag; James Zou; Mathias Niepert; Hyojin Kim; Jonathan E. Allen

Abstract

Designing and understanding molecules for biological applications requires models that can integrate rich structural information from both 2D molecular graphs and diverse 3D conformer ensembles. We introduce a fragment-aware, structure-guided graph transformer that enables scalable and expressive molecular modeling by aggregating multiple 3D conformers while incorporating fragment-level inductive biases from the 2D topology. Our approach employs a trainable attention-based fusion mechanism within a graph transformer to dynamically combine 2D and 3D representations, moving beyond static solvers and rigid fusion heuristics. This architecture enables fine-grained reasoning over chemically diverse molecules, including organocatalysts and transition-metal complexes. While originally developed for molecular property prediction, the method’s structure-aware and fragment-level modeling is readily applicable to other downstream applications in drug discovery, reaction modeling, and AI-driven biological research. The model scales to large datasets and achieves state-of-the-art results across molecular property benchmarks, demonstrating its potential as a foundational component for generative AI in molecular science.

Machine Learning, ICML

round-mode = places, round-precision = 2, [linecolor=black,linewidth=1pt,skipabove=10pt,skipbelow=10pt]boxedtheoremTheorem

1 Introduction

Machine learning has become a transformative tool in computational biology, chemistry, and drug discovery, enabling predictive and generative modeling of molecular systems (Butler et al., 2018; Vamathevan et al., 2019; Choudhary et al., 2022; Fedik et al., 2022; Batatia et al., 2023). Many existing molecular representation learning methods rely on either 2D molecular graphs, which capture topological connectivity efficiently (Kipf & Welling, 2017; Gilmer et al., 2017b; Xu et al., 2018; Veličković et al., 2018), or on 3D representations derived from a single conformer (Schütt et al., 2017, 2021b; Batzner et al., 2022; Batatia et al., 2022). While 2D graphs are computationally efficient, they lack essential geometric context - critical for understanding molecular interactions, activity, and design. Incorporating 3D conformers introduces spatial features such as bond lengths and torsion angles, but relying on a single conformation fails to reflect the intrinsic flexibility and thermodynamic diversity of real molecules.

In biological systems, molecules dynamically sample a range of conformations due to bond rotations, vibrational modes, and environmental interactions (Ramsundar et al., 2019), and many functionally relevant properties — such as solubility, binding affinity, or reactivity — emerge from this conformational ensemble (Perola & Charifson, 2004). Yet, fully modeling the conformational distribution remains computationally intensive, as quantum mechanical methods for conformer generation are expensive (Rosa et al., 2016; Wankowicz & Bonomi, 2025; Medrano Sandonas et al., 2024). This challenge has motivated hybrid learning models that combine the scalability of 2D graphs with the geometric richness of a small but representative subset of 3D conformers. These approaches are currently opening a promising path toward generative molecular design, where both topological and spatial variations are essential for modeling bioactive compounds and synthesizable drug-like molecules.

To address this, structure-aware ensemble methods based on optimal transport - especially those using fused Gromov-Wasserstein (FGW) alignment - have shown promise (Brogat-Motte et al., 2022; 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.

In this work, we propose a scalable alternative: a trainable, geometry-aware graph transformer that replaces costly 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 6× 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

Molecular representations range from fingerprints (Morgan, 1965) and 1D strings (Ahmad et al., 2022; Wang et al., 2019) to 2D graphs (Gilmer et al., 2017a; Rong et al., 2020) and 3D geometric models (Fang et al., 2021; Zhou et al., 2023). While 2D models are efficient, they lack spatial context; 3D models add geometric detail but often rely on a single conformer, overlooking structural flexibility. Recent hybrid approaches combine 2D graphs with conformer ensembles (Zhu et al., 2024b; Axelrod & Gómez-Bombarelli, 2023), using aggregation techniques like pooling or self-attention (Zaheer et al., 2017; Vaswani et al., 2017). Geometry-aware methods based on FGW alignment (Brogat-Motte et al., 2022; Nguyen et al., 2024a) better capture spatial similarity across conformers but are computationally costly and struggle to scale in generative or high-throughput settings (Zhou et al., 2023). Our method addresses this limitation by using graph transformer architectures to learn latent embeddings of 3D conformers, integrating both geometry-aware signals - akin to those used in FGW-based methods—and hierarchical features from molecular fragments. This yields a scalable and permutation-invariant framework that balances computational efficiency with high representational power, making it well-suited for accuracy-critical molecular tasks.

2.2 Scalable Optimal Transport for Graph Learning

Recent advances in learning-based Optimal Transport (OT) have introduced efficient alternatives to classical solvers. Early work leveraged differentiable Sinkhorn distances with entropic regularization to improve stability and scalability (Cuturi, 2013; Feydy et al., 2019; Genevay et al., 2018). Subsequent methods enhanced computational efficiency through structural simplifications, such as low-rank approximations (Scetbon et al., 2021; Cuturi et al., 2020) and spatially-aware geometry-based formulations (Bachmann et al., 2022; Solomon et al., 2015). Additionally, meta-learning techniques accelerated optimization by learning better initializations (Amos et al., 2023), while more recent approaches have trained neural OT models directly on data to bypass iterative solvers altogether (Courty et al., 2017; Tong et al., 2021; Haviv et al., 2024).

Despite these advances, most of this work is limited to standard OT and does not extend to structure-aware variants like Fused Gromov-Wasserstein (FGW), which account for both feature similarity and relational graph structure. To address this, we introduce the first learned approximation of FGW via a graph transformer architecture, enabling scalable and geometry-aware aggregation across conformer ensembles. By integrating fragment-level structural priors into both 2D and 3D encoders, our framework supports multi-scale reasoning that unifies topological connectivity with spatial conformational diversity—essential for rich molecular representation and downstream biological modeling.

2.3 Fragment-biases in Molecular GNN

Fragment-level molecular substructures, such as rings, functional groups, and pharmacophores, play a central role in property prediction and drug development (Merlot et al., 2003; Varnek et al., 2005). Recent studies have harnessed these motifs for scaffold-aware molecule generation (Lee et al., 2024; Li, 2020; Chan et al., 2024), fragment-centric self-supervised tasks like masking and contrastive learning (Rong et al., 2020; Zhang et al., 2021; Wen et al., 2024), and in graph neural networks that encode inductive biases at the fragment level (Fey et al., 2020; Wang et al., 2025; Wollschläger et al., 2024). These approaches consistently demonstrate improved generalization, interpretability, and data efficiency. Motivated by these findings, we take a complementary approach: embedding fragment-level priors directly into hybrid 2D–3D molecular models. Our method encodes hierarchical substructures into both 2D message passing and 3D spatial attention, supporting multi-scale reasoning across topological and geometric domains. This design enhances conformer ensemble aggregation and produces richer, geometry-aware representations for tasks that depend on molecular flexibility and spatial precision.

3 Fragment-Aware Conformer Ensemble Transformer

Notations. Let $Δ_{N} : = {𝝎 \in ℝ_{+}^{N} : 𝝎^{⊤} 𝟏_{N} = 1}$ denote the probability simplex, where $𝟏_{N}$ is the all-ones vector in $ℝ^{N}$ . For $x \in 𝛀$ , $δ_{x}$ is the Dirac measure at $x$ . We write $[K] : = {1, \dots, K}$ for $K \in ℕ$ , and use $⟨ \cdot, \cdot ⟩$ to denote the Frobenius inner product. For a tensor $𝑳 = (L_{i j k l})$ and matrix $𝑩 = (B_{k l})$ , define the contraction $𝑳 \otimes 𝑩 : = (\sum_{k l} L_{i j k l} B_{k l})_{i j}$ . A graph $G = (V, E)$ has $N : = | V |$ nodes and edges $E \subseteq {{u, v} \subseteq V : u \neq v}$ . An attributed graph is given by $𝒢 : = (𝑯, 𝑨, 𝝎)$ , where $𝑯 \in ℝ^{N \times d}$ is the node feature matrix (with row $𝑯_{v}$ for node $v$ ), $𝑨$ encodes structure (e.g., adjacency or shortest-path), and $𝝎 \in Δ_{N}$ is a node weight distribution.

Given two graphs $𝒢_{1}$ and $𝒢_{2}$ with $N_{1}$ and $N_{2}$ nodes, the Fused Gromov-Wasserstein (FGW) distance (Peyré et al., 2016; Titouan et al., 2019, 2020) is: $FGW_{p, α} (𝒢_{1}, 𝒢_{2}) : = \min_{𝝅 \in Π (𝝎_{1}, 𝝎_{2})} ⟨ (1 - α) 𝑴 + α 𝑳 (𝑨_{1}, 𝑨_{2}) \otimes 𝝅, 𝝅 ⟩,$ where $Π (𝝎_{1}, 𝝎_{2}) : = {𝝅 \in ℝ_{+}^{N_{1} \times N_{2}} : 𝝅 𝟏_{N_{2}} = 𝝎_{1}, 𝝅^{⊤} 𝟏_{N_{1}} = 𝝎_{2}}$ is the set of valid couplings, $𝑴 [i, j] = d_{f} (𝑯_{1} [i], 𝑯_{2} [j])^{p}$ is the node feature cost, $𝑳 (𝑨_{1}, 𝑨_{2}) [i, j, l, m] = | 𝑨_{1} [i, j] - 𝑨_{2} [l, m] |^{p}$ captures structural mismatch, and $α \in [0, 1]$ 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.

Refer to caption — Figure 1: FACET overview. The model integrates 2D topological features from molecular and fragment graphs using message passing neural networks (MPNNs), combines them with 3D conformer features processed through a 3D-MPNN ( $Φ$ ) and a lightweight adaptor, and feeds them into a fragment-aware graph attention for geometry-aware embedding guided by FGW distance. A final structural aggregation module fuses 2D and 3D representations into a unified embedding for downstream tasks.

3.2 Framework Overview

We propose a neural architecture as in Figure 1 composed of three components. First, a 2D MPNN captures topological features from the molecular graph, while another MPNN operates on a fragment-induced hypergraph to encode higher-order structural priors (Sec.3.3). The outputs from both are fused and passed through a lightweight adaptor module, which dynamically refines and calibrates the feature representations before feeding them into the pre-trained graph transformer (Sec.3.4). Given a set of 3D conformers sampled from an input molecule graph, we use a 3D-MPNN to extract their embedding features (Sec. 3.4.1), followed by another adaptor layer. These adaptors are crucial for handling the variability in 3D conformer and 2D molecule features extracted by the 3D-MPNN ( $Φ (.)$ ) and 2D-MPNN. Then a graph transformer is used to aggregate the conformer feature sets into a geometry-aware molecular embedding, guided by atom-level and fragment-level attention. Finally, a permutation- and E(3)-invariant fusion module combines the 2D and 3D representations into a unified embedding for downstream tasks (Sec. 3.4.5).

3.3 Fragment-Enhanced 2D Molecular Graph

Each molecule is represented as a 2D graph $G = (V, E)$ , where nodes $V$ correspond to atoms and edges $E$ to covalent bonds. Atom features $𝐡_{v}^{(0)} \in ℝ^{d}$ encode properties like atom type and valence, while bonds $(u, v)$ are annotated with features $𝒆 (u, v)$ (Scarselli et al., 2008; Gilmer et al., 2017a). We adopt a 2D message-passing neural network (MPNN) that updates node embeddings layer-wise:

𝐡_{v}^{ℓ} = 𝖴 𝖯 𝖣^{ℓ} (𝐡_{v}^{ℓ - 1}, 𝖠 𝖦 𝖦^{ℓ} (𝐌^{ℓ} (𝐡_{v}^{ℓ - 1}, 𝐡_{u}^{ℓ - 1}, 𝒆_{v, u}) ∣ u \in N (v))),

(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:

\begin{matrix} 𝐌_{v, u}^{ℓ} = α_{v, u}^{ℓ} 𝐖^{ℓ} 𝐡_{u}^{ℓ - 1}, \\ α_{v, u}^{ℓ} = s o f t m a x_{u} (L e a k y R e L U (^{⊤} [𝐖^{ℓ} 𝐡_{v}^{ℓ - 1}, |, 𝐖^{ℓ} 𝐡_{u}^{ℓ - 1}])) . \end{matrix}

(2)

After $L$ layers, we obtain final atom-level features $𝐡_{v}^{L}$ for each atom $v$ used for downstream tasks.

Fragment-Based Structural Augmentation.

To enhance atomic representations with higher-order structural context, we construct a fragment hypergraph from the input molecular graph $G$ 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. 2). Each fragment is treated as a node in a new graph $G^{frag} = (V^{frag}, E^{frag})$ , where nodes correspond to fragments and edges are induced from the connectivity in $G$ , two fragments are connected if they share an atom or are directly bonded.

We apply the same GAT formulation in Eq. (1) to the fragment graph to obtain fragment embeddings ${𝐡_{f}^{frag}}_{f \in V^{frag}}$ . Then for each atom $v$ that belongs to its fragment $f (v)$ , we fuse their atom-level representations $𝐡_{v}^{(L)}$ with ${𝐡_{f}^{frag}}$ by:

\tilde{𝐡}_{v}^{(L)} = 𝐡_{v}^{(L)} + F F N (𝐡_{f (v)}^{frag}),

(3)

where $F F N (\cdot)$ is a learnable feedforward network that projects fragment-level context into the same space as atom features. If an atom belongs to multiple fragments, its atom-level features are aggregated with the corresponding fragment structures using the shared $F F N (\cdot)$ . Finally, we define a fragment-enhanced graph-level representation that is computed by applying a readout function $𝐡_{2 D} = 𝖱 𝖤 𝖠 𝖣 𝖮 𝖴 𝖳 ({\tilde{𝐡}_{v}^{(L)} ∣ v \in V}) = \sum_{v \in V} \tilde{𝐡}_{v}^{(L)} .$ 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 $S = {𝐫_{i}, Z_{i}}_{i = 1}^{N}$ , where $N$ denotes the number of atoms, $𝐫_{i} \in ℝ^{3}$ corresponds to the 3D Cartesian coordinates of atom $i$ , and $Z_{i} \in ℕ$ indicates its atomic number.

3.4.1 3D conformer feature representation

For each conformer $S$ , we can define its graph $𝒢_{S}$ and compute its 3D feature embedding by using a geometric message-passing network SchNet (Schütt et al., 2017), though other $E (3)$ -invariant neural architectures can be readily substituted without modification. We represent the matrix of atom-level features from the final message-passing layer $L$ of SchNet as $𝐇$ , where each column $𝐇 [v]$ corresponds to the feature vector $𝐡^{(L)}_{3 d, v}$ of atom $v$ . We then compute the vector representation for a conformer $S$ as $𝐡_{3 d, S} = \sum_{v \in V} (𝐖_{3 d}) 𝐡_{3 d, v}^{(L)} + 𝐛_{3 d} \in ℝ^{d}$ with $𝐖_{3 d}$ and $𝐛_{3 d}$ are learnable vectors. Given a set of $K$ conformers ${𝕊_{k}}_{k = 1}^{K}$ , we define $𝐇_{3 d} [k] = 𝐡_{3 D, 𝕊_{k}}$ as the feature embedding for the $k$ -th conformer. The matrix $𝐇_{3 d} \in ℝ^{K \times d}$ thus summarizes the feature representations of all conformers in the set.

3.4.2 Fragment-aware Graph Former

Given the atom-wise feature matrix $𝐇$ for each conformer $S$ , we aim to learn structure-encoded latent representations using Graph Transformer architectures (Zhang et al., 2020; 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 .2). It is important to note that our framework is flexible and can incorporate alternative transformer-based models.

In particular, we compose $N$ of transformer layers (Vaswani et al., 2017). Each Transformer layer consists of a self-attention mechanism followed by a position-wise feed-forward network. Given $𝐇 = [𝐡_{1}^{⊤}, \dots, 𝐡_{n}^{⊤}]^{⊤} \in ℝ^{n \times d}$ computed in Section 3.4.1 by a 3D-MPNN, where $𝐡_{i} = 𝐡^{(L)}_{3 d, v_{i}} \in ℝ^{1 \times d}$ is the vector embedding of an atom $v_{i}$ with $d$ is the hidden size. We compute self-attention, by linearly projecting $𝐇$ into query ( $𝐐$ ), key ( $𝐊$ ), and value ( $𝐕$ ) matrices using learned weights $𝐖_{Q}, 𝐖_{K}, 𝐖_{V} \in ℝ^{d \times d}$ :

	$𝐐 = 𝐇 W_{Q}, 𝐊 = 𝐇 W_{K}, 𝐕 = 𝐇 W_{V},$
	$\tilde{𝐀} = 𝐐 𝐊^{⊤} / \sqrt{d}, Attention (𝐇) = softmax (\tilde{𝐀}) 𝐕 .$		(4)

Here, $\tilde{𝐀}$ 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) only uses feature nodes, leveraging the structural information of the 3D conformer graph is essential. Follow (Ying et al., 2021), we incorporate the (i) centrality encoding, measuring how important a node is in the graph using its degree, and (ii) spatial encoding, measuring spatial relation between two nodes $v_{i}$ and $v_{j}$ in a graph $𝒢_{S}$ by the distance of the shortest path distance (SPD) (Cormen et al., 2022; Balaban, 1985) together with a weighted learnabe value along edges of SPD between two nodes. Specifically, we incorporate (i) by:

𝐡_{i} = 𝐡_{i} + z^{-}_{deg^{-} (v_{i})} + z^{+}_{deg^{+} (v_{i})},

(5)

where $z^{-}, z^{+} \in ℝ^{d}$ are learnable embedding vectors specified by the indegree $deg^{-} (v_{i})$ and outdegree $deg^{+} (v_{i})$ of atom $v_{i}$ respectively. Assume $\tilde{𝐀}_{i j}$ as the $(i, j)$ -element of the Query-Key product matrix $\tilde{𝐀}$ , the condition (ii) extends $\tilde{𝐀}_{i j}$ as:

\tilde{𝐀}_{i j} = (𝐡_{𝐢} 𝐖_{Q}) (𝐡_{j} 𝐖_{K})^{T} / \sqrt{d} + s_{ϕ (v_{i}, v_{j})} + c_{i j},

(6)

where $s_{ϕ (v_{i}, v_{j})}$ is a learnable scalar indexed by $ϕ (v_{i}, v_{j})$ , which denotes for SPD distance between $v_{i}$ and $v_{j}$ , and shared across all layers; $c_{i j} = 𝔼 (x_{e_{n}} (w^{E}_{n})^{T})$ where $𝔼 (.)$ is the expectation operation, $x_{e_{n}}$ is the feature of the $n$ -th edge $e_{n}$ in $SPD_{i j}$ , $w_{n}^{E} \in ℝ^{d_{E}}$ is the $n$ -th weight embedding, and $d_{E}$ is the dimensionality of edge feature compute as difference in feature embeddings of two nodes belong to it.

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. 5). 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 $𝐀 (G)$ using cosine distance over the final node embeddings $\tilde{𝐡}_{v}^{(L)}$ . Specifically, for each pair of atoms $(v_{i}, v_{j})$ in the 2D molecular graph, we define

𝐀 (G)_{i j} = 1 - \frac{⟨ \tilde{𝐡}_{i}^{(L)}, \tilde{𝐡}_{j}^{(L)} ⟩}{| \tilde{𝐡}_{i}^{(L)} |_{2} \cdot | \tilde{𝐡}_{j}^{(L)} |_{2}},

(7)

which quantifies their directional dissimilarity in the embedding space. Finally, we use the attention score as:

\tilde{𝐀}_{i j} = (𝐡_{𝐢} 𝐖_{Q}) (𝐡_{j} 𝐖_{K})^{T} / \sqrt{d} + s_{ϕ (v_{i}, v_{j})} + c_{i j} + 𝐀 (G)_{i j} .

(8)

3.4.3 Learning to Approximate FGW distance

We denote by $𝒯_{θ} (.)$ be a graph transformer model that has its attention operation as Eq.(8), our goal is to train $𝒯_{θ} (\cdot)$ to map the feature representation of each conformer $S$ into a latent space where the $L_{2}$ distance between any pair $S_{i}$ , $S_{j}$ approximates their FGW distance - an effective, yet computationally expensive, geometry-aware metric (Brogat-Motte et al., 2022; Ma et al., 2023; Nguyen et al., 2024a). To this end, given a set of $Ω = {S_{i}}_{i = 1}^{K}$ of $K$ generated conformers, we sample $B$ conformers from $Ω$ , then compute their encoding features by $𝒯_{θ} (𝐇_{i})$ for each $S_{i} \in B$ . These outputs are compared with their pair-wise FGW distance to optimize the loss:

ℒ_{e n c} = \sum_{i j} [| | 𝒯_{θ} (𝐇_{i}) - 𝒯_{θ} (𝐇_{j}) | |_{2}^{2} - FGW_{p, α} (𝒢 (S_{i}), 𝒢 (S_{j}))] .

(9)

By minimizing the loss $ℒ_{enc}$ , we update the parameters of the transformation module $𝒯_{θ} (\cdot)$ using gradient descent: $θ \leftarrow θ - ϵ \nabla ℒ_{enc} .$ Once trained, we freeze $𝒯_{θ}$ and incorporate it back into the framework to compute a geometry-aware representation across $K$ conformers ${𝕊_{k}}_{k = 1}^{K}$ as follows: $\bar{𝐇} = 𝔼 ({𝒯_{θ} (𝐇_{i})}_{i = 1}^{K}),$ where $\bar{𝐇}$ denotes the aggregated structural embedding. However, the 3D conformer feature distribution, extracted by 3D-MPNN, used to train $ℒ_{enc}$ (Eq. 9) may experience a domain shift when co-trained with other components in the full framework (Sec. 3.4.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 $K$ conformers, using $\bar{𝐇}$ as the GraphTransformer (GT)-aggregated atom features. We compute the global GT representation as: $𝐡_{G T} = \sum_{v \in V} (𝐖_{G T}, \bar{𝐡}_{v} + 𝐛_{G T}),$ where $\bar{𝐡}_{v} = \bar{𝐇} [v]$ and $𝐖_{G T}, 𝐛_{G T}$ are learnable parameters. We then define $𝐇_{2 D}$ and $𝐇_{G T}$ be the matrices whose columns are, respectively, $K$ copies of the 2D feature $𝐡_{2 D}$ (Sec.3.3) and $𝐡_{G T}$ representations from previous section. We fuse those representations with the 3D conformer features $𝐇_{3 D}$ to produce the final atom-wise embedding: $𝐇_{c o m b} = \tilde{𝐖}_{2 D} 𝐇_{2 D} + \tilde{𝐖}_{3 D} 𝐇_{3 D} + \tilde{𝐖}_{G T} 𝐇_{G T},$ where each $\tilde{𝐖_{i}, i \in {2 D, 3 D, G T}}$ are trainable projection matrix. The combined embedding $𝐇_{c o m b}$ is fed into a final FFN layer to predict the target property.

4 Theoretical Bounds for Embedding Non-Euclidean FGW Distance Matrices

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 $K$ distributions as $D_{i j} : = FGW_{p, α} (𝒢 (S_{i}), 𝒢 (S_{j}))$ for all $i, j \in [K]$ , following Section 3.4.3. The empirical FGW barycenter is given by $\bar{𝒢}_{K} \in a r g m i n_{𝒢 \in 𝒫_{p} (𝛀)} \frac{1}{K} \sum_{i = 1}^{K} FGW_{p, α}^{p} (𝒢, 𝒢 (S_{i}))$ , where $𝒫_{p} (𝛀)$ denotes the space of attributed graphs with finite $p$ -th order FGW distance.

To approximate this barycenter in embedding space, we require $∥ \bar{𝒆}_{K} - 𝒆_{j} ∥_{2}^{2} \approx FGW_{p, α} (\bar{𝒢}_{K}, 𝒢 (S_{j})) : = \bar{D}_{K, j}$ for all $j \in [K]$ , where $\bar{𝒆}_{K} = \frac{1}{K} \sum_{i = 1}^{K} 𝒆_{i}$ is the mean embedding and $𝒆_{i} : = 𝒯_{θ} (𝐇_{i})$ is the learned representation. To assess how well the embeddings ${𝒆_{i}}_{i = 1}^{K} \subset ℝ^{d}$ preserve both pairwise FGW distances and barycenter structure, we define the cumulative stress: $𝒮 = \min_{𝒆_{i} \in ℝ^{d}} \sum_{i, j \in [K]} (∥ 𝒆_{i} - 𝒆_{j} ∥_{2}^{2} - D_{i j})^{2} + \sum_{j \in [K]} (∥ \bar{𝒆}_{K} - 𝒆_{j} ∥_{2}^{2} - \bar{D}_{K, j})^{2}$ . This objective enforces faithful reconstruction of both the distance structure and the barycenter alignment in the learned embedding space, as formalized in Theorem 4.1 (see proof in Appendix D).

Theorem 4.1.

Let $𝐃$ denote the pairwise $FGW_{p, α}$ distance matrix, and let ${λ_{i}, 𝐯_{i}}_{i = 1}^{K}$ represent the eigendecomposition of the associated criterion matrix $𝐅 = - 𝐂 𝐃 𝐂$ , where $𝐂 = 𝐈_{K} - \frac{1}{K} 𝟏_{K} 𝟏_{K}^{⊤}$ is the centering matrix. The optimal stress value, denoted by $𝒮^{*}$ , is bounded as follows: $ℒ \leq 𝒮^{*} \leq 𝒰$ , where $ℒ : = \sum_{i : λ_{i} < 0} λ_{i}^{2}$ , $𝒰 : = \sum_{i j} (Δ g_{i} + Δ g_{j})^{2} + ℒ + 𝒞$ , $Δ g_{i} = \frac{1}{2} \sum_{j : λ_{j} < 0} λ_{j} \cdot 𝐯_{i j}^{2}$ . Here, $𝐯_{i j}$ denotes the $j$ -th component of the $i$ -th eigenvector $𝐯_{n}$ of $𝐅$ , and $𝒞$ quantifies the approximation error between the empirical barycenter in the Euclidean embedding space and that 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 2D and 3D MPNNs to extract features from molecular graphs and conformers. These features are also used to supervise the Graph Transformer in the next stage. Stage 2: The Graph Transformer is trained independently to approximate FGW distances between conformers using the features 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 $1 e^{- 5}$ . Stage 3: We integrate all components into an end-to-end model, where only the 2D and 3D MPNNs are updated (300 epochs, learning rate $5 e^{- 4}$ ). To address feature distribution shift caused by finetuning, we apply FFN-based adaptor layers to the 2D and 3D features before feeding them to the Graph Transformer.

We use Adam for all stages. Further experimental details are provided in the Appendix.

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 $K$ conformers ( $K \in 5, 10, 15, 20$ ) using single and multi-GPU settings. Experiments are conducted on the BACE dataset and summarized in Figure 3.

It can be seen that (a) FACET exhibits strong scalability, maintaining a nearly constant runtime across varying numbers of conformers, in both single- and multi-GPU environments. In contrast, ConAN-FGW scales poorly, with runtime increasing sharply as the number of conformers grows. Although using multiple GPUs reduces the runtime compared to a single GPU, the upward trend persists, with runtimes surpassing 50 seconds for 20 conformers. (b) Secondly, FACET’s similar runtime on single- and multi-GPU setups reflects its efficiency and the small workload in this experiment. In such cases, multi-GPU overhead can outweigh speedup gains. We expect multi-GPU acceleration to be beneficial mainly for large-scale tasks, like processing a large batch size of thousands of molecules or handling memory-intensive inputs.

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 Figure 4, 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 example, 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.

5.4 State-of-the-Art Performance Comparison on Molecular Tasks

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 (pIC₅₀), lipophilicity, and hydration free energy. MARCEL consists of Drugs-75K and Kraken, where the goal is to predict the Boltzmann-averaged property $⟨ y ⟩_{k_{ℬ}}$ from sampled conformers. Drugs-75K uses quantum descriptors (IP, EA, $χ$ ), while Kraken focuses on Sterimol features (B₅, L, and their buried forms). The Boltzmann average is computed as a weighted sum over conformer-specific values $y_{i}$ with probabilities $p_{i}$ . All datasets follow the original random split settings, using the provided sampled conformers.

Table 3: Number of samples for each split on tasks of MoleculeNet and the MARCEL benchmark.

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

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., 2021a), 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

MoleculeNet. FACET achieves state-of-the-art results on ESOL, FreeSolv, and BACE, reporting the lowest MSE scores across all tasks. Its consistent outperformance of CONAN-FGW underscores the value of incorporating fragment-level substructures into geometry-aware attention, which enhances the model’s ability to capture localized chemical contexts and improves molecular property prediction.

MARCEL. FACET also demonstrates strong performance on the MARCEL benchmark, improving results across both SchNet and GemNet backbones. By combining structure-aware aggregation with hierarchical fragment information, FACET remains robust even at scale - unlike CONAN-FGW, which struggles with MARCEL’s larger dataset. Together, these results highlight FACET’s effectiveness and scalability across diverse molecular modeling tasks.

Table 4: Comparison of molecular property regression performance on the MARCEL benchmark (MAE

↓

). The results of competing methods are adapted from (Zhu et al., 2024a).

Category	Model	Drugs-75K			Kraken
	Model	IP	EA	$χ$	B₅	L	BurB₅	BurL
2D models	GIN (Xu et al., 2019)	0.4354	0.4169	0.2260	0.3128	0.4003	0.1719	0.1200
	GIN+VN (Hu et al., 2020a)	0.4361	0.4169	0.2267	0.3567	0.4344	0.2422	0.1741
	ChemProp (Yang et al., 2019b)	0.4595	0.4417	0.2441	0.4850	0.5452	0.3002	0.1948
	GraphGPS (Rampášek et al., 2022)	0.4351	0.4085	0.2212	0.3450	0.4363	0.2066	0.1500
3D models	SchNet (Schütt et al., 2017)	0.4394	0.4207	0.2243	0.3293	0.5458	0.2295	0.1861
	DimeNet++ (Klicpera et al., 2020)	0.4441	0.4233	0.2436	0.3510	0.4174	0.2097	0.1526
	GemNet (Gasteiger et al., 2021)	0.4069	0.3922	0.1970	0.2789	0.3754	0.1782	0.1635
	PaiNN (Schütt et al., 2021a)	0.4505	0.4495	0.2324	0.3443	0.4471	0.2395	0.1673
	ClofNet (Du et al., 2022)	0.4393	0.4251	0.2378	0.4873	0.6417	0.2884	0.2529
	LEFTNet (Du et al., 2023)	0.4174	0.3964	0.2083	0.3072	0.4493	0.2176	0.1486
Ensemble Strategy with DeepSets	SchNet (Schütt et al., 2017)	0.4452	0.4232	0.2243	0.2704	0.4322	0.2024	0.1443
	DimeNet++ (Klicpera et al., 2020)	0.4126	0.3944	0.2267	0.2630	0.3468	0.1783	0.1185
	GemNet (Gasteiger et al., 2021)	0.4066	0.3910	0.2027	0.2313	0.3386	0.1589	0.0947
	PaiNN (Schütt et al., 2021a)	0.4466	0.4269	0.2294	0.2225	0.3619	0.1693	0.1324
	ClofNet (Du et al., 2022)	0.4280	0.4033	0.2199	0.3228	0.4485	0.2178	0.1548
	LEFTNet (Du et al., 2023)	0.4149	0.3953	0.2069	0.2644	0.3643	0.2017	0.1386
FACET	SchNet (Schütt et al., 2017)	0.4235	0.3971	0.2155	0.2508	0.3982	0.1803	0.1245
FACET	GemNet (Gasteiger et al., 2021)	0.3891	0.3852	0.1970	0.2225	0.3402	0.1503	0.0952

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 5, 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.

Table 5: FACET ablation study.

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

6 Conclusion

We present the FACET, a scalable method that integrates 3D conformer features with fragment-level 2D graph information. Using a trainable attention mechanism, it 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.

Acknowledgement

This work was supported by Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany’s Excellence Strategy - EXC 2075 – 390740016, the DARPA ANSR program under award FA8750-23- 2-0004, the DARPA CODORD program under award HR00112590089. It was also partly funded by the Defense Threat Reduction Agency (DTRA), HDTRA1242044. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Duy M. H. Nguyen. Duy M. H. Nguyen and Daniel Sonntag are also supported by the No-IDLE project (BMBF, 01IW23002), the MASTER project (EU, 101093079), and the Endowed Chair of Applied Artificial Intelligence, Oldenburg University.

Disclaimer

The views expressed in this manuscript are those of the author(s) and do not reflect the official policy or position of the U.S. Naval Academy, Department of the Navy, the Department of Defense, or the U.S. Government.

References

Ahmad et al. (2022) Ahmad, W., Simon, E., Chithrananda, S., Grand, G., and Ramsundar, B. Chemberta-2: Towards chemical foundation models, 2022. ↩ a b c
Amos et al. (2023) Amos, B., Cohen, S., Luise, G., and Redko, I. Meta optimal transport. International Conference on Machine Learning, 2023. ↩ a
Axelrod & Gomez-Bombarelli (2022) Axelrod, S. and Gomez-Bombarelli, R. Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022. ↩ a
Axelrod & Gómez-Bombarelli (2023) Axelrod, S. and Gómez-Bombarelli, R. 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 b c
Bachmann et al. (2022) Bachmann, F., Hennig, P., and Kobak, D. Wasserstein t-sne. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 104–120. Springer, 2022. ↩ a
Balaban (1985) Balaban, A. T. Applications of graph theory in chemistry. Journal of chemical information and computer sciences, 25(3):334–343, 1985. ↩ a
Batatia et al. (2022) Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., and Csanyi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 11423–11436. Curran Associates, Inc., 2022. ↩ a
Batatia et al. (2023) Batatia, I., Benner, P., Chiang, Y., Elena, A. M., Kovács, D. P., Riebesell, J., Advincula, X. R., Asta, M., Baldwin, W. J., Bernstein, N., Bhowmik, A., Blau, S. M., Cărare, V., Darby, J. P., De, S., Pia, F. D., Deringer, V. L., Elijošius, R., El-Machachi, Z., Fako, E., Ferrari, A. C., Genreith-Schriever, A., George, J., Goodall, R. E. A., Grey, C. P., Han, S., Handley, W., Heenen, H. H., Hermansson, K., Holm, C., Jaafar, J., Hofmann, S., Jakob, K. S., Jung, H., Kapil, V., Kaplan, A. D., Karimitari, N., Kroupa, N., Kullgren, J., Kuner, M. C., Kuryla, D., Liepuoniute, G., Margraf, J. T., Magdău, I.-B., Michaelides, A., Moore, J. H., Naik, A. A., Niblett, S. P., Norwood, S. W., O’Neill, N., Ortner, C., Persson, K. A., Reuter, K., Rosen, A. S., Schaaf, L. L., Schran, C., Sivonxay, E., Stenczel, T. K., Svahn, V., Sutton, C., van der Oord, C., Varga-Umbrich, E., Vegge, T., Vondrák, M., Wang, Y., Witt, W. C., Zills, F., and Csányi, G. A foundation model for atomistic materials chemistry, 2023. ↩ a
Batzner et al. (2022) Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nat. Commun., 13(1):2453, May 2022. ISSN 2041-1723. ↩ a
Brogat-Motte et al. (2022) Brogat-Motte, L., Flamary, R., Brouard, C., Rousu, J., and d’Alché Buc, F. Learning to predict graphs with fused gromov-wasserstein barycenters. In International Conference on Machine Learning, pp. 2321–2335. PMLR, 2022. ↩ a b c
Butler et al. (2018) Butler, K. T., Davies, D. W., Cartwright, H., Isayev, O., and Walsh, A. Machine learning for molecular and materials science. Nature, 559(7715):547–555, Jul 2018. ISSN 1476-4687. ↩ a
Chan et al. (2024) Chan, B. W., Lynch, N. B., Tran, W., Joyce, J. M., Savage, G. P., Meutermans, W., Montgomery, A. P., and Kassiou, M. Fragment-based drug discovery for disorders of the central nervous system: designing better drugs piece by piece. Frontiers in Chemistry, 12:1379518, 2024. ↩ a
Cheng et al. (2025) Cheng, A. H., Sun, C., and Aspuru-Guzik, A. Scalable autoregressive 3d molecule generation. arXiv preprint arXiv:2505.13791, 2025. ↩ a
Choudhary et al. (2022) Choudhary, K., DeCost, B., Chen, C., Jain, A., Tavazza, F., Cohn, R., Park, C. W., Choudhary, A., Agrawal, A., Billinge, S. J. L., Holm, E., Ong, S. P., and Wolverton, C. 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) Chung, Y. G., Haldoupis, E., Bucior, B. J., Haranczyk, M., Lee, S., Zhang, H., Vogiatzis, K. D., Milisavljevic, M., Ling, S., Camp, J. S., 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) Cormen, T. H., Leiserson, C. E., Rivest, R. L., and Stein, C. Introduction to algorithms. MIT press, 2022. ↩ a
Courty et al. (2017) Courty, N., Flamary, R., and Ducoffe, M. Learning wasserstein embeddings. International Conference on Learning Representations (ICLR), 2017. ↩ a
Cuturi (2013) Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013. ↩ a
Cuturi et al. (2020) Cuturi, M., Teboul, O., Niles-Weed, J., and Vert, J.-P. Supervised quantile normalization for low rank matrix factorization. In International Conference on Machine Learning, pp. 2269–2279. PMLR, 2020. ↩ a
Du et al. (2022) Du, W., Zhang, H., Du, Y., Meng, Q., Chen, W., Zheng, N., Shao, B., and Liu, T.-Y. Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning, pp. 5583–5608. PMLR, 2022. ↩ a b c
Du et al. (2023) Du, Y., Wang, L., Feng, D., Wang, G., Ji, S., Gomes, C. P., Ma, Z.-M., 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 b c
Fang et al. (2021) Fang, X., Liu, L., Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H., and Wang, H. 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) Fang, X., Liu, L., Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H., and Wang, H. Geometry-enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 4(2):127–134, 2022. ↩ a b
Fedik et al. (2022) Fedik, N., Zubatyuk, R., Kulichenko, M., Lubbers, N., Smith, J. S., Nebgen, B., Messerly, R., Li, Y. W., Boldyrev, A. I., Barros, K., Isayev, O., and Tretiak, S. 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
Fey et al. (2020) Fey, M., Yuen, J.-G., and Weichert, F. Hierarchical inter-message passing for learning on molecular graphs. arXiv preprint arXiv:2006.12179, 2020. ↩ a
Feydy et al. (2019) Feydy, J., Séjourné, T., Vialard, F.-X., Amari, S.-i., Trouvé, A., and Peyré, G. 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) Gasteiger, J., Becker, F., and Günnemann, S. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790–6802, 2021. ↩ a b c d
Genevay et al. (2018) Genevay, A., Peyré, G., and Cuturi, M. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608–1617. PMLR, 2018. ↩ a
Geng et al. (2023) Geng, Z., Xie, S., Xia, Y., Wu, L., Qin, T., Wang, J., Zhang, Y., Wu, F., and Liu, T.-Y. De novo molecular generation via connection-aware motif mining. arXiv preprint arXiv:2302.01129, 2023. ↩ a
Gilmer et al. (2017a) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp. 1263–1272, 2017a. ↩ a b
Gilmer et al. (2017b) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chemistry. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th ICML, volume 70 of Proceedings of Machine Learning Research, pp. 1263–1272. PMLR, 06–11 Aug 2017b. ↩ a
Gower (1985) Gower, J. 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. URL https://www.sciencedirect.com/science/article/pii/0024379585901879. ↩ a
Haviv et al. (2024) Haviv, D., Kunes, R. Z., Dougherty, T., Burdziak, C., Nawy, T., Gilbert, A., and Pe’Er, D. Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformer. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 17697–17718. PMLR, July 2024. URL https://proceedings.mlr.press/v235/haviv24a.html. ↩ a b c
Hawkins (2017) Hawkins, P. C. Conformation generation: the state of the art. Journal of chemical information and modeling, 57(8):1747–1756, 2017. ↩ a
Higham (1988) Higham, N. J. 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. URL https://www.sciencedirect.com/science/article/pii/0024379588902236. ↩ a
Hu et al. (2020a) Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118–22133, 2020a. ↩ a b
Hu et al. (2020b) Hu, W., Liu, B., Gomes, J., Zitnik, M., Liang, P., Pande, V., and Leskovec, J. Strategies for pre-training graph neural networks. In International Conference on Learning Representations, 2020b. ↩ a b
Kipf & Welling (2017) Kipf, T. N. and Welling, M. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017. ↩ a
Klicpera et al. (2020) Klicpera, J., Groß, J., Günnemann, S., et al. Directional message passing for molecular graphs. In International Conference on Learning Representations, pp. 1–13, 2020. ↩ a b c
Kong et al. (2022) Kong, X., Huang, W., Tan, Z., and Liu, Y. Molecule generation by principal subgraph mining and assembling. Advances in Neural Information Processing Systems, 35:2550–2563, 2022. ↩ a
Kreuzer et al. (2021) Kreuzer, D., Beaini, D., Hamilton, W., Létourneau, V., and Tossou, P. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34:21618–21629, 2021. ↩ a
Landrum (2016) Landrum, G. Rdkit: open-source cheminformatics http://www.rdkit.org, 2016. ↩ a
Lee et al. (2024) Lee, S., Lee, S., Kawaguchi, K., and Hwang, S. J. Drug discovery with dynamic goal-aware fragments. International Conference on Machine Learning, 2024. ↩ a
Li (2020) Li, Q. Application of fragment-based drug discovery to versatile targets. Frontiers in molecular biosciences, 7:180, 2020. ↩ a
Liu et al. (2021) Liu, M., Fu, C., Zhang, X., Wang, L., Xie, Y., Yuan, H., Luo, Y., Xu, Z., Xu, S., and Ji, S. Fast quantum property prediction via deeper 2d and 3d graph networks. arXiv preprint arXiv:2106.08551, 2021. ↩ a b
Liu et al. (2015) Liu, Z., Li, Y., Han, L., Li, J., Liu, J., Zhao, Z., Nie, W., Liu, Y., and Wang, R. Pdb-wide collection of binding data: current status of the pdbbind database. Bioinformatics, 31(3):405–412, 2015. ↩ a
Luo et al. (2024) Luo, Y., Li, H., Shi, L., and Wu, X.-M. Enhancing graph transformers with hierarchical distance structural encoding. Advances in Neural Information Processing Systems, 37:57150–57182, 2024. ↩ a
Ma et al. (2023) Ma, X., Chu, X., Wang, Y., Lin, Y., Zhao, J., Ma, L., and Zhu, W. Fused gromov-wasserstein graph mixup for graph-level classifications. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. ↩ a b
Medrano Sandonas et al. (2024) Medrano Sandonas, L., Van Rompaey, D., Fallani, A., Hilfiker, M., Hahn, D., Perez-Benito, L., Verhoeven, J., Tresadern, G., Kurt Wegner, J., Ceulemans, H., 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) Merlot, C., Domine, D., Cleva, C., and Church, D. J. Chemical substructures in drug discovery. Drug Discovery Today, 8(13):594–602, 2003. ↩ a
Morgan (1965) Morgan, H. L. 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) Nguyen, D. M., Lukashina, N., Nguyen, T., Le, A. T., Nguyen, T., Ho, N., Peters, J., Sonntag, D., Zaverkin, V., and Niepert, M. Structure-aware e (3)-invariant molecular conformer aggregation networks. International Conference on Machine Learning, 2024a. ↩ a b c d
Nguyen et al. (2024b) Nguyen, D. M. H., Lukashina, N., Nguyen, T., Le, A. T., Nguyen, T., Ho, N., Peters, J., Sonntag, D., Zaverkin, V., and Niepert, M. 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 h i
Otsuka et al. (2011) Otsuka, S., Kuwajima, I., Hosoya, J., Xu, Y., and Yamazaki, M. 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) Perola, E. and Charifson, P. S. 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) Peyré, G., Cuturi, M., and Solomon, J. Gromov-Wasserstein Averaging of Kernel and Distance Matrices. In Balcan, M. F. and Weinberger, K. Q. (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) Peyré, G., Cuturi, M., 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) Rampášek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35:14501–14515, 2022. ↩ a b
Ramsundar et al. (2019) Ramsundar, B., Eastman, P., Walters, P., and Pande, V. 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) Raza, A., Henle, E. A., and Fern, X. Non-equilibrium molecular geometries in graph neural networks. arXiv preprint arXiv:2203.04697, 2022. ↩ a
Rong et al. (2020) Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W., and Huang, J. Self-supervised graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems, 33:12559–12571, 2020. ↩ a b c d
Rosa et al. (2016) Rosa, M., Micciarelli, M., Laio, A., and Baroni, S. Sampling molecular conformers in solution with quantum mechanical accuracy at a nearly molecular-mechanics cost. Journal of Chemical Theory and Computation, 12(9):4385–4389, 2016. ↩ a
Ross et al. (2022) Ross, J., Belgodere, B., Chenthamarakshan, V., Padhi, I., Mroueh, Y., and Das, P. Large-scale chemical language representations capture molecular structure and properties. Nature Machine Intelligence, 4(12):1256–1264, 2022. ↩ a b
Scarselli et al. (2008) Scarselli, F., Gori, M., Tsoi, A. C., Hagenbuchner, M., and Monfardini, G. The graph neural network model. IEEE transactions on neural networks, 20(1):61–80, 2008. ↩ a
Scetbon et al. (2021) Scetbon, M., Cuturi, M., and Peyré, G. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pp. 9344–9354. PMLR, 2021. ↩ a
Schütt et al. (2017) Schütt, K., Kindermans, P., Felix, H. E. S., Chmiela, S., Tkatchenko, A., and Müller, K. 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) Schütt, K., Kindermans, P.-J., Sauceda Felix, H. E., Chmiela, S., Tkatchenko, A., and Müller, K.-R. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. ↩ a b c d e f g h
Schütt et al. (2021a) Schütt, K., Unke, O., and Gastegger, M. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377–9388. PMLR, 2021a. ↩ a b c
Schütt et al. (2021b) Schütt, K. T., Unke, O. T., and Gastegger, M. Equivariant message passing for the prediction of tensorial properties and molecular spectra. ICML, pp. 1–13, 2021b. ↩ a
Solomon et al. (2015) Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1–11, 2015. ↩ a
Sonthalia et al. (2021) Sonthalia, R., Van Buskirk, G., Raichel, B., and Gilbert, A. How can classical multidimensional scaling go wrong? In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P. S., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 12304–12315. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/66121d1f782d29b62a286909165517bc-Paper.pdf. ↩ a
Titouan et al. (2019) Titouan, V., Courty, N., Tavenard, R., Laetitia, C., and Flamary, R. Optimal Transport for structured data with application on graphs. In Chaudhuri, K. and Salakhutdinov, R. (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) Titouan, V., Chapel, L., Flamary, R., Tavenard, R., and Courty, N. Fused Gromov-Wasserstein Distance for Structured Objects. Algorithms, 13(9):212, August 2020. ISSN 1999-4893. doi: 10.3390/a13090212. URL https://www.mdpi.com/1999-4893/13/9/212. ↩ a
Tong et al. (2021) Tong, A. Y., Huguet, G., Natik, A., MacDonald, K., Kuchroo, M., Coifman, R., Wolf, G., and Krishnaswamy, S. Diffusion earth mover’s distance and distribution embeddings. In International Conference on Machine Learning, pp. 10336–10346. PMLR, 2021. ↩ a
Tong et al. (2025) Tong, V., Trung-Dung, H., Liu, A., Broeck, G. V. d., and Niepert, M. Learning to discretize denoising diffusion odes. International Conference on Learning Representations (ICLR)), 2025. ↩ a
Torgerson (1952) Torgerson, W. S. Multidimensional Scaling: I. Theory and Method. Psychometrika, 17(4):401–419, 1952. ISSN 0033-3123. doi: 10.1007/BF02288916. URL https://www.cambridge.org/core/product/5A026426380B7639E7EA6D92D5DF19AB. Edition: 2025/01/01 Publisher: Cambridge University Press & Assessment. ↩ a
Vamathevan et al. (2019) Vamathevan, J., Clark, D., Czodrowski, P., Dunham, I., Ferran, E., Lee, G., Li, B., Madabhushi, A., Shah, P., Spitzer, M., and Zhao, S. Applications of machine learning in drug discovery and development. Nat. Rev. Drug Discov., 18(6):463–477, Jun 2019. ISSN 1474-1784. ↩ a
Varnek et al. (2005) Varnek, A., Fourches, D., Hoonakker, F., and Solov’ev, V. P. 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) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017. ↩ a b c
Veličković et al. (2017) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y. Graph attention networks. arXiv preprint arXiv:1710.10903, 2017. ↩ a
Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph attention networks. In International Conference on Learning Representations, 2018. ↩ a b
Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph attention networks. In ICLR, 2018. ↩ a
Wang et al. (2025) Wang, J., Min, Y., Li, M., and Wu, J. Fragformer: A fragment-based representation learning framework for molecular property prediction. Transactions on Machine Learning Research, 2025. ↩ a
Wang et al. (2019) Wang, S., Guo, Y., Wang, Y., Sun, H., and Huang, J. 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
Wankowicz & Bonomi (2025) Wankowicz, S. A. and Bonomi, M. From possibility to precision in macromolecular ensemble prediction. arXiv preprint arXiv:2505.01919, 2025. ↩ a
Wen et al. (2024) Wen, Q., Ju, M., Ouyang, Z., Zhang, C., and Ye, Y. 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) Wollschläger, T., Kemper, N., Hetzel, L., Sommer, J., and Günnemann, S. Expressivity and generalization: Fragment-biases for molecular gnns. International Conference on Machine Learning, 2024. ↩ a b
Wu et al. (2018) Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Geniesse, C., Pappu, A. S., Leswing, K., and Pande, V. MoleculeNet: A benchmark for molecular machine learning. Chemical Science, pp. 513–530, 2018. ↩ a b
Xiong et al. (2019) Xiong, Z., Wang, D., Liu, X., Zhong, F., Wan, X., Li, X., Li, Z., Luo, X., Chen, K., Jiang, H., 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 b
Xu et al. (2018) Xu, K., Li, C., Tian, Y., Sonobe, T., Kawarabayashi, K.-i., and Jegelka, S. Representation learning on graphs with jumping knowledge networks. In Dy, J. and Krause, A. (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) Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? International Conference on Learning Representations, 2019. ↩ a b
Yang et al. (2019a) Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Guzman-Perez, A., Hopper, T., Kelley, B., Mathea, M., Palmer, A., Settels, V., Jaakkola, T., Jensen, K., and Barzilay, R. 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) Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Guzman-Perez, A., Hopper, T., Kelley, B., Mathea, M., et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8):3370–3388, 2019b. ↩ a b
Ying et al. (2021) Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. 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) Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R. R., and Smola, A. J. Deep sets. Advances in neural information processing systems, 30, 2017. ↩ a b
Zhang et al. (2020) Zhang, J., Zhang, H., Xia, C., and Sun, L. Graph-bert: Only attention is needed for learning graph representations. arXiv preprint arXiv:2001.05140, 2020. ↩ a
Zhang et al. (2021) Zhang, Z., Liu, Q., Wang, H., Lu, C., and Lee, C.-K. 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) Zhou, G., Gao, Z., Ding, Q., Zheng, H., Xu, H., Wei, Z., Zhang, L., and Ke, G. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023. ↩ a b c d
Zhu et al. (2023) Zhu, Y., Hwang, J., Adams, K., Liu, Z., Nan, B., Stenfors, B., Du, Y., Chauhan, J., Wiest, O., Isayev, O., Coley, C. W., Sun, Y., and Wang, W. Learning over molecular conformer ensembles: Datasets and benchmarks, 2023. ↩ a
Zhu et al. (2024a) Zhu, Y., Hwang, J., Adams, K., Liu, Z., Nan, B., Stenfors, B., Du, Y., Chauhan, J., Wiest, O., Isayev, O., et al. Learning over molecular conformer ensembles: Datasets and benchmarks. In International Conference on Learning Representations, 2024a. ↩ a b c
Zhu et al. (2024b) Zhu, Y., Hwang, J., Adams, K., Liu, Z., Nan, B., Stenfors, B., Du, Y., Chauhan, J., Wiest, O., Isayev, O., et al. Learning over molecular conformer ensembles: Datasets and benchmarks. International Conference on Learning Representations (ICLR), 2024b. ↩ a

Supplementary Materials for
“From Fragments to Geometry: A Unified Graph Transformer for Molecular Representation from Conformer Ensembles”

Appendix A Limitations of FACET

A.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.

A.2 Limitations in Scope

A.2.1 Focusing 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.

A.2.2 Limitation in Generative Capabilities.

While FACET demonstrates strong performance on discriminative tasks such as molecular property prediction across MoleculeNet and MARCEL benchmarks, its current formulation and evaluation are limited to regression settings where the goal is to predict properties from given molecular structures. As a result, the model’s potential for generative applications such as de novo molecule generation, scaffold decoration, or fragment-based drug design remains unexplored. This limits our understanding of how well FACET can serve as a foundational model for inverse molecular problems, where structural creativity and diversity are critical. Future work should explore extensions of FACET with generative decoding mechanisms, such as auto-regressive sampling (Cheng et al., 2025), diffusion models (Tong et al., 2025), or variational objectives, to fully leverage its design for structure-conditioned generation.

Appendix B 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 $𝐇_{2 d - 3 d} = \tilde{𝐖}_{2 D} 𝐇_{2 D} + \tilde{𝐖}_{3 D} 𝐇_{3 D}$ , we then train for 150 epochs and set the learning rate to $1 e^{- 3}$ . to optimize target property tasks $ℒ_{p r e d} = | | \hat{𝒚}_{2 d - 3 d} - \tilde{𝐲} | |_{2}^{2}$ where $\tilde{𝐲}$ be the ground-truth value and $\hat{𝒚}$ be our predicted value defined by:

\hat{𝒚}_{2 d - 3 d} = 𝐖^{𝒢} (\frac{1}{K} \sum_{k = 1}^{K} 𝐇_{2 d - 3 d} [k]) + 𝐛^{𝒢},

(10)

with $𝐖^{𝒢}$ and $𝐛^{𝒢}$ are learnable parameters.

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 $1 e^{- 5}$ and train for 1000 epochs with the following loss function:

ℒ_{e n c} = \sum_{i j} [| | 𝒯_{θ} (𝐇_{i}) - 𝒯_{θ} (𝐇_{j}) | |_{2}^{2} - FGW_{p, α} (𝒢 (S_{i}), 𝒢 (S_{j}))] .

(11)

Stage 3. Training Fragment-aware Graph Transformer. In the final stage, we freeze the trained GraphTransformer $𝒯 θ (\cdot)$ 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 $𝒯 θ (\cdot)$ , 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 $ℒ_{p r e d} = | | \hat{𝒚} - \tilde{𝐲} | |_{2}^{2}$ where

\hat{𝒚} = 𝐖^{𝒢} (\frac{1}{K} \sum_{k = 1}^{K} 𝐇_{c o m b} [k]) + 𝐛^{𝒢} .

(12)

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 $32$ molecules, each with different conformers ranging from $3$ , $5$ , $10$ , $15$ , and $20$ , and another experiment with four GPUs on the same batch size, i.e., 8 molecules per GPU.

Figure 5 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 6 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.

Appendix D Proof of Theorem 4.1

Recall that we aim to establish the following novel theoretical bounds: Let $𝑫$ denote the pairwise $FGW_{p, α}$ distance matrix, and let ${λ_{k}, 𝒗_{k}}_{k = 1}^{K}$ represent the eigendecomposition of the associated criterion matrix $𝑭 = - 𝑪 𝑫 𝑪$ , where $𝑪 = 𝑰_{K} - \frac{1}{K} 𝟏_{K} 𝟏_{K}^{⊤}$ is the centering matrix. The optimal stress value, denoted by $𝒮^{*}$ , is bounded as follows: $ℒ \leq 𝒮^{*} \leq 𝒰$ , where

ℒ : = \sum_{k : λ_{k} < 0} λ_{k}^{2}, 𝒰 : = \sum_{k l} (Δ g_{k} + Δ g_{l})^{2} + ℒ + 𝒞, Δ g_{k} = \frac{1}{2} \sum_{l : λ_{l} < 0} λ_{l} \cdot 𝒗_{k l}^{2}, \forall k \in [K] .

Here, $𝒗_{k l}$ denotes the $l$ -th component of the $k$ -th eigenvector $𝒗_{k}$ 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 $𝒆 : = {𝒆_{k}}_{k \in [K]} \in ℝ^{d \times K}$ , our objective is to derive lower and upper bounds for the following cumulative stress:

$𝒮^{*}$	$= \min_{𝒆 \in ℝ^{d \times K}} 𝒮 (𝒆), 𝒮 (𝒆) = 𝒮_{1} (𝒆) + 𝒮_{2} (𝒆),$	(13)
$𝒮^{*}_{1}$	$: = \min_{𝒆 \in ℝ^{d \times K}} 𝒮_{1} (𝒆), 𝒮_{1} (𝒆) : = \sum_{k, l \in [K]} (∥ 𝒆_{k} - 𝒆_{l} ∥_{2}^{2} - D_{k l})^{2},$	(14)
$𝒮^{*}_{2}$	$: = \min_{𝒆 \in ℝ^{d \times K}} 𝒮_{2} (𝒆), 𝒮_{2} (𝒆) : = \sum_{l \in [K]} (∥ \bar{𝒆}_{K} - 𝒆_{l} ∥_{2}^{2} - \bar{D}_{K, l})^{2} .$	(15)

To this end, we begin by specifying and formally defining the following important concepts in Section D.1.

D.1 Non-Euclidean Nature of Pairwise FGW Distance Matrix

Definition D.1 (Euclidean Distance Matrix).

A $K \times K$ distance matrix $𝑫$ is said to be Euclidean if there exists a set of points $𝒆 = {𝒆_{k}}_{k = 1}^{K}$ in some Euclidean space $ℝ^{d}$ such that

\forall k, l \in [K], D_{k l} = ∥ 𝒆_{k} - 𝒆_{l} ∥_{2}^{2} .

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: $D_{k l} \geq 0$ for all $k, l \in [K]$ ,
(ii)

Hollow diagonal: $D_{k k} = 0$ for all $k \in [K]$ ,
(iii)

Positive semidefiniteness: the associated double-centered matrix $𝑭 : = - 𝑪 𝑫 𝑪$ is positive semidefinite (PSD), where $𝑪 = 𝑰_{K} - \frac{1}{K} 𝟏_{K} 𝟏_{K}^{⊤}$ is the centering matrix, and $𝟏_{K}$ denotes the $K$ -dimensional vector of ones.

Recall that the pairwise FGW distance matrix $𝑫$ for a collection of $K$ distributions is defined entry-wise by $D_{k l} : = FGW_{p, α} (𝒢 (𝕊_{k}), 𝒢 (S_{l}))$ for all $k, l \in [K]$ , as introduced in Section 3. The following result establishes that this matrix does not correspond to a Euclidean distance matrix:

Lemma D.2 (Non-Euclidean Nature of Pairwise FGW Distance Matrix).

Consider the case where $d_{f} = ∥ \cdot ∥_{2}$ . Then the FGW distance matrix $𝐃$ , whose entries are given by

FGW_{p, α} (𝒢_{1}, 𝒢_{2}) : = \min_{𝝅 \in Π (𝝎_{1}, 𝝎_{2})} ⟨ (1 - α) 𝑴 + α 𝑳 (𝑨_{1}, 𝑨_{2}) \otimes 𝝅, 𝝅 ⟩,

with $α \in [0, 1]$ , does not define a Euclidean distance matrix.

As established in Lemma D.2, which is proved in Section D.4, the distance $FGW_{p, α}$ 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 D.2 and D.3, respectively.

D.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 $\hat{D}_{k l} : = ∥ 𝒆_{k} - 𝒆_{l} ∥_{2}^{2}$ and $\bar{\hat{D}}_{K l} : = ∥ \bar{𝒆}_{K} - 𝒆_{l} ∥_{2}^{2}$ induced by the embedding, and by incorporating the necessary conditions to ensure that $\hat{𝑫}$ corresponds to a valid Euclidean distance matrix, the reformulated problem becomes convex for $𝒮_{1}$ . Note that we can prove that $𝒮$ has a lower bound at $\hat{𝑳}^{*}$ , where $\hat{𝑳}^{*}$ is a minimizer of $𝒮_{1}$ , that is,

	$𝒮^{*}$	$= \min_{\hat{𝑫} \in ℰ} [𝒮_{1} (\hat{𝑫}) + 𝒮_{2} (\hat{𝑫})], 𝒮_{2} (\hat{𝑫}) : = \sum_{l \in [K]} (\bar{\hat{D}}_{K l} - \bar{D}_{K, l})^{2},$		(16)
	$𝒮_{1} (\hat{𝑳}^{*})$	$= \min_{\hat{𝑫} \in ℰ} 𝒮_{1} (\hat{𝑫}), 𝒮_{1} (\hat{𝑫}) : = \sum_{k, l \in [K]} (\hat{D}_{k l} - D_{k l})^{2} .$		(17)

Indeed, given the previous reformulation of $𝒮$ , we can establish the following lower bound via Proposition D.3. Notably, to simplify the problem, in Proposition D.3, 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 D.3.

Proposition D.3 (Error Lower Bound of $𝒮^{*}$ ).

The lower bound of $𝒮$ is provided as follows:

$𝒮^{*}$	$= \min_{\hat{𝑫} \in ℰ} [𝒮_{1} (\hat{𝑫}) + 𝒮_{2} (\hat{𝑫})] \geq 𝒮_{1} (\hat{𝑳}^{}) + 𝒮_{2} (\hat{𝑳}^{}) \geq ℒ_{1} + ℒ_{2} = : ℒ,$	(18)
$𝒮_{1} (\hat{𝑳}^{*})$	$= \min_{\hat{𝑫} \in ℰ_{ℒ}} 𝒮_{1} (\hat{𝑫}) \geq \sum_{k : λ_{k} < 0} λ_{k}^{2} = : ℒ_{1},$	(19)
$𝒮_{2} (\hat{𝑳}^{*})$	$= \min_{\hat{𝑫} \in ℰ_{ℒ}} 𝒮_{2} (\hat{𝑫}) = 0 = : ℒ_{2} .$	(20)

Here $ℰ_{ℒ}$ contains $ℰ$ by keeping only the PSD property from the EDM definition in Fact 1.

Proof of Proposition D.3.

Note that if $𝒮_{1}$ is minimized at $\hat{𝑳}^{*}$ , that is,

𝒮_{1} (\hat{𝑳}^{*})

= \min_{\hat{𝑫} \in ℰ} 𝒮_{1} (\hat{𝑫}), 𝒮_{1} (\hat{𝑫}) : = \sum_{k, l \in [K]} (\hat{D}_{k l} - D_{k l})^{2} .

(21)

We then can find the lower bound of $𝒮^{*} = \min_{\hat{𝑫} \in ℰ} [𝒮_{1} (\hat{𝑫}) + 𝒮_{2} (\hat{𝑫})]$ via the minimizer $\hat{𝑳}^{*}$ .

Using the definition of Frobenius norm and $ℰ_{ℒ}$ , we can obtain:

𝒮_{1} (\hat{𝑳}^{*}) : = \min_{\hat{𝑫} \in ℰ} 𝒮_{1} (\hat{𝑫}) \geq \min_{\hat{𝑫} \in ℰ_{ℒ}} 𝒮_{1} (\hat{𝑫}), 𝒮_{1} (\hat{𝑫}) = ∥ \hat{𝑫} - 𝑫 ∥^{2}_{F},

We then obtain the following decomposition:

	$∥ \hat{𝑫} - 𝑫 ∥^{2}_{F}$	$= ∥ 𝑨 ∥^{2}_{F} + ∥ 𝑩 ∥^{2}_{F}, 𝑨 : = 𝑪 \hat{𝑫} 𝑪 - 𝑪 𝑫 𝑪,$
	$𝑩$	$: = \frac{1}{K} 𝑶 \hat{𝑫} 𝑪 + \frac{1}{K} 𝑪 \hat{𝑫} 𝑶 + \frac{1}{K^{2}} 𝑶 \hat{𝑫} 𝑶 - (\frac{1}{K} 𝑶 𝑫 𝑪 + \frac{1}{K} 𝑪 𝑫 𝑶 + \frac{1}{K^{2}} 𝑶 𝑫 𝑶),$

where $𝑪 = 𝑰_{K} - \frac{1}{K} 𝑶$ is the centering matrix and $𝑶 = 𝟏_{K} 𝟏_{K}^{⊤}$ is the all-ones matrix. Indeed, using the definition of the centering matrix $𝑪 = 𝑰_{K} - \frac{1}{K} 𝑶$ , we have $𝑰_{K} = 𝑪 + \frac{1}{K} 𝑶$ .

∥ \hat{𝑫} - 𝑫 ∥^{2}_{F}

= ∥ 𝑰_{K} \hat{𝑫} 𝑰_{K} - 𝑰_{K} 𝑫 𝑰_{K} ∥^{2}_{F} = ∥ 𝑨 + 𝑩 ∥^{2}_{F} = ∥ 𝑨 ∥^{2}_{F} + ∥ 𝑩 ∥^{2}_{F} + 2 T r (𝑨 𝑩) = ∥ 𝑨 ∥^{2}_{F} + ∥ 𝑩 ∥^{2}_{F},

Here we used the fact that the matrix product is invariant under cyclic permutation:

T r (𝑨 𝑩) = T r (𝑪 (\hat{𝑫} - 𝑫) 𝑪 (\hat{𝑫} - 𝑫) \frac{1}{K} 𝑶) = T r (\frac{1}{K} 𝑶 𝑪 (\hat{𝑫} - 𝑫) 𝑪 (\hat{𝑫} - 𝑫)) = 0,

and

\frac{1}{K} 𝑶 𝑪 = \frac{1}{K} 𝑶 (𝑰_{K} - \frac{1}{K} 𝑶) = \frac{1}{K} 𝑶 - \frac{1}{K^{2}} 𝑶 𝑶 = 0 .

Under only the PSD constraint, the optimal solution $\hat{𝑳}^{*}$ that minimizes $𝒮_{1} (\hat{𝑫})$ can be decomposed as:

\hat{𝑳}^{*} = \hat{𝑳}^{*}_{𝑨} + \hat{𝑳}^{*}_{𝑩},

where $\hat{𝑳}^{*}_{𝑨}$ and $\hat{𝑳}^{*}_{𝑩}$ respectively minimize the terms $∥ 𝑨 ∥_{F}^{2}$ and $∥ 𝑩 ∥_{F}^{2}$ independently.

In particular, using the definition of the centering matrix $𝑪 = 𝑰_{K} - \frac{1}{K} 𝑶$ , the entries of $\hat{𝑳}^{*}_{𝑩}$ are given by:

	$\hat{𝑳}^{*}_{𝑩, k l}$	$: = [\frac{1}{K} 𝑶 𝑫 𝑪 + \frac{1}{K} 𝑪 𝑫 𝑶 + \frac{1}{K^{2}} 𝑶 𝑫 𝑶]_{k l}$
		$= [\frac{1}{K} 𝑶 𝑫 + \frac{1}{K} (𝑶 𝑫)^{⊤} - \frac{1}{K^{2}} 𝑶 𝑫 𝑶]_{k l} = \bar{𝑫}_{k} + \bar{𝑫}_{l} - \bar{𝑫},$

where $\bar{𝑫}_{k}$ denotes the mean of the $k$ -th row (or column) of $𝑫$ , and $\bar{𝑫}$ is the global mean of all elements in $𝑫$ . Therefore, the rows/columns mean of $\hat{𝑳}^{*}_{𝑩}$ equal those of $𝑫$ itself, and hence

\hat{𝑳}^{*}_{𝑩} = a r g m i n_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑩 ∥^{2}_{F}, \min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑩 ∥^{2}_{F} = 0 .

Therefore,

\min_{\hat{𝑫} \in ℰ_{ℒ}} 𝒮_{2} (\hat{𝑫}) = \min_{\hat{𝑫} \in ℰ_{ℒ}} \sum_{l \in [K]} (\bar{\hat{D}}_{K l} - \bar{D}_{K, l})^{2} = 0 .

Here we used the fact that the matrix $𝑫$ is given by $D_{k l} : = FGW_{p, α} (𝒢 (𝕊_{k}), 𝒢 (S_{l}))$ for all $k, l \in [K]$ and the empirical FGW barycenter is given by

	$\bar{𝒢}_{K}$	$\in a r g m i n_{𝒢 \in 𝒫_{p} (𝛀)} \frac{1}{K} \sum_{l = 1}^{K} FGW_{p, α}^{p} (𝒢, 𝒢 (S_{l})) = a r g m i n_{𝒢 \in 𝒫_{p} (𝛀)} \frac{1}{K} \sum_{l = 1}^{K} FGW_{p, α} (𝒢, 𝒢 (S_{l})),$
	$\bar{D}_{K, l}$	$: = FGW_{p, α} (\bar{𝒢}_{K}, 𝒢 (S_{l})) = \min_{𝒢 \in 𝒫_{p} (𝛀)} \frac{1}{K} \sum_{l = 1}^{K} FGW_{p, α} (𝒢, 𝒢 (S_{l})) (=: column l -th means of 𝑫),$

where $𝒫_{p} (𝛀)$ denotes the space of attributed graphs with finite $p$ -th order FGW distance. To approximate this barycenter in embedding space, we require

∥ \bar{𝒆}_{K} - 𝒆_{l} ∥_{2}^{2} \approx FGW_{p, α} (\bar{𝒢}_{K}, 𝒢 (S_{l})) : = \bar{D}_{K, l} for all l \in [K],

where $\bar{𝒆}_{K} = \frac{1}{K} \sum_{k = 1}^{K} 𝒆_{k}$ is the mean embedding and $𝒆_{k} : = 𝒯_{θ} (𝐇_{k})$ is the learned representation.

Now we would like to find a local analytic solution $\hat{𝑳}^{*}_{𝑨}$ minimizing $∥ 𝑨 ∥_{F}^{2}$ such that the global solution $\hat{𝑳}^{*} = \hat{𝑳}^{*}_{𝑨} + \hat{𝑳}^{*}_{𝑩}$ minimizes both terms $∥ 𝑨 ∥_{F}^{2}$ and $∥ 𝑩 ∥_{F}^{2}$ simultaneously. That is,

	$\min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑨 ∥_{F}^{2}$	$= \min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑪 (\hat{𝑳}_{𝑨} + \hat{𝑳}_{𝑩}) 𝑪 - 𝑪 𝑫 𝑪 ∥_{F}^{2}$
		$= ∥ 𝑪 (\hat{𝑳}^{}_{𝑨} + \hat{𝑳}^{}_{𝑩}) 𝑪 - 𝑪 𝑫 𝑪 ∥_{F}^{2} = ∥ 𝑪 \hat{𝑳}^{*}_{𝑨} 𝑪 - 𝑪 𝑫 𝑪 ∥_{F}^{2} .$

Here we used the fact that by definition of $\hat{𝑳}^{*}_{𝑩}$ , it holds that $𝑪 \hat{𝑳}^{*}_{𝑩} 𝑪 = 0$ . Hence, the optimization becomes:

\min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑪 \hat{𝑳}_{𝑨} 𝑪 - 𝑪 𝑫 𝑪 ∥_{F}^{2} .

This is in fact the problem of computing the nearest PSD approximation $𝑪 \hat{𝑳}_{𝑨} 𝑪$ to a symmetric matrix $𝑪 𝑫 𝑪$ . Using the result from (Higham, 1988), we find the analytic solution as follows:

\hat{𝑳}^{*}_{𝑨} = - \sum_{k : λ_{k} > 0} λ_{k} 𝒗_{k} 𝒗_{k}^{⊤} .

(22)

Here ${λ_{k}, 𝒗_{k}}_{k \in [K]}$ are the eigenvalues and eigenvectors of $𝑭 = - 𝑪 𝑫 𝑪$ . Because $𝑪 𝑫 𝑪$ has rows/columns means $0$ , the ones vector $𝟏_{K}$ is an eigenvector of $𝑪 𝑫 𝑪$ with eigenvalue $0$ . This leads to $𝟏_{K}$ is also in the null space $\hat{𝑳}^{*}_{𝑨}$ and:

\hat{𝑳}^{*}_{𝑨} = 𝑪 \hat{𝑳}^{*}_{𝑨} 𝑪, \frac{1}{K} 𝑶 \hat{𝑳}^{*}_{𝑨} = \frac{1}{K} (𝑶 \hat{𝑳}^{*}_{𝑨})^{⊤} = 0 .

Therefore,

∥ \hat{𝑳}^{*} - 𝑫 ∥^{2}_{F} = ∥ \hat{𝑳}^{*}_{𝑨} + \hat{𝑳}^{*}_{𝑩} - 𝑫 ∥^{2}_{F} = \sum_{k : λ_{k} < 0} λ_{k}^{2} .

Combining all together, Proposition D.3 is derived as follows:

𝒮^{*} \geq \min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑨 ∥^{2}_{F} + \min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑩 ∥^{2}_{F} + \min_{\hat{𝑫} \in ℰ_{ℒ}} 𝒮_{2} (\hat{𝑫}) = \sum_{k : λ_{k} < 0} λ_{k}^{2} + 0 + 0 = \sum_{k : λ_{k} < 0} λ_{k}^{2} = : ℒ .

∎

D.3 Upper Bounds on Embedding of Pairwise Empirical FGW Barycenter Distances

As discussed in Section D.2, the lower bound stated in Proposition D.3 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 D.4, 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 D.4 (Error Upper Bound of $𝒮^{*}$ ).

There exists a matrix $\hat{𝐔}^{*} \in ℰ$ such that the following upper bounds hold:

$𝒮^{*}$	$= \min_{\hat{𝑫} \in ℰ} [𝒮_{1} (\hat{𝑫}) + 𝒮_{2} (\hat{𝑫})] \leq 𝒮_{1} (\hat{𝑼}^{}) + 𝒮_{2} (\hat{𝑼}^{}) \leq 𝒰_{1} + 𝒰_{2} = : 𝒰,$	(23)
$𝒮_{1} (\hat{𝑼}^{*})$	$= \min_{\hat{𝑫} \in ℰ} 𝒮_{1} (\hat{𝑫}) \leq 𝒰_{1} : = \sum_{k : λ_{k} < 0} λ_{k}^{2} + \sum_{k l} (Δ p_{k} + Δ p_{l})^{2},$
$Δ p_{k}$	$= \frac{1}{2} \sum_{l : λ_{l} < 0} λ_{l} \cdot 𝒗_{k l}^{2}, \forall k \in [K]$	(24)
$𝒮_{2} (\hat{𝑼}^{*})$	$= \min_{\hat{𝑫} \in ℰ} 𝒮_{2} (\hat{𝑫}) \leq \sum_{l} (Δ \bar{𝒑}_{l})^{2} = : 𝒰_{2},$	(25)

where the aggregated error term is defined as:

Δ \bar{𝒑}_{l} : = \frac{1}{2 K} \sum_{k = 1}^{K} \sum_{l : λ_{l} < 0} λ_{l} \cdot 𝒗_{k l}^{2} .

We aim to exploit the information derived from the truncation of the negative eigenspace of the matrix $𝑪 𝑫 𝑪$ , specifically the matrix introduced in Equation 22, defined as:

\hat{𝑳}^{*}_{𝑨} = - \sum_{k : λ_{k} > 0} λ_{k} 𝒗_{k} 𝒗_{k}^{⊤},

where ${λ_{k}, 𝒗_{k}}_{k \in [K]}$ denote the eigenvalues and corresponding eigenvectors of the matrix $𝑭 = - 𝑪 𝑫 𝑪$ .

Recall that the entries of $\hat{𝑳}^{*}_{𝑩}$ are given by:

\hat{𝑳}^{*}_{𝑩, k l}

= [\frac{1}{K} 𝑶 𝑫 + \frac{1}{K} (𝑶 𝑫)^{⊤} - \frac{1}{K^{2}} 𝑶 𝑫 𝑶]_{k l} = \bar{𝑫}_{k} + \bar{𝑫}_{l} - \bar{𝑫} .

As a consequence, the sum $\hat{𝑳}^{*}_{𝑨} + \hat{𝑳}^{*}_{𝑩}$ 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 $\hat{𝑳}^{*}_{𝑨}$ , resulting in the matrix $\hat{𝑼}^{*} : = \hat{𝑳}^{*}_{𝑨} + 𝑷$ , which is both hollow and PSD. This adjustment is designed to avoid any additional penalty on the term $∥ 𝑨 ∥^{2}_{F}$ , though it may introduce some approximation errors in $∥ 𝑩 ∥^{2}_{F}$ and in the quantity $𝒮_{2}$ . 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 $∥ 𝑨 ∥_{F}^{2}$ . This can be ensured by imposing the constraint $𝑪 𝑷 𝑪 = 0$ . Under this condition, the matrix $\hat{𝑼}^{*}$ remains a minimizer of $∥ 𝑨 ∥_{F}^{2}$ , as demonstrated below:

	$\min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑨 ∥_{F}^{2}$	$= \min_{\hat{𝑫} \in ℰ_{ℒ}} ∥ 𝑪 (\hat{𝑳}_{𝑨} + \hat{𝑳}_{𝑩}) 𝑪 - 𝑪 𝑫 𝑪 ∥_{F}^{2}$
		$= ∥ 𝑪 (\hat{𝑳}^{}_{𝑨} + 𝑷 + \hat{𝑳}^{}_{𝑩}) 𝑪 - 𝑪 𝑫 𝑪 ∥_{F}^{2}$
		$= ∥ 𝑪 \hat{𝑳}^{*}_{𝑨} 𝑪 - 𝑪 𝑫 𝑪 ∥_{F}^{2},$

where the final equality holds due to the constraint $𝑪 𝑷 𝑪 = 0$ .

This leads to the condition $(𝑪 𝑷) 𝑪 = 𝑪 (𝑷 𝑪) = 0$ , 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:

𝟏_{K} 𝒄^{⊤} = 𝑷 𝑪 = 𝑷 (𝑰_{K} - \frac{1}{K} 𝑶) or 𝑷 = 𝟏_{K} 𝒄^{⊤} + 𝑷 \frac{1}{K} 𝑶,

where $𝒄$ is a column vector to be defined subsequently. Here, we have used the fact that $𝑪$ is the centering matrix defined by $𝑪 = 𝑰_{K} - \frac{1}{K} 𝑶$ .

Multiplying both sides on the left by $\frac{1}{K} 𝑶$ yields:

\frac{1}{K} 𝑶 𝑷 = \frac{1}{K} 𝑶 𝟏_{K} 𝒄^{⊤} + \frac{1}{K} 𝑶 (\frac{1}{K} 𝑷 𝑶) = 𝟏_{K} 𝒄^{⊤} + \frac{1}{K^{2}} 𝑶 𝑷 𝑶 .

This leads to

𝒄^{⊤} = \frac{1}{K} 𝟏_{K}^{⊤} 𝑷 - \frac{1}{K^{2}} 𝟏_{K}^{⊤} 𝑶 𝑷 𝑶 .

Indeed, via the definition of $𝑶 = 𝟏_{K} 𝟏_{K}^{⊤}$ , we can verify this as follows:

	$𝟏_{K} 𝒄^{⊤} + \frac{1}{K^{2}} 𝑶 𝑷 𝑶$	$= 𝟏_{K} (\frac{1}{K} 𝟏_{K}^{⊤} 𝑷 - \frac{1}{K^{2}} 𝟏_{K}^{⊤} 𝑶 𝑷 𝑶) + \frac{1}{K^{2}} 𝑶 𝑷 𝑶$
		$= \frac{1}{K} 𝟏_{K} 𝟏_{K}^{⊤} 𝑷 - \frac{1}{K^{2}} 𝟏_{K} 𝟏_{K}^{⊤} 𝑶 𝑷 𝑶 + \frac{1}{K^{2}} 𝑶 𝑷 𝑶$
		$= \frac{1}{K} 𝟏_{K} 𝟏_{K}^{⊤} 𝑷 - \frac{1}{K^{2}} 𝑶 𝑶 𝑷 𝑶 + \frac{1}{K^{2}} 𝑶 𝑷 𝑶$
		$= \frac{1}{K} 𝑶 𝑷 .$

Hence,

	$𝑷$	$= 𝟏_{K} (\frac{1}{K} 𝟏_{K}^{⊤} 𝑷 - \frac{1}{K^{2}} 𝟏_{K}^{⊤} 𝑶 𝑷 𝑶) + 𝑷 \frac{1}{K} 𝑶$
		$= \frac{1}{K} 𝟏_{K} (𝟏_{K}^{⊤} 𝑷) + \frac{1}{K} (𝑷 𝟏_{K}) 𝟏_{K}^{⊤} - \frac{1}{K^{2}} 𝟏_{K} 𝟏_{K}^{⊤} 𝑶 𝑷 𝑶$

Since $𝑷 𝟏_{K}$ is a column vector, to satisfy this constraint, $𝑷$ must be of the form:

𝑷 = 𝟏_{K} \frac{𝒑^{⊤}}{K} + \frac{𝒑}{K} 𝟏_{K}^{⊤} - \hat{𝒑} \frac{𝟏_{K} 𝟏_{K}^{⊤}}{K},

where $𝒑 \in ℝ^{K}$ is a vector of free parameters, and $\hat{𝒑}$ denotes its average. This construction implies that $𝑷$ has only $K$ degrees of freedom. However, to ensure that $\hat{𝑳}^{*}_{𝑨} + 𝑷$ has zero diagonal (i.e., the resulting matrix is hollow), the diagonal entries of $𝑷$ must satisfy the following $K$ linear constraints:

𝒑_{k} - \frac{1}{2} \hat{𝒑} = - \frac{1}{2} [\hat{𝑳}^{*}_{𝑨}]_{k k}, \forall k \in [K] .

Solving this linear system yields:

	$𝒑_{k}$	$= \frac{1}{2} (\sum_{l : λ_{l} > 0} λ_{l} \cdot 𝒗_{k l}^{2} + \frac{1}{K} \hat{𝒑}),$
	$\hat{𝒑}$	$= \frac{1}{K} \sum_{k = 1}^{K} 𝒑_{k} = \frac{1}{K} \sum_{k = 1}^{K} \sum_{l : λ_{l} > 0} λ_{l} \cdot 𝒗_{k l}^{2},$

where we have used the fact that $\hat{𝑳}^{*}_{𝑨} = - \sum_{l : λ_{l} > 0} λ_{l} 𝒗_{l} 𝒗_{l}^{⊤}$ , and hence its diagonal entries are given by $[\hat{𝑳}^{*}_{𝑨}]_{k k} = - \sum_{l : λ_{l} > 0} λ_{l} \cdot 𝒗_{k l}^{2}$ .

Consequently, the resulting matrix $𝑷$ can be expressed element-wise as:

𝑷_{k, l} = - \frac{[\hat{𝑳}^{*}_{𝑨}]_{k k} + [\hat{𝑳}^{*}_{𝑨}]_{l l}}{2} \geq 0,

where the inequality follows from the fact that $\hat{𝑳}^{*}_{𝑨}$ is negative semi-definite.

In summary, the matrix $\hat{𝑼}^{*} : = \hat{𝑳}^{*}_{𝑨} + 𝑷$ satisfies all three constraints specified in Definition D.1.

Although $\hat{𝑼}^{*}$ preserves the value of $∥ 𝑨 ∥_{F}^{2}$ , it differs from $\hat{𝑳}^{*}_{𝑨}$ and introduces approximation errors in the $∥ 𝑩 ∥_{F}^{2}$ term and the $𝒮_{2}$ term. Note that the sum of the untruncated version of $𝑪 𝑫 𝑪$ and the matrix

\frac{1}{K} 𝑶 𝑫 𝑪 + \frac{1}{K} 𝑪 𝑫 𝑶 + \frac{1}{K^{2}} 𝑶 𝑫 𝑶

is equal to $𝑫$ and remains hollow. Recall the decomposition:

	$∥ \hat{𝑫} - 𝑫 ∥_{F}^{2}$	$= ∥ 𝑨 ∥_{F}^{2} + ∥ 𝑩 ∥_{F}^{2}, 𝑨 : = 𝑪 \hat{𝑫} 𝑪 - 𝑪 𝑫 𝑪,$
	$𝑩$	$: = \frac{1}{K} 𝑶 \hat{𝑫} 𝑪 + \frac{1}{K} 𝑪 \hat{𝑫} 𝑶 + \frac{1}{K^{2}} 𝑶 \hat{𝑫} 𝑶$
		$- (\frac{1}{K} 𝑶 𝑫 𝑪 + \frac{1}{K} 𝑪 𝑫 𝑶 + \frac{1}{K^{2}} 𝑶 𝑫 𝑶),$

where $𝑪 = 𝑰_{K} - \frac{1}{K} 𝑶$ is the centering matrix and $𝑶 = 𝟏_{K} 𝟏_{K}^{⊤}$ is the all-ones matrix.

The matrix

\frac{1}{K} 𝑶 𝑫 𝑪 + \frac{1}{K} 𝑪 𝑫 𝑶 + \frac{1}{K^{2}} 𝑶 𝑫 𝑶

can be written similarly to $𝑷$ by including the contributions from the negative eigenvalues, resulting in the matrix $\tilde{𝑷}$ , parameterized by:

	$\tilde{𝒑}_{k}$	$= \frac{1}{2} (\sum_{l} λ_{l} \cdot 𝒗_{k l}^{2} + \frac{1}{K} \tilde{\hat{𝒑}}),$
	$\tilde{\hat{𝒑}}$	$= \frac{1}{K} \sum_{k = 1}^{K} \tilde{𝒑}_{k} = \frac{1}{K} \sum_{k = 1}^{K} \sum_{l} λ_{l} \cdot 𝒗_{k l}^{2} .$

Define the correction due to negative eigenvalues as:

Δ 𝒑_{k} : = \frac{1}{2} \sum_{l : λ_{l} < 0} λ_{l} \cdot 𝒗_{k l}^{2}, \forall k \in [K] .

The resulting error in the $∥ 𝑩 ∥_{F}^{2}$ term is given by:

∥ 𝑩 ∥_{F}^{2}

= ∥ \tilde{𝑷} - 𝑷 ∥_{F}^{2} = \sum_{k, l} (Δ 𝒑_{k} + Δ 𝒑_{l})^{2} .

Furthermore, the contribution to $𝒮_{2}$ is bounded as:

𝒮_{2}

= \min_{\hat{𝑫} \in ℰ} 𝒮_{2} (\hat{𝑫}) = \sum_{l \in [K]} (\bar{\hat{𝑫}}_{K, l} - \bar{𝑫}_{K, l})^{2} \leq \sum_{l} (Δ \bar{𝒑}_{l})^{2} = : 𝒰_{2},

where the aggregated error term is defined as:

Δ \bar{𝒑}_{l} : = \frac{1}{2 K} \sum_{k = 1}^{K} \sum_{l : λ_{l} < 0} λ_{l} \cdot 𝒗_{k l}^{2} .

D.4 Proof of Lemma D.2

The proof is proved via leveraging Proposition 8.2 from (Peyré et al., 2019), applied to the specific case $α = 0$ , and relies on the relationships among FGW, Wasserstein (W), and Gromov-Wasserstein (GW) distances.

The FGW cost $FGW_{p, α} (𝒢_{1}, 𝒢_{2})$ is defined via two components: the node feature cost matrix $𝑴 [i, j] = d_{f} (𝑯_{1} [i], 𝑯_{2} [j])^{p}$ , and the structural discrepancy tensor $𝑳 (𝑨_{1}, 𝑨_{2}) [i, j, l, m] = | 𝑨_{1} [i, j] - 𝑨_{2} [l, m] |^{p}$ .

Let $𝒢_{1} = (𝑯_{1}, 𝑨_{1}, 𝝎_{1})$ and $𝒢_{2} = (𝑯_{2}, 𝑨_{2}, 𝝎_{2})$ be two attributed graphs with $N_{1}$ and $N_{2}$ nodes, respectively. Their associated probability measures are

μ_{1} = \sum_{k} ω_{1 k} δ_{(𝒙_{1 k}, 𝒂_{1 k})}, μ_{2} = \sum_{l} ω_{2 l} δ_{(𝒙_{2 l}, 𝒂_{2 l})} .

We define the marginals $μ_{𝑯_{1}} = \sum_{k} ω_{k} δ_{𝒙_{k}}$ and $μ_{𝑨_{1}} = \sum_{k} ω_{k} δ_{𝒂_{k}}$ (and analogously for $μ_{𝑯_{2}}$ and $μ_{𝑨_{2}}$ ) as projections of $μ_{1}$ and $μ_{2}$ onto the feature and structural spaces, respectively.

Using these definitions, we introduce the following notation:

$J_{p} (𝑨_{1}, 𝑨_{2}, 𝝅)$	$= \sum_{i j k l} L_{i j k l} (𝑨_{1}, 𝑨_{2})^{p} 𝝅_{i j} 𝝅_{k l},$	(26)
$GW_{p} (μ_{𝑯_{1}}, μ_{𝑯_{2}})^{p}$	$= \min_{𝝅 \in 𝚷 (𝝎_{1}, 𝝎_{2})} J_{p} (𝑨_{1}, 𝑨_{2}, 𝝅),$	(27)
$H_{p} (𝑴, 𝝅)$	$= \sum_{k l} d_{f} (𝒙_{1 k}, 𝒙_{2 l})^{p} 𝝅_{k l},$	(28)
$W_{p} (μ_{𝑨_{1}}, μ_{𝑨_{2}})^{p}$	$= \min_{𝝅 \in 𝚷 (𝝎_{1}, 𝝎_{2})} H_{p} (𝑴, 𝝅) .$	(29)

Let $𝝅 \in 𝚷 (𝝎_{1}, 𝝎_{2})$ be any admissible coupling. If both $μ_{1}$ and $μ_{2}$ are defined over a common metric space $(𝛀, 𝑨, μ)$ , then the FGW distance is given by:

FGW_{p, α} (𝒢_{1}, 𝒢_{2}) : = \min_{𝝅 \in Π (𝝎_{1}, 𝝎_{2})} ⟨ (1 - α) 𝑴 + α 𝑳 (𝑨_{1}, 𝑨_{2}) \otimes 𝝅, 𝝅 ⟩ .

(30)

We now derive the following fundamental identity:

	$𝔼_{p, α} (𝑴, 𝑨_{1}, 𝑨_{2}, 𝝅)$	$: = \sum_{i j k l} [(1 - α) d_{f} (𝒙_{1 k}, 𝒙_{2 l})^{p} + α \| 𝑨_{1} (i, k) - 𝑨_{2} (j, l) \|^{p}] 𝝅_{i j} 𝝅_{k l}$
		$= (1 - α) H_{p} (𝑴, 𝝅) + α J_{p} (𝑨_{1}, 𝑨_{2}, 𝝅) .$		(31)

Moreover, let $𝝅_{α}$ denote the optimal coupling that minimizes the FGW objective $𝔼_{p, α} (𝑴, 𝑨_{1}, 𝑨_{2}, \cdot)$ . Then the FGW distance admits the following decomposition:

$FGW_{p, α}^{p} (μ_{1}, μ_{2})$	$= \min_{𝝅 \in 𝚷 (𝝎_{1}, 𝝎_{2})} 𝔼_{p, α} (𝑴, 𝑨_{1}, 𝑨_{2}, 𝝅) = 𝔼_{p, α} (𝑴, 𝑨_{1}, 𝑨_{2}, 𝝅_{α})$
	$= (1 - α) H_{p} (𝑴, 𝝅_{α}) + α J_{p} (𝑨_{1}, 𝑨_{2}, 𝝅_{α})$
	$\geq (1 - α) W_{p}^{p} (μ_{𝑨_{1}}, μ_{𝑨_{2}}) + α GW_{p}^{p} (μ_{𝑯_{1}}, μ_{𝑯_{2}}) .$	(32)

This inequality follows from the optimality of the W and GW distances with respect to the cost functions $H_{p}$ and $J_{p}$ , respectively, and highlights the interpolation nature of the FGW distance between these two metrics as governed by the parameter $α$ .

The generalized FGW cost $𝔼_{p, α} (𝑴, 𝑨_{1}, 𝑨_{2}, 𝝅)$ admits the following explicit formulation:

	$𝔼_{p, α} (𝑴, 𝑨_{1}, 𝑨_{2}, 𝝅)$	$= ⟨ (1 - α) 𝑴^{p} + α 𝑳 (𝑨_{1}, 𝑨_{2})^{p} \otimes 𝝅, 𝝅 ⟩$
		$= \sum_{i, j, k, l} [(1 - α) d_{f} (𝒙_{1 k}, 𝒙_{2 l})^{p} + α \| 𝑨_{1} (i, k) - 𝑨_{2} (j, l) \|^{p}] 𝝅_{i j} 𝝅_{k l} .$

Based on the formulation above, we obtain the following upper bound on the FGW distance:

	$FGW_{p, α} (G_{1}, G_{2})$	$\leq ⟨ (1 - α) 𝑴 + α 𝑳 (𝑨_{1}, 𝑨_{2}) \otimes 𝝅, 𝝅 ⟩$
		$\leq \sum_{k, l} [(1 - α) d_{f} (𝒙_{1 k}, 𝒙_{2 l}) + 2^{p - 1} α 𝑨 [k, l]]^{p} 𝝅_{k l},$		(33)

where the second inequality follows from the convexity of the function $x \mapsto x^{p}$ for $p \geq 1$ and an application of Minkowski-type bounds on the structural term. Importantly, inequality in equation (33) holds for any admissible coupling $𝝅 \in 𝚷 (𝝎_{1}, 𝝎_{2})$ , and in particular, it remains valid when $𝝅 = \bar{𝝅}$ , the optimal coupling associated with the Wasserstein distance $W_{p} (μ_{1}, μ_{2})$ over the product metric space $(𝛀, \bar{d})$ . Here, the effective distance $\bar{d}$ between structured nodes $(𝒙_{1}, 𝒂_{1})$ and $(𝒙_{2}, 𝒂_{2})$ is defined as:

\bar{d} ((𝒙_{1}, 𝒂_{1}), (𝒙_{2}, 𝒂_{2})) = (1 - α) d_{f} (𝒙_{1}, 𝒙_{2}) + 2^{p - 1} α 𝑨 (𝒂_{1}, 𝒂_{2}) .

Combining this with the Wasserstein formulation in equation (29), we observe the following inequality:

FGW_{p, α} (𝒢_{1}, 𝒢_{2}) \leq W_{p} (μ_{𝑨_{1}}, μ_{𝑨_{2}}), and FGW_{p, α} (𝒢_{1}, 𝒢_{2}) = W_{p} (μ_{𝑨_{1}}, μ_{𝑨_{2}}) when α = 0 .

(34)

[bib.bib1] Ahmad et al. (2022) Ahmad, W., Simon, E., Chithrananda, S., Grand, G., and Ramsundar, B. Chemberta-2: Towards chemical foundation models, 2022. ↩ a b c

[bib.bib2] Amos et al. (2023) Amos, B., Cohen, S., Luise, G., and Redko, I. Meta optimal transport. International Conference on Machine Learning, 2023. ↩ a

[bib.bib3] Axelrod & Gomez-Bombarelli (2022) Axelrod, S. and Gomez-Bombarelli, R. Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022. ↩ a

[bib.bib4] Axelrod & Gómez-Bombarelli (2023) Axelrod, S. and Gómez-Bombarelli, R. 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 b c

[bib.bib5] Bachmann et al. (2022) Bachmann, F., Hennig, P., and Kobak, D. Wasserstein t-sne. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 104–120. Springer, 2022. ↩ a

[bib.bib6] Balaban (1985) Balaban, A. T. Applications of graph theory in chemistry. Journal of chemical information and computer sciences, 25(3):334–343, 1985. ↩ a

[bib.bib7] Batatia et al. (2022) Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., and Csanyi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 11423–11436. Curran Associates, Inc., 2022. ↩ a

[bib.bib8] Batatia et al. (2023) Batatia, I., Benner, P., Chiang, Y., Elena, A. M., Kovács, D. P., Riebesell, J., Advincula, X. R., Asta, M., Baldwin, W. J., Bernstein, N., Bhowmik, A., Blau, S. M., Cărare, V., Darby, J. P., De, S., Pia, F. D., Deringer, V. L., Elijošius, R., El-Machachi, Z., Fako, E., Ferrari, A. C., Genreith-Schriever, A., George, J., Goodall, R. E. A., Grey, C. P., Han, S., Handley, W., Heenen, H. H., Hermansson, K., Holm, C., Jaafar, J., Hofmann, S., Jakob, K. S., Jung, H., Kapil, V., Kaplan, A. D., Karimitari, N., Kroupa, N., Kullgren, J., Kuner, M. C., Kuryla, D., Liepuoniute, G., Margraf, J. T., Magdău, I.-B., Michaelides, A., Moore, J. H., Naik, A. A., Niblett, S. P., Norwood, S. W., O’Neill, N., Ortner, C., Persson, K. A., Reuter, K., Rosen, A. S., Schaaf, L. L., Schran, C., Sivonxay, E., Stenczel, T. K., Svahn, V., Sutton, C., van der Oord, C., Varga-Umbrich, E., Vegge, T., Vondrák, M., Wang, Y., Witt, W. C., Zills, F., and Csányi, G. A foundation model for atomistic materials chemistry, 2023. ↩ a

[bib.bib9] Batzner et al. (2022) Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nat. Commun., 13(1):2453, May 2022. ISSN 2041-1723. ↩ a

[bib.bib10] Brogat-Motte et al. (2022) Brogat-Motte, L., Flamary, R., Brouard, C., Rousu, J., and d’Alché Buc, F. Learning to predict graphs with fused gromov-wasserstein barycenters. In International Conference on Machine Learning, pp. 2321–2335. PMLR, 2022. ↩ a b c

[bib.bib11] Butler et al. (2018) Butler, K. T., Davies, D. W., Cartwright, H., Isayev, O., and Walsh, A. Machine learning for molecular and materials science. Nature, 559(7715):547–555, Jul 2018. ISSN 1476-4687. ↩ a

[bib.bib12] Chan et al. (2024) Chan, B. W., Lynch, N. B., Tran, W., Joyce, J. M., Savage, G. P., Meutermans, W., Montgomery, A. P., and Kassiou, M. Fragment-based drug discovery for disorders of the central nervous system: designing better drugs piece by piece. Frontiers in Chemistry, 12:1379518, 2024. ↩ a

[bib.bib13] Cheng et al. (2025) Cheng, A. H., Sun, C., and Aspuru-Guzik, A. Scalable autoregressive 3d molecule generation. arXiv preprint arXiv:2505.13791, 2025. ↩ a

[bib.bib14] Choudhary et al. (2022) Choudhary, K., DeCost, B., Chen, C., Jain, A., Tavazza, F., Cohn, R., Park, C. W., Choudhary, A., Agrawal, A., Billinge, S. J. L., Holm, E., Ong, S. P., and Wolverton, C. Recent advances and applications of deep learning methods in materials science. npj Comput. Mater., 8(1):59, Apr 2022. ISSN 2057-3960. ↩ a

[bib.bib15] Chung et al. (2019) Chung, Y. G., Haldoupis, E., Bucior, B. J., Haranczyk, M., Lee, S., Zhang, H., Vogiatzis, K. D., Milisavljevic, M., Ling, S., Camp, J. S., 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

[bib.bib16] Cormen et al. (2022) Cormen, T. H., Leiserson, C. E., Rivest, R. L., and Stein, C. Introduction to algorithms. MIT press, 2022. ↩ a

[bib.bib17] Courty et al. (2017) Courty, N., Flamary, R., and Ducoffe, M. Learning wasserstein embeddings. International Conference on Learning Representations (ICLR), 2017. ↩ a

[bib.bib18] Cuturi (2013) Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013. ↩ a

[bib.bib19] Cuturi et al. (2020) Cuturi, M., Teboul, O., Niles-Weed, J., and Vert, J.-P. Supervised quantile normalization for low rank matrix factorization. In International Conference on Machine Learning, pp. 2269–2279. PMLR, 2020. ↩ a

[bib.bib20] Du et al. (2022) Du, W., Zhang, H., Du, Y., Meng, Q., Chen, W., Zheng, N., Shao, B., and Liu, T.-Y. Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning, pp. 5583–5608. PMLR, 2022. ↩ a b c

[bib.bib21] Du et al. (2023) Du, Y., Wang, L., Feng, D., Wang, G., Ji, S., Gomes, C. P., Ma, Z.-M., 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 b c

[bib.bib22] Fang et al. (2021) Fang, X., Liu, L., Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H., and Wang, H. Chemrl-gem: Geometry enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 2021. doi: 10.48550/ARXIV.2106.06130. ↩ a

[bib.bib23] Fang et al. (2022) Fang, X., Liu, L., Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H., and Wang, H. Geometry-enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 4(2):127–134, 2022. ↩ a b

[bib.bib24] Fedik et al. (2022) Fedik, N., Zubatyuk, R., Kulichenko, M., Lubbers, N., Smith, J. S., Nebgen, B., Messerly, R., Li, Y. W., Boldyrev, A. I., Barros, K., Isayev, O., and Tretiak, S. 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

[bib.bib25] Fey et al. (2020) Fey, M., Yuen, J.-G., and Weichert, F. Hierarchical inter-message passing for learning on molecular graphs. arXiv preprint arXiv:2006.12179, 2020. ↩ a

[bib.bib26] Feydy et al. (2019) Feydy, J., Séjourné, T., Vialard, F.-X., Amari, S.-i., Trouvé, A., and Peyré, G. 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

[bib.bib27] Gasteiger et al. (2021) Gasteiger, J., Becker, F., and Günnemann, S. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790–6802, 2021. ↩ a b c d

[bib.bib28] Genevay et al. (2018) Genevay, A., Peyré, G., and Cuturi, M. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608–1617. PMLR, 2018. ↩ a

[bib.bib29] Geng et al. (2023) Geng, Z., Xie, S., Xia, Y., Wu, L., Qin, T., Wang, J., Zhang, Y., Wu, F., and Liu, T.-Y. De novo molecular generation via connection-aware motif mining. arXiv preprint arXiv:2302.01129, 2023. ↩ a

[bib.bib30] Gilmer et al. (2017a) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp. 1263–1272, 2017a. ↩ a b

[bib.bib31] Gilmer et al. (2017b) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chemistry. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th ICML, volume 70 of Proceedings of Machine Learning Research, pp. 1263–1272. PMLR, 06–11 Aug 2017b. ↩ a

[bib.bib32] Gower (1985) Gower, J. 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. URL https://www.sciencedirect.com/science/article/pii/0024379585901879. ↩ a

[bib.bib33] Haviv et al. (2024) Haviv, D., Kunes, R. Z., Dougherty, T., Burdziak, C., Nawy, T., Gilbert, A., and Pe’Er, D. Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformer. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 17697–17718. PMLR, July 2024. URL https://proceedings.mlr.press/v235/haviv24a.html. ↩ a b c

[bib.bib34] Hawkins (2017) Hawkins, P. C. Conformation generation: the state of the art. Journal of chemical information and modeling, 57(8):1747–1756, 2017. ↩ a

[bib.bib35] Higham (1988) Higham, N. J. 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. URL https://www.sciencedirect.com/science/article/pii/0024379588902236. ↩ a

[bib.bib36] Hu et al. (2020a) Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118–22133, 2020a. ↩ a b

[bib.bib37] Hu et al. (2020b) Hu, W., Liu, B., Gomes, J., Zitnik, M., Liang, P., Pande, V., and Leskovec, J. Strategies for pre-training graph neural networks. In International Conference on Learning Representations, 2020b. ↩ a b

[bib.bib38] Kipf & Welling (2017) Kipf, T. N. and Welling, M. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017. ↩ a

[bib.bib39] Klicpera et al. (2020) Klicpera, J., Groß, J., Günnemann, S., et al. Directional message passing for molecular graphs. In International Conference on Learning Representations, pp. 1–13, 2020. ↩ a b c

[bib.bib40] Kong et al. (2022) Kong, X., Huang, W., Tan, Z., and Liu, Y. Molecule generation by principal subgraph mining and assembling. Advances in Neural Information Processing Systems, 35:2550–2563, 2022. ↩ a

[bib.bib41] Kreuzer et al. (2021) Kreuzer, D., Beaini, D., Hamilton, W., Létourneau, V., and Tossou, P. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34:21618–21629, 2021. ↩ a

[bib.bib42] Landrum (2016) Landrum, G. Rdkit: open-source cheminformatics http://www.rdkit.org, 2016. ↩ a

[bib.bib43] Lee et al. (2024) Lee, S., Lee, S., Kawaguchi, K., and Hwang, S. J. Drug discovery with dynamic goal-aware fragments. International Conference on Machine Learning, 2024. ↩ a

[bib.bib44] Li (2020) Li, Q. Application of fragment-based drug discovery to versatile targets. Frontiers in molecular biosciences, 7:180, 2020. ↩ a

[bib.bib45] Liu et al. (2021) Liu, M., Fu, C., Zhang, X., Wang, L., Xie, Y., Yuan, H., Luo, Y., Xu, Z., Xu, S., and Ji, S. Fast quantum property prediction via deeper 2d and 3d graph networks. arXiv preprint arXiv:2106.08551, 2021. ↩ a b

[bib.bib46] Liu et al. (2015) Liu, Z., Li, Y., Han, L., Li, J., Liu, J., Zhao, Z., Nie, W., Liu, Y., and Wang, R. Pdb-wide collection of binding data: current status of the pdbbind database. Bioinformatics, 31(3):405–412, 2015. ↩ a

[bib.bib47] Luo et al. (2024) Luo, Y., Li, H., Shi, L., and Wu, X.-M. Enhancing graph transformers with hierarchical distance structural encoding. Advances in Neural Information Processing Systems, 37:57150–57182, 2024. ↩ a

[bib.bib48] Ma et al. (2023) Ma, X., Chu, X., Wang, Y., Lin, Y., Zhao, J., Ma, L., and Zhu, W. Fused gromov-wasserstein graph mixup for graph-level classifications. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. ↩ a b

[bib.bib49] Medrano Sandonas et al. (2024) Medrano Sandonas, L., Van Rompaey, D., Fallani, A., Hilfiker, M., Hahn, D., Perez-Benito, L., Verhoeven, J., Tresadern, G., Kurt Wegner, J., Ceulemans, H., et al. Dataset for quantum-mechanical exploration of conformers and solvent effects in large drug-like molecules. Scientific Data, 11(1):742, 2024. ↩ a

[bib.bib50] Merlot et al. (2003) Merlot, C., Domine, D., Cleva, C., and Church, D. J. Chemical substructures in drug discovery. Drug Discovery Today, 8(13):594–602, 2003. ↩ a

[bib.bib51] Morgan (1965) Morgan, H. L. 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

[bib.bib52] Nguyen et al. (2024a) Nguyen, D. M., Lukashina, N., Nguyen, T., Le, A. T., Nguyen, T., Ho, N., Peters, J., Sonntag, D., Zaverkin, V., and Niepert, M. Structure-aware e (3)-invariant molecular conformer aggregation networks. International Conference on Machine Learning, 2024a. ↩ a b c d

[bib.bib53] Nguyen et al. (2024b) Nguyen, D. M. H., Lukashina, N., Nguyen, T., Le, A. T., Nguyen, T., Ho, N., Peters, J., Sonntag, D., Zaverkin, V., and Niepert, M. 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 h i

[bib.bib54] Otsuka et al. (2011) Otsuka, S., Kuwajima, I., Hosoya, J., Xu, Y., and Yamazaki, M. Polyinfo: Polymer database for polymeric materials design. In 2011 International Conference on Emerging Intelligent Data and Web Technologies, pp. 22–29. IEEE, 2011. ↩ a

[bib.bib55] Perola & Charifson (2004) Perola, E. and Charifson, P. S. 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

[bib.bib56] Peyré et al. (2016) Peyré, G., Cuturi, M., and Solomon, J. Gromov-Wasserstein Averaging of Kernel and Distance Matrices. In Balcan, M. F. and Weinberger, K. Q. (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

[bib.bib57] Peyré et al. (2019) Peyré, G., Cuturi, M., 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

[bib.bib58] Rampášek et al. (2022) Rampášek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35:14501–14515, 2022. ↩ a b

[bib.bib59] Ramsundar et al. (2019) Ramsundar, B., Eastman, P., Walters, P., and Pande, V. Deep learning for the life sciences: applying deep learning to genomics, microscopy, drug discovery, and more. O’Reilly Media, 2019. ↩ a

[bib.bib60] Raza et al. (2022) Raza, A., Henle, E. A., and Fern, X. Non-equilibrium molecular geometries in graph neural networks. arXiv preprint arXiv:2203.04697, 2022. ↩ a

[bib.bib61] Rong et al. (2020) Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W., and Huang, J. Self-supervised graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems, 33:12559–12571, 2020. ↩ a b c d

[bib.bib62] Rosa et al. (2016) Rosa, M., Micciarelli, M., Laio, A., and Baroni, S. Sampling molecular conformers in solution with quantum mechanical accuracy at a nearly molecular-mechanics cost. Journal of Chemical Theory and Computation, 12(9):4385–4389, 2016. ↩ a

[bib.bib63] Ross et al. (2022) Ross, J., Belgodere, B., Chenthamarakshan, V., Padhi, I., Mroueh, Y., and Das, P. Large-scale chemical language representations capture molecular structure and properties. Nature Machine Intelligence, 4(12):1256–1264, 2022. ↩ a b

[bib.bib64] Scarselli et al. (2008) Scarselli, F., Gori, M., Tsoi, A. C., Hagenbuchner, M., and Monfardini, G. The graph neural network model. IEEE transactions on neural networks, 20(1):61–80, 2008. ↩ a

[bib.bib65] Scetbon et al. (2021) Scetbon, M., Cuturi, M., and Peyré, G. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pp. 9344–9354. PMLR, 2021. ↩ a

[bib.bib66] Schütt et al. (2017) Schütt, K., Kindermans, P., Felix, H. E. S., Chmiela, S., Tkatchenko, A., and Müller, K. 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

[bib.bib67] Schütt et al. (2017) Schütt, K., Kindermans, P.-J., Sauceda Felix, H. E., Chmiela, S., Tkatchenko, A., and Müller, K.-R. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. ↩ a b c d e f g h

[bib.bib68] Schütt et al. (2021a) Schütt, K., Unke, O., and Gastegger, M. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377–9388. PMLR, 2021a. ↩ a b c

[bib.bib69] Schütt et al. (2021b) Schütt, K. T., Unke, O. T., and Gastegger, M. Equivariant message passing for the prediction of tensorial properties and molecular spectra. ICML, pp. 1–13, 2021b. ↩ a

[bib.bib70] Solomon et al. (2015) Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1–11, 2015. ↩ a

[bib.bib71] Sonthalia et al. (2021) Sonthalia, R., Van Buskirk, G., Raichel, B., and Gilbert, A. How can classical multidimensional scaling go wrong? In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P. S., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 12304–12315. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/66121d1f782d29b62a286909165517bc-Paper.pdf. ↩ a

[bib.bib72] Titouan et al. (2019) Titouan, V., Courty, N., Tavenard, R., Laetitia, C., and Flamary, R. Optimal Transport for structured data with application on graphs. In Chaudhuri, K. and Salakhutdinov, R. (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

[bib.bib73] Titouan et al. (2020) Titouan, V., Chapel, L., Flamary, R., Tavenard, R., and Courty, N. Fused Gromov-Wasserstein Distance for Structured Objects. Algorithms, 13(9):212, August 2020. ISSN 1999-4893. doi: 10.3390/a13090212. URL https://www.mdpi.com/1999-4893/13/9/212. ↩ a

[bib.bib74] Tong et al. (2021) Tong, A. Y., Huguet, G., Natik, A., MacDonald, K., Kuchroo, M., Coifman, R., Wolf, G., and Krishnaswamy, S. Diffusion earth mover’s distance and distribution embeddings. In International Conference on Machine Learning, pp. 10336–10346. PMLR, 2021. ↩ a

[bib.bib75] Tong et al. (2025) Tong, V., Trung-Dung, H., Liu, A., Broeck, G. V. d., and Niepert, M. Learning to discretize denoising diffusion odes. International Conference on Learning Representations (ICLR)), 2025. ↩ a

[bib.bib76] Torgerson (1952) Torgerson, W. S. Multidimensional Scaling: I. Theory and Method. Psychometrika, 17(4):401–419, 1952. ISSN 0033-3123. doi: 10.1007/BF02288916. URL https://www.cambridge.org/core/product/5A026426380B7639E7EA6D92D5DF19AB. Edition: 2025/01/01 Publisher: Cambridge University Press & Assessment. ↩ a

[bib.bib77] Vamathevan et al. (2019) Vamathevan, J., Clark, D., Czodrowski, P., Dunham, I., Ferran, E., Lee, G., Li, B., Madabhushi, A., Shah, P., Spitzer, M., and Zhao, S. Applications of machine learning in drug discovery and development. Nat. Rev. Drug Discov., 18(6):463–477, Jun 2019. ISSN 1474-1784. ↩ a

[bib.bib78] Varnek et al. (2005) Varnek, A., Fourches, D., Hoonakker, F., and Solov’ev, V. P. Substructural fragments: an universal language to encode reactions, molecular and supramolecular structures. Journal of computer-aided molecular design, 19:693–703, 2005. ↩ a

[bib.bib79] Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017. ↩ a b c

[bib.bib80] Veličković et al. (2017) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y. Graph attention networks. arXiv preprint arXiv:1710.10903, 2017. ↩ a

[bib.bib81] Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph attention networks. In International Conference on Learning Representations, 2018. ↩ a b

[bib.bib82] Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph attention networks. In ICLR, 2018. ↩ a

[bib.bib83] Wang et al. (2025) Wang, J., Min, Y., Li, M., and Wu, J. Fragformer: A fragment-based representation learning framework for molecular property prediction. Transactions on Machine Learning Research, 2025. ↩ a

[bib.bib84] Wang et al. (2019) Wang, S., Guo, Y., Wang, Y., Sun, H., and Huang, J. 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

[bib.bib85] Wankowicz & Bonomi (2025) Wankowicz, S. A. and Bonomi, M. From possibility to precision in macromolecular ensemble prediction. arXiv preprint arXiv:2505.01919, 2025. ↩ a

[bib.bib86] Wen et al. (2024) Wen, Q., Ju, M., Ouyang, Z., Zhang, C., and Ye, Y. 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

[bib.bib87] Wollschläger et al. (2024) Wollschläger, T., Kemper, N., Hetzel, L., Sommer, J., and Günnemann, S. Expressivity and generalization: Fragment-biases for molecular gnns. International Conference on Machine Learning, 2024. ↩ a b

[bib.bib88] Wu et al. (2018) Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Geniesse, C., Pappu, A. S., Leswing, K., and Pande, V. MoleculeNet: A benchmark for molecular machine learning. Chemical Science, pp. 513–530, 2018. ↩ a b

[bib.bib89] Xiong et al. (2019) Xiong, Z., Wang, D., Liu, X., Zhong, F., Wan, X., Li, X., Li, Z., Luo, X., Chen, K., Jiang, H., 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 b

[bib.bib90] Xu et al. (2018) Xu, K., Li, C., Tian, Y., Sonobe, T., Kawarabayashi, K.-i., and Jegelka, S. Representation learning on graphs with jumping knowledge networks. In Dy, J. and Krause, A. (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

[bib.bib91] Xu et al. (2019) Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? International Conference on Learning Representations, 2019. ↩ a b

[bib.bib92] Yang et al. (2019a) Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Guzman-Perez, A., Hopper, T., Kelley, B., Mathea, M., Palmer, A., Settels, V., Jaakkola, T., Jensen, K., and Barzilay, R. 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

[bib.bib93] Yang et al. (2019b) Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Guzman-Perez, A., Hopper, T., Kelley, B., Mathea, M., et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8):3370–3388, 2019b. ↩ a b

[bib.bib94] Ying et al. (2021) Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877–28888, 2021. ↩ a b c d e f

[bib.bib95] Zaheer et al. (2017) Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R. R., and Smola, A. J. Deep sets. Advances in neural information processing systems, 30, 2017. ↩ a b

[bib.bib96] Zhang et al. (2020) Zhang, J., Zhang, H., Xia, C., and Sun, L. Graph-bert: Only attention is needed for learning graph representations. arXiv preprint arXiv:2001.05140, 2020. ↩ a

[bib.bib97] Zhang et al. (2021) Zhang, Z., Liu, Q., Wang, H., Lu, C., and Lee, C.-K. Motif-based graph self-supervised learning for molecular property prediction. Advances in Neural Information Processing Systems, 34:15870–15882, 2021. ↩ a

[bib.bib98] Zhou et al. (2023) Zhou, G., Gao, Z., Ding, Q., Zheng, H., Xu, H., Wei, Z., Zhang, L., and Ke, G. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023. ↩ a b c d

[bib.bib99] Zhu et al. (2023) Zhu, Y., Hwang, J., Adams, K., Liu, Z., Nan, B., Stenfors, B., Du, Y., Chauhan, J., Wiest, O., Isayev, O., Coley, C. W., Sun, Y., and Wang, W. Learning over molecular conformer ensembles: Datasets and benchmarks, 2023. ↩ a

[bib.bib100] Zhu et al. (2024a) Zhu, Y., Hwang, J., Adams, K., Liu, Z., Nan, B., Stenfors, B., Du, Y., Chauhan, J., Wiest, O., Isayev, O., et al. Learning over molecular conformer ensembles: Datasets and benchmarks. In International Conference on Learning Representations, 2024a. ↩ a b c

[bib.bib101] Zhu et al. (2024b) Zhu, Y., Hwang, J., Adams, K., Liu, Z., Nan, B., Stenfors, B., Du, Y., Chauhan, J., Wiest, O., Isayev, O., et al. Learning over molecular conformer ensembles: Datasets and benchmarks. International Conference on Learning Representations (ICLR), 2024b. ↩ a

Model	Lipo	ESOL	FreeSolv	BACE
2D-GAT (Veličković et al., 2018)	1.387 $\pm$ 0.206	2.288 $\pm$ 0.017	8.564 $\pm$ 1.345	1.844 $\pm$ 0.33
D-MPNN (Yang et al., 2019a)	0.534 $\pm$ 0.022	0.923 $\pm$ 0.045	4.213 $\pm$ 0.068	0.723 $\pm$ 0.021
Attentive FP (Xiong et al., 2019)	0.520 $\pm$ 0.001	0.771 $\pm$ 0.026	4.197 $\pm$ 0.193	-
PretrainGNN (Hu et al., 2020b)	0.545 $\pm$ 0.003	1.210 $\pm$ 0.005	6.392 $\pm$ 0.003	-
GROVER_large (Rong et al., 2020)	0.676 $\pm$ 0.012	0.798 $\pm$ 0.018	5.162 $\pm$ 0.047	-
ChemBERTa-2* (Ahmad et al., 2022)	0.639 $\pm$ 0.006	0.795 $\pm$ 0.033	-	1.858 $\pm$ 0.029
ChemRL-GEM (Fang et al., 2022)	0.486 $\pm$ 0.008	0.706 $\pm$ 0.061	3.924 $\pm$ 0.436	-
MolFormer (Ross et al., 2022)	0.492 $\pm$ 0.012	0.766 $\pm$ 0.026	5.485 $\pm$ 0.045	1.091 $\pm$ 0.021
ConfNet (Liu et al., 2021)	1.360 $\pm$ 0.038	2.115 $\pm$ 0.484	-	1.329 $\pm$ 0.042
UniMol (Zhou et al., 2023)	0.374 $\pm$ 0.012	0.741 $\pm$ 0.014	2.867 $\pm$ 0.186	-
SchNet-scalar (Schütt et al., 2017)	0.704 $\pm$ 0.032	0.672 $\pm$ 0.027	1.608 $\pm$ 0.158	0.723 $\pm$ 0.100
SchNet-emb (Schütt et al., 2017)	0.589 $\pm$ 0.022	0.635 $\pm$ 0.057	1.587 $\pm$ 0.136	0.692 $\pm$ 0.028
ChemProp3D (Axelrod & Gómez-Bombarelli, 2023)	0.602 $\pm$ 0.035	0.681 $\pm$ 0.023	2.014 $\pm$ 0.182	0.815 $\pm$ 0.170
ConAN (Nguyen et al., 2024b)	0.556 $\pm$ 0.013	0.571 $\pm$ 0.019	1.496 $\pm$ 0.158	0.635 $\pm$ 0.051
ConAN-FGW (Nguyen et al., 2024b)	0.422 $\pm$ 0.016	0.529 $\pm$ 0.022	1.068 $\pm$ 0.083	0.549 $\pm$ 0.016
FACET	0.424 $\pm$ 0.009	0.516 $\pm$ 0.044	0.967 $\pm$ 0.082	0.495 $\pm$ 0.115

From Fragments to Geometry: A Unified Graph Transformer for Molecular Representation from Conformer Ensembles

Abstract

1 Introduction

2 Related Work

2.1 Conformer Ensemble Learning in Molecular Representations

2.2 Scalable Optimal Transport for Graph Learning

2.3 Fragment-biases in Molecular GNN

3 Fragment-Aware Conformer Ensemble Transformer

3.1 Conformer Generation

3.2 Framework Overview

3.3 Fragment-Enhanced 2D Molecular Graph

Fragment-Based Structural Augmentation.

3.4 Learning Graph Transformer for 3D Molecule Aggregations

3.4.1 3D conformer feature representation

3.4.2 Fragment-aware Graph Former

3.4.3 Learning to Approximate FGW distance

3.4.4 Invariant Aggregation of 2D and 3D Representation

4 Theoretical Bounds for Embedding Non-Euclidean FGW Distance Matrices

Theorem 4.1.

5 Experiments

5.1 Implementation Details

5.2 Approximation of FGW Distance via Graph Transformer

5.3 Scaling Fragment Geometry-Aware Aggregation

5.4 State-of-the-Art Performance Comparison on Molecular Tasks

5.4.1 Results

5.5 Ablation study

6 Conclusion

Acknowledgement

Disclaimer

References

Appendix A Limitations of FACET

A.1 FACET Operates on a Predefined Set of 3D Conformers.

A.2 Limitations in Scope

A.2.1 Focusing on Small Molecules

A.2.2 Limitation in Generative Capabilities.

Appendix B Implementation Details

Stage 1. Learning 2D and 3D features.

Stage 2. Training Graph Transformer to approximate FGW distance.

Appendix C Additional Analysis of FACET’s Scalability and Performance with More 3D Conformers

C.1 Inference Time when Increasing the Number of 3D Conformers for Each Molecule.

C.2 Average Training Time per Epoch as a Function of Dataset Size.

Appendix D Proof of Theorem 4.1

D.1 Non-Euclidean Nature of Pairwise FGW Distance Matrix

Definition D.1 (Euclidean Distance Matrix).

Fact 1 (Conditions for Euclidean Distance Matrix, see, e.g., (Gower, 1985)).

Lemma D.2 (Non-Euclidean Nature of Pairwise FGW Distance Matrix).

D.2 Lower Bounds on Embedding non-Euclidean FGW Distances

Proposition D.3 (Error Lower Bound of 𝒮∗).

Proof of Proposition D.3.

D.3 Upper Bounds on Embedding of Pairwise Empirical FGW Barycenter Distances

Proposition D.4 (Error Upper Bound of 𝒮∗).

D.4 Proof of Lemma D.2

Cite this paper

Related work

Proposition D.3 (Error Lower Bound of $𝒮^{*}$ ).

Proposition D.4 (Error Upper Bound of $𝒮^{*}$ ).