Changelog

All notable changes to TorchGW are documented in this file.

[Unreleased]

New Features

  • Two-sided unbalanced transport — sampled_gw accepts rho_a and rho_b (KL penalty weights for the source and target marginals) when semi_relaxed=True. The legacy rho= keyword still works and maps to rho_b; passing both raises ValueError.

  • Implicit differentiation for semi-relaxed transport — grad_mode="implicit" now handles the KL-relaxed fixed point (tau_a, tau_b enter the adjoint system). Previously semi_relaxed=True silently fell back to unrolled autograd.

Refactor

  • _solver.py split into three layers: _sinkhorn.py (inner OT solvers, all returning (T, log_u, log_v) instead of stashing potentials on the tensor), _loop.py (the sampled GW iteration with the OT step injected as step(Lambda, reg)), and _solver.py (input handling and the two public entry points, parameters grouped by purpose). Public API and defaults unchanged; rho compatibility handled once at the API boundary.

  • Removed dead code: unused torch.compile path, maybe_gc, the NumPy anchor sampler sample_pairs_from_plan (the torch sampler runs on CPU too), and the "autograd" alias of grad_mode="unrolled".

  • Internal working notes moved from docs/ to notes/; optimization-log no longer published on the docs site.

Bug Fixes

  • Source-side relaxation on CUDA — rho_a < 1 with rho_b = 1 was dispatched to the single-tau Triton kernel, which silently ran balanced Sinkhorn. The Triton path is now used only when the source side is balanced.

  • Multiscale coarse solve no longer builds an autograd graph when differentiable=True (the coarse plan is only a warm start).

Bug Fixes (differentiable mode)

  • Cold-start Sinkhorn — the differentiable wrapper dropped the warm-start potentials, so every inner Sinkhorn restarted from zero with a 100-iteration cap. At the default epsilon=0.001 the returned plan violated its marginals by ~50% and implicit gradients were meaningless. Warm start is now passed through; differentiable=True produces the same forward plan as differentiable=False.

  • Gradient scaled by momentum — the blend (1-alpha) T_prev + alpha T_new detached T_prev, so every gradient was multiplied by alpha (0.9 by default). The blend is now straight-through: value unchanged, gradient equal to the fixed-point derivative.

  • Memory growth in FGW differentiable mode — reusing the augmented cost buffer chained all iterations into one autograd graph and retained an N x K cost matrix per GW iteration (O(max_iter * N * K)). The buffer is now allocated per iteration when differentiable=True.

  • alpha=1.0 raised ZeroDivisionError in the in-place momentum update.

  • grad_mode is validated upfront in sampled_gw instead of on first use.

  • Warning for differentiable=True with epsilon < 0.01 — at that regularization Sinkhorn cannot reach a fixed point, so no gradient mode is reliable. Recommended range for differentiable training: epsilon >= 0.01.

[0.4.2] — 2026-04-09

Performance

  • ~40% lower peak GPU memory — in-place momentum update, in-place convergence error, cost matrix kept in the Sinkhorn dtype, and in-place log-kernel conversion. Removes four N x K temporaries. Max scale on L40S 48GB rises from 25k x 30k to 35k x 40k.

Bug Fixes

  • int32 overflow for large plans — <Lambda, T> now uses row-wise bmm instead of torch.dot on the flattened plan, which overflows at N x K > 2^31. Output cast to float64 falls back to float32 on OOM.

Benchmarks

  • H100 results up to 45k x 45k added to examples/benchmark_scale.py and the docs.

[0.4.1] — 2026-04-09

Exact differentiable gradients via implicit differentiation. Fixes a correctness bug where the previous “envelope theorem” backward produced gradients with up to 30x error (cosine similarity as low as 0.07).

Breaking Changes

  • Default gradient computation for differentiable=True is now implicit differentiation (exact) instead of the old frozen-potentials approximation. No API change needed — the new default is strictly better.

