Wasserstein DRO Optimization

• Wasserstein-based DRO is a robust optimization framework that minimizes loss functions under uncertainty using transportation metrics.
• It employs a linear coupling of gradient and mirror descent steps, achieving accelerated convergence rates of O(1/T²) for smooth objectives.
• The framework integrates probabilistic ambiguity sets and Bregman divergence to enhance stability and performance against adversarial data shifts.

Wasserstein-based distributionally robust optimization (DRO) is a framework in convex optimization and machine learning for minimizing a loss function under the assumption that the underlying data distribution is only approximately known. This approach typically models adversarial uncertainty by defining an ambiguity set containing probability measures within a certain Wasserstein distance from a reference empirical distribution. The objective is to find solutions that perform well under the worst-case distribution within this set, thereby ensuring robustness against distributional shifts and data perturbations. Wasserstein-based DRO connects foundational ideas in mirror descent, smooth convex minimization, and accelerated gradient methods via notions such as Bregman divergence and dual progress.

1. Problem Setup and Wasserstein Ambiguity

Fundamental to DRO is the formulation

$$\min_{x \in X} \sup_{Q \in \mathcal{U}_\varepsilon(\widehat{P}_n)} \mathbb{E}_Q[f(x; \xi)],$$

where $X \subset \mathbb{R}^n$ is a closed convex set, $f(x; \xi)$ is an objective function parameterized by the random variable $\xi$, $\widehat{P}_n$ is the empirical distribution, and $\mathcal{U}_\varepsilon(\widehat{P}_n)$ is the set of distributions within Wasserstein radius $\varepsilon$ of $\widehat{P}_n$. The Wasserstein metric quantifies the minimal "cost" of transporting one probability distribution to another, based on an underlying ground metric on the sample space.

This setup places the optimization in the domain of large-scale machine learning and robust statistics. The function $f(x) = \mathbb{E}_P[f(x; \xi)]$ inherits convexity and often smoothness from its integrand. When $f$ is further assumed to be $L$-smooth with respect to a norm $\|\cdot\|$, that is,

$$\forall x, y \in X,\,\, f(x) \leq f(y) + \langle \nabla f(y), x-y \rangle + \frac{L}{2}\|x-y\|^2,$$

gradient-based first-order methods become applicable (Allen-Zhu et al., 2014).

2. Algorithmic Framework: Gradient and Mirror Descent Coupling

Key advances in first-order methods reflect the complementary roles of primal (gradient) and dual (mirror) updates (Allen-Zhu et al., 2014). For Wasserstein-based DRO, these principles are particularly relevant when the ambiguity set induces nonsmoothness or complicates direct gradient computation, motivating hybrid schemes.

The "linear coupling" framework maintains three sequences, initialized at $x_0 = y_0 = z_0 \in X$:

• The gradient descent step updates

$x_{k+1} = y_k - \frac{1}{L} \nabla f(y_k),$

promoting primal progress in $X$.

• The mirror ascent step solves

$$z_{k+1} = \arg\min_{z \in X} \{ \langle \nabla f(y_k), z \rangle + \frac{1}{\beta_k} D_h(z; z_k) \},$$

where $h: X \to \mathbb{R}$ is a 1-strongly convex "mirror map" inducing the Bregman divergence

$$D_h(u; v) = h(u) - h(v) - \langle \nabla h(v), u-v \rangle \geq \frac{1}{2}\|u-v\|^2.$$

• The linear coupling combines the iterates:

$$y_{k+1} = \tau_{k+1} z_{k+1} + (1-\tau_{k+1}) x_{k+1}.$$

Parameters are chosen as $\beta_k = \frac{(k+1)(k+2)}{4L}$ and $\tau_{k+1} = \frac{2}{k+2}$. These sequences balance dual and primal progress for optimal convergence (Allen-Zhu et al., 2014).

3. Convergence Properties

The potential function

$$\Phi_k = (k+1)(k+2)[f(y_k) - f(x^*)] + 4L D_h(x^*; z_k)$$

