1 Introduction↩︎

Distributionally robust optimization (DRO) addresses inference and decision-making problems under uncertainty of the data distribution for stochastic optimization models. Since parameterized uncertainty set constructions, e.g., moments [1], [2] and deviations [3] ambiguity sets, may encompass an overly broad range of distributions significantly different from the reference, recent research has increasingly focused on non-parametric discrepancy-based ambiguity sets [4]–[9], which have a tunable radius that quantifies the uncertainty level. In this paper, we focus on this category of discrepancy-based DRO: \[\label{Problem} \inf_{f \in \mathcal{F}}\sup_{\mathbb{Q} \in \mathcal{B}_{\delta}(\mathbb{P})}\mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f,\xi)].\tag{1}\] Here \(\mathbb{P}\) is a \(d\)-dimensional reference distribution with a finite second moment and \(\mathcal{B}_{\delta}(\mathbb{P}):=\{\mathbb{Q}\in\mathcal{P}_{2}: \mathcal{D}(\mathbb{Q},\mathbb{P})\leq\delta\}\) for some \(\delta>0\) is the uncertainty set defined by the discrepancy function \(\mathcal{D}\). Meanwhile, \(\mathcal{F}\) is a measurable decision function class and \(\ell:\mathcal{F}\times\mathbb{R}^{d}\rightarrow\mathbb{R}\) is the loss function.

Although the theory and applications of DRO models have been extensively studied (see, e.g., [10]–[12]) and efficient solvers have been recently developed [13], comparatively less attention has been given to the explicit generation of the worst-case distribution (also called the least favorable distribution (LFD)). One of the key differences between DRO and traditional optimization is that it seeks models that perform best under the LFD in the ambiguity set (i.e., the maximization problem in 1 ), which makes it an infinite-dimensional optimization problem. To construct worst-case samples from 1 , existing methods often rely on dual formulations or perturbations of nominal support points. In particular, by leveraging semi-infinite duality [14], many infinite-dimensional minimax problems can be reformulated as finite tractable programs when the reference distribution is discrete [5], [7]. Such strong duality results have also been extended to general optimal transport costs [8], [9], allowing the LFD to be obtained by solving the dual problem. On the other hand, [15], [16] propose to solve a Lagrangian regularization formulation of the optimal transport based distributional problem from point-to-point perturbation by stochastic gradient methods in the discrete case.

However, the above-mentioned approaches that focus on discrete worst-case distributions face several challenges when directly applied to generate the LFD. First, there is a significant computational challenge. The dual formulation of DRO is not scalable to large datasets; for example, constructing the LFD under the Wasserstein metric requires solving a linear program with \(\mathcal{O}(N^{2})\) decision variables, where \(N\) is the total number of training data points, and the complexity of solving a linear program is typically quadratic in the number of decision variables. Such computational complexity for problems with thousands of training data points can be prohibitive, which means that current DRO formulations can usually only be used to find discrete LFDs in small-sample settings. Second, discrete LFDs limit the generalization capability of the resulting algorithm. Since the LFD is discrete and supported solely on the training dataset, the corresponding optimal detector is confined to training points. This also creates a challenge when generating samples from the LFD, as the existing point-wise perturbation method [15] requires retraining for each unseen data point to obtain the corresponding LFD sample.

Given the concerns outlined above, we are motivated to explore the DRO problem 1 as a minimax optimization in the continuous probability space for LFD generation. This is closely related to optimization problems on the probability space equipped with Wasserstein geometry, e.g., [17]–[20]. In particular, from a computational perspective, [19] studied the Frank–Wolfe algorithm in Wasserstein space, whose subproblem is insightfully designed as a linear program over a Wasserstein ball, a problem well studied in the DRO literature [7], [9], [21]. However, these methods typically rely on particle-based approximations, which require costly gradient computations to perturb each new data point in the DRO setting. When the inference system, such as a classifier, involves a complex neural network architecture, these methods incur additional overhead due to the need for auto-differentiation with respect to the inputs. Moreover, generating new samples from the LFD still requires repeating the expensive training process, limiting their practicality in real-time or large-scale applications.

