PaperlayerView exact revision on arXiv

Source-linked AI summary

Light Entropic Optimal Transport on Riemannian Manifolds

Xavier Aramayo-Carrasco, Petr Mokrov, Alexander Korotin

arXiv:2610.03085v1cs.LG

TL;DR

Existing neural EOT methods are formulated for Euclidean spaces, while manifold extensions can rely on approximations. ManifoldLightOT builds geometry-specific kernel-potential pairs for manifold EOT, and the paper reports strong empirical performance alongside simpler training and sampling. Extending the approach to a new geometry requires deriving a tractable pair, which can require additional analytical work.

  • Problem

    Neural EOT methods are formulated for Euclidean spaces, while manifold-aware methods rely on approximations that can add optimization error, sampling bias, or target-support sensitivity in high-dimensional settings.

    Source: §1 · pp. 1–2

  • Method

    ManifoldLightOT constructs compatible geometry-specific Gibbs kernels and potentials for common manifolds and their products, enabling closed-form normalization, direct conditional sampling, and Monte Carlo training from samples.

    Source: §1 · p. 2, §3 · p. 4, §3.2 · pp. 4, 6

  • Results

    ManifoldLightOT combines strong empirical performance with simpler training and sampling than existing manifold-aware EOT solvers, adapting across spheres, tori, SO(3), and SE(3).

    Source: §6 · p. 9

  • Takeaways & Limitations

    The induced geometry-aware costs can also, in principle, be used by other continuous OT/EOT solvers that support general costs.

    Source: §6 · p. 9

    Extending ManifoldLightOT to a new geometry requires deriving a kernel-potential pair that preserves tractable normalization and direct conditional sampling, which can require additional analytical work.

    Source: §6 · p. 9

Abstract

from arXiv · show

Entropic Optimal Transport (EOT) has become a practical framework for learning stochastic couplings between complex distributions, with applications in generative modeling and domain adaptation. However, most EOT solvers are designed for Euclidean spaces, while manifold extensions remain limited and often rely on costly iterative methods, simulated dynamics, or generic neural models that do not fully exploit the underlying geometry. We introduce ManifoldLightOT, a light approach for learning kernel-induced EOT couplings directly on common manifolds. Using the kernel form of the EOT solution, we construct geometry-specific Gibbs kernels together with compatible potential parameterizations for spheres, tori, $\mathrm{SO}(3)$, and $\mathrm{SE}(3)$. These choices yield closed-form normalization and directly sampleable conditional distributions. Our formulation naturally extends to products of manifolds, making it applicable to more complex geometries. The parameters of the potentials are optimized directly from samples using Monte Carlo estimates of the learning objective. Through synthetic and real-world experiments, we show that ManifoldLightOT often outperforms existing manifold OT methods while retaining direct sampling.

1 INTRODUCTION

Continuous EOT solvers support out-of-sample transport, but existing neural methods are formulated for Euclidean data and may not transfer directly to curved manifolds. ManifoldLightOT uses geometry-specific kernels and potentials to provide tractable manifold EOT couplings, and the paper evaluates it on synthetic and real-world data.

1 INTRODUCTION

  • EOT seeks a coupling minimizing expected transport cost with an entropy regularizer, and continuous solvers learn conditional distributions or transport maps from samples.The paper focuses on continuous solvers rather than discrete solvers that represent couplings as sample-indexed matrices.Source: §1 · p. 1
  • Existing neural EOT and Schrödinger bridge methods are formulated in Euclidean spaces, while manifold-aware alternatives use approximations that can add optimization error, sampling bias, or target-support sensitivity.The stated sensitivity concerns high-dimensional settings.Source: §1 · p. 1
  • ManifoldLightOT pairs geometry-specific Gibbs kernels with compatible potentials to give closed-form normalization and directly sampleable conditional distributions on spheres, tori, SO(3), and SE(3).The construction also preserves these tractability properties on Cartesian products and trains from unpaired samples using Monte Carlo objective estimates, without additional inner optimization.Source: §1 · p. 2
  • The paper evaluates ManifoldLightOT on synthetic transport problems across the considered manifolds and on real-world manifold-valued data against existing OT methods.Source: §1 · p. 2

2 BACKGROUND

The background introduces manifold EOT through its Gibbs-kernel characterization and then frames continuous EOT learning as estimating an out-of-sample coupling from independent, unpaired samples. It distinguishes this setting from discrete or Euclidean-specific approaches and reviews the assumptions and parameterizations used in the paper.