New Features

  • grad_mode parameter for sampled_gw — controls how gradients are computed when differentiable=True:

    • "implicit" (default): exact gradient via adjoint system at the Sinkhorn fixed point. Solved via Schur complement on the Sinkhorn Jacobian. Memory: O(NK + K^2). Same speed as the old approximate mode.

    • "unrolled": exact gradient via unrolled PyTorch autograd. Memory: O(NK * sinkhorn_iters). Useful as fallback at extremely small epsilon.

Bug Fixes

  • Gradient correctness — The old backward formula grad_C = -grad_T * T / reg treated Sinkhorn potentials as constants (“frozen-potentials”). This is a first-order approximation that ignores how the potentials depend on C through the Sinkhorn iterations. The new implicit differentiation backward solves the adjoint system derived from the implicit function theorem at the Sinkhorn fixed point, giving exact gradients.

  • Adjoint solver stability — Initial implementation used fixed-point iteration for the adjoint system, which diverges when the spectral radius >= 1 (common at small epsilon). Replaced with Schur complement direct solve on the Sinkhorn Jacobian J^T (eigenvalues in [0,2], well-conditioned). Null space from potential constant ambiguity removed via rank-1 correction (11^T/K).

  • Warning for non-differentiable pure GW — differentiable=True with fgw_alpha=0 now emits a warning, since gradients cannot flow through precomputed graph distances.

Internal

  • _SinkhornAutograd renamed to _SinkhornApproximate (frozen-potentials, used internally for semi-relaxed only)

  • New _SinkhornImplicit class (exact VJP via adjoint)

  • New _sinkhorn_unrolled function (exact VJP via autograd)

  • _sinkhorn_differentiable rewritten as dispatcher for the three backends

Tests

  • 8 new gradient correctness tests in tests/test_sinkhorn_grad.py: implicit vs unrolled (rel_err < 2%), implicit vs finite differences, descent direction, non-uniform marginals, approximate formula verification, grad_mode validation

  • Integration tests for grad_mode through sampled_gw

  • 111 total tests, all passing

Documentation

  • docs/algorithm.rst: full derivation of implicit differentiation (Jacobian structure, adjoint equation, Schur complement, null-space handling)

  • README: updated news, differentiable mode example with grad_mode

[0.4.0] — 2026-04-07

Major performance and robustness release. 3-6x faster on typical workloads, with Triton GPU kernel acceleration, mixed precision support, and comprehensive numerical stability fixes.

Performance

  • Triton fused Sinkhorn kernels — Custom GPU kernels for the Sinkhorn row/column logsumexp updates, reducing kernel launches from ~6 to 1 per half-step. 2-5x speedup on the Sinkhorn portion (5001×6001 fp32: 261ms → 50ms). Includes fused transport plan materialization and fused marginal error check. Falls back to PyTorch automatically when Triton is unavailable. (_triton_sinkhorn.py)

  • Sinkhorn warm-start — Reuse log-domain potentials (log_u, log_v) from the previous GW iteration as initial values. Reduces Sinkhorn convergence from ~10 to ~3-5 steps.

  • GPU sampling — Replace CPU numpy sampling with torch.multinomial on GPU. Transfers 2×M integers instead of the full N×K transport plan per iteration.

  • Mixed precision — New mixed_precision=True parameter runs Sinkhorn in float32 (safe in log domain) while keeping marginals and output in float64. Up to 1.7x faster on A100/L40S; larger gains expected on consumer GPUs.

  • Dijkstra caching — DijkstraProvider caches per-node SSSP results across iterations with FIFO eviction (max 2000 rows per side). Avoids redundant computation when the same anchor nodes are re-sampled.

  • Cost plateau early stopping — GW cost EMA + patience-based convergence detection. Stops when the smoothed cost stops improving, rather than waiting for the noisy ||T - T_prev|| to drop below tol (which may never happen due to sampling noise). Example: dijkstra 1000×1200 stops at 97 iters instead of running all 500.

  • Parallel all-pairs Dijkstra — PrecomputedProvider runs source and target graph Dijkstra in parallel via process-based parallelism (scipy holds the GIL). 1.2-1.5x speedup on large graphs (≥2000 total nodes).

  • Reduced CUDA sync points — Convergence checks batched every 5 iterations; augmented penalty computed on GPU without .item() sync.

  • Sinkhorn convergence check via logsumexp — Avoids materializing full N×K matrix for the marginal error computation.

  • Pre-allocated augmented cost matrix — Reused across iterations instead of re-allocated each step.