In this paper, we design an algorithmic framework for distributionally robust inference that enables efficient worst-case generation. Assuming the solution of the inner maximization problem is attainable, we focus on the following regularized formulation of 1 : \[\label{Problem-r} \min_{\phi \in \Phi}\max_{\mathbb{Q}\in \mathcal{P}_2}\mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f_{\phi}, \xi)]-\lambda\cdot \mathcal{D}(\mathbb{Q},\mathbb{P}).\tag{2}\] Here, we incorporate a parameterized decision function \(f_\phi\) with \(\phi \in \Phi \subseteq \mathbb{R}^m\) and a predetermined Lagrangian multiplier \(\lambda \ge 0\). When the discrepancy is the Wasserstein distance, the inner maximization reduces to a Wasserstein proximal problem, as commonly studied in score-based models [22]. Different from traditional discrete DRO approaches and sampling-based methods, we leverage the existence of an optimal transport (OT) map guaranteed by Brenier’s theorem in the continuous setting. This allows us to derive an equivalent optimization problem over the transport map: \[\min_{\phi \in \Phi}\max_{T \in L^2(\mathbb{P})}\mathbb{E}_{\xi \sim \mathbb{P}}[\ell(f_{\phi}, T(\xi ))]-\lambda\cdot \mathcal{D}(T_{\#}\mathbb{P},\mathbb{P}).\] Once the minimax-optimal transport map is learned for problem 2 , worst-case samples can be efficiently generated by passing inputs through the learned map. This reformulation enables a practical implementation via flow-based neural network parameterizations [23], which rely on the change-of-variables formula and explicitly model data distributions through transport maps.

With this transport map reformulation in hand, we focus on the convergence guarantees of iterative algorithms for solving 2 in infinite-dimensional spaces. We propose an iterative minimax algorithmic framework (Algorithm 1) operating in Wasserstein space, which for step \(k\ge0\) performs: \[\left\{ \begin{align} &\mathbb{Q}_{k}:=\underset{\mathbb{Q} \in \mathcal{P}_2} {\mathop{\mathrm{argmax}}}\left\{\mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f_{\phi_k}, \xi )]-\lambda\cdot \mathcal{D}(\mathbb{Q},\mathbb{P})\right\}; \\ &\phi_{k+1}:=\mathop{\mathrm{proj}} _{\Phi}(\phi_{k}- \eta_k\cdot\partial_{\phi}\mathbb{E}_{\xi \sim \mathbb{Q}_k}[\ell(f_{\phi_{k}} ,\xi)]). \end{align} \right.\] We prove that this framework converges globally to a stationary point of 2 under mild assumptions that only require weak convexity of the loss function, covering a wide range of modern machine learning models.

To further enhance performance, we develop two algorithmic variants. The first is a proximal regularized algorithm inspired by the Jordan–Kinderlehrer–Otto (JKO) scheme (Algorithm 2), which iteratively computes the LFD in the inner maximization problem. The proximal regularization stabilizes the neural transport-map update by discouraging large distributional jumps between successive iterates, which improves the exploration of multimodal LFDs. When the discrepancy function \(\mathcal{D}\) satisfies strong convexity along generalized geodesics (a.g.g.) with a sufficiently large regularization parameter \(\lambda\), we establish exponential convergence of the inner loop. The second variant is an alternating algorithm (Algorithm 3) that, under an additional smoothness assumption, updates the decision function and the distribution in turn, performing only a single JKO step per outer iteration and thereby significantly reducing the computational burden of the inner-loop updates. The proposed algorithms allow us to quantify the total number of (sub)gradient evaluations and inexact JKO iterations needed to achieve approximate stationary points.

Our analysis builds on the regularization structure of the DRO model 2 but is also of independent interest for optimization over distributions, as it provides a novel convergence analysis of proximal-type methods in Wasserstein space using geometric properties along fixed generalized geodesics. This requirement is much weaker than strong convexity along all generalized geodesics assumed in the existing literature and constitutes the only geometric structure exploited in our framework, exemplified by Wasserstein-2 regularization. Another recent work [24] studies a related problem from a different perspective, using a transport-map formulation based on the Monge formulation of optimal transport and Brenier’s theorem. In contrast, our work adopts a regularized Wasserstein formulation in which the optimization variable is a continuous distribution in the \(\mathcal{W}_2\) space. As a result, the two approaches lead to distinct analytical frameworks and rely on different technical assumptions.

Numerical experiments illustrate the benefit of neural transport maps for worst-case generation. On a two-dimensional Gaussian-mixture task, both the elimination-based and alternating variants are implemented as neural transport maps and benchmarked against particle-based solvers that optimize moved samples directly. The JKO-regularized neural update produces particle-like exploration of multimodal least favorable distributions while enabling fast inference through a learned map. On CIFAR-10 feature representations, learned neural maps further provide effective initializations for particle-based optimization, reducing the cost of high-dimensional worst-case inference.

1.1 Organization↩︎

The rest of the paper is organized as follows. In Section 2, we propose a minimax optimization framework to solve the DRO problem 2 and establish its global convergence rate. Section 3 presents the JKO scheme for computing the LFD as a subproblem solver within the proposed framework, and Section 4 establishes improved complexity results for its alternating variant under smoothness assumptions. Implementation details and numerical results on both synthetic and real datasets are presented in Sections 5 and 6. Finally, we end with some closing remarks in Section 7. Key definitions, preliminary results, and detailed proofs are given in the Appendices.

1.2 Notation↩︎

Throughout this paper, we use the standard notation. Let \(\mathcal{X}\subseteq\mathbb{R}^{d}\) and denote the space \(\mathcal{P}(\mathcal{X})\) of Borel probability measures over \(\mathcal{X}\). Let \(\mathcal{P}_{2}(\mathcal{X}):=\{\mathbb{P}\in\mathcal{P}(\mathcal{X}): \int_{\mathcal{X}}\|x\|^{2}\operatorname{d}\mathbb{P}(x)<\infty\}\) be the space of all distributions on domain \(\mathcal{X}\) that have a finite second moment. Define \(\mathcal{P}_{2}^{r}(\mathcal{X})\) as all distributions in \(\mathcal{P}_{2}(\mathcal{X})\) that also have densities (absolutely continuous with respect to the Lebesgue measure). We use the notation \(\mathcal{P}_{2}\) and \(\mathcal{P}_{2}^{r}\) when \(\mathcal{X}=\mathbb{R}^d\). Given a (measurable) map \(T: \mathbb{R}^d\rightarrow\mathbb{R}^d\) and \(\mathbb{P}\) a distribution on \(\mathbb{R}^d\), its \(L^{2}\) norm is denoted as \(\|T\|_{\mathbb{P}}:=(\int_{\mathbb{R}^d}\|T(x)\|^{2}\operatorname{d\mathbb{P}}(x))^{1/2}\). For \(\mathbb{P}\in \mathcal{P}_{2}\), we denote by \(L^{2}(\mathbb{P})\) the \(L^{2}\) space of vector fields, that is, \(L^{2}(\mathbb{P}):=\{T:\|T\|_{\mathbb{P}}^2<\infty\}\). For \(T_{1}, T_{2}\in L^{2}(\mathbb{P})\), define \(\langle T_{1}, T_{2}\rangle_{\mathbb{P}}:=\int_{\mathbb{R}^d}T_{1}(x)^{\top} T_{2}(x) \operatorname{d \mathbb{P}}(x)\). For \(T: \mathbb{R}^d\rightarrow\mathbb{R}^d\), the pushforward of a distribution \(\mathbb{P}\) is denoted as \(T_{\#}\mathbb{P}\), such that \(T_{\#}\mathbb{P}(A)=\mathbb{P}(T^{-1}(A))\) for any measurable set \(A\). If the density \(p\) exists, the above mentioned notation related to \(\mathbb{P}\) can also be replaced by \(p\). We use \(\iota_{\mathcal{X}}: \mathbb{R}^{d}\rightarrow \{0, +\infty\}\) to denote the indicator function associated with \(\mathcal{X}\). For an extended-valued function \(h:\mathbb{R}^d\to\mathbb{R}\cup\{+\infty\}\) and \(\epsilon\ge0\), the set of all \(\epsilon\)-optimal solutions of \(h\) is defined by \(\epsilon\operatorname{-}\mathop{\mathrm{argmin}}h:= \{x\in\mathbb{R}^d: h(x)\leq \inf_{y\in\mathbb{R}^d}h(y)+\epsilon\}\), which is reduced to the optimal solution set \(\mathop{\mathrm{argmin}}h\) when \(\epsilon=0\). We denote by \(\partial h\) the Fréchet subdifferential. We refer to other key concepts and their properties for our analysis in Appendix 8.

2 Minimax Algorithmic Framework↩︎

In this section, we present the main algorithmic framework for worst-case generation by solving the DRO problem 2 in probability space and analyze its convergence. The detailed iterative procedure is given in Algorithm 1.

From a computational perspective, the key challenge for practical implementation is solving the inexact maximization step in updating the probability measure \(\mathbb{Q}\). In fact, with fixed decision function \(f_{\phi}\), we can consider the following equivalent transport map maximization problem for non-atomic \(\mathbb{P}\) (from the similar argument in [25]) that \[\label{eq:max-subproblem} \max_{T \in L^2(\mathbb{P})}\mathbb{E}_{\xi \sim \mathbb{P}}[\ell(f_{\phi}, T(\xi))]-\lambda\cdot \mathcal{D}(T_{\#}\mathbb{P},\mathbb{P}).\tag{3}\] Moreover, if \(\mathcal{D}=\frac{1}{2}\mathcal{W}_{2}^{2}\) and \(\mathbb{P}\in\mathcal{P}_2^r\), then we can show that the problem 3 is further equivalent to the \(\mathcal{W}_{2}\)-proximal problem by Brenier’s theorem (cf. [23]) that \[\label{eq:ot-prox} \max_{T \in L^2(\mathbb{P})}\mathbb{E}_{\xi \sim \mathbb{P}}\left[\ell(f_{\phi}, T(\xi))-\frac{\lambda}{2}\|T (\xi)-\xi \|^{2}\right].\tag{4}\] Hence, we can solve the inexact maximization step in Algorithm 1 using the above mentioned equivalent form by parameterizing the transport map.

A fundamental question regarding Algorithm 1 is its global convergence, the characterization of its limit points, and its oracle complexity. These results are established in the following subsection.

2.1 Convergence Guarantees↩︎

We begin by introducing the key notation for the distributionally robust inference model:

• Minimax loss function: \(\mathcal{H}(\phi,\mathbb{Q}):=\mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f_{\phi}, \xi)]-\lambda\cdot \mathcal{D}(\mathbb{Q},\mathbb{P})\);

• Worst-case value function: \(\mathcal{V}(\phi):=\max_{\mathbb{Q}\in \mathcal{P}_2}\mathcal{H}(\phi,\mathbb{Q})\).

The following mild assumption is required to establish the convergence guarantees of Algorithm 1.

In the remainder of this paper, we focus on weakly convex functions, which are not necessarily convex. This class encompasses a wide variety of models with Lipschitz continuous gradients, and it even includes certain non-smooth functions, both of which are commonly encountered in modern machine learning. Specifically, our analysis covers traditional convex models (e.g., classification with logistic loss, or linear regression with ridge and sparse \(\ell_{1}\)-regularizations). Additionally, it extends to more advanced non-convex scenarios involving nonlinear regression, support vector machines (SVM) equipped with nonlinear kernels, non-convex regularizers such as smoothly clipped absolute deviation (SCAD) and minimax concave penalty (MCP), and neural networks parameterized by smooth activation functions. We provide further illustration through two concrete examples.

Before presenting the main convergence results, we first introduce a suitable stationarity measure to assess the algorithm’s performance. This requires establishing the weakly convex property of the worst-case value function \(\mathcal{V}\).

Thanks to Lemma 1, we can adopt a stationary notion based on the Moreau envelope, which is commonly used in optimization for weakly convex functions. To be more precise, we define the Moreau envelope of \(\mathcal{V}\) with \(r>\rho\) for any \(\phi\in\Phi\), denoted by \(\mathcal{V}_{1/r}(\cdot)\), that \[\mathcal{V}_{1/r}(\phi):=\min_{\phi^{\prime}\in\Phi}\left\{\mathcal{V}(\phi ^{\prime})+\frac{r}{2}\|\phi-\phi^{\prime}\|^{2}\right\}.\] The Moreau envelope has the following important property by directly adopting the proof in [26].

The above lemma indicates that a small \(\|\nabla \mathcal{V}_{1/r}(\phi)\|\) implies \(\phi\) is close to a point \(\mathop{\mathrm{prox}}_{\mathcal{V}/r}(\phi)\), which is approximately a stationary point of the original function \(\mathcal{V}\). In this context, the approximate stationarity of \(\mathcal{V}\) can be evaluated using the norm of the gradient of its Moreau envelope. We are now ready to present the convergence guarantee for Algorithm 1.

Theorem 2 characterizes the number of iterations required by Algorithm 1 to reach an approximate stationary point with a desired accuracy. However, compared to optimization algorithms in Euclidean space, the primary distinction of our proposed algorithm framework lies in the approach to solving the inexact maximization step on the probability space, or equivalently, the transport map maximization problem 3 . This involves finding the LFD at each iterative step of the algorithm. Thus, the dominant computational cost arises from solving the \(\mathbb{Q}_k\)-subproblem, which is generally more expensive than the \(\phi_k\)-update. Moreover, solving this problem directly by parameterization may lead to instability and mode collapse, phenomena often observed in probability density optimization.

3 JKO Scheme for Least Favorable Distribution↩︎

To address these challenges, we draw inspiration from the JKO scheme [27] and derive the oracle complexity of our algorithm in terms of the number of finer-grained JKO steps required in the iterative updates over probability space. The JKO scheme provides a stable approximation of the backward discrete-time gradient flow along generalized geodesics in Wasserstein space. In our framework, we adopt this scheme to approximately solve the inexact inner maximization step in Algorithm 1.

An approximate solution to the JKO scheme can be obtained by solving the following parameterized transport map maximization problem, initialized with a map \(T_0\) that pushes forward a continuous distribution \(\nu\in\mathcal{P}_2^r\) to \(\mathbb{P}\):

\[\label{eq:practical95problem} T_{i+1}:=\mathop{\mathrm{argmax}}_{T\in L^{2}(\nu)}\left\{\mathbb{E}_{\xi \sim\nu}\left[\ell(f_{\phi}, T(\xi))-\frac{1}{2\gamma}\|T(\xi)-T_{i}(\xi)\|^{2}\right]-\lambda\cdot \mathcal{D}(T_{\#}\nu,\mathbb{P})\right\}.\tag{5}\]

In the remaining part of this section, we show that with a sufficiently large regularization coefficient \(\lambda\) associated with common discrepancy functions over a weakly concave function \(\ell(f_{\phi}, \cdot)\), the JKO scheme produces a sequence of transported distributions \(\{\mathbb{Q}_{i}\}\) with desirable convergence properties. Specifically, \(\{\mathbb{Q}_{i}\}\) achieves exponential convergence in both the Wasserstein-2 distance \(\mathcal{W}_{2}(\mathbb{Q}_{i}, \mathbb{Q}^{*})\) and the objective gap \(\mathcal{H}(\phi, \mathbb{Q}^{*})-\mathcal{H}(\phi, \mathbb{Q}_{i})\), where \(\mathbb{Q}^{*}\) is the global maximizer of \(\mathcal{H}(\phi, \cdot)\). Consequently, for any desired accuracy in the stationarity measure, we can characterize the required number of subgradient calls and JKO steps in our proposed Algorithm 1, combined with the subproblem solver Algorithm 2.

3.1 Strong Convexity along Generalized Geodesics↩︎

As a starting point for analyzing the convergence rate of the JKO scheme, we aim to establish exponential convergence to significantly reduce the number of costly JKO steps required in the algorithm. This is motivated by the favorable structural properties of our minimax loss function: \[\label{eq:robustloss} \mathcal{H}(\phi,\cdot)=\mathbb{E}_{\xi \sim (\cdot)}[\ell(f_{\phi} , \xi)]-\lambda \cdot \mathcal{D}(\cdot,\mathbb{P}),\tag{6}\] which may exhibit desirable growth behavior similar to that of a strongly concave function. Strong convexity (or strong concavity in this context) plays a central role in ensuring fast convergence of iterative optimization methods. Before we proceed, we discuss how strong convexity is defined in the Wasserstein space as it is a geodesic space and does not share the usual linear structure. We should clarify the definition of strong/weak convexity along the geodesics and more general curves.

By examining the structure of 6 , we can observe that the mapping \(\mathbb{Q} \mapsto \mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f_{\phi}, \xi)]\) is “weakly concave”. If the discrepancy term \(\mathcal{D}(\mathbb{Q}, \mathbb{P})\) is “strongly convex”, then for sufficiently large \(\lambda\), the overall objective \(\mathcal{H}(\phi, \cdot)\) becomes “strongly concave”. However, whether such properties hold depends heavily on which curves we consider. For example, from [28] we know that when \(\mathcal{D}=\frac{1}{2}\mathcal{W}_{2}^{2}\), the squared Wasserstein distance \(\mathbb{Q}\mapsto \frac{1}{2}\mathcal{W}_{2}^{2}(\mathbb{Q},\mathbb{P})\) is \(1\)-weakly concave (i.e., not convex) along geodesics. Moreover, in dimensions greater than one, [28] provides a counterexample showing that strong convexity does not generally hold.

Consequently, we need a more general notion than standard geodesics between probability measures. Here, we introduce convexity along generalized geodesics as mentioned in 7 to address this gap. Different from the geodesics, the generalized geodesics involve a third continuous distribution \(\nu\) and are defined using interpolation of the two OT maps from \(\nu\) to \(\mu_{1}\) and \(\mu_{2}\), respectively. Specifically, let \(\nu \in \mathcal{P}_{2}^{r}\) and \(\mu_{1}, \mu_{2}\in \mathcal{P}_{2}\), then a generalized geodesic joining \(\mu_{1}\) to \(\mu_{2}\) centered at \(\nu\) is a curve of type \[\label{eq:gene95geodesic}\mu_{t}^{1 \rightarrow 2}:=\left((1-t) T_{\nu}^{\mu_1} +t T_{\nu}^{\mu_2}\right)_{\#}\nu, \quad t \in[0,1] .\tag{7}\] This will also lead to the generalized definition of the Wasserstein distance (i.e., linearized OT distance): \[\mathcal{W}_{\nu}^{2}(\mu_{1}, \mu_{2}):=\int_{\mathbb{R}^d}\|T_{\nu}^{\mu_1} (x)-T_{\nu}^{\mu_2}(x)\|^{2}\operatorname{d \nu}(x).\]

Fortunately, we can identify several important cases where strong convexity holds along certain different generalized geodesics.

3.2 Oracle Complexity↩︎

Leveraging the previously discussed strongly convex a.g.g. property, we are now prepared to derive the oracle complexity of our proposed algorithm. We begin by outlining the following necessary assumptions.

A sufficiently large regularization parameter \(\lambda\), corresponding to a small ambiguity set radius, induces a favorable optimization landscape as demonstrated by the following lemma.

Since \(\mathcal{H}(\phi, \cdot)\) is \((\lambda - \rho)\)-strongly concave a.g.g. in \(\mathcal{P}_{2}\) by Lemma 3, we can derive the exponential convergence of the inexact proximal point method along the generalized geodesics in Wasserstein space, as proposed in Algorithm 2, to the optimal solution of the subproblem in Algorithm 1.

Using the aforementioned assumption on the accuracy of the inexact JKO step in Algorithm 2 (which can also be satisfied by the parameterized transport map maximization problem thanks to ?? ), we obtain the following results.

The above proposition allows us to establish nonasymptotic convergence results for our proposed Algorithm 1, specifically regarding the number of subgradient evaluations and the required number of inexact JKO steps.

One remaining issue lies in the small step size \(\eta = \mathcal{O}(\varepsilon^{2})\) required in Theorems 2 and 11, which compensates for the nonsmooth nature of the problem. In the next section, we show that under a smoothness condition, the complexity bounds can be further improved and the double-loop structure can be avoided.

4 Smoothness and Alternating Scheme↩︎

The nonsmoothness arises from two sources in the problem 2 :

1. The nonsmoothness of \(\ell(f_{(\cdot)}, \xi)\) for any \(\xi \in \mathbb{R}^d\), which is typically from nonsmooth modeling elements, e.g., regularizers and parameterizations;

2. The nonunique solutions of the LFD in the maximization problem.

For (i), in this section we additionally assume that \(\ell(f _{(\cdot)}, \xi)\) is gradient Lipschitz continuous. For (ii), we impose Assumption 2 to ensure that the problem is strongly concave a.g.g., which guarantees the uniqueness of the LFD.

Although this assumption introduces an additional requirement for the model, it does not significantly restrict its applicability. For example, in practice, we can use smooth activation functions in neural networks, such as the softplus function, to replace the highly nonsmooth ReLU activation function. Then, with the help of the generalized Danskin’s theorem, we derive the smoothness of the worst-case value function \(\mathcal{V}\).

By applying Danskin’s theorem as established in Proposition 12, we derive key Lipschitz continuity properties of the solution mapping, which in turn imply the Lipschitz gradient of the worst-case value function \(\mathcal{V}\).

Next, we derive improved oracle complexity results under smoothness for Algorithm 1, using the subproblem solver Algorithm 2 introduced in the previous section.

Furthermore, we propose the following alternating scheme for the fully smooth case, i.e., the parameter space is \(\Phi = \mathbb{R}^m\). Unlike the eliminating approach in Algorithm 1, it avoids the time-consuming JKO iterations in the subproblem by performing only a single update per outer iteration.

Together, Theorems 13 and 14 quantify the computational cost of the elimination-based and alternating schemes in terms of gradient oracle calls and inexact JKO oracle calls. In practice, each inexact JKO step can be implemented by optimizing a parameterized neural transport map in 5 ; the resulting oracle error \(\epsilon'\) captures the combined effect of optimization, approximation, and sampling errors. This oracle viewpoint separates the Wasserstein-space convergence analysis from the specific neural parameterization, while retaining the computational advantage that a learned map can generate worst-case samples by a forward evaluation.

5 Implementation of Wasserstein-2 Inner Update↩︎

The convergence results above use inexact JKO steps as the basic computational primitive over probability measures. To apply the algorithms, these steps must be approximated by finite-dimensional computations. The \(\mathcal{W}_2\)-regularized DRO problem is a natural setting for this approximation because it is closely related to the Wasserstein robust training formulation based on pointwise perturbations [15]: the least-favorable distribution update maximizes expected loss penalized by a quadratic transport cost. Because this quadratic cost can be written through transport maps, the abstract update over \(\mathbb{Q}\) can be approximated in finite dimensions either by optimizing transported samples directly or by optimizing a parameterized map.

At the population level, following the transport-map formulation in 3 –5 , let \(\nu\) be a continuous base distribution, let \(T_0{}_{\#}\nu=\mathbb{P}=\mathbb{Q}_0\), and let \(T_k{}_{\#}\nu=\mathbb{Q}_k\). For fixed \(f_{\phi_k}\), a new distribution is represented as \(T_{\#}\nu\), and the map objective for the JKO step is \[\begin{align} \max_T \mathbb{E}_{\zeta\sim\nu} \left[ \ell(f_{\phi_k},T(\zeta)) -\frac{\lambda}{2}\|T(\zeta)-T_0(\zeta)\|^2 -\frac{1}{2\gamma}\|T(\zeta)-T_k(\zeta)\|^2 \right]. \label{eq:flow-ot} \end{align}\tag{8}\] Setting \(\gamma=\infty\) removes the proximal term and recovers the direct inner maximization in Algorithm 1; finite \(\gamma\) gives the proximal update used by the alternating scheme.

A common implementation is the direct-sampling specialization of this formulation. When samples can be drawn directly from \(\mathbb{P}\), one may take \(\nu=\mathbb{P}\) and \(T_0=I_d\). Then the base sample \(\zeta\) above is a reference sample, written as \(\xi_i^0\sim\mathbb{P}\) in a finite batch \(i = 1, 2\ldots, n\). The current map gives \(\xi_i^k=T_k(\xi_i^0)\), and the empirical inner problem chooses updated samples \(\xi_i\) by maximizing \[\max_{\xi_1,\ldots,\xi_n} \frac{1}{n}\sum_{i=1}^n \left[ \ell(f_{\phi_k},\xi_i) -\frac{\lambda}{2}\|\xi_i-\xi_i^0\|^2 -\frac{1}{2\gamma}\|\xi_i-\xi_i^k\|^2 \right].\] This empirical problem leads to two implementation families. Particle methods optimize the updated samples \(\xi_i\) directly and provide finite-sample baselines for the direct inner maximization, especially in the elimination case \(\gamma=\infty\). Neural transport maps instead impose the shared form \(\xi_i=T_\theta(\xi_i^0)\), which amortizes the update across samples and can be applied immediately to newly drawn data. This is particularly useful for the alternating scheme, where the proximal term can be evaluated on new samples by comparing \(T_\theta(\xi_i^0)\) with \(T_k(\xi_i^0)\).

For supervised learning, a sample can be written as \(\xi=(x,y)\). The transport may keep the label fixed and move only the feature vector, so the quadratic costs are computed in feature space. If \(\xi_i^0=(x_i,y_i)\), we write the updated sample as \(\xi_i=(z_i,y_i)\). Particle methods optimize the feature vectors \(z_i\) directly, while neural maps can use the label-conditioned form \(z_i=T_{\theta,y_i}(x_i)\).

6 Numerical Experiments↩︎

We now evaluate the implementation choices above on two tasks. The two-dimensional Gaussian mixture isolates the inner \(\mathbb{Q}\)-update and visualizes the transported samples, while the CIFAR-10 feature-space experiment tests the same inner maximization in a higher-dimensional representation. The comparisons focus on the effect of the JKO proximal term for neural maps and, in the CIFAR-10 experiment, on whether neural maps are better used directly or as initializers for particle L-BFGS. The cross-entropy losses used below satisfy the assumptions in the theoretical analysis under the boundedness and smoothness conditions induced by the compact parameter/feature domains and the smooth SiLU networks. Together with the Wasserstein-2 regularizer, whose a.g.g. strong convexity is established, this verifies the additional assumptions used by the JKO and alternating schemes. The synthetic experiment verifies that neural transport maps can qualitatively match particle-based least-favorable movements while preserving the fast inference advantage of a learned map. The CIFAR-10 experiment tests the same mechanism in a higher-dimensional setting and shows that neural maps are useful as standalone solvers, while being particularly effective as initializers for particle methods.

6.1 Two-Dimensional Gaussian Mixture↩︎

6.1.0.1 Setup.

We instantiate the Wasserstein-2 DRO problem on a binary two-dimensional classification task. Class 0 follows \(\mathbb{P}(x|{y=0})=\mathcal{N}(0,I_2)\), and class 1 follows a 25-component Gaussian mixture \(\mathbb{P}(x|{y=1})\) with means on \(\{-10,-5,0,5,10\}^2\) and covariance \(0.1I_2\). For \(\xi=(x,y)\), \(y\in\{0,1\}\), we use the binary cross-entropy \(\ell(f_\phi,(x,y))=-\log\frac{\exp(f_{\phi,y}(x))}{\exp(f_{\phi,0}(x))+\exp(f_{\phi,1}(x))}\). For numerical illustrations, we consider the example where only one class sample is perturbed. The inner maximization is applied only to class-0 samples; class-1 samples remain fixed. Thus below \(T_\theta\) denotes \(T_{\theta,0}\). At outer iteration \(k\), we draw class-0 samples \(\{x_i\}_{i=1}^n\) and class-1 samples \(\{x'_j\}_{j=1}^n\). We denote the moved class-0 points by \(z_i\): particle methods optimize \(z_i\) directly, while neural-map methods set \(z_i=T_\theta(x_i)\). After the inner update, the classifier is updated using \(\widehat{\mathbb{Q}}_k = \frac{1}{2n}\sum_{i=1}^{n}\delta_{(z_i,0)} + \frac{1}{2n}\sum_{j=1}^{n}\delta_{(x'_j,1)}\). The corresponding empirical minimax objective is \[\widehat{\mathcal{H}}_k = \frac{1}{2n}\sum_{i=1}^{n}\ell(f_{\phi_k},(z_i,0)) + \frac{1}{2n}\sum_{j=1}^{n}\ell(f_{\phi_k},(x'_j,1)) - \frac{\lambda}{4n}\sum_{i=1}^{n}\|z_i-x_i\|^2 .\] The first two terms are the empirical classifier loss. For the inner stationarity metric, define the particle objective \(h_i(z;\phi_k):=\ell(f_{\phi_k},(z,0))-\frac{\lambda}{2}\|z-x_i\|^2\). For a fixed batch, the particle inner problem is \(\max_{z_1,\ldots,z_n}\sum_i h_i(z_i;\phi_k)\). The full setup is given in Appendix Table 1.

6.1.0.2 Methods.

We compare particle updates with neural transport maps. The particle baselines are L-BFGS, implemented as one batched solve over the stacked vector \((z_1,\ldots,z_n)\), and GD, which follows the rescaled update of [15] by minimizing \(-\lambda^{-1}\ell(f_\phi,(z,y))+\frac{1}{2}\|z-x\|^2\) with diminishing steps and the gradient warm start from [15]. We also report an unrescaled GD variant in Figure 9. The neural methods use residual maps \(T_\theta(x)=x+R_\theta(x)\). Elim-\(B\) performs \(B\in\{5,10,15\}\) optimizer steps on \[\max_\theta \frac{1}{n}\sum_{i=1}^{n} \left[ \ell(f_{\phi_k},(T_\theta(x_i),0)) - \frac{\lambda}{2}\|T_\theta(x_i)-x_i\|^2 \right],\] which corresponds to 8 with \(\gamma=\infty\). Alt-\(B\) uses the same \(B\) values and adds the JKO proximal term \(-\frac{1}{2\gamma}\|T_\theta(x_i)-T_{\theta_k}(x_i)\|^2\) with \(\gamma=5.0\). All neural maps are warm-started from the previous outer iteration.

6.1.0.3 Metrics.

We report two convergence metrics: the classifier gradient norm \(\|\nabla_\phi\widehat{\mathcal{H}}\|\), computed before gradient clipping, and the RMS particle-gradient norm \[\|\nabla_z h\|_{\mathrm{RMS}} := \left( \frac{1}{n}\sum_{i=1}^{n} \|\nabla_z h_i(z_i;\phi_k)\|^2 \right)^{1/2},\] which measures stationarity of the finite-particle inner problem.

6.1.0.4 Results.

The two-dimensional experiment supports the first claim: the JKO proximal term stabilizes the neural \(\mathbb{Q}\)-update. Figure 4 shows the moved samples: the particle methods and Alt-5 move class-0 samples toward high-loss regions near the class-1 modes, while Elim-5 is less aligned with the particle updates. Figure 5 shows the same pattern in the convergence metrics. With the same neural-map update budget, Alt-\(B\) ends with smaller classifier gradients than Elim-\(B\) (\(0.085\)–\(0.107\) versus \(0.53\)–\(1.61\)) and smaller \(\|\nabla_z h\|_{\mathrm{RMS}}\) (\(5.0\times10^{-3}\)–\(6.7\times10^{-3}\) versus \(1.0\times10^{-2}\)–\(4.9\times10^{-2}\)). Additional loss and objective curves are given in Figure 8.

6.2 CIFAR-10 Image Classification↩︎

6.2.0.1 Setup.

We next consider CIFAR-10 feature vectors. Images are encoded by a fixed ViT-base model, and the classifier is trained in the resulting feature space. For \(\xi=(x,y)\) with \(y\in\{0,\ldots,9\}\), we use the multiclass cross-entropy \(\ell(f_\phi,(x,y))=-\log\frac{\exp(f_{\phi,y}(x))}{\sum_{c=0}^{9}\exp(f_{\phi,c}(x))}\). For a mini-batch \(\{(x_i,y_i)\}_{i=1}^{n}\), the particle objective is \[\max_{z_1,\ldots,z_n} \frac{1}{n}\sum_{i=1}^{n} \left[ \ell(f_{\phi_k},(z_i,y_i)) - \frac{\lambda}{2}\|z_i-x_i\|^2 \right].\]

6.2.0.2 Methods.

We compare ERM, GD, L-BFGS, Neural, and Neural+L-BFGS. GD and L-BFGS optimize the particles, i.e., feature vectors, directly. Neural uses a label-conditioned transport map \(T_{\theta,y_i}(x_i)\). Neural+L-BFGS uses this map output, \(T_{\theta,y_i}(x_i)\), to initialize the L-BFGS solve. We sweep \(\lambda\in\{10^{-3},10^{-2},10^{-1},1\}\) over five seeds. Robustness is evaluated by 50-step \(\ell_2\) PGD in feature space with budget \(0.2\) times the average feature norm. The full setup is given in Appendix Table 2.

6.2.0.3 Metrics.

We report clean test accuracy, feature-space PGD-50 accuracy, and CIFAR-10-C [30] mean accuracy. The CIFAR-10-C mean averages the Weather, Blur, Noise, and Digital corruption groups. For computational cost, we report L-BFGS iterations and function evaluations. An L-BFGS iteration is one accepted quasi-Newton update of the stacked particles, while a function evaluation is one batched objective and gradient evaluation. These counts differ because line search can evaluate several candidate step lengths before accepting an update.

6.2.0.4 Results.

Figure 6 shows the same pattern on both robustness benchmarks. Particle GD and L-BFGS give the strongest direct inner updates, while Neural alone is consistently weaker in the high-dimensional feature space. Using the neural map as an initialization largely closes this gap: Neural+L-BFGS is within \(0.7\) percentage points of GD and L-BFGS on PGD-50 and is essentially tied with them on CIFAR-10-C at the best \(\lambda\) values. Clean accuracy remains near \(98\%\), so the robustness improvement does not come from sacrificing nominal performance. Figure 7 further shows that the neural initialization reduces both L-BFGS iterations and function evaluations under the same stopping criteria. Overall, these results suggest that neural maps are most useful here as amortized warm starts for particle optimization, rather than as standalone replacements for direct particle updates.

7 Conclusions↩︎

We have investigated iterative algorithms for solving the DRO problem from the perspective of optimization in Wasserstein space. Unlike prior works that analyze dynamics in continuous time, we adopt a viewpoint based on transforming probability distributions via iterative steps. In addition to being a robust classifier, the solution derived from our algorithm can also be utilized to efficiently generate samples from the LFD. A future direction involves extending our framework to the general nonlinear DRO problems [31], which are receiving increasing attention due to their relevance in real-world applications.

Acknowledgment↩︎

This work is partially supported by an NSF CAREER CCF-1650913, NSF DMS-2134037, CMMI-2015787, CMMI-2112533, DMS-1938106, DMS-1830210, ONR N000142412278, and the Coca-Cola Foundation. The authors thank Professors Xiuyuan Cheng and Johannes Milz for helpful discussions.

8 Definitions and Basic Facts↩︎

Let us first introduce weakly convex functions, the class of nonconvex functions studied in this paper.

Next, we will introduce the discrepancy functions considered in the paper. There are two natural ways to model changes in distributions to build distributional ambiguity sets. The first way is to assume that the likelihood of the baseline model (often uniformly distributed over an observed data set) is corrupted. In this case, we may account for model misspecification using likelihood ratios. This leads to \(\varphi\)-divergence-based uncertainty sets. The second way is to consider integral probability metrics, e.g., Wasserstein distance, to account for potential corruptions in the baseline model by considering perturbations in the actual outcomes.

When \(\mathbb{P}\in \mathcal{P}_{2}^{r}\), Brenier’s theorem [32] guarantees a well-defined and unique optimal transport (OT) map from \(\mathbb{P}\) to any \(\mathbb{Q}\in \mathcal{P}_{2}\). Specifically, let \(T_{\mathbb{P}}^{\mathbb{Q}}\) denote the OT map from \(\mathbb{P}\) to \(\mathbb{Q}\), which is defined \(\mathbb{P}\)-almost everywhere and satisfies \((T_{\mathbb{P}}^{\mathbb{Q}})_{\#}\mathbb{P}= \mathbb{Q}\). For any map \(T: \mathbb{R}^{d}\to \mathbb{R}^{d}\) such that \(T_{\#}\mathbb{P}= \mathbb{Q}\), the measure \((I_{d}, T)_{\#}\mathbb{P}\) forms a coupling of \(\mathbb{P}\) and \(\mathbb{Q}\). Consequently, we know that \(\mathcal{W}_{2}^{2}(\mathbb{Q}, \mathbb{P}) \leq \mathbb{E}_{\xi \sim \mathbb{P}}\|\xi - T(\xi )\|^{2}\). Minimizing the right-hand side over all \(T\) that push forward \(\mathbb{P}\) to \(\mathbb{Q}\) is known as the Monge problem. By Brenier’s theorem, when \(\mathbb{P}\in \mathcal{P}_{2}^{r}\), the OT map \(T_{\mathbb{P}}^{\mathbb{Q}}\) solves the Monge problem, leading to the following important identity for the Wasserstein distance in this case: \[\mathcal{W}_{2}^{2}(\mathbb{Q}, \mathbb{P}) = \mathbb{E}_{\xi \sim \mathbb{P}}\|\xi - T_{\mathbb{P}}^{\mathbb{Q}} (\xi)\|^{2}.\]

Finally, because our approach involves optimization over a probability space equipped with the Wasserstein metric, it is natural to introduce first-order concepts such as subdifferentials. Consider a proper, lower semi-continuous functional \(\mathcal{V}: \mathcal{P}_{2}\to (-\infty, \infty]\), and let its domain be \(\operatorname{dom}(\mathcal{V}) := \{\mu \in \mathcal{P}_{2}: \mathcal{V}(\mu ) < \infty\}.\) We adopt the Fréchet subdifferential of \(\mathcal{V}\) from [28] Definition 10.1.1, focusing on its strong form. Under this definition, the Wasserstein subgradients provide linear approximations even if the perturbations do not arise from the optimal transport [33] Proposition 2.12.

9 Proofs for Section 2↩︎

Proof of Lemma 1↩︎

Since for each \(\xi\in\mathbb{R}^d\) the function \(\ell(f_{(\cdot)}, \xi)\) is \(\rho\)-weakly convex on \(\Phi\) by Assumption 1, we can directly check by definition that the function \(\mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f_{(\cdot)}, \xi)]=\int_{\mathbb{R}^d} \ell(f_{(\cdot)}, \xi) \operatorname{d \mathbb{Q}}(\xi)\) is a \(\rho\)-weakly convex function on \(\Phi\) for each \(\mathbb{Q}\in \mathcal{P}_{2}\). Then we know that \(\mathcal{H}(\cdot,\mathbb{Q})\) is \(\rho\)-weakly convex on \(\Phi\), which implies for any \(\phi,\psi\in\Phi\) and \(\alpha \in[0,1]\) that \[\begin{align} & \mathcal{H}(\alpha \phi+(1-\alpha) \psi,\mathbb{Q}) \\ \leq\; & \alpha \mathcal{H}(\phi,\mathbb{Q})+(1-\alpha) \mathcal{H}(\psi,\mathbb{Q})+\frac{\rho \alpha(1-\alpha)}{2}\|\phi-\psi\|^{2} \\ \leq\; & \max_{\mathbb{Q}'}\left\{\alpha \mathcal{H}(\phi,\mathbb{Q}')+(1-\alpha) \mathcal{H}(\psi,\mathbb{Q}')+\frac{\rho \alpha(1-\alpha)}{2}\|\phi-\psi\|^{2}\right\} \\ \leq\; & \alpha\max_{\mathbb{Q}'}\mathcal{H}(\phi,\mathbb{Q}')+(1-\alpha) \max_{\mathbb{Q}'}\mathcal{H}(\psi,\mathbb{Q}')+\frac{\rho\alpha(1-\alpha)}{2}\|\phi-\psi\|^{2} \\ =\; & \alpha \mathcal{V}(\phi)+(1-\alpha) \mathcal{V}(\psi)+\frac{\rho\alpha(1-\alpha)}{2}\|\phi-\psi\|^{2}, \end{align}\] where the first inequality is by the definition of the weakly convex function. Taking the supremum of the left hand side over \(\mathbb{Q}\) gives the desired result: \[\mathcal{V}(\alpha \phi+(1-\alpha) \psi) \leq \alpha \mathcal{V}(\phi )+(1-\alpha) \mathcal{V}(\psi)+\frac{\rho\alpha(1-\alpha)}{2}\|\phi-\psi \|^{2}.\] The proof is complete.

Proof of Theorem 2↩︎

From the proof of Lemma 1, we know that \(\mathcal{H}(\cdot ,\mathbb{Q})\) is \(\rho\)-weakly convex. Furthermore, since \(\mathbb{Q}_{k}\) is an \(\epsilon\)-approximate maximizer of \(\max_{\mathbb{Q}\in \mathcal{P}_2}\mathcal{H}(\phi_{k},\mathbb{Q})\), it follows for any \(\phi \in \Phi\) and \(\zeta(\phi_{k},\mathbb{Q}_{k})\in \partial_{\phi}\mathbb{E}_{\xi \sim \mathbb{Q}_k}[\ell (f_{(\cdot)},\xi)](\phi_{k})\) that \[\label{eq:max-property} \begin{align} \mathcal{V}(\phi) \geq \mathcal{H}(\phi, \mathbb{Q}_{k}) & =\mathbb{E}_{\xi \sim \mathbb{Q}_k}[\ell(f_{\phi}, \xi)]-\lambda\cdot \mathcal{D}(\mathbb{Q}_{k},\mathbb{P}) \\ & \geq \mathbb{E}_{\xi \sim \mathbb{Q}_k}[\ell(f_{\phi_k}, \xi)+\langle\partial_{\phi}\ell(f_{\phi_k}, \xi), \phi-\phi_{k}\rangle]-\frac{\rho}{2}\|\phi-\phi_{k}\|^{2}-\lambda\cdot \mathcal{D}(\mathbb{Q}_{k},\mathbb{P}) \\ & = \mathcal{H}(\phi_{k}, \mathbb{Q}_{k})+\langle\zeta(\phi_{k},\mathbb{Q}_{k}), \phi-\phi_{k}\rangle-\frac{\rho}{2}\|\phi-\phi_{k}\|^{2} \\ & \geq \mathcal{V}(\phi_{k})-\epsilon+\langle\zeta(\phi_{k},\mathbb{Q}_{k}), \phi-\phi_{k}\rangle-\frac{\rho}{2}\|\phi-\phi_{k}\|^{2}, \end{align}\tag{9}\] where the first inequality is from the \(\rho\)-weakly convexity of the function \(\ell(f_{(\cdot)}, \xi)\). On the other hand, we know from the iterate of Algorithm 1 that \[\label{eq:sub-key00} \begin{align} \|\phi_{k+1}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2} & = \|\mathop{\mathrm{proj}}_{\Phi}(\phi_{k}-\eta_k\cdot\zeta(\phi_{k},\mathbb{Q}_{k}))-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2} \\ & \leq \|\phi_{k}-\eta_k\cdot\zeta(\phi_{k},\mathbb{Q}_{k})-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2} \\ & \leq \|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2}+2 \eta_k\langle\zeta(\phi_{k},\mathbb{Q}_{k}), \mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})-\phi_{k}\rangle+\eta_k^{2}L^{2}, \end{align}\tag{10}\] where the first inequality is from the nonexpansiveness of projection onto convex sets, and the second one is from \(\|\zeta(\phi_{k},\mathbb{Q}_{k})\|\leq L\) by \(L\)-Lipschitz continuity of \(\ell(f_{(\cdot)}, \xi)\) on \(\Phi\) for any \(\xi\in\mathbb{R}^d\). Combining 10 and 9 with \(\phi=\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\) indicates that \[\label{eq:sub-key1} \begin{align} \mathcal{V}_{1/2\rho}(\phi_{k+1}) & \leq \mathcal{V}(\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k}))+\rho\|\phi_{k+1}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2} \\ & \leq \mathcal{V}_{1/2\rho}(\phi_{k})+2 \eta_k \rho\langle\zeta(\phi_{k},\mathbb{Q}_{k}), \mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})-\phi_{k}\rangle+\eta_k^{2}\rho L^{2} \\ & \leq \mathcal{V}_{1/2\rho}(\phi_{k})+2 \eta_k \rho\left(\mathcal{V}(\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k}))-\mathcal{V}(\phi_{k})+\epsilon+\frac{\rho}{2}\|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2}\right)+\eta_k^{2}\rho L^{2}, \end{align}\tag{11}\] Taking the sum over \(k=0,\ldots,K-1\) of 11 we obtain \[\begin{align} \mathcal{V}_{1/2\rho}(\phi_{K}) -\mathcal{V}_{1/2\rho}(\phi_{0}) \leq\;&2 \rho \sum_{k=0}^{K-1}\eta_k\left(\mathcal{V}(\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k}))-\mathcal{V}(\phi_{k})+\epsilon+\frac{\rho}{2}\|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2}\right)\\ &+\rho L^{2}\sum_{k=0}^{K-1}\eta_k^2, \end{align}\] which implies that \[\label{eq:key-convergence} \begin{align} &\frac{1}{\sum_{k=0}^{K-1}\eta_k}\sum_{k=0}^{K-1}\eta_k\left(\mathcal{V} (\phi_{k})-\mathcal{V}(\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k}))-\frac{\rho}{2} \|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2}\right) \\ \leq\;& \frac{\mathcal{V}_{1/2\rho}(\phi_{0})-\min_{\phi\in\Phi}\mathcal{V}(\phi)}{2 \rho \sum_{k=0}^{K-1}\eta_k}+\frac{ L^{2}\sum_{k=0}^{K-1}\eta_k^2}{2\sum_{k=0}^{K-1}\eta_k}+\epsilon. \end{align}\tag{12}\] Moreover, since \(\mathcal{V}(\cdot)+\rho\|\cdot-\phi_{k}\|^{2}\) is \(\rho\)-strongly convex, we have \[\begin{align} & \mathcal{V}(\phi_{k})-\mathcal{V}(\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k}))-\frac{\rho}{2}\|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2} \\ =\; & \mathcal{V}(\phi_{k})+\rho\|\phi_{k}-\phi_{k}\|^{2}-\mathcal{V}(\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k}))-\rho\|\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})-\phi_{k}\|^{2}+\frac{\rho}{2}\|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2} \\ =\; & \mathcal{V}(\phi_{k})+\rho\|\phi_{k}-\phi_{k}\|^{2}-\min_{\phi\in\Phi}\left\{ \mathcal{V}(\phi)+\rho\|\phi-\phi_{k}\|^{2}\right\}+\frac{\rho}{2}\|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2} \\ \geq\; & \rho\|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2}=\frac{1}{4 \rho}\|\nabla \mathcal{V}_{1/2\rho}(\phi_{k})\|^{2}, \end{align}\] where the last equality is from the explicit form of the gradient of the Moreau envelope [26] that \(\nabla \mathcal{V}_{1/2\rho} (\phi)=2\rho(\phi-\mathop{\mathrm{prox}}_{\frac{1}{2\rho}\mathcal{V}+\iota_{\Phi}}(\phi))\). Plugging this in 12 , it implies that there exists a \(k'\in\{0,1,\ldots,K-1\}\) such that \[\begin{align} \|\nabla \mathcal{V}_{1/2\rho}(\phi_{k'})\|^{2} &\leq\frac{1}{\sum_{k=0}^{K-1}\eta_k}\sum_{k=0}^{K-1}\eta_k\|\nabla \mathcal{V}_{1/2\rho}(\phi_{k})\|^{2} \\ & \leq \frac{4\rho}{\sum_{k=0}^{K-1}\eta_k}\sum_{k=0}^{K-1}\eta_k\left(\mathcal{V}(\phi_{k})-\mathcal{V}(\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k}))-\frac{\rho}{2}\|\phi_{k}-\mathop{\mathrm{prox}}_{\mathcal{V}/2\rho}(\phi_{k})\|^{2}\right) \\ & \leq \frac{2(\mathcal{V}_{1/2\rho}(\phi_0)-\min_{\phi\in\Phi} \mathcal{V}(\phi))+2 \rho L^{2}\sum_{k=0}^{K-1}\eta_k^2}{ \sum_{k=0}^{K-1}\eta_k}+4\rho\epsilon\\ &= \frac{2(\mathcal{V}_{1/2\rho}(\phi_0)-\min_\phi \mathcal{V}(\phi)+\rho L^2) }{\sqrt{K}}+4 \rho \epsilon, \end{align}\] where the equality is derived from \(\eta_k=1 / \sqrt{K}\). The proof is complete.

