DoublyStochasticQuadraticAffinity

class torchdr.DoublyStochasticQuadraticAffinity(eps: float = 1.0, init_dual: Tensor | None = None, tol: float = 1e-05, max_iter: int = 1000, optimizer: str = 'Adam', lr: float = 1.0, base_kernel: str = 'gaussian', tolog: bool = False, metric: str = 'sqeuclidean', zero_diag: bool = True, device: str = 'auto', keops: bool = False, verbose: bool = False)[source]

Bases: Affinity

Compute the symmetric doubly stochastic affinity.

Implement the doubly stochastic normalized matrix with controlled global \(\ell_2\) norm.

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

\[\mathbf{P}^{\star} \mathbf{1} = \mathbf{1} \quad \text{where} \quad \forall (i,j), \: P^{\star}_{ij} = \left[f^\star_i + f^\star_j - C_{ij} / \varepsilon \right]_{+} \:.\]

\(\mathbf{f}^\star\) is computed by performing dual ascent.

Convex problem. Consists in solving the following symmetric quadratic optimal transport problem [Z23]:

\[\mathop{\arg\min}_{\mathbf{P} \in \mathcal{DS}} \: \langle \mathbf{C}, \mathbf{P} \rangle + \varepsilon \| \mathbf{P} \|_2^2\]

where :

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

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

  • \(\varepsilon\): quadratic regularization parameter.

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

Bregman projection. Another way to write this problem is to consider the \(\ell_2\) projection of \(- \mathbf{C} / \varepsilon\) onto the set of doubly stochastic matrices \(\mathcal{DS}\), as follows:

\[\mathrm{Proj}_{\mathcal{DS}}^{\ell_2}(- \mathbf{C} / \varepsilon) := \mathop{\arg\min}_{\mathbf{P} \in \mathcal{DS}} \: \| \mathbf{P} + \mathbf{C} / \varepsilon \|_2 \:.\]
Parameters:

  • eps (float, optional) – Regularization parameter.

  • init_dual (torch.Tensor of shape (n_samples), optional) – Initialization for the dual variable (default None).

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

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

  • optimizer ({"Adam", "SGD", "NAdam"}, optional) – Optimizer to use for the dual ascent.

  • lr (float, optional) – Learning rate for the optimizer.

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

  • tolog (bool, optional) – Whether to store intermediate result in a dictionary.

  • metric (str, optional)

  • "sqeuclidean"). (Metric to use for computing distances (default)

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

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

  • keops (bool, optional) – Whether to use KeOps for computation.

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

References

[Z23]

Stephen Zhang, Gilles Mordant, Tetsuya Matsumoto, Geoffrey Schiebinger (2023). Manifold Learning with Sparse Regularised Optimal Transport. arXiv preprint.