CHANGELOG.md

Changelog

All notable changes to TorchGW are documented in this file.

[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)