10 Proofs for Section 3↩︎

Proof of Proposition 4↩︎

Denote \(F_{i+1}\) the mapping such that \(T_{i+1}= F_{i+1}\circ T_{i}\) since \(T_{i}\) is non-degenerate. Then we know that \[\begin{align} &\max_{T\in L^{2}(\nu)}\left\{\mathbb{E}_{\xi \sim \nu}\left[\ell(f_{\phi}, T(\xi))-\frac{1}{2\gamma}\|T (\xi)-T_{i}(\xi)\|^{2}\right]-\lambda\cdot \mathcal{D}(T _{\#}\nu,\mathbb{P})\right\}\\ =\;&\max_{F\in L^2(\mathbb{Q}_i)}\Bigg\{ \mathbb{E}_{\xi \sim \nu}\left[ \ell(f_{\phi}, F\circ T_{i}(\xi))\right] -\lambda\cdot \mathcal{D}((F\circ T_{i} )_{\#}\nu,\mathbb{P})- \frac{1}{2\gamma}\mathbb{E}_{\xi \sim \nu}\left[\|F\circ T_{i}(\xi) - T_{i}(\xi)\|^{2}\right] \Bigg\}\\ =\;&\max_{F\in L^2(\mathbb{Q}_i)}\left\{\mathcal{H}(\phi,(F\circ T_{i})_{\#}\nu)-\frac{1}{2\gamma}\mathbb{E}_{\xi \sim \nu}\left[\|(F-I)\circ T_i(\xi)\|^{2}\right]\right\}\\ =\;&\max_{F\in L^2(\mathbb{Q}_i)}\left\{\mathcal{H}(\phi,F_{\#}((T_{i})_{\#}\nu))-\frac{1}{2\gamma}\mathbb{E}_{\xi \sim (T_{i})_{\#}\nu}\left[\|F(\xi)-\xi\|^{2}\right]\right\}\\ =\;&\max_{F\in L^2(\mathbb{Q}_i)}\left\{\mathcal{H}(\phi,F_{\#}\mathbb{Q}_i)-\frac{1}{2\gamma}\mathbb{E}_{\xi \sim \mathbb{Q}_i}\left[\|F(\xi)-\xi\|^{2}\right]\right\}\\ =\;&\max_{\mathbb{Q}\in\mathcal{P}_2}\left\{\mathcal{H}(\phi,\mathbb{Q})-\frac{1}{2\gamma} \cdot\mathcal{W}_{2}^{2}(\mathbb{Q}, \mathbb{Q}_{i})\right\}, \end{align}\] where the third equality is from the chain rule for pushforward map and the last one is from Brenier’s theorem (cf. [23]).

Next, from the optimality condition of the problem \(\max_{\mathbb{Q}\in\mathcal{P}_2}\{\mathcal{H}(\phi,\mathbb{Q})-\frac{1}{2\gamma}\cdot\mathcal{W}_{2}^{2}(\mathbb{Q}, \mathbb{Q}_{i})\}\) with the calculus rule in [33], we know that \[0\in -\partial_{\mathcal{W}_2}\mathcal{H}(\phi, (T_{{i+1}})_{\#} \nu))+ \frac{1}{\gamma}(I- T^{i}_{(T_{{i+1}})_{\#}\nu}).\] Then we have \[\begin{align} 0 &\in -\partial_{\mathcal{W}_2}\mathcal{H}(\phi, (T_{{i+1}})_{\#}\nu)+ \frac{1}{\gamma}(I- T^{i}_{(T_{{i+1}})_{\#}\nu})\\ &= -\partial_{\mathcal{W}_2}\mathcal{H}(\phi, (T_{{i+1}})_{\#}\nu)\circ T_{\nu}^{(T_{{i+1}})_{\#}\nu}+ \frac{1}{\gamma}(T_{\nu}^{(T_{{i+1}})_{\#}\nu}-T^{i}_{(T_{{i+1}})_{\#}\nu}\circ T_{\nu}^{(T_{{i+1}})_{\#}\nu} ). \end{align}\] The proof is complete.

Proof of Proposition 7↩︎

The first part about the \(1\)-strong convexity a.g.g. of the function \(\frac{1}{2}\mathcal{W}_{2}^{2}(\cdot,\mathbb{P})\) is directly from [33] Corollary 2.19(v) and [28]. For the proof of the second part, we know from \(\mathbb{Q}\ll \mathbb{P}\) that \(\mathbb{Q}\) has the density \(q\) and \[\label{eq:KL95equi}\operatorname{KL}(\mathbb{Q},\mathbb{P})=\int_{\mathbb{R}^d} \frac{\operatorname{d\mathbb{Q}}}{\operatorname{d\mathbb{P}}}\log \left(\frac{\operatorname{d\mathbb{Q}}}{\operatorname{d\mathbb{P}}} \right) \operatorname{d\mathbb{P}}=\int_{\mathbb{R}^d}(q\log (q)-q\log (p))=c+\int_{\mathbb{R}^d}(q\log (q)+G q),\tag{13}\] where \(c\) is a constant and the last equality is by \(p\propto \exp(-G)\).

Since \(s \mapsto \log(s^{-d})= s^{d}F(s^{-d})\) with \(F(s)= s \log s\) is convex and non-increasing on \((0,+\infty)\), we know from [28] that \(q\mapsto \int_{\mathbb{R}^d}F(q)=\int_{\mathbb{R}^d}q\log (q)\) is convex a.g.g. in \(\mathcal{P}_{2}^{r}\). On the other hand, since \(G\) is \(1\)-strongly convex and also proper, lower semi-continuous, bounded from below, by [28] Proposition 9.3.2(i) it follows that \(q\mapsto\int_{\mathbb{R}^d} G q\) is \(1\)-convex along any interpolation curve, which implies \(1\)-strong convexity a.g.g. in \(\mathcal{P}_{2}^{r}\). Hence, taking these into 13 we have the conclusion that \(\operatorname{KL}(\cdot,\mathbb{P})\) is \(1\)-convex a.g.g. in \(\mathcal{P}_{2}^{r}\).

Proof of Lemma 3↩︎

By Assumption 2 we know that \(\mathbb{E}_{\xi \sim \mathbb{Q}}\left [\ell(f_{\phi}, \xi)\right]\) is \(\rho\)-weakly concave a.g.g. along all generalized geodesics with respect to \(\mathbb{Q}\) from [28]. This together with the fact that \(\lambda\mathcal{D}(\cdot,\mathbb{P})\) is \(\lambda\)-strongly convex a.g.g. centered at \(\nu\) with the regularization parameter \(\lambda> \rho\) implies that the function \(\mathcal{H}(\phi,\cdot)=\mathbb{E}_{\xi \sim (\cdot)}\left[\ell(f_{\phi} , \xi)\right]-\lambda\mathcal{D}(\cdot,\mathbb{P})\) is \((\lambda-\rho)\)-strongly concave a.g.g. centered at \(\nu\) in \(\mathcal{P}_{2}\) by definition. The proof is complete.

Proof of Proposition 9↩︎

The OT map from \(\nu \in \mathcal{P}_{2}^{r}\) to \(\mathbb{Q}_{i}\) is uniquely defined \(\nu\)-a.e. by Brenier’s theorem. From the calculus rule of the Wasserstein subdifferential with the assumption on \(\xi_{i+1}\) by Assumption 3, we have for each \(i\) that there exists \(\eta_{i+1}\in \partial_{\mathcal{W}_2}(-\mathcal{H}(\phi,\cdot))(\mathbb{Q}_{i+1} )\) such that \[\label{eq:sub95bridge}\gamma (\xi_{i+1}- \eta_{i+1}\circ T_{\nu}^{i+1})+T^{i}_{i+1}\circ T_{\nu}^{i+1}-T_{\nu}^i =T_{\nu}^{i+1}-T_{\nu}^{i}, \quad \nu \text{-a.e. }\tag{14}\] Let \(T_{\nu}^{*}\) denote the unique OT map from \(\nu\) to \(\mathbb{Q}^{*}\). Then \[\begin{align} \|T_{\nu}^{i+1}-T_{\nu}^{*}\|_{\nu}^{2} & =\|T_{\nu}^{i+1}-T_{\nu}^{i}+T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2} \\ & =\|T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2}+2\langle T_{\nu}^{i}-T_{\nu}^{*}, T_{\nu}^{i+1}-T_{\nu}^{i}\rangle_{\nu}+\|T_{\nu}^{i+1}-T_{\nu}^{i}\|_{\nu}^{2} \\ & \leq\|T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2}+2\langle T_{\nu}^{i+1}-T_{\nu}^{*}, T_{\nu}^{i+1}-T_{\nu}^{i}\rangle_{\nu}, \end{align}\] where the last inequality is from \(\|T_{\nu}^{i+1}-T_{\nu}^{i}\|_{\nu}^{2} \geq 0\). This together with 14 implies that \[\label{eq:sub95key95ineq} \begin{align} & \|T_{\nu}^{i+1}-T_{\nu}^{*}\|_{\nu}^{2}\leq \|T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2}+2\gamma\left\langle T_{\nu}^{i+1}-T_{\nu}^{*}, \xi_{i+1}- \eta_{i+1}\circ T_{\nu}^{i+1}+\frac{1}{\gamma}(T^{i}_{i+1}\circ T_{\nu}^{i+1}-T_{\nu}^i)\right\rangle_{\nu}. \end{align}\tag{15}\] On the other hand, since from Assumption 2 we know that \(\lambda\mathcal{D}(\cdot,\mathbb{P})\) is \(\lambda\)-strongly convex a.g.g. centered at \(\nu\) with the regularization parameter \(\lambda> \rho\), it follows from the proof of [29] (with \(\mathcal{W}_{\nu}\) replacing \(\mathcal{W}_{2}\)) that \[\label{eq:sub95sc95key} \begin{align} & -\mathcal{H}(\phi,\mathbb{Q}^{*})+\mathcal{H}(\phi,\mathbb{Q}_{i+1}) \geq\langle T_{\nu}^{*}-T_{\nu}^{i+1}, \eta_{i+1}\circ T_{\nu}^{i+1}\rangle_{\nu}+\frac{\lambda-\rho}{2}\cdot\mathcal{W}_{\nu}^{2}(\mathbb{Q}_{i+1}, \mathbb{Q}^{*}). \end{align}\tag{16}\] Meanwhile, by Cauchy–Schwarz inequality, \[\label{eq:sub95innerproduct} \begin{align} &\left|\left\langle T_{\nu}^{*}-T_{\nu}^{i+1}, \xi_{i+1}+\frac{1}{\gamma}(T^{i}_{i+1}\circ T_{\nu}^{i+1}-T_{\nu}^i)\right\rangle_{\nu}\right|\\ \leq\;&\|T_{\nu}^{*}-T_{\nu}^{i+1}\|_{\nu}\cdot\left\|\xi_{i+1}+\frac{1}{\gamma}(T^{i}_{i+1}\circ T_{\nu}^{i+1}-T_{\nu}^i)\right\|_{\nu} \\ \leq\;& \|T_{\nu}^{*}-T_{\nu}^{i+1}\|_{\nu}\cdot \left(\epsilon'+\frac{1}{\gamma}\|T^{i}_{i+1}\circ T_{\nu}^{i+1}-T_{\nu}^i\|_{\nu}\right)\\ \leq\;&\frac{\lambda-\rho}{4}\|T_{\nu}^{*}-T_{\nu}^{i+1}\|_{\nu}^{2}+\frac{2}{\lambda-\rho}\left(\epsilon'^2+\frac{1}{\gamma^2}\|T^{i}_{i+1}\circ T_{\nu}^{i+1}-T_{\nu}^i\|_{\nu}^2\right), \end{align}\tag{17}\] where the second inequality is by \(\|\xi_{i+1}\|_{\nu}\leq \epsilon'\) from Assumption 3. On the other hand, we know that \[\label{eq:threepoint} \begin{align} \frac{1}{\gamma^2}\|T^{i}_{i+1}\circ T_{\nu}^{i+1}-T^{i}_{\nu}\|_\nu^2 &= \frac{1}{\gamma^2}\|(T^{i}_{i+1}-I)\circ T_{\nu}^{i+1}+T_{\nu}^{i+1}-T_{\nu}^*+T_{\nu}^*-T^{i}_{\nu}\|_\nu^2\\ &\leq \frac{3}{\gamma^2}(\mathcal{W}_2^2(\mathbb{Q}_i,\mathbb{Q}_{i+1})+\|T_{\nu}^{i+1}-T_{\nu}^*\|_{\nu}^2+\|T_{\nu}^{i}-T_{\nu}^*\|_{\nu}^2)\\ &\leq \frac{9}{\gamma^2}(\|T_{\nu}^{i+1}-T_{\nu}^*\|_{\nu}^2+\|T_{\nu}^{i}-T_{\nu}^*\|_{\nu}^2), \end{align}\tag{18}\] where the last inequality is from \[\mathcal{W}_2^2(\mathbb{Q}_i,\mathbb{Q}_{i+1})\leq 2\mathcal{W}_2^2(\mathbb{Q}_{i+1},\mathbb{Q}^*)+2\mathcal{W}_2^2(\mathbb{Q}_{i},\mathbb{Q}^*)\leq 2(\|T_{\nu}^{i+1}-T_{\nu}^*\|_{\nu}^2+\|T_{\nu}^{i}-T_{\nu}^*\|_{\nu}^2).\] Inserting 16 , 17 and 18 into 15 gives \[\label{eq:ppm95EVI} \begin{align} & \left(1-\frac{\gamma (\lambda-\rho)}{2}-\frac{36}{\gamma (\lambda-\rho)}\right)\|T_{\nu}^{i+1}-T_{\nu}^{*}\|_{\nu}^{2} \\ \leq\; &\left(1+\frac{36}{\gamma (\lambda-\rho)}\right) \|T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2}+2 \gamma\left(-\mathcal{H}(\phi,\mathbb{Q}^{*})+\mathcal{H}(\phi,\mathbb{Q}_{i+1})-\frac{\lambda-\rho}{2}\cdot W_{\nu}^{2}(\mathbb{Q}_{i+1}, \mathbb{Q}^{*})\right)+\frac{4 \gamma\epsilon'^2 }{\lambda-\rho}, \end{align}\tag{19}\] which implies that \[\begin{align} \left(1+\frac{\gamma (\lambda-\rho)}{2}-\frac{36}{\gamma (\lambda-\rho)}\right)\|T_{\nu}^{i+1}-T_{\nu}^{*}\|_{\nu}^{2}\leq \left(1+\frac{36}{\gamma (\lambda-\rho)}\right)\|T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2}+\frac{4 \gamma\epsilon'^2}{\lambda-\rho}. \end{align}\] Since \(\gamma(\lambda-\rho)> 36\), we know that \[\tau:=\frac{\gamma (\lambda-\rho)+36}{\gamma^2 (\lambda-\rho)^2/2+\gamma (\lambda-\rho)-36}=\frac{1+\frac{36}{\gamma (\lambda-\rho)}}{1+\frac{\gamma (\lambda-\rho)}{2}-\frac{36}{\gamma (\lambda-\rho)}}<1\] By recursively applying the procedure from \(i = 0\) to \(i = I-1\), we obtain that \[\label{eq:jko95linear} \begin{align} \|T_{\nu}^{I}-T_{\nu}^{*}\|_{\nu}^{2} & \leq \tau^{I}\cdot\|T_{\nu}^{0}-T_{\nu}^{*}\|_{\nu}^{2}+\frac{4 \gamma\epsilon'^2}{\lambda-\rho}\cdot\frac{\left(1+\frac{\gamma (\lambda-\rho)}{2}-\frac{36}{\gamma (\lambda-\rho)}\right)^{-1}}{1-\tau} \\ & = \tau^{I}\cdot\|T_{\nu}^{0}-T_{\nu}^{*}\|_{\nu}^{2}+\frac{8\epsilon'^2}{(\lambda-\rho)^2-\frac{144}{\gamma^2 }}\\ & \leq \tau^{I}\cdot\|T_{\nu}^{0}-T_{\nu}^{*}\|_{\nu}^{2}+\frac{9\epsilon'^2}{(\lambda-\rho)^2}, \end{align}\tag{20}\] which is the desired result ?? . From equation ?? , it follows that to ensure \[\|T_{\nu}^{i}- T_{\nu}^{*}\|_{\nu}^{2}\leq \frac{10\epsilon'^{2}}{(\lambda - \rho)^{2}},\] we require \[i \geq \frac{2}{\log(\tau^{-1})}\left( \log \mathcal{W} _{\nu}(\mathbb{Q}_{0}, \mathbb{Q}^{*}) + \log \left( \frac{\lambda - \rho}{\epsilon'} \right) \right).\] Additionally, from 19 and the condition \(\|T_{\nu}^{i}- T _{\nu}^{*}\|_{\nu}^{2}\leq \frac{10\epsilon'^{2}}{(\lambda - \rho)^{2}}\) and \(\gamma(\lambda-\rho) > 36\), we have \[\begin{align} 2 \gamma (\mathcal{H}(\phi,\mathbb{Q}^{*})-\mathcal{H}(\phi,\mathbb{Q}_{i+1})) &\leq \left(1+\frac{36}{\gamma (\lambda-\rho)}\right)\|T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2}+\frac{4 \gamma\epsilon'^2}{\lambda-\rho}\\ &<2\|T_{\nu}^{i}-T_{\nu}^{*}\|_{\nu}^{2}+\frac{4 \gamma\epsilon'^2}{\lambda-\rho}\\ &\leq \frac{20\epsilon'^{2}}{(\lambda - \rho)^{2}}+\frac{4 \gamma\epsilon'^2}{\lambda-\rho}\leq \frac{5\gamma\epsilon'^2}{\lambda-\rho}. \end{align}\] The proof is complete.