2 BACKGROUND

  • Manifold EOT couples source and target measures under a transport cost and entropy regularization, with the paper assuming absolute continuity, finite entropy, and finite expected cost.For SE(3), it additionally assumes finite translational second moments; its formulation uses geometry-specific kernel-induced costs rather than squared geodesic cost.Source: §2 · p. 2
  • The optimal coupling has a Gibbs-kernel form whose positive scaling functions enforce the marginals; its conditional distribution uses a normalization function, and the formulation matches static Schrödinger bridges for suitable reference measures.Source: §2 · pp. 2–3
  • LightSB learns a parameterized coupling by minimizing a KL divergence equivalent up to constants to a surrogate loss, using Gaussian-mixture potentials with closed-form normalization and conditional coupling in Euclidean space.That construction relies on Gaussian distributions and Euclidean identities, so it is not directly applicable to manifold-supported distributions.Source: §2 · p. 3

2.1 CONTINUOUS EOT LEARNING SETUP

  • Continuous EOT learning uses independent, unpaired samples from unknown measures to approximate a conditional coupling that can generate outputs for unseen source points.Unlike discrete EOT, which returns a sample-indexed coupling matrix, the paper's learning setup targets out-of-sample transport; existing map, barycentric, and smoothing approaches are specialized to quadratic Euclidean costs, while continuous SB solvers typically use Euclidean reference processes and kernels.Source: §2.1 · p. 3

3 MANIFOLD ENTROPIC OPTIMAL TRANSPORT SOLVER

ManifoldLightOT constructs kernel-induced EOT couplings by matching geometry-specific Gibbs kernels with compatible potential families. This design supports closed-form normalization and direct conditional sampling, including on products of manifolds.

3 MANIFOLD ENTROPIC OPTIMAL TRANSPORT SOLVER

  • Kernel–potential compatibility defines tractable kernel-induced EOT problems, with closed-form normalization and conditional sampling rather than a universal squared-geodesic objective.The kernels induce costs cε(x, y) = −ε log Kε(x, y), up to additive constants.Source: §3 · pp. 3–4, §3.2 · p. 4

3.1 LEARNING OBJECTIVE

  • The learned coupling is induced by a positive potential and fixed Gibbs kernel, and its objective can be optimized from independent, unpaired source and target samples.The objective has the same minimizers as the forward-KL learning objective, and does not require access to the unknown optimal coupling.Source: §3.1 · p. 4

3.2 TRACTABLE MANIFOLD-AWARE PARAMETERIZATIONS

  • On spheres, compatible von Mises–Fisher potentials yield closed-form normalization and direct conditional samples under the induced cost c(x, y) = −x⊤y.Samples are drawn from vMF(fk + x/ε).Source: §3.2 · p. 5
  • On tori, coordinate-wise circular von Mises distributions support direct sampling under the induced cost c(x, y) = −Σ_j[1 − cos(x_j − y_j)], up to an additive constant.Sampling is independent across coordinates.Source: §3.2 · p. 5
  • For SO(3), antipodally symmetric quaternion-based components use a hyperbolic-cosine kernel, then sample a sign and a von Mises–Fisher rotation.The induced cost is c(x, y) = −ε log cosh(λR⟨x, y⟩/ε).Source: §3.2 · p. 5
  • For SE(3), the construction combines Gaussian translations with the SO(3) rotation construction, sampling translational and rotational components independently.Its induced cost combines a squared translation term and a hyperbolic-cosine rotation term.Source: §3.2 · pp. 5–6
  • Cartesian products preserve closed-form normalization and direct conditional sampling, enabling combinations such as SO(3)^J, SE(3)^J, and heterogeneous products without additional procedures.The general construction also yields component-wise normalizers and conditionals for product spaces.Source: §3.2 · p. 6

3.3 TRAINING AND SAMPLING

  • Training estimates the objective from independent minibatches and uses standard gradient methods, without paired samples, trajectory simulation, or inner optimization; trained conditionals are sampled with geometry-specific samplers.The sampling procedure selects a mixture component using its normalized weight before drawing from that component’s manifold-specific conditional.Source: §3.3 · pp. 6–7

4 RELATED WORKS

The related-work section reviews continuous OT/EOT solvers and manifold-aware variants. It contrasts their approaches to transport maps, couplings, and geometric structure.

4 RELATED WORKS

  • The review organizes relevant continuous solvers into general OT/EOT approaches and methods that explicitly incorporate manifold geometry.Both categories are described as learning transport representations from samples and supporting out-of-sample inference.Source: §4 · p. 7

4.1 CONTINOUS OT SOLVERS

  • Continuous OT methods learn transport from samples; NeuralOT is used as a representative general continuous OT baseline.The review also covers regularized OT, score-based coupling sampling, and energy-based conditional plans.Source: §4.1 · p. 7

