Conformer Aggregation by Fused Gromov-Wasserstein Barycenter

Each conformer is a small 2D point cloud standing in for a molecule (n = 6 atoms, each drawn in its type color). From conformer k the method reads the paper's exact ingredients: an E(3)-invariant structure matrix A_k, the atom features H_k, and a uniform mass 1/n. (The released code builds A_k as a binary radius-graph adjacency, A_k = 1[dist ≤ cutoff] with a SchNet cutoff of 10 Å / ViSNet 5 Å, and H_k as E(3)-invariant SchNet/ViSNet embeddings min-max normalized into [0.1, 2.0]; this toy uses the real-valued distance matrix and one scalar per-atom feature so the 2D geometry stays legible.) The FGW barycenter is the single consensus graph Ḡ = (Ā, H̄) minimizing the mean Fused Gromov-Wasserstein distance to the K conformers, FGW = α · (Gromov-Wasserstein structural cost on the structure matrices) + (1 − α) · (Wasserstein feature cost on the atom features). The released code hardcodes α = 0.1, so aggregation is feature-dominated (0.9) and structure-light (0.1); here α = 0.1 is used exactly. It is solved live by block-coordinate descent: an entropic Sinkhorn coupling π_k per conformer (computed detached, so the barycenter stays differentiable through A_k and H_k only), then two closed-form updates, the structure update (Eq. 14) Ā = (∑_k λ_k π_k A_k π_k^⊤) / (ω ⊗ ω) and the feature update (Eq. 15) H̄ = diag(1/ω) ∑_k λ_k π_k H_k, with uniform masses ω = 1/n and λ_k = 1/K; the barycenter is finally read out by a sum over the n nodes (ViSNet L2-normalizes first). The middle panel shows a naive per-atom coordinate mean: because the conformers sit at different orientations the atoms pile toward the centroid and get red halos where they overlap, the paper's stated failure mode. (The code's actual ConAN ablation is mean-pooling of the per-conformer invariant embeddings, not coordinate averaging; the coordinate mean is shown here because it makes the E(3) failure vivid.) The right panel shows the FGW barycenter, which preserves the pairwise geometry (its structure-matrix error is many times smaller). Now drag Rotate: the highlighted conformer spins, its structure-matrix heatmap on the right stays pixel-identical, the naive mean swings (the readout tracks its drift and its wobbling radius of gyration) and the FGW barycenter does not move at all (drift 0.00), because it depends on the conformer only through invariant quantities. That is the E(3)-invariance the paper proves. (The released code hardcodes α = 0.1 and ε = 0.1 inside _compute_barycenter, so the YAML/CLI values are dead code; here α = 0.1 is used exactly, but a sharper Sinkhorn ε is the toy default so the 2D geometry preservation is visible: drag ε up to the code's 0.1 to watch the coupling blur, and the E(3)-invariance holds at every ε. Because the toy feature is a single scalar shared across conformers, the barycenter features are warm-started rather than initialized to 0 as in the code; with the real distinct 64-d features that difference vanishes.) (No neural network runs here: there is no SchNet/ViSNet message passing and no trained property head, and the molecules are 2D toys. What is genuinely computed live is the aggregation math, the invariant structure matrices, the entropic-Sinkhorn FGW couplings, and the exact Eq. 14/15 barycenter, so the invariance you watch is a real consequence of that computation. In the full model the barycenter is one of three summed branches, x = x_3d (SchNet invariant backbone) + x_covalent (a 2D covalent-graph GAT) + 0.2 · x_bary, mean-pooled across the K conformers and passed to a linear property head, trained in two stages. On MoleculeNet it reports RMSE (lower is better) for ConAN-FGW of 0.454 (Lipo, K=3), 0.514 (ESOL, K=3), 1.423 (FreeSolv, K=10) and 0.654 (BACE, K=5), against 0.531 / 0.591 / 1.548 / 0.816 for ConAN without the FGW branch, and ROC-AUC / PRC-AUC on the CoV-2 and CoV-2-3CL classification tasks.)