Proof of Theorem 11↩︎

With the help of Theorem 2 and \(\eta=\mathcal{O}(\varepsilon^{2})\), it suffices to obtain \[\label{eq:complexityrequire}\|\nabla \mathcal{V}_{1/2\rho}(\phi_{k})\| ^{2}\lesssim \frac{2(\mathcal{V}_{1/2\rho}(\phi_{0})-\min_{\phi}\mathcal{V}(\phi)+\rho L^{2}K\varepsilon^4)}{K\varepsilon^2}+4 \rho \epsilon\leq\varepsilon^{2}.\tag{21}\] We divide this requirement into two components. For the first term \[\frac{2(\mathcal{V}_{1/2\rho}(\phi_{0})-\min_{\phi}\mathcal{V}(\phi)+\rho L^{2}K\varepsilon^4)}{K\varepsilon^2}\leq\frac{\varepsilon^{2}}{2},\] it leads to \(K= \Omega(\varepsilon^{-4})\), which provides the lower bound for the number of outer iterations required by Algorithm 1.

On the other hand, for the second term in 21 , which concerns the inner JKO steps by Algorithm 2 that approximately solve the subproblem, Proposition 9 indicates that if \(I-1 \geq \frac{2}{\log(\tau^{-1})}\left(\log \mathcal{W} _{\nu}(\mathbb{Q}_{0},\mathbb{Q}^{*})+\log ((\lambda-\rho) / \epsilon')\right)\), then \[\mathcal{H}(\phi, \mathbb{Q}^{*})-\mathcal{H}(\phi, \mathbb{Q}_{I}) \leq \frac{5\epsilon'^2}{2(\lambda - \rho)}.\] Since the accuracy \(\epsilon\) of the subproblem in Algorithm 1 is also defined as \(\mathcal{H}(\phi, \mathbb{Q}^{*})-\mathcal{H}(\phi, \mathbb{Q}_{I}) \leq \epsilon\), to ensure that \(4\rho\epsilon \leq \frac{\varepsilon^{2}}{2}\), it suffices to require \[4\rho\cdot \frac{5\epsilon'^2}{2(\lambda - \rho)}\leq\frac{\varepsilon^{2}}{2}.\] This leads us to choose \(\epsilon' = \mathcal{O}(\varepsilon)\). By combining these two components, we obtain the desired oracle complexity results. Specifically, the number of subgradient evaluations is equal to the number of outer iterations, and the inexact JKO steps require a product of the number of outer iterations \(K\) and the number of inner iterations \(I\).

11 Proofs for Section 4↩︎

Proof of Proposition 12↩︎

Since \(\mathcal{V}(\phi)=\max_{\mathbb{Q}\in \mathcal{P}_2}\mathcal{H}(\phi,\mathbb{Q})=\max_{\mathbb{Q}\in \mathcal{P}_2(\mathcal{X})}\mathcal{H}(\phi,\mathbb{Q})\) for any \(\phi\in\Phi\) from Assumption 4, we focus on the function \(\mathcal{H}(\phi,\cdot)\) with the space \(\mathcal{P}_2(\mathcal{X})\). By Assumption 2 and Lemma 3 we know that \(\mathcal{H}(\phi, \cdot)\) is strongly concave a.g.g. centered at \(\nu \in \mathcal{P}_{2}^{r}\) for each \(\phi \in \Phi\). Consequently, by definition, the set \(\mathop{\mathrm{argmax}}_{\mathbb{Q}\in \mathcal{P}_2}\mathcal{H}(\phi, \mathbb{Q})\) is a singleton. On the other hand, from Assumption 4 we know that \(\mathcal{H}(\cdot,\mathbb{Q})=\mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f_{(\cdot)}, \xi)]-\lambda \mathcal{D}(\mathbb{Q},\mathbb{P})\) is differentiable for every \(\mathbb{Q}\in\mathcal{P}_{2}\), and \(\nabla_{\phi}\mathcal{H} (\phi,\mathbb{Q})=\mathbb{E}_{\xi \sim \mathbb{Q}}[\nabla_{\phi}\ell(f_{\phi}, \xi)]\) is Lipschitz continuous with respect to \(\phi\times \mathbb{Q}\in\Phi\times \mathcal{P} _{2}\). We then focus on applying Danskin’s theorem [34], which extends Danskin’s theorem to general compact topological spaces. In order to derive the desired results, it suffices to verify that the space \(\mathcal{P} _{2}(\mathcal{X})\) is a compact topological space.

