SinkhornAffinity#

class torchdr.SinkhornAffinity(eps: float = 1.0, tol: float = 1e-05, max_iter: int = 1000, base_kernel: str = 'gaussian', metric: str = 'sqeuclidean', zero_diag: bool = True, device: str = 'auto', backend: str | None = None, verbose: bool = False, with_grad: bool = False)[source]#

Bases: LogAffinity

Compute the symmetric doubly stochastic affinity matrix.

The algorithm computes the doubly stochastic matrix ${\mathbf{P}}^{\mathrm{ds}}$ with controlled global entropy using the symmetric Sinkhorn algorithm [Sinkhorn and Knopp, 1967].

The algorithm computes the optimal dual variable ${\mathbf{f}}^{\star }\in {\mathbb{R}}^n$ such that

$${\mathbf{P}}^{\mathrm{ds}}\mathbf{1}=\mathbf{1}\text{where}\mathrm{\forall }(i,j),P_{ij}^{\mathrm{ds}}=\exp (f_i^{\star }+f_j^{\star }-C_{ij}/\epsilon ).$$

where :

• $\mathbf{C}$: symmetric pairwise distance matrix between the samples.

• $\epsilon$: entropic regularization parameter.

• $\mathbf{1}:=(1,...,1{)}^{\mathrm{\top }}$: all-ones vector.

${\mathbf{f}}^{\star }$ is computed by performing dual ascent via the Sinkhorn fixed-point iteration (eq. 25 in [Feydy et al., 2019]).

Convex problem. Consists in solving the following symmetric entropic optimal transport problem [Cuturi, 2013]:

$${\mathbf{P}}^{\mathrm{ds}}\in \underset{\mathbf{P}\in \mathcal{DS}}{\arg min}⟨\mathbf{C},\mathbf{P}⟩+\epsilon \mathrm{H}(\mathbf{P})$$

where :

• $\mathcal{DS}:=\{\mathbf{P}\in {\mathbb{R}}_{+}^{n\times n}:\mathbf{P}={\mathbf{P}}^{\mathrm{\top }},\mathbf{P}\mathbf{1}=\mathbf{1}\}$: set of symmetric doubly stochastic matrices.

• $\mathrm{H}$: (global) Shannon entropy such that $\mathrm{H}(\mathbf{P}):=-\sum _{ij}{P}_{ij}(\log P_{ij}-1)$.

Bregman projection. Another way to write this problem is to consider the KL projection of the Gaussian kernel ${\mathbf{K}}_{\epsilon }=\exp (-\mathbf{C}/\epsilon )$ onto the set of doubly stochastic matrices:

$${\mathbf{P}}^{\mathrm{ds}}={\mathrm{Proj}}_{\mathcal{DS}}^{\mathrm{KL}}({\mathbf{K}}_{\epsilon }):=\underset{\mathbf{P}\in \mathcal{DS}}{\arg min}\mathrm{KL}(\mathbf{P}\Vert {\mathbf{K}}_{\epsilon })$$

where $\mathrm{KL}(\mathbf{P}\Vert \mathbf{Q}):=\sum _{ij}{P}_{ij}(\log (Q_{ij}/P_{ij})-1)+Q_{ij}$ is the Kullback Leibler divergence between $\mathbf{P}$ and $\mathbf{Q}$.

Parameters:

• eps (float, optional) – Regularization parameter for the Sinkhorn algorithm.

• tol (float, optional) – Precision threshold at which the algorithm stops.

• max_iter (int, optional) – Number of maximum iterations for the algorithm.

• base_kernel ({"gaussian", "student"}, optional) – Which base kernel to normalize as doubly stochastic.

• metric (str, optional) – Metric to use for computing distances (default “sqeuclidean”).

• zero_diag (bool, optional) – Whether to set the diagonal of the distance matrix to 0.

• device (str, optional) – Device to use for computation.

• backend ({"keops", "faiss", None}, optional) – Which backend to use for handling sparsity and memory efficiency. Default is None.

• verbose (bool, optional) – Verbosity. Default is False.

• with_grad (bool, optional (default=False)) – If True, the Sinkhorn iterations are done with gradient tracking. If False, torch.no_grad() is used for the iterations.