New Features

  • mixed_precision parameter for sampled_gw and sampled_lowrank_gw

  • lambda_ema_beta parameter for cost matrix EMA smoothing (variance reduction)

  • Verbose Sinkhorn output (verbose=True prints per-iteration marginal errors)

  • sample_pairs_gpu() — GPU-native weighted sampling function

  • sample_pairs_from_plan() now accepts optional rng parameter for reproducibility

Bug Fixes

  • Numerical stability

    • torch.log(a + 1e-300) replaced with .clamp(min=1e-30) — the 1e-300 constant vanishes in float32, providing no protection against log(0)

    • Regularization decay capped at 10x to prevent instability with large epsilon values

    • Low-rank mirror descent: enforce gamma * reg >= 1 to prevent exponential overflow

    • Handle all-inf distance matrices (fully disconnected subgraphs) without crashing

    • sample_pairs_gpu casts to float32 before torch.multinomial (required on some PyTorch versions/devices)

  • Correctness

    • kNN graph symmetrized via .maximum(.T) — kneighbors_graph returns directed edges

    • Semi-relaxed Sinkhorn: correct KL proximal blend tau * new + (1-tau) * old instead of tau * new which discarded history

    • differentiable=True + semi_relaxed=True now raises NotImplementedError (envelope theorem gradient is invalid for unbalanced Sinkhorn)

    • Detach T_prev in differentiable mode to prevent computation graph accumulation across GW iterations (OOM after many iterations)

    • joint_embedding: prevent index out-of-bounds when out_dim > k_svds

    • lambda_ema_beta=0.0 now disables EMA (previously locked to first iteration’s cost)

    • Dijkstra cache eviction safety: never evict keys needed by the current request

  • Compatibility

    • scipy.sparse.linalg.cg: auto-detect tol vs rtol parameter name for SciPy 1.10-1.17+ compatibility

    • sampled_lowrank_gw: semi_relaxed validation moved to function start (fail-fast)

  • API consistency

    • sampled_lowrank_gw now accepts mixed_precision parameter

    • Removed unused semi_relaxed/rho/**kwargs from sinkhorn_lowrank signature

    • sample_pairs_from_plan returns (rows, cols) arrays instead of list[tuple]

Tests

  • 72 tests covering all solver modes, mixed precision, early stopping, Dijkstra cache, differentiable gradients, boundary values, and semi-relaxed mode

  • Test suite runs in ~18s (down from ~68s before optimizations)

Documentation

  • docs/optimization-log.md — Detailed optimization history with benchmarks

  • docs/improvements.md — Updated future directions (torch.compile, cuGraph, Triton extensions)

[0.3.0] — 2026-04-03

Initial public release.

  • Sampled Gromov-Wasserstein solver (sampled_gw) with log-domain Sinkhorn

  • Low-rank solver (sampled_lowrank_gw) via mirror descent + Dykstra

  • Three distance modes: dijkstra, precomputed, landmark

  • Fused Gromov-Wasserstein (fgw_alpha blending)

  • Multiscale warm-start via farthest-point sampling

  • Differentiable transport plans (differentiable=True)

  • Semi-relaxed mode for unbalanced transport

  • Joint manifold embedding (joint_embedding)

  • kNN graph construction with component stitching (build_knn_graph)