To proceed, we prove that the space of probability measures \(\mathcal{P}_{2} (\mathcal{X})\) equipped with the Wasserstein-2 metric \(\mathcal{W}_{2}\) is compact. Let \(\{\mu_{i}\}_{i \geq 1}\) be an arbitrary sequence in \(\mathcal{P}_{2}(\mathcal{X})\). Since \(\mathcal{X}\) is compact, every sequence of probability measures on \(\mathcal{X}\) is tight. By Prokhorov’s theorem [35], the sequence \(\{\mu_{i}\}\) has a convergent subsequence that converges weakly to some probability measure \(\mu \in \mathcal{P}_{2}(\mathcal{X})\). Then there exists a probability space \((\Omega, \mathcal{F}, \nu)\) and random variables \(X, X_{1}, X_{2}, \ldots\) with distribution \(\mu,\mu_{1},\mu_{2},\ldots\) such that \(\lim_{i \rightarrow \infty}\|X_{i}- X\| = 0\) almost surely under \(\nu\) by Skorokhod’s representation theorem. Since \(\mathcal{X}\) is compact, there exists a constant \(c > 0\) such that \(\|X_{i}- X\| \leq c\) for all \(i \geq 1\). Hence, we know from the dominated convergence theorem that \(\lim_{i\rightarrow\infty}\mathbb{E} [\|X_{i}-X\|^{2}]= 0.\) Consequently, for the couplings \(\pi_{i}\in \Pi(\mu_{i}, \mu)\) one has that \[\lim_{i\rightarrow\infty}\int_{\mathcal{X} \times \mathcal{X}}\|x_i-x\| ^{2}\mathrm{~d}\pi_i(x_i, x) = 0.\] Taking the infimum of the left-hand side over \(\Pi(\mu_{i}, \mu)\) and recalling the definition of the Wasserstein distance, we deduce that \(\lim_{i\rightarrow\infty}\mathcal{W}_{2}(\mu_{i}, \mu) = 0\), which means that \(\mu\) is also a limit point of \(\{\mu_{i}\}_{i \geq 1}\) for the Wasserstein-\(2\) metric. Then \((\mathcal{P}_{2}(\mathcal{X}), \mathcal{W} _{2})\) is compact, which completes the proof.

