Fast Sampled Gromov-Wasserstein Optimal Transport

TorchGW aligns two point clouds by matching their internal distance structures, so the two inputs may live in different dimensions. Instead of the full O(NK(N+K)) Gromov-Wasserstein cost, it samples M anchor pairs per iteration and approximates the cost in O(NKM). Together with Triton-fused Sinkhorn kernels, this aligns 40k x 50k point clouds on a single GPU in seconds.

Typical uses: single-cell multi-omics integration, cross-domain graph matching, shape correspondence, manifold alignment.

Installation

pip install torchgw

Requires Python 3.10+ and torch>=2.0. Triton ships with PyTorch and is used automatically on GPU. POT is not required.

Quick start

import torch
from torchgw import sampled_gw

X = torch.randn(500, 3)   # source: 500 points in 3D
Y = torch.randn(600, 5)   # target: 600 points in 5D

T = sampled_gw(X, Y, epsilon=0.005, M=80, max_iter=200)
# T is a (500, 600) transport plan; T[i, j] couples X[i] with Y[j]

For large inputs, use landmark distances and mixed precision:

T = sampled_gw(X, Y, distance_mode="landmark", mixed_precision=True)

Choosing a configuration

Data size	Call	Why
N, K below 5k	sampled_gw(X, Y, distance_mode="precomputed")	All-pairs graph distances computed once
5k and above	sampled_gw(X, Y, distance_mode="landmark", mixed_precision=True)	O(Nd) memory, runs fully on GPU
N x K exceeds GPU memory	sampled_lowrank_gw(X, Y, rank=30, distance_mode="landmark")	Stores O((N+K) rank) instead of O(NK)

The default distance_mode="dijkstra" works at any size but is slower than "landmark" on GPU. To bring your own distance matrices, pass dist_source and dist_target with distance_mode="precomputed".

Recipes

Fused Gromov-Wasserstein. Blend graph structure with a feature cost matrix. fgw_alpha=1.0 gives pure Wasserstein.

C = torch.cdist(F_src, F_tgt)                    # (N, K) feature cost
T = sampled_gw(X, Y, fgw_alpha=0.5, C_linear=C)

Unbalanced transport. Relax the marginals with KL penalties, for example when a cell type exists in only one sample. A smaller rho relaxes that side more.

T = sampled_gw(X, Y, semi_relaxed=True, rho_a=1.0, rho_b=0.1)

Differentiable transport. Exact gradients via implicit differentiation flow back through the feature cost. Use epsilon >= 0.01 so the inner Sinkhorn reaches a fixed point.

C = torch.cdist(encoder(X), encoder(Y))
T = sampled_gw(fgw_alpha=1.0, C_linear=C, differentiable=True, epsilon=0.01)
loss = (C.detach() * T).sum()
loss.backward()

Multi-scale warm start. Solve a coarse problem first, then refine.

T = sampled_gw(X, Y, multiscale=True, n_coarse=200)

Diagnostics. Return the convergence log alongside the plan.

T, info = sampled_gw(X, Y, log=True)

Full parameter reference: API docs.

Performance

Spiral (2D) to Swiss roll (3D) on an NVIDIA H100, distance_mode="landmark", mixed_precision=True:

Scale	Time	Spearman	GPU memory
4,000 x 5,000	0.8 s	0.999	0.7 GB
20,000 x 25,000	4.6 s	0.999	16 GB
45,000 x 45,000	18 s	0.999	65 GB

At 4,000 x 5,000 TorchGW finishes in 1 s where POT needs 183 s. Full results, L40S numbers and reproduction scripts are in the benchmark docs.

How it works

Gromov-Wasserstein compares distances within each space rather than across spaces, so a 2D spiral can be matched to a 3D Swiss roll, or gene expression to chromatin accessibility. Standard solvers rebuild the full O(NK(N+K)) cost tensor every iteration. TorchGW instead samples M anchor pairs from the current plan, computes graph distances only from those anchors, and assembles a low-variance O(NKM) cost estimate. A Triton-fused log-domain Sinkhorn step with warm-started potentials projects it back to a transport plan, and cost plateau detection stops the loop early.

Derivations, the semi-relaxed extension and the implicit-differentiation backward are in the algorithm docs.

Development

git clone https://github.com/chansigit/torchgw.git
cd torchgw
pip install -e ".[dev]"
pytest

Citation

@software{torchgw,
  author  = {Sijie Chen},
  title   = {TorchGW: Fast Sampled Gromov-Wasserstein Optimal Transport},
  url     = {https://github.com/chansigit/torchgw},
  version = {0.4.2},
  year    = {2026},
}

License

Free for academic and non-commercial use under the terms in LICENSE. Commercial use requires a separate license; see COMMERCIAL_LICENSE.md or contact Stanford OTL at otl@stanford.edu.

Copyright (c) 2026 The Board of Trustees of the Leland Stanford Junior University.