SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Meyer Scutbon*
t-mscetbon@microsoft.com
Giovanni Palla
Helmholtz Center Munich
giovanni.palla@helmholtz-muenchen.de
Michael Klein*
michalk@apple.com
Marco Cuturi
cUTILi@apple.com
Two salient limitations have long hindered the relevance of optimal transport methods to machine learning. First, the  $O(n^3)$  computational cost of standard sample-based solvers (when used on batches of  $n$  samples) is prohibitive. Second, the mass conservation constraint makes OT solvers too rigid in practice: because they must match all points from both measures, their output can be heavily influenced by outliers. A flurry of recent works has addressed these computational and modeling limitations, but has resulted in two separate strains of methods: While the computational outlook was much improved by entropic regularization, more recent  $O(n)$  linear-time low-rank solvers hold the promise to scale up OT further. In terms of modeling flexibility, the rigidity of mass conservation has been eased for entropic regularized OT, thanks to unbalanced variants of OT that can penalize couplings whose marginals deviate from those specified by the source and target distributions. The goal of this paper is to merge these two strains, low-rank and unbalanced, to achieve the promise of solvers that are both scalable and versatile. We propose custom algorithms to implement these extensions for the linear OT problem and its fused-Gromov-Wasserstein generalization, and demonstrate their practical relevance to challenging spatial transcriptomics matching problems. These algorithms are implemented in the ott-jax toolbox [Cuturi et al., 2022].
Recent machine learning (ML) works have witnessed a flurry of activity around optimal transport (OT) methods. The OT toolbox provides convenient, intuitive and versatile ways to quantify the difference between two probability measures, either to quantify a distance (the Wasserstein and Gromov-Wasserstein distances), or, in more elaborate scenarios, by computing a push-forward map that can transform one measure into the other [Peyre and Cuturi, 2019]. Recent examples include, e.g., single-cell omics [Bunne et al., 2021, 2022, Demetci et al., 2020, Nitzan et al., 2019, Cang et al., 2023, Klein et al., 2023], attention mechanisms [Tay et al., 2020, Sander et al., 2022], self-supervised learning[Caron et al., 2020, Oquab et al., 2023], and learning on graphs [Vincent-Cuaz et al., 2023].
On the challenges of using OT. Despite their long history in ML [Rubner et al., 2000], OT methods have long suffered from various limitations, that arise from their statistical, computational, and modelling aspects. The statistical argument is commonly referred to as the curse-of-dimensionality of OT estimators: the Wasserstein distance between two probability densities, and its associated optimal Monge map, is poorly approximated using samples as the dimension  $d$  of observation grows [Dudley et al., 1966, Boissard and Le Gouic, 2014]. On the computational side, computing OT between a pair of  $n$  samples involves solving a (generalized) matching problem, with a price of  $O(n^{3})$  and above [Kuhn, 1955, Ahuja et al., 1993]. Finally, the original model for OT rests on a
mass conservation constraint: all observations from either samples must be accounted for, including outliers that are prevalent in machine learning datasets. Combined, these weaknesses have long hindered the use of OT, until a more recent generation of solvers addressed these three crucial issues.
The Entropic Success Story. The winning approach, so far, to carry out that agenda has been entropic regularization methods [Cuturi, 2013]. The computational virtues of the Sinkhorn algorithm when solving OT [Altschuler et al., 2017, Peyre et al., 2016, Solomon et al., 2016, Le et al., 2021] come with statistical efficiency [Genevay et al., 2019, Mena and Niles-Weed, 2019, Chizat et al., 2020], and can also be seamlessly combined with unbalanced formulations by penalizing - rather than constraint - mass conservation, both for the linear [Frogner et al., 2015, Chizat et al., 2018, Séjourne et al., 2022, Fatras et al., 2021, Pham et al., 2020] and quadratic [Séjourne et al., 2021] problems. These developments have all been implemented in popular OT packages [Feydy et al., 2019, Flamary et al., 2021, Cuturi et al., 2022].
The Low-Rank Alternative. A recent strain of solvers relies instead on low-rank (LR) properties of cost and coupling matrices [Forrow et al., 2018, Scetbon and Cuturi, 2020, Scetbon et al., 2021]. Much like entropic solvers, these LR solvers have a better statistical outlook [Scetbon and Cuturi, 2022] and extend to GW problems [Scetbon et al., 2022]. In stark contrast to entropic solvers, however, LR solvers benefit from linear complexity  $O(nrd)$  w.r.t sample size  $n$  (using rank  $r$  and cost dimension  $d$ ) that can scale to ambitious tasks where entropic solvers fail [Klein et al., 2023].
The Need for Unbalanced Low-Rank Solvers. LR solvers do suffer, however, from a major practical limitation: their inability to handle unbalanced problems. Yet, unbalancedness is a crucial ingredient for OT to be practically relevant. This is exemplified by the fact that unbalancedness played a crucial role in the seminal reference [Schiebinger et al., 2019], where it is used to model cell birth and death.
Our Contributions We propose in this work to lift this last limitation for LR solvers to:
- Incorporate unbalanced regularizers to define a LR linear solver (§ 3.1);
- Provide accelerated algorithms, inspired by some of the recent corrections proposed by [Séjourné et al., 2022], to isolate translation terms that appear in dual subroutines (§ 3.2);
- Carry over and adapt these approaches to the GW (§ 3.3) and Fused-GW problems (§ 3.4);
- Carry out an exhaustive hyperparameter selection procedure within large scale OT tasks (spatial transcriptomics, brain imaging), and demonstrate the benefits of our approach ( $\S 4$ ).
2 Reminders on Low-Rank Transport and Unbalanced Transport
We consider two metric spaces  $(\mathcal{X},d_{\mathcal{X}})$  and  $(\mathcal{Y},d_{\mathcal{Y}})$ , as well as a cost function  $c:\mathcal{X}\times \mathcal{Y}\to [0, + \infty[$ . The simplex  $\Delta_n^+$  holds all positive  $n$ -vectors summing to 1. For  $n,m\geq 1,a\in \Delta_n^+$ , and  $b\in \Delta_m^+$ , given points  $x_{1},\ldots ,x_{n}\in \mathcal{X}$  and  $y_{1},\ldots ,y_{m}\in \mathcal{V}$ , we define two discrete probability measures  $\mu$  and  $\nu$  as  $\mu \coloneqq \sum_{i = 1}^{n}a_{i}\delta_{x_{i}}$ ,  $\nu \coloneqq \sum_{j = 1}^{m}b_{j}\delta_{y_{j}}$  where  $\delta_z$  is the Dirac mass at  $z$ .
Cost matrices. For  $q \geq 1$ , consider first two square pairwise cost matrices, each encoding the geometries of points within  $\mu$  and  $\nu$ , and a rectangular matrix that studies that across their support:
$$
A:= [d^{q}_{\mathcal{X}}(x_{i},x_{i^{\prime}})]_{1\leq i,i^{\prime}\leq n}, B:= [d^{q}_{\mathcal{Y}}(y_{j},y_{j^{\prime}})]_{1\leq j,j^{\prime}\leq m}, C:= [c(x_{i},y_{j})]_{\substack{1\leq i \leq n,\\ 1\leq j\leq m}}.
$$
The Kantorovich Formulation of OT is defined as the following parameterized linear program:
$$
\mathrm {O T} (\mu , \nu) := \min  _ {P \in \Pi_ {a, b}} \langle C, P \rangle , \quad \text {w h e r e} \quad \Pi_ {a, b} := \left\{P \in \mathbb {R} _ {+} ^ {n \times m}, \text {s . t .} P \mathbf {1} _ {m} = a, P ^ {T} \mathbf {1} _ {n} = b \right\}. \tag {1}
$$
The Low-Rank Formulation of OT is best understood as a variant of (1) that rests on a low-rank property for cost matrix  $C$ , and low-rank constraints for couplings  $P$ . More precisely, Scutbon et al. [2021] propose to constraint the set of admissible couplings to those, within  $\Pi_{a,b}$ , that have a non-negative rank of  $r \geq 1$ . That set can be equivalently reparameterized as
$$
\Pi_ {a, b} (r) = \left\{P \in \mathbb {R} _ {+} ^ {n \times m} | P = Q \operatorname {d i a g} (1 / g) R ^ {T}, Q \in \Pi_ {a, g}, R \in \Pi_ {b, g}, \text {a n d} g \in \Delta_ {r} ^ {+} \right\}.
$$
The low-rank optimal transport (LOT) problem simply uses that restriction in (1) to define :
$$
\operatorname {L O T} _ {r} (\mu , \nu) := \min  _ {P \in \Pi_ {a, b} (r)} \langle C, P \rangle = \min  _ {Q \in \Pi_ {a, g}, R \in \Pi_ {a, g}, g \in \Delta_ {r} ^ {+}} \langle C, Q \operatorname {d i a g} (1 / g) R \rangle . \tag {2}
$$
Scetbon et al. [2021] propose and prove the convergence of a mirror-descent scheme to solve (2), and obtain linear time and memory complexities with respect to the number of samples, where each iteration in that descent scales as  $(n + m)rd$ , where  $d$  is the rank of  $C$ .
The Unbalanced Formulation of OT starts from (1) as well, but proposes to do without  $\Pi_{a,b}$  and its marginal constraints [Frogner et al., 2015, Chizat et al., 2018], and rely instead on two regularizers:
$$
\mathrm {U O T} (\mu , \nu) := \min  _ {P \in \mathbb {R} _ {+} ^ {n \times m}} \langle C, P \rangle + \tau_ {1} \mathrm {K L} (P \mathbf {1} _ {m} | a) + \tau_ {2} \mathrm {K L} (P ^ {T} \mathbf {1} _ {n} | b) \tag {3}
$$
where  $\tau_{1},\tau_{2} > 0$  and  $\mathrm{KL}(p|q)\coloneqq \sum_{i}p_{i}\log (p_{i} / q_{i}) + q_{i} - p_{i}$ . This formulation is solved using entropic regularization, with modified Sinkhorn updates [Frogner et al., 2015]. Proposing an efficient algorithm able to merge (2) with (3) is the first goal of this paper.
Gromov-Wasserstein (GW) Considerations. The GW problem [Mémoli, 2011] is a generalization of (1) where the energy  $\mathcal{Q}_{A,B}$  is a quadratic function of  $P$  defined through inner cost matrices  $A$ ,  $B$ :
$$
\mathcal {Q} _ {A, B} (P) := \sum_ {i, j, i ^ {\prime}, j ^ {\prime}} \left(A _ {i i ^ {\prime}} - B _ {j j ^ {\prime}}\right) ^ {2} P _ {i j} P _ {i ^ {\prime} j ^ {\prime}} = \mathbf {1} _ {m} ^ {T} P ^ {T} A ^ {\odot 2} P \mathbf {1} _ {m} + \mathbf {1} _ {n} ^ {T} P B ^ {\odot 2} P ^ {T} \mathbf {1} _ {n} - 2 \langle A P B, P \rangle \tag {4}
$$
where  $\odot$  is the Hadamard product. To minimize (4), the default approach rests on entropic regularization [Solomon et al., 2016, Peyre et al., 2016] and variants [Sato et al., 2020, Blumberg et al., 2020, Xu et al., 2019, Li et al., 2023]. Scutbon et al. [2022] adapted the low-rank framework to minimize  $Q_{A,B}$  over low-rank matrices  $P$ , achieving a linear-time complexity when  $A$  and  $B$  are themselves low-rank. Independently, [Sejourné et al., 2021] proposed an unbalanced generalization that also applies to GW and which can be implemented practically using entropic regularization. Finally, the minimization of a composite objective involving the sum of  $Q_{A,B}$  with  $\langle C,\cdot\rangle$  is known as the fused GW problem [Vayer et al., 2018].
3 Unbalanced Low-Rank Transport
3.1 Unbalanced Low-rank Linear Optimal Transport
We incorporate unbalancedness to low-rank solvers [Scetbon et al., 2021, 2022], moving gradually from the linear problem to the more involved GW and FGW problem. Using the framework of [Frogner et al., 2015, Chizat et al., 2018], we extend first the definition of LOT, introduced in (2), to the unbalanced case by considering the following optimization problem:
$$
\operatorname {U L O T} _ {r} (\mu , \nu) := \min  _ {P: \operatorname {r k} _ {+} (P) \leq r} \langle C, P \rangle + \tau_ {1} \mathrm {K L} \left(P \mathbf {1} _ {m} | a\right) + \tau_ {2} \mathrm {K L} \left(P ^ {T} \mathbf {1} _ {n} | b\right), \tag {5}
$$
where  $\mathrm{rk}_{+}(P)$  denotes the non-negative rank of  $P$ . Therefore by denoting  $\Pi_r \coloneqq \{(Q, R, g) \in \mathbb{R}_+^{n \times r} \times \mathbb{R}_+^{m \times r} \times \mathbb{R}_+^r : Q^T \mathbf{1}_n = R^T \mathbf{1}_m = g\}$ , and using the reparameterization of low-rank couplings, we obtain the following equivalent formulation of ULOT:
$$
\mathrm {U L O T} _ {r} (\mu , \nu) = \min  _ {(Q, R, g) \in \Pi_ {r}} \underbrace {\langle C , Q \operatorname {d i a g} (1 / g) R ^ {T} \rangle} _ {\mathcal {L} _ {C} (Q, R, g)} + \underbrace {\tau_ {1} \mathrm {K L} (Q \mathbf {1} _ {r} | a) + \tau_ {2} \mathrm {K L} (R \mathbf {1} _ {r} | b)} _ {\mathcal {G} _ {a, b} (Q, R, g)}. \tag {6}
$$
We introduce the more compact notation  $\mathcal{G}_{a,b}(Q,R,g) := F_{\tau_1,a}(Q\mathbf{1}_r) + F_{\tau_2,b}(R\mathbf{1}_r)$ , where  $F_{\tau ,z}(s)\coloneqq \tau \mathrm{KL}(s|z)$  for  $\tau >0$  and  $z\geq 0$  coordinate-wise. To solve (6), and using this split, we move away from mirror-descent and apply instead proximal gradient-descent for the KL divergence. At each iteration, we consider a linear approximation of  $\mathcal{L}_C$  where a KL penalization is added to the objective (as in the classical mirror descent scheme). However, we leave  $\mathcal{G}_{a,b}$  intact at each iteration. Borrowing notations from [Scetbon et al., 2021], we must solve at each iteration the convex optimization problem:
$$
\left(Q _ {k + 1}, R _ {k + 1}, g _ {k + 1}\right) := \underset {\zeta \in \Pi_ {r}} {\operatorname {a r g m i n}} \frac {1}{\gamma_ {k}} \mathrm {K L} (\boldsymbol {\zeta}, \boldsymbol {\xi} _ {k}) + \tau_ {1} \mathrm {K L} \left(Q \mathbf {1} _ {r} | a\right) + \tau_ {2} \mathrm {K L} \left(R \mathbf {1} _ {r} | b\right), \tag {7}
$$
where  $(Q_0,R_0,g_0)\in \Pi_r$  is the initialization, and the triplet  $\xi_{k}\coloneqq (\xi_{k}^{(1)},\xi_{k}^{(2)},\xi_{k}^{(3)})$  holds synthetic costs matrices that are re-computed at each iteration  $k$ :
$$
\xi_ {k} ^ {(1)} := Q _ {k} \odot e ^ {- \gamma_ {k} C R _ {k} \operatorname {d i a g} (1 / g _ {k})}, \xi_ {k} ^ {(2)} := R _ {k} \odot e ^ {- \gamma_ {k} C ^ {T} Q _ {k} \operatorname {d i a g} (1 / g _ {k}))}, \xi_ {k} ^ {(3)} := g _ {k} \odot e ^ {\gamma_ {k} \omega_ {k} / g _ {k} ^ {2}},
$$
with  $[\omega_k]_i \coloneqq [Q_k^T CR_k]_{i,i} \in \mathbb{R}^r$ , and  $(\gamma_k)_{k \geq 0}$  is a sequence of positive step sizes.
Reformulation using Duality. To solve (7), we apply Dykstra's algorithm [1983], whose iterations correspond to an alternating maximization on the dual formulation of (7):
Proposition 1. The convex optimization problem defined in (7) admits the following dual:
$$
\begin{array}{l} \sup  _ {f _ {1}, h _ {1}, f _ {2}, h _ {2}} \mathcal {D} _ {k} \left(f _ {1}, h _ {1}, f _ {2}, h _ {2}\right) := - F _ {\tau_ {1}, a} ^ {\star} (- f _ {1}) - \frac {1}{\gamma_ {k}} \left\langle e ^ {\gamma_ {k} \left(f _ {1} \oplus h _ {1}\right)} - 1, \xi_ {k} ^ {(1)} \right\rangle \tag {8} \\ - F _ {\tau_ {2}, b} ^ {\star} (- f _ {2}) - \frac {1}{\gamma_ {k}} \langle e ^ {\gamma_ {k} (f _ {2} \oplus h _ {2})} - 1, \xi_ {k} ^ {(2)} \rangle - \frac {1}{\gamma_ {k}} \langle e ^ {- \gamma_ {k} (h _ {1} + h _ {2})} - 1, \xi_ {k} ^ {(3)} \rangle \\ \end{array}
$$
where  $h_1, h_2 \in \mathbb{R}^r$ ,  $f_1 \in \mathbb{R}^n$ ,  $f_2 \in \mathbb{R}^m$ ,  $F_{\tau,z}^\star(\cdot) \coloneqq \sup_y \{\langle y, \cdot \rangle - F_{\tau,z}(y)\}$  is the convex conjugate of  $F_{\tau,z}$ . In addition strong duality holds and the primal problem admits a unique minimizer.
Remark 1. While we stick to KL regularizers in this work for simplicity, it is worth noting that this can be extended to more generic regularizers  $F_{\tau_1,a}$  and  $F_{\tau_2,b}$ , as considered by Chizat et al. [2018].
We use an alternating maximization scheme to solve (8). Starting from  $h_1^{(0)} = h_2^{(0)} = \mathbf{0}_r$ , we apply for  $\ell \geq 0$  the following updates (dropping iteration number  $k$  in (7) for simplicity):
$$
\begin{array}{l} f _ {1} ^ {(\ell + 1)} := \arg \sup  _ {z} \mathcal {D} (z, h _ {1} ^ {(\ell)}, f _ {2} ^ {(\ell)}, h _ {2} ^ {(\ell)}), f _ {2} ^ {(\ell + 1)} := \arg \sup  _ {z} \mathcal {D} (f _ {1} ^ {(\ell + 1)}, h _ {1} ^ {(\ell)}, z, h _ {2} ^ {(\ell)}), \\ \left(h _ {1} ^ {(\ell + 1)}, h _ {2} ^ {(\ell + 1)}\right) := \operatorname * {a r g   \sup  _ {z _ {1}, z _ {2}} \mathcal {D}} \left(f _ {1} ^ {(\ell + 1)}, z _ {1}, f _ {2} ^ {(\ell + 1)}, z _ {2}\right). \\ \end{array}
$$
These maximizations can all be obtained in closed form, to result in the closed-form updates:
$$
\begin{array}{l} \exp (\gamma f _ {1} ^ {(\ell + 1)}) = \left(\frac {a}{\xi^ {(1)} \exp (\gamma h _ {1} ^ {(\ell)})}\right) ^ {\frac {\tau_ {1}}{\tau_ {1} + 1 / \gamma}}, \exp (\gamma f _ {2} ^ {(\ell + 1)}) = \left(\frac {b}{\xi^ {(2)} \exp (\gamma h _ {2} ^ {(\ell)})}\right) ^ {\frac {\tau_ {2}}{\tau_ {2} + 1 / \gamma}}, \\ g _ {\ell + 1} := \left(\xi^ {(3)} \odot (\xi^ {(1)}) ^ {T} \exp (\gamma f _ {1} ^ {(\ell + 1)}) \odot (\xi^ {(2)}) ^ {T} \exp (\gamma f _ {2} ^ {(\ell + 1)})\right) ^ {1 / 3} \\ \exp (\gamma h _ {1} ^ {(\ell + 1)}) = \frac {g _ {\ell + 1}}{(\xi^ {(1)}) ^ {T} \exp (\gamma f _ {1} ^ {(\ell + 1)})}, \quad \exp (\gamma h _ {2} ^ {(\ell + 1)}) = \frac {g _ {\ell + 1}}{(\xi^ {(2)}) ^ {T} \exp (\gamma f _ {2} ^ {(\ell + 1)})} \\ \end{array}
$$
When using "scaling" representations for these dual variables,  $\ell \geq 0$ ,  $u_{i}^{(\ell)} := \exp (\gamma f_{i}^{(\ell)})$  and  $v_{i}^{(\ell)} := \exp (\gamma h_{i}^{(\ell)})$  for  $i \in \{1,2\}$ , we obtain a simple update, provided in the appendix (Alg. 5).
Initialization and Termination. We use the stopping criterion proposed in [Scetbon et al., 2021] to terminate the algorithm,  $\Delta (\zeta ,\tilde{\zeta},\gamma)\coloneqq \frac{1}{\gamma^2} (\mathrm{KL}(\zeta ,\tilde{\zeta}) + \mathrm{KL}(\tilde{\zeta},\zeta))$ . Finding an efficient initialization for that problem is challenging, and various choices have been implemented for instance in [Cuturi et al., 2022]. We adopt the practical choices proposed in [Scetbon and Cuturi, 2022], using either random subcoupling matrices or a  $k$ -means approach, and also follow them in adapting the choice of  $\gamma_{k}$  at each iteration  $k$  of the outer loop. We summarize our proposal in Algorithm 1, which can be seen as an extension of [Scetbon et al., 2021, Alg.2].
Convergence. The convergence proof for Dykstra's algorithm, as implemented in Alg. 5 (see appendix), follows from [Bauschke and Combettes, 2008]). Scutbon et al. [2021] show the convergence of their scheme towards a stationary point, w.r.t to the criterion  $\Delta (\cdot ,\cdot ,\gamma)$ , for fixed  $\gamma$ . The stationary convergence of our proposed algorithm can be directly derived from their result.
Complexity. Given  $\xi$ , solving Eq. (7) requires a time and memory complexity of  $\mathcal{O}((n + m)r)$ . However computing  $\xi$  requires in general  $\mathcal{O}((n^2 + m^2)r)$  time and  $\mathcal{O}(n^2 + m^2)$  memory. Scutbon et al. [2021] propose to consider low-rank factorizations of the cost matrix  $C$  of the form  $C \simeq C_1C_2^T$  where  $C_1 \in \mathbb{R}^{n \times d}$  and  $C_2 \in \mathbb{R}^{m \times d}$ . In that case computing  $\xi$  can be done in  $\mathcal{O}((n + m)rd)$  time and  $\mathcal{O}((n + m)(r + d))$  memory. Such factorizations are either known explicitly (e.g. when using squared-Euclidean distances) or can be obtained as approximations using the algorithm in [Indyk et al., 2019], which guarantees that for any distance matrix  $C \in \mathbb{R}^{n \times m}$  and  $\alpha > 0$  it can output matrices  $C_1 \in \mathbb{R}^{n \times d}$ ,  $C_2 \in \mathbb{R}^{m \times d}$  in  $\mathcal{O}((m + n)\mathrm{poly}(\frac{d}{\alpha}))$  algebraic operations such that with probability at least 0.99,  $\|C - C_1C_2^T\|_F^2 \leq \|C - C_d\|_F^2 + \alpha \|C\|_F^2$ , where  $C_d$  denotes the best rank- $d$  approximation to  $C$ .
Algorithm 1 ULOT  $(C, a, b, r, \gamma_0, \tau_1, \tau_2, \delta)$
Inputs:  $C,a,b,r,\gamma_0,\tau_1,\tau_2,\delta$ $Q,R,g\gets$  Initialization as proposed in [Scetbon and Cuturi, 2022]   
repeat   
 $\tilde{Q} = Q,\tilde{R} = R,\tilde{g} = g,$ $\nabla_{Q} = CR\mathrm{diag}(1 / g),\nabla_{R} = C^{\top}Q\mathrm{diag}(1 / g),$ $\omega \leftarrow \mathcal{D}(Q^{T}CR),\nabla_{g} = -\omega /g^{2},$ $\gamma \leftarrow \gamma_0 / \max (\| \nabla_Q\|_\infty^2,\| \nabla_R\|_\infty^2,\| \nabla_g\|_\infty^2),$ $\xi^{(1)}\gets Q\odot \exp (-\gamma \nabla_Q),\xi^{(2)}\gets R\odot \exp (-\gamma \nabla_R),\xi^{(3)}\gets g$ $Q,R,g\gets$  ULR-Dykstra(a,b,  $\pmb {\xi},\gamma ,\tau_1,\tau_2,\delta)$  (Alg. 5)   
until  $\Delta ((Q,R,g),(\tilde{Q},\tilde{R},\tilde{g}),\gamma) <   \delta ;$    
Result:  $Q,R,g$
3.2 Improvements on the Unbalanced Dykstra Algorithm
A well documented source of instability of unbalanced formulations of OT lies in the fact that the total mass of the optimal unbalanced coupling is not known beforehand. Séjourne et al. [2022] have proposed a technique to address this issue, with the benefit of reduced computational costs. They propose first a dual objective that is translation invariant (TI). We take inspiration from this strategy and adapt it to our problem, to propose the following variant of (8):
$$
\sup  _ {\tilde {f} _ {1}, \tilde {h} _ {1}, \tilde {f} _ {2}, \tilde {h} _ {2}} \left(\mathcal {D} _ {\mathrm {T I}} \left(\tilde {f} _ {1}, \tilde {h} _ {1}, \tilde {f} _ {2}, \tilde {h} _ {2}\right) := \sup  _ {\lambda_ {1}, \lambda_ {2} \in \mathbb {R}} \mathcal {D} \left(\tilde {f} _ {1} + \lambda_ {1}, \tilde {h} _ {1} - \lambda_ {1}, \tilde {f} _ {2} + \lambda_ {2}, \tilde {h} _ {2} - \lambda_ {2}\right)\right) \tag {9}
$$
It is clear from the reparameterization that both problems (8) and (9) have the same value and also that  $(\tilde{f}_1,\tilde{h}_1,\tilde{f}_2,\tilde{h}_2)$  is solution of (9) if and only if  $(\tilde{f}_1 + \lambda_1^\star ,\tilde{h}_1 - \lambda_1^\star ,\tilde{f}_2 + \lambda_2^\star ,\tilde{h}_2 - \lambda_2^\star)$  is solution of (8) where  $(\lambda_1^\star ,\lambda_2^\star)$  solves  $\mathcal{D}_{\mathrm{TI}}(\tilde{f}_1,\tilde{h}_1,\tilde{f}_2,\tilde{h}_2)$ . To solve (9), we show that the variational formulation of the translation invariant dual objective targeted inside (9) can be obtained in closed form.
Proposition 2. Let  $\tilde{f}_1 \in \mathbb{R}^n$ ,  $\tilde{f}_2 \in \mathbb{R}^m$  and  $\tilde{h}_1, \tilde{h}_2 \in \mathbb{R}^r$ , then the inner problem defined in (9) by  $\mathcal{D}_{TI}(\tilde{f}_1, \tilde{h}_1, \tilde{f}_2, \tilde{h}_2)$  admits a unique solution  $(\lambda_1^\star, \lambda_2^\star)$  and we have that
$$
\lambda_ {1} ^ {\star} := \left(1 - \frac {\tau_ {1} \tau_ {2}}{(1 / \gamma + \tau_ {1}) (1 / \gamma + \tau_ {2})}\right) ^ {- 1} \left(\frac {\tau_ {1} / \gamma}{1 / \gamma + \tau_ {1}} c _ {1} - \frac {\tau_ {1} / \gamma}{1 / \gamma + \tau_ {1}} \frac {\tau_ {2}}{1 / \gamma + \tau_ {2}} c _ {2}\right) \tag {10}
$$
$$
\lambda_ {2} ^ {\star} := \left(1 - \frac {\tau_ {1} \tau_ {2}}{(1 / \gamma + \tau_ {1}) (1 / \gamma + \tau_ {2})}\right) ^ {- 1} \left(\frac {\tau_ {2} / \gamma}{1 / \gamma + \tau_ {2}} c _ {2} - \frac {\tau_ {1} / \gamma}{1 / \gamma + \tau_ {1}} \frac {\tau_ {2}}{1 / \gamma + \tau_ {2}} c _ {1}\right) \tag {11}
$$
$$
c _ {1} := \log \left(\frac {\langle \exp (- \tilde {f} _ {1} / \tau_ {1}) , a \rangle}{\langle \exp (- \gamma (\tilde {h} _ {1} + \tilde {h} _ {2})) , \xi^ {(3)} \rangle}\right), \quad a n d \quad c _ {2} := \log \left(\frac {\langle \exp (- \tilde {f} _ {2} / \tau_ {2}) , a \rangle}{\langle \exp (- \gamma (\tilde {h} _ {1} + \tilde {h} _ {2})) , \xi^ {(3)} \rangle}\right).
$$
Using Proposition 2, we perform an alternate maximization scheme on the TI formulation of the dual  $\mathcal{D}_{\mathrm{TI}}$ . Indeed using Danskin's theorem (under the assumption that  $\lambda_1^{\star}, \lambda_2^{\star}$  do not diverge), one obtains a variant of Algorithm 5. That TI approach is summarized below in Algorithm 3, using Algorithm 2 as a subroutine. We show in the experiments section (Exp. 1) that the TI approach has better computational performance on a simple task.
Algorithm 2 compute-lambdas  $(a,b,\xi^{(3)},u_1,v_1,u_2,v_2,\gamma ,\tau_1,\tau_2)$
Inputs:  $a, b, \xi^{(3)}, u_1, v_1, u_2, v_2, \gamma, \tau_1, \tau_2$ $\tilde{u}_1 \gets u_1^{-1/\gamma/\tau_1}$ ,  $\tilde{u}_2 \gets u_2^{-1/\gamma/\tau_2}$ $c_1 \gets \log(\langle \tilde{u}_1, a \rangle) - \log(\langle \xi^{(3)}, v_1^{-1} \odot v_2^{-1} \rangle)$ ,  $c_2 \gets \log(\langle \tilde{u}_2, b \rangle) - \log(\langle \xi^{(3)}, v_1^{-1} \odot v_2^{-1} \rangle)$   
Result:  $\lambda_1^{\star}, \lambda_2^{\star}$  as in (10), (11)