Proof of Lemma 4↩︎

From Lemma 3 we know that the function \(\mathcal{H}(\phi,\cdot )\) is \((\lambda-L)\)-strongly concave a.g.g. centered at \(\nu \in \mathcal{P}_{2}^{r}\) in \(\mathcal{P}_{2}\). Then \[\begin{align} & -\mathcal{H}(\phi_{1},\mathbb{Q}^{*}(\phi_{2}))+\mathcal{H}(\phi_{1},\mathbb{Q}^{*}(\phi_{1}))\geq \frac{\lambda-L}{2}\cdot\mathcal{W}_{\nu}^{2}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2})),\tag{22} \\ & -\mathcal{H}(\phi_{2},\mathbb{Q}^{*}(\phi_{1}))+\mathcal{H}(\phi_{2},\mathbb{Q}^{*}(\phi_{2}))\geq \frac{\lambda-L}{2}\cdot\mathcal{W}_{\nu}^{2}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2})).\tag{23} \end{align}\] Moreover, by the \(L\)-Lipschitz continuity of \(\nabla_{\phi}\mathcal{H}(\cdot, \mathbb{Q})\) for any \(\mathbb{Q}\in \mathcal{P}_{2}\) we have \[\begin{align} \label{q-yconcave}-\mathcal{H}(\phi_{2},\mathbb{Q}^{*}(\phi_{1}))+\mathcal{H}(\phi_{1},\mathbb{Q}^{*}(\phi_{1}))\leq \langle \nabla_{\phi}\mathcal{H}(\phi_{1},\mathbb{Q}^{*}(\phi_{1})), \phi_{1}-\phi_{2}\rangle+\frac{L}{2}\|\phi_{2}-\phi_{1}\|^{2} \end{align}\tag{24}\] and \[\begin{align} \label{q-ylip}-\mathcal{H}(\phi_{1},\mathbb{Q}^{*}(\phi_{2}))+\mathcal{H}(\phi_{2},\mathbb{Q}^{*}(\phi_{2}))\leq \langle \nabla_{\phi}\mathcal{H}(\phi_{2},\mathbb{Q}^{*}(\phi_{2})), \phi_{2}-\phi_{1}\rangle+\frac{L}{2}\|\phi_{2}-\phi_{1}\|^{2}. \end{align}\tag{25}\] Incorporating 22 –25 , we obtain \[\label{q-perturbe-key} \begin{align} & (\lambda-L)\cdot\mathcal{W}_{\nu}^{2}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2})) \leq \langle\nabla_{\phi}\mathcal{H}(\phi_{1},\mathbb{Q}^{*}(\phi_{1}))-\nabla_{\phi}\mathcal{H}(\phi_{2},\mathbb{Q}^{*}(\phi_{2})), \phi_{1}-\phi_{2}\rangle +L\|\phi_{2}-\phi_{1}\|^{2}. \end{align}\tag{26}\] Recall that \(\mathcal{H}(\phi,\mathbb{Q})=\mathbb{E}_{\xi \sim \mathbb{Q}}[\ell(f_{\phi}, \xi )]-\lambda\cdot \mathcal{D}(\mathbb{Q},\mathbb{P})\), it follows for any transport map \(T\) satisfying \(T_{\#}\nu=\mathbb{Q}\) that \[\nabla_{\phi}\mathcal{H}(\phi,\mathbb{Q})=\mathbb{E}_{\xi \sim \mathbb{Q}}[\nabla _{\phi}\ell(f_{\phi}, \xi)]=\mathbb{E}_{\xi \sim \nu}[\nabla_{\phi} \ell(f_{\phi}, T(\xi))].\] Since \(\nabla_{\phi}\ell(f_{\phi}, \cdot)\) is \(L\)-Lipschitz continuous, we have for the OT maps \(T_{1}\), \(T_{2}\) satisfy \((T_{1})_{\#}\nu=\mathbb{Q}^{*}( \phi_{1})\) and \((T_{2})_{\#}\nu=\mathbb{Q}^{*}(\phi_{2})\) that \[\|\nabla_{\phi}\ell(f_{\phi}, T_{1}(\xi))-\nabla_{\phi}\ell(f_{\phi}, T_{2}(\xi))\|\leq L\|T_{1}(\xi)-T_{2}(\xi)\|\quad \text{for } \nu\text{-a.e.}\;\xi.\] Integrating with respect to \(\nu\), we derive \[\label{eq:key95distri95perturb} \begin{align} & \|\nabla_{\phi}\mathcal{H}(\phi,\mathbb{Q}^{*}(\phi_{1}))-\nabla_{\phi}\mathcal{H}(\phi,\mathbb{Q}^{*}(\phi_{2}))\| \\ =\; & \|\mathbb{E}_{\xi\sim\nu}[\nabla_{\phi}\ell(f_{\phi}, T_{1}(\xi))-\nabla_{\phi}\ell(f_{\phi}, T_{2}(\xi))]\| \\ \leq\; & \mathbb{E}_{\xi\sim\nu}[\|\nabla_{\phi}\ell(f_{\phi}, T_{1}(\xi))-\nabla_{\phi}\ell(f_{\phi}, T_{2}(\xi))\|] \\ \leq\; & L\cdot\mathbb{E}_{\xi\sim\nu}[\|T_{1}(\xi)-T_{2}(\xi)\|]\leq L\cdot\mathcal{W}_{\nu}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2})), \end{align}\tag{27}\] where the last inequality is from Jensen’s inequality. Then it follows that \[\begin{align} \|\nabla_{\phi}\mathcal{H}(\phi_{1},\mathbb{Q}^{*}(\phi_{1}))-\nabla_{\phi}\mathcal{H}(\phi_{2},\mathbb{Q}^{*}(\phi_{2}))\|\leq L(\|\phi_{1}-\phi_{2}\|+\mathcal{W}_{\nu}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2}))). \end{align}\] This together with 26 implies that \[\begin{align} (\lambda-L)\cdot\mathcal{W}_{\nu}^{2}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2})) \leq L\cdot\mathcal{W}_{\nu}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2}))\cdot\|\phi_{1}-\phi_{2}\|+2L\|\phi_{1}-\phi_{2}\|^{2}. \end{align}\] Let \(\beta:=\frac{\mathcal{W}_{\nu}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2}))}{\|\phi_{1}-\phi_{2}\|}\). Then, the above inequality gives \[\begin{align} \beta^{2}\leq \frac{L}{\lambda-L}\beta+\frac{2L}{\lambda-L}\leq \frac{1}{2}\beta^{2}+\frac{L^{2}}{2(\lambda-L)^{2}}+\frac{2L}{\lambda-L} & \leq \frac{1}{2}\beta^{2}+\frac{L^{2}+4L(\lambda-L)}{2(\lambda-L)^{2}} \\ & \leq \frac{1}{2}\beta^{2}+\frac{(L+2(\lambda-L))^{2}}{2(\lambda-L)^{2}}, \end{align}\] where the second inequality holds because \(ab\leq \frac{1}{2}(a^{2}+b^{2} )\) for any \(a,b\in\mathbb{R}\). Thus, we get \[\mathcal{W}_{2}(\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2}))\leq \mathcal{W}_{\nu} (\mathbb{Q}^{*}(\phi_{1}),\mathbb{Q}^{*}(\phi_{2})) \leq \frac{2\lambda-L}{\lambda-L} \cdot \|\phi_{1}-\phi_{2}\|.\] Furthermore, with the help of the Danskin’s theorem, i.e., Proposition 12, we know that \[\|\nabla \mathcal{V}(\phi_{1})-\nabla \mathcal{V}(\phi_{2})\|=\|\nabla _{\phi}\mathcal{H}(\phi_{1}, \mathbb{Q}^{*}(\phi_{1}))-\nabla_{\phi}\mathcal{H} (\phi_{2}, \mathbb{Q}^{*}(\phi_{2}))\| \leq L\left(1+\frac{2\lambda-L}{\lambda-L} \right)\|\phi_{1}-\phi_{2}\|.\] The proof is complete.

Proof of Theorem 13↩︎

First we observe from Lemma 4 that \[\label{eq:smooth95des} \begin{align} \mathcal{V}(\phi_{k+1}) & \leq \mathcal{V}(\phi_{k})+\langle\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q}^{*}_{k}), \phi_{k+1}-\phi_{k}\rangle+\kappa_2L\|\phi_{k+1}-\phi_{k}\|^{2}, \end{align}\tag{28}\] where \(\mathbb{Q}_{k}^{*}:=\mathop{\mathrm{argmax}}_{\mathbb{Q}\in\mathcal{P}_2}\mathcal{H}(\phi_{k},\mathbb{Q})\). The update \(\phi_{k+1}=\mathop{\mathrm{proj}}_{\Phi}(\phi_{k}-\eta_k\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q} _{k}))\) together with its optimality by projection implies that \[\langle\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q}_{k}), \phi_{k+1}-\phi_{k} \rangle+\frac{1}{2\eta_k}\|\phi_{k+1}-\phi_{k}\|^{2}\leq0.\] Then combining this with 28 we have that \[\begin{align} \label{eq:smooth95key} & \mathcal{V}(\phi_{k+1}) \notag \\ \leq\; & \mathcal{V}(\phi_{k})+\langle\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q}^{*}_{k})-\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q}_{k})+\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q}_{k}), \phi_{k+1}-\phi_{k}\rangle +\kappa_2L\|\phi_{k+1}-\phi_{k}\|^{2}\notag \\ \leq\; & \mathcal{V}(\phi_{k})+\langle\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q}^{*}_{k})-\nabla_{\phi}\mathcal{H}(\phi_{k},\mathbb{Q}_{k}), \phi_{k+1}-\phi_{k}\rangle+\left(\frac{(3\lambda-2L)L}{2(\lambda-L)}-\frac{1}{2\eta_k}\right)\|\phi_{k+1}-\phi_{k}\|^{2}\notag \\ \leq\; & \mathcal{V}(\phi_{k})-\left(\frac{1}{2\eta_k}-\frac{(3\lambda-2L)L}{2(\lambda-L)}-\frac{L}{2}\right)\|\phi_{k+1}-\phi_{k}\|^{2}+\frac{L}{2}\mathcal{W}_{\nu}^{2}(\mathbb{Q}_{k}^{*},\mathbb{Q}_{k}), \end{align}\tag{29}\] where the last inequality is from similar argument for deriving gradient Lipschitz property 27 .