serves as a Lyapunov function certifying convergence at $O(1/k^2)$ for smooth convex objectives. Monotonicity of $\Phi_k$ is established under the coupling framework via three key properties:

1. Smoothness yields a quadratic descent at the gradient step.
2. Mirror descent optimality at $z_{k+1}$ ensures shrinking Bregman gaps.
3. The convexity of $f$ and the structure of $y_{k+1}$ aggregate progress optimally.

Explicitly,

$$(k+2)(k+3)[f(y_{k+1})-f(x^*)] + 4L D_h(x^*; z_{k+1}) \leq (k+1)(k+2)[f(y_k)-f(x^*)] + 4L D_h(x^*; z_k).$$

This structure yields the accelerated rate $O(1/T^2)$ after $T$ steps, improving upon non-coupled gradient or mirror descent methods, both of which admit only $O(1/T)$ convergence (Allen-Zhu et al., 2014).

4. Special Cases and Methodological Connections

Linear coupling reveals relations to classical optimization:

• Gradient Descent (drop mirror step): recovers $O(1/T)$ rate.
• Mirror Descent (drop gradient step): also $O(1/T)$.
• Nesterov Accelerated Gradient: Euclidean mirror map $h(x) = \frac{1}{2}\|x\|^2$ yields Bregman divergence as squared norm and precise recovery of Nesterov sequences and rates.
• Composite and Proximal Methods: By embedding nonsmooth regularizers $\psi$ in the mirror step, accelerated proximal-gradient methods equivalent to FISTA are derived.
• Stochastic and Coordinate Variants: Replacing $\nabla f(y_k)$ by unbiased estimators or partial gradients maintains acceleration properties in expected value (Allen-Zhu et al., 2014).

These results imply that Wasserstein-based DRO can leverage modular and extendable algorithms from the linear coupling paradigm, extending also to non-Euclidean geometries and strongly convex functions—with exponential convergence $O(\exp(-\mu k / L))$ for $\mu$-strongly convex $f$.

5. Role of Bregman Divergence and Mirror Maps

The choice of mirror map $h$ directly influences computational efficiency and the geometry of iterates in DRO settings. The induced Bregman divergence offers a natural framework for non-Euclidean updates, which is especially significant in Wasserstein balls constructed under various ground metrics.

When $h$ is 1-strongly convex with respect to the chosen norm, $D_h(u; v)$ lower bounds $\frac{1}{2}\|u-v\|^2$, enforcing stability of mirror updates. This property is critical in designing algorithms that robustly adapt to uncertainty in the data distribution, as encoded by the Wasserstein radius and the ambiguity set (Allen-Zhu et al., 2014).

6. Extensions and Practical Considerations

The linear coupling framework is agnostic to norm choice and supports direct adaptation to Wasserstein-based DRO where ambiguity is defined with respect to domain-specific metrics. Composite objective structures, stochastic approximations, and coordinate descent variants are accommodated by minor modifications to the step definitions.

A plausible implication is that, in DRO applications, the above methodology allows for accelerated robust optimization even when the adversarial distributional uncertainty induces complex nonsmoothness or constraints. The modularity further supports application to empirical risk minimization, robust machine learning, and risk-sensitive control, where ambiguity sets in Wasserstein space express natural robustness constraints.

7. Summary Table: Methodological Connections

Method	Mirror Map $h$	Convergence Rate
Gradient Descent	Not used	$O(1/T)$
Mirror Descent	General $h$, 1-strongly convex	$O(1/T)$
Nesterov Acceleration	$h(x) = \frac{1}{2}\|x\|^2$	$O(1/T^2)$
Linearly Coupled DRO	Any 1-strongly convex $h$ w.r.t. chosen norm	$O(1/T^2)$

In Wasserstein-based DRO, linear coupling of gradient and mirror steps unifies and generalizes classical first-order optimization methodologies, delivering optimal rates while ensuring robustness to distributional shift via principled ambiguity set construction (Allen-Zhu et al., 2014).