4.2 CONTINOUS RIEMMANIAN OT SOLVERS

  • RNOT learns continuous manifold transport maps via c-transform potentials, whereas ERNOT learns neural Schrödinger potentials defining conditional couplings.ERNOT can summarize its conditional distribution using a Riemannian barycenter or the mode of a heat-smoothed distribution.Source: §4.2 · p. 7

5 EXPERIMENTAL RESULTS

The evaluation covers synthetic transport across four manifolds and real-world continental-drift and crystallographic-texture tasks. It compares methods on shared source and target distributions using manifold-aware metrics, with the cost choices made explicit.

5 EXPERIMENTAL RESULTS

  • Experiments cover synthetic transport on spheres, tori, SO(3), and SE(3), plus continental drift on S2 and crystallographic texture transfer on SO(3).Performance is measured with manifold-aware metrics comparing transported source samples with target samples.Source: §5 · p. 7
  • RNOT and ERNOT use squared geodesic cost, while NeuralOT uses a manifold-specific induced cost.The protocol compares these methods on the same source and target distributions while making their cost choices explicit.Source: §5 · p. 8

Synthetic Experiments

  • Synthetic sources and targets are independently drawn from known manifold distributions, and Table 2 reports squared MMD across the experiments.The table reports mean ± standard deviation over five seeds, in units of 10^-2; lower is better.Source: §5 · p. 8, Table 2 · p. 8

Continental Drift (S2)

  • The continental-drift task transports unpaired samples from continental boundaries 150 million years ago to present-day boundaries on S2.Figure 1 shows ManifoldLightOT's prediction.Source: §5 · p. 8, Figure 1 · p. 8

Crystallographic Texture Transfer

  • The EBSD task transports Ti-6Al-4V orientations from starting material to material rolled at 825°C with 75% reduction, evaluating crystallographic metrics separately for rolling-direction (RD) and transverse-direction (TD) conditions.Metrics are shown in Table 2. Figure 2 shows {0001} pole-density plots for RD and TD, with colors indicating orientation concentration in MRD.Source: §5 · pp. 8–9, Figure 2 · p. 9

6 DISCUSSION

The discussion reports strong empirical performance with simpler training and sampling than existing manifold-aware EOT solvers. It presents tractable kernel-potential pairs as broadly reusable, while noting that deriving them for new geometries can require analytical work.

6 DISCUSSION

  • ManifoldLightOT combines strong empirical performance with simpler training and sampling than existing manifold-aware EOT solvers.The discussion attributes avoidance of inner optimization and Monte Carlo training to manifold-specific kernel-potential pairs.Source: §6 · p. 9
  • The pipeline, learning objective, and sampling procedure are not tied to a particular manifold, but each new geometry needs an appropriate tractable kernel-potential pair.Deriving that pair is the main manifold-specific challenge and can require additional analytical work.Source: §6 · p. 9
  • The induced costs can define geometry-aware EOT objectives for other continuous OT/EOT solvers that support general costs.The discussion presents the kernels both as a practical solver and as a set of cost functions for manifold-valued transport.Source: §6 · p. 9
  • The authors provide implementations, configurations, and appendix documentation intended to enable reproduction of the reported results and comparisons.The materials include baseline implementations and details on setup, datasets, preprocessing, metrics, and implementation.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 PROOFS

The proofs establish tractable component conditionals for the considered manifolds, then show that finite mixtures and Cartesian products preserve closed-form normalization and direct sampling.

Mixture construction

  • A finite mixture conditional is sampled by first selecting a component with its normalized weight, then sampling from that component conditional.This construction preserves closed-form normalization and direct conditional sampling.Source: §A.1 · p. 13

Cartesian products

  • For Cartesian products, each component conditional factorizes across manifolds, so independent component-wise draws yield a directly sampled product conditional.Closed-form normalization and direct conditional sampling are preserved under Cartesian products.Source: §A.1 · pp. 13–14

Sphere

  • On the sphere, each kernel-potential product is an unnormalized von Mises–Fisher density, giving a closed-form normalizer and a directly sampleable conditional.The component's natural parameter is fk + x/ε.Source: §A.1 · p. 14

Torus

  • On the torus, the kernel-potential product factorizes across coordinates, enabling coordinate-wise normalization and independent sampling.Each factor's normalizer is 2πI0(ρkj(x)).Source: §A.1 · p. 15