On the other hand, by the definition of \(\mathbb{Q}^{*}_{k}\) and the subproblem of Algorithm 1, we know that \[\mathcal{H}(\phi_{k},\mathbb{Q}_{k})\geq \mathcal{H}(\phi_{k},\mathbb{Q}_{k}^{*}) -\epsilon .\] This together with the a.g.g. strong concavity with respect to \(\mathbb{Q}\) centered at \(\nu \in \mathcal{P}_{2}^{r}\) implies that \[\epsilon\geq \mathcal{H}(\phi_{k},\mathbb{Q}^{*}_{k}) -\mathcal{H}(\phi_{k},\mathbb{Q} _{k})\ge\frac{\lambda-L}{2}\cdot\mathcal{W}_{\nu}^{2}(\mathbb{Q}_{k},\mathbb{Q}^{*}_{k} ),\] and consequently we know from 29 that \[\mathcal{V}(\phi_{k+1}) \leq \mathcal{V}(\phi_{k})-\left(\frac{1}{2\eta_k} -\frac{(3\lambda-2L)L}{2(\lambda-L)}-\frac{L}{2}\right)\|\phi_{k+1}-\phi _{k}\|^{2}+\frac{L\epsilon}{\lambda-L}.\] Taking the sum over \(k\), we obtain \[\mathcal{V}(\phi_{K}) \leq \mathcal{V}(\phi_{0})-\left(\frac{1}{2\eta_k} -\frac{(4\lambda-3L)L}{2(\lambda-L)}\right) \sum_{k=0}^{K-1}\|\phi_{k+1} -\phi_{k}\|^{2}+\frac{L\epsilon}{\lambda-L}\cdot K.\] Rearranging this, then we have \[\frac{1}{K}\sum_{k=0}^{K-1}\|\phi_{k+1}-\phi_{k}\|^{2}\leq \frac{2\eta_k(\lambda-L)\cdot(\mathcal{V}(\phi_{0})-\min_{\phi}\mathcal{V}(\phi))}{(\lambda-L-\eta_k L(4\lambda-3L))K} +\frac{2\eta_k L\epsilon}{\lambda-L-\eta_k L(4\lambda-3L)}.\] In the meantime, by the optimality condition of the projection step \(\phi_{k+1}=\mathop{\mathrm{proj}}_\Phi(\phi_k-\eta_k \nabla_\phi\mathcal{H}(\phi_k,\mathbb{Q}_k))\), we have \[\frac{\phi_k-\phi_{k+1}}{\eta_k}-\nabla_\phi\mathcal{H}(\phi_k,\mathbb{Q}_k) \in \partial\iota_\Phi(\phi_{k+1}).\] Therefore, \[\begin{align} & \mathop{\mathrm{dist}}^{2}\left(0,\nabla \mathcal{V}(\phi_{k+1}) +\partial\iota_{\Phi}(\phi_{k+1})\right) \\ \leq\;& \left\| \nabla\mathcal{V}(\phi_{k+1}) +\frac{\phi_k-\phi_{k+1}}{\eta_k} -\nabla_\phi\mathcal{H}(\phi_k,\mathbb{Q}_k) \right\|^2 \\ \leq\;& 2\left(\frac{1}{\eta_k}+2\kappa_2L\right)^2 \|\phi_{k+1}-\phi_k\|^2 + 2\|\nabla_\phi\mathcal{H}(\phi_k,\mathbb{Q}_k)-\nabla_\phi\mathcal{H}(\phi_k,\mathbb{Q}_k^*)\|^2 \\ \leq\;& 2\left(\frac{1}{\eta_k}+2\kappa_2L\right)^2 \|\phi_{k+1}-\phi_k\|^2 + \frac{4L^2\epsilon}{\lambda-L}. \end{align}\] Then letting \(\eta_k=\frac{\lambda-L}{2L(4\lambda-3L)}\) we have \[\begin{align} & \min_{k\in[K-1]} \left\{ \mathop{\mathrm{dist}}^{2}\left( 0,\nabla \mathcal{V}(\phi_{k+1}) +\partial\iota_{\Phi}(\phi_{k+1}) \right) \right\} \\ \leq\;& 4\left(\frac{1}{\eta_k}+2\kappa_2 L\right)^{2} \frac{ \eta_k(\lambda-L) \left(\mathcal{V}(\phi_0)-\min_{\phi}\mathcal{V}(\phi)\right) }{ \left(\lambda-L-\eta_k L(4\lambda-3L)\right)K } + \frac{ 4\left(\frac{1}{\eta_k}+2\kappa_2 L\right)^2\eta_k L\epsilon }{ \lambda-L-\eta_k L(4\lambda-3L) } + \frac{4L^2\epsilon}{\lambda-L}\\ \leq\;& \frac{ 4L(11\lambda-8L)^2 \left(\mathcal{V}(\phi_0)-\min_{\phi}\mathcal{V}(\phi)\right) }{ (\lambda-L)(4\lambda-3L)K } + \left( \frac{(11\lambda-8L)L}{\lambda-L} \right)^2 \frac{4\epsilon}{4\lambda-3L} + \frac{4L^2\epsilon}{\lambda-L}. \end{align}\] The oracle complexity results follow from a similar analysis as in Theorem 11. The proof is complete.

