Fragment-Aware, Fast Conformer Ensemble Transformer
FACET extends the Fused Gromov-Wasserstein conformer-aggregation idea with two moves. First, fragment awareness: a RingPath decomposition splits a molecule's 2D graph into SSSR rings and maximal path linkers that meet at junction atoms, producing a per-atom 30-dim fragment-type count prior that is projected (Linear 30→64), concatenated to the 9 base atom features, run through a 2-layer GAT, and folded back as an additive bias on each conformer's intra-conformer Graphormer atom-atom attention. Second, fast aggregation: the predecessor pays an O(N²) price to align all N conformers by pairwise FGW; FACET instead trains a Graphormer to regress the true pairwise FGW distance into a metric encoding, then aggregates the frozen per-conformer encodings by a plain mean (enc.mean, the latent FGW barycenter). Attention is atom-atom inside a conformer; across conformers the aggregate is an average, which is permutation-invariant. That turns N(N−1)/2 optimal-transport solves into an O(N) encode-then-mean pass, a measured over 6x faster aggregation. Rebuild the ensemble and watch the encoding track FGW while the mean aggregate stays put under reordering.
The side panel runs the RingPath decomposition on a fixed 2D molecule graph (a 6-ring and a 5-ring joined by a short path, plus one degree-1 leaf). Rings are the SSSR (smallest set of smallest rings, the minimum cycle basis): every bond's shortest cycle is enumerated, deduplicated, and reduced to a GF(2)-independent set whose size equals the cycle rank. Paths are grown by extending each non-ring bond into a maximal linear chain, stopping at ring atoms and at atoms with more than two neighbors; with the default cut_leafs=False the leaf atom folds into its incident path rather than forming a separate class. An atom shared by exactly two fragments creates only a fragment-fragment edge (drawn between super-nodes in Fragment view), while an atom shared by more than two fragments would be promoted to its own junction node. The 30-dim bars are the per-atom fragment-type count prior (ring bins 0 to 14, path bins 15 to 29); in the full model that vector is Linear(30→64), concatenated to the 9 base atom features, passed through a 2-layer GAT, and added as a cosine bias on the Graphormer attention logits. The main panel tells the efficiency story. Each conformer is the same molecule under a seeded deformation; from every pair the demo computes a genuine small Fused Gromov-Wasserstein distance (an entropic-Sinkhorn optimal-transport solve on the E(3)-invariant distance matrices; the code uses alpha = 0.1, epsilon = 0.1, sinkhorn_log, square_loss and a binary radius-graph adjacency), giving the N×N heatmap and its N(N−1)/2 solve count. Classical MDS of those distances is the latent encoding: a stand-in for the trained Graphormer Tθ that in the paper minimizes the FGW-regression loss Lenc = ∑ij ( ‖Tθ(Hi) − Tθ(Hj)‖² − FGW(G(Si), G(Sj)) )², with Theorem 1 giving the cumulative-stress bound L ≤ S* ≤ U that justifies the metric embedding. The bottom-left scatter plots OT distance (FGW) against encoding distance with the y = x reference and a live Pearson rho anchored near the selected dataset's Figure-2 value (ESOL 0.871, FreeSolv 0.964, BACE 0.852, Lipo 0.812; drag Encoding fidelity to push it below the anchor). The star is the mean aggregate, ŷ = W (1/K ∑k Hk) + b (Eq. 12), so Shuffle conformers scrambles the heatmap yet the star's drift stays 0.00 by averaging. The bars contrast N(N−1)/2 pairwise FGW solves against N+1 encode-plus-mean operations, with the paper's measured aggregation speedup (over 6x; MARCEL Drugs-75K training 1107.58 → 214 GPU-hours over 300 epochs, MoleculeNet 2.28x to 3.17x per epoch, inference nearly constant in the conformer count while CONAN-FGW grows past 30 s at 20 conformers). Backbones: a SchNet E(3)-invariant 3D encoder (3 interaction blocks, output 64) feeds a 12-layer, 8-head Graphormer aggregator (hidden 64, 372k params, centrality + shortest-path + fragment biases), swappable to ViSNet or GemNet. (No neural network runs here: there is no SchNet/ViSNet encoder, no fragment-aware Graphormer, no training and no property head, and the molecules are 2D toys. The RingPath decomposition is genuine graph work (real SSSR rings, maximal path growth, junction detection, the 30-dim prior), but the latent encoding is classical MDS computed live on the demo's own toy FGW distances, a stand-in for the paper's trained frozen transformer. The rho anchors are the paper's Figure-2 values and the speedups are the paper's measured numbers, not a benchmarked wall-clock here. Reported MoleculeNet MSE (SchNet, lower is better) FACET vs CONAN-FGW: ESOL 0.516/0.529, FreeSolv 0.967/1.068, BACE 0.495/0.549, Lipo 0.424/0.422; MARCEL Drugs-75K and Kraken best with the GemNet backbone.)
Run the experiments
Every animation runs live in your browser. Click a button to play that experiment on the demo (it scrolls up and starts); drag any control to take over. Nothing is downloaded.
Number of conformers
Add conformers and watch the fragment-aware embedding and its FGW-fidelity track them at linear cost.
Permutation-invariance
Shuffle the conformer order; the mean-aggregated representation does not move.
Encoding fidelity
Sweep how faithfully the embedding preserves the FGW geometry.
Fragment prior
Sweep the fragment-count prior strength.
FGW α
Trade off structure and features in the FGW cost.
The idea in three steps
A molecule flexes into many conformers, and a predictor must fuse them cheaply and in a way that ignores both their labeling and their pose. FACET keeps the invariance of FGW aggregation while removing its quadratic cost, and adds a fragment-level prior on top.
RingPath fragments
Split the molecular graph into SSSR rings and maximal path linkers meeting at junction atoms, build a fragment graph, and fold a 30-dim fragment-type count prior back into the atoms as an additive bias on the intra-conformer Graphormer attention.
Regress FGW into a metric
Instead of aligning every pair of conformers by optimal transport, train a Graphormer whose Euclidean encoding distance regresses the true pairwise FGW distance (reported Pearson rho of 0.81 to 0.96). Now each conformer is one point.
Permutation-invariant, O(N)
Aggregate the frozen encodings by a plain mean, the latent FGW barycenter, which ignores conformer order. Encode-then-mean costs O(N) rather than the O(N²) of pairwise FGW, a measured aggregation speedup of over 6x.
For the RingPath decomposition and fragment-aware encoder, the FGW-regression training objective and its metric-embedding guarantee, the mean-aggregation and the property-prediction experiments, see FACET: A Fragment-Aware Conformer Ensemble Transformer (ICLR 2026), which extends Structure-Aware E(3)-Invariant Molecular Conformer Aggregation Networks (Nguyen et al., ICML 2024).