SO(3)

  • On SO(3), a density proportional to cosh(g⊤y) is sampled by drawing a quaternion from vMF(g) on S3 and mapping it to its rotation.The quaternion pair y and −y represents the same rotation, and the components are directly sampleable.Source: §A.1 · p. 15

Special Euclidean Group

  • For SE(3), the rotational factor is tractable as on SO(3), while completing the square makes the translation conditional Gaussian.Together, the factors satisfy the mixture construction's tractability conditions.Source: §A.1 · p. 16

A.2.1 SYNTHETIC DATASETS

The synthetic datasets pair independently generated source and target distributions on spheres, tori, SO(3), and SE(3). Training uses unpaired samples, with held-out samples reserved for evaluation.

A.2.1 SYNTHETIC DATASETS

  • Each synthetic experiment uses 20,000 training samples and 5,120 test samples for each marginal, with training based only on unpaired source and target samples.Held-out test sets are used for evaluation.Source: §A.2.1 · p. 16

Synthetic Sphere

  • For the sphere task, the source is uniform and the target is a wrapped normal centered at a fixed point, with σ = 0.3 in all dimensions.Target samples are mapped from tangent-space Gaussian perturbations through the spherical exponential map.Source: §A.2.1 · p. 16

Synthetic Torus

  • For the torus task, the source angles are uniform and the target is a coordinate-wise wrapped Gaussian centered at (π, . . . , π), with σ = 0.3.Angles are represented on the unit circle when an ambient Euclidean representation is required.Source: §A.2.1 · p. 16

Synthetic SO(3)

  • For SO(3), source and target are wrapped normal distributions around different center rotations, with perturbations drawn using σ = 0.8.The source center is the identity; the target center is a 2.5-radian rotation around the z-axis.Source: §A.2.1 · pp. 16–17

Synthetic SE(3)

  • For SE(3), the source combines Haar-uniform rotations and uniform translations in [-4, 4]^3, while the target uses a concentrated rotation and Gaussian translations.Target translations outside [-4, 4]^3 are rejected and resampled to keep translational support identical to the source.Source: §A.2.1 · p. 17

A.2.2 CONTINENTAL DRIFT

The continental-drift benchmark transports between spherical samples of continental regions reconstructed for 150 million years ago and the present. The samples are independently generated from the two time-specific polygon maps.

A.2.2 CONTINENTAL DRIFT

  • The source and target are continental regions at 150 million years ago and the present, respectively, represented by paleogeographic polygon boundaries.The benchmark follows Cohen et al. (2021) and uses reconstructions from GPlates (Müller et al., 2018).Source: §A.2.2 · p. 17
  • Candidate points are sampled uniformly by surface area on S2 and rejected unless they fall within a continental polygon for the corresponding time period.Candidates are generated from uniform longitude and z coordinates, with latitude set to arcsin z.Source: §A.2.2 · p. 17
  • The independently sampled source and target are unpaired empirical distributions, with 50,000 training points drawn from each.Accepted samples are represented as three-dimensional unit vectors.Source: §A.2.2 · p. 17

A.2.3 CRYSTALLOGRAPHIC TEXTURE DATASET

The appendix describes an unpaired crystallographic-texture transport task, its evaluation metrics and implementation, and additional synthetic sphere and torus results. The added results are reported as strong relative performance, without values in the supplied passages.

A.2.3 CRYSTALLOGRAPHIC TEXTURE DATASET

  • Texture transport compares orientations in starting Ti-6Al-4V material with those after rolling at 825°C and 75% thickness reduction, using separate RD and TD tasks.The separate EBSD maps make the source and target unpaired; orientations are normalized as quaternions and canonicalized so q0 ≥ 0.Source: §A.2.3 · pp. 17–18

A.3 METRICS

  • Texture agreement is assessed with texture-index error, normalized ODF difference, and M-index error, which measure complementary aspects of distribution agreement.The M-index error uses crystallographic misorientation-angle distributions; lower values indicate better agreement for all three metrics.Source: §A.3 · pp. 19–20

A.4 IMPLEMENTATION DETAILS

  • ManifoldLightOT directly optimizes geometry-specific mixture-potential parameters without an additional neural network, using Adam.For crystallographic texture, it uses 1,000 components and trains for 200,000 steps.Source: §A.4 · p. 20

A.5 ADDITIONAL RESULTS

  • Table 4 reports MMD2 for additional high-dimensional sphere and torus cases, and the text describes ManifoldLightOT as remaining strong against other methods.The table reports mean ± standard deviation over five seeds in units of 10^-2, with lower values better; no cell values are supplied here.Source: §A.5 · p. 22, Table 4 · p. 22
Loading 2610.03085v1…