Federated Sinkhorn

TL;DR

This paper introduces Federated Sinkhorn, a distributed algorithm for entropy-regularized optimal transport that efficiently exchanges scaled vectors across multiple workers without centralizing the kernel, validated through theoretical analysis and experiments.

Contribution

It proposes two exact synchronous protocols for distributed Sinkhorn iterations with minimal communication, including a privacy layer, and provides detailed cost analysis and empirical validation.

Findings

Distributed Sinkhorn matches centralized results under positivity assumptions.

Communication becomes the bottleneck as the number of workers increases.

The protocols are validated on multi-node CPU/GPU setups.

Abstract

We study distributed Sinkhorn iterations for entropy-regularized optimal transport when the Gibbs kernel operator is row-partitioned across c workers and cannot be centralized. We present Federated Sinkhorn, two exact synchronous protocols that exchange only scaling-vector slices: (i) an All-to-All scheme implemented by Allgather, and (ii) a Star (parameter-server) scheme implemented by client to server sends and server to client broadcasts. For both, we derive closed-form per-iteration compute, communication, and memory costs under an alpha-beta latency--bandwidth model, and show that the distributed iterates match centralized Sinkhorn under standard positivity assumptions. Multi-node CPU/GPU experiments validate the model and show that repeated global scaling exchange quickly becomes the dominant bottleneck as c increases. We also report an optional bounded-delay asynchronous schedule and an optional privacy measurement layer for communicated log-scalings.