SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xinyu Ma $^{1}$  Xu Chu $^{2*}$  Yasha Wang $^{1,3}$  Yang Lin $^{1}$  Junfeng Zhao $^{1}$  Liantao Ma $^{1,3}$  Wenwu Zhu $^{2}$
$^{2}$ Department of Computer Science and Technology, Tsinghua University  
 $^{3}$ National Research and Engineering Center of Software Engineering, Peking University maxinyu@pku.edu.cn, chu_xu@mail.tsinghua.edu.cn
Graph data augmentation has shown superiority in enhancing generalizability and robustness of GNNs in graph-level classifications. However, existing methods primarily focus on the augmentation in the graph signal space and the graph structure space independently, neglecting the joint interaction between them. In this paper, we address this limitation by formulating the problem as an optimal transport problem that aims to find an optimal inter-graph node matching strategy considering the interactions between graph structures and signals. To solve this problem, we propose a novel graph mixup algorithm called FGWMixup, which seeks a "midpoint" of source graphs in the Fused Gromov-Wasserstein (FGW) metric space. To enhance the scalability of our method, we introduce a relaxed FGW solver that accelerates FGWMixup by improving the convergence rate from  $\mathcal{O}(t^{-1})$  to  $\mathcal{O}(t^{-2})$ . Extensive experiments conducted on five datasets using both classic (MPNNs) and advanced (Graphormers) GNN backbones demonstrate that FGWMixup effectively improves the generalizability and robustness of GNNs. Codes are available at https://github.com/ArthurLeoM/FGWMixup.
In recent years, Graph Neural Networks (GNNs) [1, 2] have demonstrated promising capabilities in graph-level classifications, including molecular property prediction [3, 4], social network classification [5], healthcare prediction [6, 7], etc. Nevertheless, similar to other successfully deployed deep neural networks, GNNs also suffer from data insufficiency and perturbation, requiring the application of regularization techniques to improve generalizability and robustness [8]. Data augmentation is widely adopted for regularization in deep learning. It involves creating new training data by applying various semantic-invariant transformations to the original data, such as cropping or rotating images in computer vision [9], randomly inserting and rephrasing words in natural language [10, 11], etc. The augmented data fortify deep neural networks (DNNs) against potential noise and outliers underlying insufficient samples, enabling DNNs to learn more robust and representative features.
Data augmentation for GNNs requires a unique design due to the distinctive properties of attributed graphs [12], such as irregular sizes, misaligned nodes, and diverse topologies, which are not encountered when dealing with data in the Euclidean spaces such as images and tabular data. Generally, the augmentation for GNNs necessitates the consideration of two intertwined yet complementary input spaces, namely the graph signal space  $\mathcal{X}$  and the graph structure space  $\mathcal{A}$ , which are mapped to an aligned latent space  $\mathcal{H}$  with GNNs. The graph signal space  $\mathcal{X}$  consists of node features. Current
GNNs rely on parameterized nonlinear transformations to process graph signals, which can be efficiently encoded and serve as crucial inputs for downstream predictions. Therefore, the augmentation of graph signals is significant for regularizing the GNN parameter space. On the other hand, the graph structure space  $\mathcal{A}$  consists of information about graph topology. Traditional Message Passing Neural Networks (MPNNs), such as GCN and GIN, perform feature aggregation based on edge connectivity. Great efforts [13, 14, 15, 16] have been further exerted on enhancing the expressive power of GNNs [17], which carry out new GNN architectures with stronger sensitivity to graph topology compared with MPNNs. Hence, a good augmentation method should also consider the graph structure space.
Recently, huge efforts have been made to design graph data augmentation methods based on the two spaces. Mainstream research [18, 19, 20, 21, 22] considers the augmentation in the graph signal space and the graph structure space independently. For instance, ifMixup [19] conducts Euclidean mixup in the graph signal space, yet fails to preserve key topologies of the original graphs.  $\mathcal{G}$  -Mixup [20] realizes graph structure mixup based on the estimated graphons, yet fails to assign semantically meaningful graph signals. In fact, the graph signal and structure spaces are not isolated from each other, and there are strong entangling relations between them [23]. Therefore, a joint modeling of the interaction between the two spaces is essential for conducting graph data augmentation (joint modeling problem for short).
Aiming at graph data augmentation and addressing the joint modeling problem, we design a novel graph mixup [24] method considering the interaction of the two spaces during the mixup procedure. The key idea is to solve a proper inter-graph node matching strategy in a metric space that measures the distance with respect to both graph signals and structures. We propose to compute the distance metric by solving an optimal transport (OT) problem [25]. The OT problem solves the optimal coupling between nodes across graphs in the fused Gromov-Wasserstein metric space [26], wherein the distance between points takes the interaction between the graph signals and structures into account. Specifically, following [26], graphs can be modeled as probability distributions embedded in the product metric space  $\mathcal{X} \times \mathcal{A}$ . Our objective is to solve an augmented graph that minimizes the weighted sum of transportation distances between the distributions of the source graphs and the objective graph in this metric space. Developing from the Gromov-Wasserstein metric [27], Fused Gromov-Wasserstein (FGW) distance [26] has been designed to calculate the transportation distance between two unregistered probability distributions defined on different product metric spaces comprising two components, such as graph signals and structures, which defines a proper distance metric for attributed graphs. In short words, the solved graph is the augmented mixup graph in a space considering the interaction between the graph signals and structures.
However, trivially adopting FGW distance solvers [28, 26] is not scalable to large graph datasets due to a heavy computation burden. The computational bottleneck of FGW-based solvers is the nested triple-loop optimization [29], mainly due to the polytope constraint for the coupling matrix in the FGW solver. Inspired by [30], we disentangle the polytope constraint into two simplex constraints for rows and columns respectively, thence executing mirror descent with projections on the two simplexes in an alternating manner to approximate the original constraint. We prove that with a bounded gap with the ideal optimum, we may optimize the entire algorithm into a double-loop structure at convergence rate  $\mathcal{O}(t^{-2})$ , improving the  $\mathcal{O}(t^{-1})$  rate of traditional FGW solvers.
In summary, we highlight the contributions of this paper: We address the challenge of enhancing the generalizability and robustness of GNNs by proposing a novel graph data augmentation method that models the interaction between the graph signal and structure spaces. We formulate our objective as an OT problem and propose a novel graph mixup algorithm dubbed FGWMixup that seeks a "midpoint" of two graphs defined in the graph structure-signal product metric space. We employ FGW as the distance metric and speed up FGWMixup by relaxing the polytope constraint into disentangled simplex constraints, reducing the complexity from nested triple loops to double loops and meanwhile improve the convergence rate of FGWMixup. Extensive experiments are conducted on five datasets and four classic (MPNNs) and advanced (Graphormers) backbones. The results demonstrate that our method substantially improves the performance of GNNs in terms of their generalizability and robustness.
In this section, we formally introduce the proposed graph data augmentation method dubbed FGWMixup. We first introduce Fused Gromov-Wasserstein (FGW) distance that presents a distance metric between graphs considering the interaction of graph signal and structure spaces.
Then we propose our optimization objective of graph mixup based on the FGW metric and the algorithmic solutions. Finally, we present our acceleration strategy. In the following we denote  $\Delta_{n} := \{\pmb{\mu} \in \mathbb{R}_{+}^{n} | \sum_{i} \mu_{i} = 1\}$  as the probability simplex with  $n$ -bins, and  $\mathbb{S}_n(\mathbb{A})$  as the set of symmetric matrices of size  $n$  taking values in  $\mathbb{A} \subset \mathbb{R}$ .
2.1 Fused Gromov-Wasserstein Distance
In OT problems, an undirected attributed graph  $G$  with  $n$  nodes is defined as a tuple  $(\pmb{\mu}, \pmb{X}, \pmb{A})$ .  $\pmb{\mu} \in \Delta_{n}$  denotes the probability measure of nodes within the graph, which can be modeled as the relative importance weights of graph nodes. Common choices of  $\pmb{\mu}$  are uniform distributions [31]  $(\pmb{\mu} = \mathbf{1}_n / n)$  and degree distributions [32]  $(\pmb{\mu} = [\deg(v_i)]_i / (\sum_i \deg(v_i)))$ .  $\pmb{X} = (\pmb{x}^{(1)}, \dots, \pmb{x}^{(n)})^\top \in \mathbb{R}^{n \times d}$  denotes the node feature matrix with  $d$ -dimensional feature on each node.  $\pmb{A} \in \mathbb{S}_n(\mathbb{R})$  denotes a matrix that encodes structural relationships between nodes, which can be selected from adjacency matrix, shortest path distance matrix or other distance metrics based on the graph topologies. Given two graphs  $G_1 = (\pmb{\mu}_1, \pmb{X}_1, \pmb{A}_1)$  and  $G_2 = (\pmb{\mu}_2, \pmb{X}_2, \pmb{A}_2)$  of sizes  $n_1$  and  $n_2$  respectively, Fused Gromov-Wasserstein distance can be defined as follows:
$$
\operatorname {F G W} _ {q} \left(G _ {1}, G _ {2}\right) = \min  _ {\boldsymbol {\pi} \in \Pi \left(\boldsymbol {\mu} _ {1}, \boldsymbol {\mu} _ {2}\right)} \sum_ {i, j, k, l} \left((1 - \alpha) d \left(\boldsymbol {x} _ {1} ^ {(i)}, \boldsymbol {x} _ {2} ^ {(j)}\right) ^ {q} + \alpha \left| \boldsymbol {A} _ {1} (i, k) - \boldsymbol {A} _ {2} (j, l) \right| ^ {q}\right) \pi_ {i, j} \pi_ {k, l}, \tag {1}
$$
where  $\Pi (\pmb {\mu}_1,\pmb {\mu}_2)\coloneqq \{\pmb {\pi}\in \mathbb{R}_{+}^{n_1\times n_2}|\pmb {\pi}\mathbf{1}_{n_2} = \pmb {\mu}_1,\pmb {\pi}^\top \mathbf{1}_{n_1} = \pmb {\mu}_2\}$  is the set of all valid couplings between node distributions  $\pmb{\mu}_{1}$  and  $\pmb{\mu}_{2}$ ,  $d(\cdot ,\cdot)$  is the distance metric in the feature space, and  $\alpha \in [0,1]$  is the weight that trades off between the Gromov-Wasserstein cost on graph structure and Wasserstein cost on graph signals. The FGW distance is formulated as an optimization problem that aims to determine the optimal coupling between nodes in a fused metric space considering the interaction of graph structure and node features. The optimal coupling matrix  $\pi^*$  serves as a soft node matching strategy that tends to match two pairs of nodes from different graphs that have both similar graph structural properties (such as  $k$  -hop connectivity, defined by  $\mathbf{A}$ ) and similar node features (such as Euclidean similarity, defined by  $d(\cdot ,\cdot)$ ). In fact, FGW also defines a strict distance metric on the graph space  $(\Delta ,\mathcal{X},\mathcal{A})$  when  $q = 1$  and  $\mathbf{A}s$  are distance matrices, and defines a semi-metric whose triangle inequality is relaxed by  $2^{q - 1}$  for  $q > 1$ . In practice, we usually choose  $q = 2$  and Euclidean distance for  $d(\cdot ,\cdot)$  to calculate FGW distance.
Solving FGW Distance Solving FGW distance is a non-convex optimization problem, whose non-convexity comes from the quadratic term of the GW distance. There has been a line of research contributing to solving this problem. [26] proposes to apply the conditional gradient (CG) algorithm to solve this problem, and [28] presents an entropic approximation of the problem and solves the optimization through Mirror Descent (MD) algorithm according to KL-divergence.
2.2 Solving Graph Mixup in the FGW Metric Space
Building upon the FGW distance and its properties, we propose a novel graph mixup method dubbed FGWMixup. Formally, our objective is to solve a synthetic graph  $\tilde{G} = (\tilde{\mu},\tilde{X},\tilde{A})$  of size  $\tilde{n}$  that minimizes the weighted sum of FGW distances between  $\tilde{G}$  and two source graphs  $G_{1} = (\pmb{\mu}_{1},\pmb{X}_{1},\pmb{A}_{1})$  and  $G_{2} = (\pmb{\mu}_{2},\pmb{X}_{2},\pmb{A}_{2})$  respectively. The optimization objective is as follows:
$$
\arg \min  _ {\tilde {G} \in \left(\Delta_ {\tilde {n}}, \mathbb {R} ^ {\tilde {n} \times d}, \mathrm {S} _ {\tilde {n}} (\mathbb {R})\right)} \lambda \operatorname {F G W} \left(\tilde {G}, G _ {1}\right) + (1 - \lambda) \operatorname {F G W} \left(\tilde {G}, G _ {2}\right), \tag {2}
$$
where  $\lambda \in [0,1]$  is a scalar mixing ratio, usually sampled from a  $\mathrm{Beta}(k,k)$  distribution with hyperparameter  $k$ . This optimization problem formulates the graph mixup problem as an OT problem that aims to find the optimal graph  $\tilde{G}$  at the "midpoint" of  $G_{1}$  and  $G_{2}$  in terms of both graph signals and structures. When the optimum  $\tilde{G}^*$  is reached, the solutions of FGW distances between  $\tilde{G}^*$  and the source graphs  $G_{1}, G_{2}$  provide the node matching strategies that minimize the costs of aligning graph structures and graph signals. The label of  $\tilde{G}$  is then assigned as  $y_{\tilde{G}} = \lambda y_{G_1} + (1 - \lambda)y_{G_2}$ .
In practice, we usually fix the node probability distribution of  $\tilde{G}$  with a uniform distribution (i.e.,  $\tilde{\mu} = 1_{\tilde{n}} / \tilde{n}$ ) [26]. Then, Eq.2 can be regarded as a nested bi-level optimization problem, which composes the upper-level optimization w.r.t. the node feature matrix  $\tilde{X}$  and the graph structure  $\tilde{A}$ ,
and the lower-level optimization w.r.t. the couplings between current  $\tilde{G}$  and  $G_{1}, G_{2}$  denotes as  $\pi_{1}, \pi_{2}$ . Inspired by [28, 26, 33], we propose to solve Eq.2 using a Block Coordinate Descent algorithm, which iteratively minimizes the lower-level and the upper-level with a nested loop framework. The algorithm is presented in Alg.1. The inner loop solves the lower-level problem (i.e., FGW distance and the optimal couplings  $\pi_{1}^{(k)}, \pi_{2}^{(k)}$ ) based on  $\tilde{\pmb{X}}^{(k)}, \tilde{\pmb{A}}^{(k)}$  at the  $k$ -th outer loop iteration, which is non-convex and requires optimization algorithms introduced in Section 2.1. The outer loop solves the upper-level problem (i.e., minimizing the weighted sum of FGW distances), which is a convex quadratic optimization problem w.r.t.  $\tilde{\pmb{X}}^{(k)}$  and  $\tilde{\pmb{A}}^{(k)}$  with exact analytical solution (i.e., Line 7,8).
Algorithm 1 FGWMixup: Solving Eq.2 with BCD Algorithm
1: Input:  $\tilde{\mu}, G_1 = (\pmb{\mu}_1, \pmb{X}_1, \pmb{A}_1), G_2 = (\pmb{\mu}_2, \pmb{X}_2, \pmb{A}_2)$   
2: Optimizing:  $\tilde{\pmb{X}} \in \mathbb{R}^{\tilde{n} \times d}$ ,  $\tilde{\pmb{A}} \in \mathbb{S}_{\tilde{n}}(\mathbb{R})$ ,  $\pmb{\pi}_1 \in \Pi(\tilde{\pmb{\mu}}, \pmb{\mu}_1)$ ,  $\pmb{\pi}_2 \in \Pi(\tilde{\pmb{\mu}}, \pmb{\mu}_2)$ .  
3: for  $k$  in outer iterations and not converged do:  
4:  $\tilde{G}^{(k)} := (\tilde{\pmb{\mu}}, \tilde{X}^{(k)}, \tilde{A}^{(k)})$   
5: Solve arg  $\min_{\pmb{\pi}_1^{(k)}}$  FGW  $(\tilde{G}^{(k)}, G_1)$  with MD or CG (inner iterations)  
6: Solve  $\arg \min_{\pmb{\pi}_2^{(k)}}$  FGW  $(\tilde{G}^{(k)}, G_2)$  with MD or CG (inner iterations)  
7: Update  $\tilde{A}^{(k+1)} \leftarrow \frac{1}{\tilde{\pmb{\mu}}\tilde{\pmb{\mu}}^{\top}} (\lambda \pmb{\pi}_1^{(k)} \pmb{A}_1 \pmb{\pi}_1^{(k)^{\top}} + (1 - \lambda) \pmb{\pi}_2^{(k)} \pmb{A}_2 \pmb{\pi}_2^{(k)^{\top}})$   
8: Update  $\tilde{X}^{(k+1)} \leftarrow \lambda \mathrm{diag}(1 / \tilde{\pmb{\mu}}) \pmb{\pi}_1^{(k)} \pmb{X}_1 + (1 - \lambda) \mathrm{diag}(1 / \tilde{\pmb{\mu}}) \pmb{\pi}_2^{(k)} \pmb{X}_2$   
9: end for  
10: return  $\tilde{G}^{(k)}, y_{\tilde{G}} = \lambda y_{G_1} + (1 - \lambda) y_{G_2}$
2.3 Accelerating FGwMixup
Algorithm 1 provides a practical solution to optimize Eq.2, whereas the computation complexity is relatively high. Specifically, Alg.1 adopts a nested double-loop framework, where the inner loop is the FGW solver that optimizes the couplings  $\pi$ , and the outer loop updates the optimal feature matrix  $\tilde{\pmb{X}}$  and graph structure  $\tilde{\pmb{A}}$  accordingly. However, the most common FGW solvers, such as MD and CG, require another nested double-loop algorithm. This algorithm invokes (Bregman) projected gradient descent type methods to address the non-convex optimization of FGW, which involves taking a gradient step in the outer loop and projecting to the polytope-constrained feasible set w.r.t. the couplings in the inner loop (e.g., using Sinkhorn [30] or Dykstra [34] iterations). Consequently, this makes the entire mixup algorithm a triple-loop framework, resulting in a heavy computation burden.
Therefore, we attempt to design a method that efficiently accelerates the algorithm. There lie two efficiency bottlenecks of Alg.1: 1) the polytope constraints of couplings makes the FGW solver a nested double-loop algorithm, 2) the strict projection of couplings to the feasible sets probably modifies the gradient step to another direction that deviates from the original navigation of the gradient, possibly leading to a slower convergence rate. In order to alleviate the problems, we are motivated to slightly relax the feasibility constraints of couplings to speed up the algorithm. Inspired by [29, 30], we do not strictly project  $\pi$  to fit the polytope constraint  $\Pi (\pmb {\mu}_i,\pmb {\mu}_j)\coloneqq \{\pmb {\pi}\in \mathbb{R}_{+}^{n_{1}\times n_{2}}|\pmb{\pi}\mathbf{1}_{n_{2}} = \pmb{\mu}_{1},\pmb{\pi}^{\top}\mathbf{1}_{n_{1}} = \pmb{\mu}_{2}\}$  after taking each gradient step. Instead, we relax the constraint into two simplex constraints of rows and columns respectively (i.e.,  $\Pi_1\coloneqq \{\pi \in \mathbb{R}_+^{n_1\times n_2}|\pi \mathbf{1}_{n_2} = \pmb {\mu}_1\} ,\Pi_2\coloneqq \{\pi \in \mathbb{R}_+^{n_1\times n_2}|\pi^\top \mathbf{1}_{n_1} = \pmb {\mu}_2\} ,$  ) and project  $\pi$  to the relaxed constraints  $\Pi_1$  and  $\Pi_2$  in an alternating fashion. The accelerated algorithm is presented in Alg.2, dubbed FGWMixup*.
Alg.2 mainly substitutes Line 5 and Line 6 of Alg.1 with a single-loop FGW distance solver (Line 7-12 in Alg.2) that relaxes the joint polytope constraints of the couplings. Specifically, we remove the constant term in FGW distance, and denote the equivalent metric as  $\overline{\mathrm{FGW}}(G_1, G_2)$ . The optimization objective of  $\overline{\mathrm{FGW}}(G_1, G_2)$  can be regarded as a function of  $\pi$ , and we name it the FGW function  $f(\pi)$  (See Appendix A.1 for details). Then we employ entropic regularization on  $f(\pi)$  and select Mirror Descent as the core algorithm for the FGW solver. With the negative entropy  $\phi(x) = \sum_{i} x_i \log x_i$  as the Bregman projection, the MD update takes the form of:
$$
\pi \leftarrow \pi \odot \exp (- \gamma \nabla_ {\pi} f (\pi)), \quad \pi \leftarrow \operatorname {P r o j} _ {\pi \in \Pi_ {i}} (\pi) := \arg \min  _ {\pi^ {*} \in \Pi_ {i}} \| \pi^ {*} - \pi \|, \tag {3}
$$
where  $\gamma$  is the step size. The subgradient of FGW w.r.t. the coupling  $\pi$  can be calculated as:
$$
\nabla_ {\boldsymbol {\pi}} f (\boldsymbol {\pi}) = (1 - \alpha) \boldsymbol {D} - 4 \alpha \boldsymbol {A} _ {1} \boldsymbol {\pi} \boldsymbol {A} _ {2}, \tag {4}
$$
Algorithm 2 FGWMixup*: Accelerated FGWMixup
1: Input:  $\tilde{\mu}, G_1 = (\pmb{\mu}_1, \pmb{X}_1, \pmb{A}_1), G_2 = (\pmb{\mu}_2, \pmb{X}_2, \pmb{A}_2)$   
2: Optimizing:  $\tilde{\pmb{X}} \in \mathbb{R}^{\tilde{n} \times d}$ ,  $\tilde{\pmb{A}} \in \mathbb{S}_{\tilde{n}}(\mathbb{R})$ ,  $\pi_1 \in \Pi(\tilde{\mu}, \pmb{\mu}_1)$ ,  $\pi_2 \in \Pi(\tilde{\mu}, \pmb{\mu}_2)$ .  
3: for  $k$  in outer iterations and not converged do:  
4:  $\tilde{G}^{(k)} := (\tilde{\mu}, \tilde{X}^{(k)}, \tilde{A}^{(k)})$   
5:  $D_1^{(k)} := (d(\tilde{X}^{(k)}[i], X_1[j]))_{\tilde{n} \times n_1}$ ,  $D_2^{(k)} := (d(\tilde{X}^{(k)}[i], X_2[j]))_{\tilde{n} \times n_2}$   
6: for  $i$  in  $\{1, 2\}$  do  
7: while not convergence do:  
8:  $\pi_i^{(k)} \gets \pi_i^{(k)} \odot \exp\left(\gamma(4\alpha\tilde{A}^{(k)}\pi_i^{(k)}A_i - (1 - \alpha)D_i^{(k)})\right)$   
9:  $\pi_i^{(k)} \gets \mathrm{diag}(\tilde{\mu}. / \pi_i^{(k)}1_{\tilde{n}})\pi_1^{(k)}$   
10:  $\pi_i^{(k)} \gets \pi_i^{(k)} \odot \exp\left(\gamma(4\alpha\tilde{A}^{(k)}\pi_i^{(k)}A_i - (1 - \alpha)D_i^{(k)})\right)$   
11:  $\pi_i^{(k)} \gets \pi_i^{(k)}\mathrm{diag}(\pmb{\mu}_i. / \pmb{\pi}_i^{(k)})^\top 1_{n_i}$   
12: end while  
13: end for  
14: Update  $\tilde{A}^{(k+1)} \gets \frac{1}{\tilde{\mu}\tilde{\mu}^\top} (\lambda\pi_1^{(k)}A_1\pi_1^{(k)})^\top + (1 - \lambda)\pi_2^{(k)}A_2\pi_2^{(k)})^\top$   
15: Update  $\tilde{X}^{(k+1)} \gets \lambda\mathrm{diag}(1/\tilde{\mu})\pi_1^{(k)}X_1 + (1 - \lambda)\mathrm{diag}(1/\tilde{\mu})\pi_2^{(k)}X_2$   
16: end for  
17: return  $\tilde{G}^{(k)}, y_{\tilde{G}} = \lambda y_{G_1} + (1 - \lambda)y_{G_2}$
where  $D = (d(\pmb{X}_1[i],\pmb{X}_2[j]))_{n_1\times n_2}$  is the distance matrix of node features between two graphs. The detailed derivation can be found in Appendix A.2.
Our relaxation is conducted in Line 9 and 11, where the couplings are projected to the row and column simplexes alternately instead of directly to the strict polytope. Although this relaxation may sacrifice some feasibility due to the relaxed projection, the efficiency of the algorithm has been greatly promoted. On the one hand, noted that  $\Pi_1$  and  $\Pi_2$  are both simplexes, the Bregman projection w.r.t. the negative entropy of a simplex can be extremely efficiently conducted without invoking extra iterative optimizations (i.e., Line 9, 11), which simplifies the FGW solver from a nested double-loop framework to a single-loop one. On the other hand, the relaxed constraint may also increase the convergence efficiency due to a closer optimization path to the unconstrained gradient descent.
We also provide some theoretical results to justify our algorithm. Proposition 1 presents a convergence rate analysis on our algorithm. Taking  $1 / \gamma$  as the entropic regularization coefficient, our FGW solver can be formulated as Sinkhorn iterations, whose convergence rate can be optimized from  $\mathcal{O}(t^{-1})$  to  $\mathcal{O}(t^{-2})$  by conducting marginal constraint relaxation. Proposition 2 presents a controlled gap between the optima given by the relaxed single-loop FGW solver and the strict FGW solver.
Proposition 1. Let  $(\mathcal{X},\mu),(\mathcal{Y},\nu)$  be Polish probability spaces, and  $\pi \in \Pi (\mu ,\nu)$  the probability measure on  $\mathcal{X}\times \mathcal{Y}$  with marginals  $\mu ,\nu$  .Let  $\pi_t$  be the Sinkhorn iterations  $\pi_{2t} =$  arg  $\min_{\Pi (*,\nu)}H(\cdot |\pi_{2t - 1}),\pi_{2t + 1} = \arg \min_{\Pi (\mu ,*)}H(\cdot |\pi_{2t})$  where  $H(p|q) = -\sum_{i}p_{i}\log \frac{q_{i}}{p_{i}}$  is the Kullback-Leibler divergence, and  $\Pi (*,\nu)$  is the set of measures with second marginal  $\nu$  and arbitrary first marginal  $(\Pi (\mu ,*)$  is defined analogously). Let  $\pi^{*}$  be the unique optimal solution. We have the convergence rate as follows:
$$
H \left(\pi_ {t} \mid \pi^ {*}\right) + H \left(\pi^ {*} \mid \pi_ {t}\right) = \mathcal {O} \left(t ^ {- 1}\right), \tag {5}
$$
$$
H \left(\mu_ {t} | \mu\right) + H \left(\mu | \mu_ {t}\right) + H \left(\nu_ {t} | \nu\right) + H \left(\nu | \nu_ {t}\right) = \mathcal {O} \left(t ^ {- 2}\right), \tag {6}
$$