Auxiliary Lemmas for Convergence Analysis of Algorithm 3↩︎

The convergence proof of Algorithm 3 relies primarily on a basic descent property, along with an exponential bound on the growth of the associated increasing term.

Proof of Theorem 14↩︎

Define \(\delta:=(1+\tau)/ 2 +4 \kappa_1^2 L^2 \eta^2/(1-\tau)\). Using the inequality in Lemma 6 recursively yields that \[\begin{align} \|T_{\nu}^{k}-T_{\nu}^*(\phi_{k})\|_{\nu}^2 &\leq \delta^k C^2+\frac{4 \kappa_1^2 \eta^2}{1-\tau}\left(\sum_{j=0}^{k-1} \delta^{k-1-j}\left\|\nabla \mathcal{V}(\phi_{j})\right\|^2\right)+\frac{4\gamma\epsilon'^2}{\lambda-\rho}\cdot\sum_{j=0}^{k-1} \delta^{k-1-j}, \end{align}\] where \(C:=\|T_{\nu}^{0}-T_{\nu}^*(\phi_{0})\|_{\nu}<\infty\).

Now we substitute this into the descent inequality of Lemma 5. From the JKO contraction bound we have \[\|T_{\nu}^{k+1}-T_{\nu}^*(\phi_{k})\|_{\nu}^2\leq\tau\left(\|T_{\nu}^{k}-T_{\nu}^*(\phi_{k})\|_{\nu}^2+\frac{4 \gamma\epsilon'^2}{\lambda-\rho}\right),\] and therefore Lemma 5 gives \[\begin{align} \mathcal{V}(\phi_{k+1}) \leq\;& \mathcal{V}(\phi_{k})-\frac{7 \eta}{16}\left\|\nabla \mathcal{V}(\phi_{k})\right\|^2+\frac{9 \eta L^2 \tau}{16}\|T_{\nu}^{k}-T_{\nu}^*(\phi_{k})\|_{\nu}^2 +\frac{9 \eta L^2\tau}{4}\cdot\frac{\gamma\epsilon'^2}{\lambda-\rho}. \end{align}\] Substituting the recursive bound for \(\|T_{\nu}^{k}-T_{\nu}^*(\phi_{k})\|_{\nu}^2\) and noting \(\tau\leq 1\), we obtain \[\begin{align} \mathcal{V}(\phi_{k+1}) \leq\;& \mathcal{V}(\phi_{k})-\frac{7 \eta}{16}\left\|\nabla \mathcal{V}(\phi_{k})\right\|^2+\frac{9 \eta L^2 \delta^{k} C^2}{16}+\frac{9 \eta^3 L^2 \kappa_1^2}{4(1-\tau)}\left(\sum_{j=0}^{k-1} \delta^{k-1-j}\left\|\nabla \mathcal{V}(\phi_{j})\right\|^2\right)\\ &+ \frac{9 \eta L^2}{16} \cdot\frac{4\gamma\epsilon'^2}{\lambda-\rho}\cdot\left(1+\sum_{j=0}^{k-1} \delta^{k-1-j}\right). \end{align}\] Summing up over \(k=0,1, \ldots, K-1\) and rearranging the terms yields that \[\begin{align} \mathcal{V}(\phi_{K}) \leq\;& \mathcal{V}(\phi_{0})-\frac{7 \eta}{16}\left(\sum_{k=0}^{K-1}\left\|\nabla \mathcal{V}(\phi_{k})\right\|^2\right)+\frac{9 \eta L^2 C^2}{16}\left(\sum_{k=0}^{K-1} \delta^k\right)+\frac{9 \eta^3 L^2 \kappa_1^2}{4(1-\tau)}\left(\sum_{k=0}^{K-1} \sum_{j=0}^{k-1} \delta^{k-1-j}\left\|\nabla \mathcal{V}(\phi_{j})\right\|^2\right)\\ &+ \frac{9 \eta L^2}{4}\cdot \frac{\gamma\epsilon'^2}{\lambda-\rho}\cdot\sum_{k=0}^{K-1}\left(1+\sum_{j=0}^{k-1} \delta^{k-1-j}\right). \end{align}\] Since \(\eta\leq (1-\tau)/ 16\kappa_1 L\), we have \(\delta \leq \frac{3+\tau}{4}\) and \(\frac{9 \eta^3 L^2 \kappa_1^2}{4(1-\tau)} \leq \frac{9 \eta(1-\tau)}{1024}\). This implies that \(\sum_{k=0}^{K-1} \delta^k \leq \frac{4}{1-\tau}\), \[\sum_{k=0}^{K-1} \sum_{j=0}^{k-1} \delta^{k-1-j}\left\|\nabla \mathcal{V}(\phi_{j})\right\|^2 \leq \frac{4}{1-\tau}\left(\sum_{k=0}^{K-1}\left\|\nabla \mathcal{V}(\phi_{k})\right\|^2\right)\] and \[\sum_{k=0}^{K-1}\left(1+\sum_{j=0}^{k-1} \delta^{k-1-j}\right) \leq K\left(1+\frac{4}{1-\tau}\right)\leq \frac{5 K}{1-\tau}.\] Putting these pieces together yields that \[\mathcal{V}(\phi_{K}) \leq \mathcal{V}(\phi_{0})-\frac{103 \eta}{256}\left(\sum_{k=0}^{K-1}\left\|\nabla \mathcal{V}(\phi_{k})\right\|^2\right)+\frac{9 \eta L^2 C^2}{4(1-\tau)} + \frac{45 \eta L^2}{4(1-\tau)}\cdot \frac{\gamma\epsilon'^2}{\lambda-\rho}\cdot K.\] Rearranging the terms, we have \[\frac{1}{K}\left(\sum_{k=0}^{K-1}\left\|\nabla \mathcal{V}(\phi_{k})\right\|^2\right) \leq \frac{256\left(\mathcal{V}(\phi_{0})-\min_{\phi}\mathcal{V}(\phi)\right)}{103 \eta K}+\frac{576 L^2 C^2}{103(1-\tau)K}+ \frac{2880 L^2}{103(1-\tau)} \cdot\frac{\gamma\epsilon'^2}{\lambda-\rho}.\] Then the oracle complexity results can be directly derived. The proof is complete.

12 Additional Experimental Details and Figures↩︎

Additional Figures↩︎

12.0.0.1 Additional two-dimensional Gaussian mixture results.

Figure 8 reports the empirical classifier loss \(\mathbb{E}_{\widehat{\mathbb{Q}}_k}[\ell]\) and empirical minimax objective \(\widehat{\mathcal{H}}_k\). These curves complement the stationarity curves in Figure 5.

Experimental Details↩︎

12.0.0.2 Two-dimensional Gaussian mixture.

The detailed hyperparameter setup can be found in Table 1. At each outer iteration, we draw a fresh batch from the two distributions described in Section 6. The inner maximization is solved only for the class-0 samples; the class-1 samples are kept fixed and are used only in the classifier update. The particle L-BFGS baseline is used for this inner finite-dimensional problem, while the classifier \(f_\phi\) is trained by Adam for all methods. For the neural transport maps, we use the residual parameterization \(T_\theta(x)=x+R_\theta(x)\) and initialize the map as the identity by setting the final layer of \(R_\theta\) to zero.

12.0.0.3 CIFAR-10 feature-space experiment.

The detailed hyperparameter setup can be found in Table 2. We first encode CIFAR-10 images with a fixed ViT-base encoder, and all classifiers and transport maps are trained in the feature space. A labeled sample is written as \((x,y)\), where \(x\) is the feature vector and \(y\) is the class label. The transport map changes only \(x\) and leaves \(y\) unchanged. We implement the family \(\{T_{\theta,y}\}_{y=0}^{9}\) by one shared neural network: the label \(y\) is mapped to a learned embedding, the embedding is concatenated with \(x\), and the network outputs a residual vector denoted by \(R_{\theta,y}(x)\). Thus, \(T_{\theta,y}(x)=x+R_{\theta,y}(x)\). The classifier \(f_\phi\) takes only the feature vector as input. We initialize \(T_{\theta,y}\) as the identity map by setting the final residual layer to zero. All CIFAR-10 experiments use \(\gamma=\infty\); in Neural+L-BFGS, the output of \(T_{\theta,y}\) is used to initialize the L-BFGS particle solve.

References↩︎

---

1. llzzhu@gatech.edu↩︎

2. yzhu812@gatech.edu↩︎

3. yao.xie@isye.gatech.edu↩︎