Learning the Ground Cost: SE(3)-Equivariant Entropic Optimal Transport with an Anisotropic Riemannian Metric

Neural optimal transport for structure prediction between molecular partners fixes the ground cost in advance, so a residue translation or rotation is charged the same regardless of its local environment. We ask whether the transport geometry itself can be learned. We introduce Geometry-Aware Optimal Transport (GOAT), an entropic optimal-transport framework on the configuration manifold SE(3)^N in which a network shared across residues maps invariant local features to one positive-definite $6 \times 6$ block per mobile residue. The resulting block-diagonal metric defines a source-local anisotropic quadratic cost between an unbound configuration and candidate bound configurations. Combined with a neural semidual potential, the cost defines an entropic weighting over a set of observed bound structures, whose weighted average is the prediction. Because features, costs and weights are invariant under a joint global rigid motion, the predictor is $\mathrm{SE}(3)$-equivariant by construction while still permitting per-residue induced-fit deformation. On a controlled synthetic benchmark spanning refinement, docking-refinement, and general-docking regimes, GOAT improves mean iRMSD in four of five setting--feature combinations, including gains of 9.9% in refinement and 11.6% in atom-level general docking. In general docking, it also reduces the configuration-level 2-Wasserstein distance by up to 16.0% and the energy distance by up to 12.1%. We present this as evidence that the ground cost is a usable learnable component of equivariant transport, and as a foundation for extending GOAT to experimentally derived complexes.