source_v1/main.tex

\documentclass{article}


\usepackage{arxiv}

\usepackage[english]{babel}
\usepackage[utf8]{inputenc} %
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{float}
\usepackage{booktabs} %
\usepackage{multirow} %
\usepackage{graphicx}
\usepackage{subcaption}
\usepackage[colorinlistoftodos]{todonotes}
\usepackage[authoryear]{natbib}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{pgfplots}
\pgfplotsset{compat=1.12}
\usepackage{hyperref} %

             
\newcommand{\Qrob}{\mathbb{Q}}
\newcommand{\calP}{\mathcal{P}}
\newcommand{\calM}{\mathcal{M}}
\newcommand{\calZ}{\mathcal{Z}}
\newcommand{\calB}{\mathcal{B}}
\newcommand{\W}{\mathbb{W}}
\newcommand{\X}{\mathcal{X}}
\newcommand{\calK}{\mathcal{K}}
\newcommand{\esssup}{\operatornamewithlimits{ess\,sup}}
\newcommand{\norm}[1]{\left\lVert#1\right\rVert}

\newtheorem{theorem}{Theorem}
\newtheorem{lemma}{Lemma}
\newtheorem{Proposition}{Proposition}
\newtheorem{corollary}{Corollary}
\newtheorem{definition}{Definition}
\newtheorem{assumption}{Assumption}
\newtheorem{condition}{Condition}
\newtheorem{remark}{Remark}
\newenvironment{proof}{\paragraph{Proof:}}{\hfill$\square$}




\usepackage{preamble-JZ}

\title{Gradient Flow Sampler-based\\
Distributionally Robust Optimization}

\author{
Zusen Xu\\
Weierstrass Institute for \\
Applied Analysis and Stochastics\\
Berlin, Germany\\
\texttt{zusen.xu@wias-berlin.de} \\
    \And
Jia-Jie Zhu\\
Weierstrass Institute for \\
Applied Analysis and Stochastics\\
Berlin, Germany\\
\texttt{zplusj@gmail.com}\\
}

\begin{document}

\maketitle



\begin{abstract}
We propose a mathematically principled PDE gradient flow framework for distributionally robust optimization (DRO). Exploiting the recent advances in the intersection of Markov Chain Monte Carlo sampling and gradient flow theory, we show that our theoretical framework can be implemented as practical algorithms for sampling from worst-case distributions and, consequently, DRO. While numerous previous works have proposed various reformulation techniques and iterative algorithms, we contribute a sound gradient flow view of the distributional optimization that can be used to construct new algorithms.
As an example of applications, we solve a class of Wasserstein and Sinkhorn DRO problems using the recently-discovered Wasserstein Fisher-Rao and Stein variational gradient flows. Notably, we also show some simple reductions of our framework recover exactly previously proposed popular DRO methods, and provide new insights into their theoretical limit and optimization dynamics. Numerical studies based on stochastic gradient descent provide empirical backing for our theoretical findings.

\end{abstract}



\section{Introduction}
Distributionally robust optimization (DRO)~\citep{delageDistributionallyRobustOptimization2010,kuhnDistributionallyRobustOptimization2024}
is a framework that aims to enhance the robustness of the solution to optimization problems.
Since the original Wasserstein distributionally robust optimization (DRO) works by \citep{mohajerin2018data,zhaoDatadrivenRiskaverseStochastic2018,gaoDistributionallyRobustStochastic2016}, many subsequent works have presented variations of problems with various numerical solutions.
In particular, the entropy-regularized Wasserstein DRO problem can be formulated as the penalized optimization problem:
\begin{align}
    \min_{\theta\in\Theta} \max_{\rho\in\calP}
    \int \ell(\theta, z)\dd\rho{(z)} - \tfrac1{2\tau} W^2_{\epsilon}(\rho,\widehat{\rho_N})
    \label{eq:main-ent-dro}
\end{align}
where
$W_\epsilon$ is the
entropy-regularized optimal transport (OT) divergence to be defined later,
$\epsilon > 0$ is the entropy regularization parameter, and $\tau > 0$ is a temperature parameter (Lagrange multiplier); see later discussions.
This DRO problem was considered by \citep{wang2021sinkhorn}.
If we set $\epsilon = 0 $, we recover the penalized version of the Wasserstein DRO as considered by \citep{sinha2020certifyingdistributionalrobustnessprincipled}.
For the mathematical analysis, we will temporarily focus on the interesting choice of $W_\epsilon$, while other choices such as the KL divergences can also be adapted to our framework.
There is a thread of works that considered the (D)RO problems via the gradient descent-ascent dynamics, e.g.,
\citep{wang2022exponentially,yuFastDistributionallyRobust2022,trillosAdversarialRobustnessUse2023,congerStrategicDistributionShift2023}.
For illustration, let $V:=-\ell(\theta, x)$ and $\rho_0=\widehat{\rho_N}$. 
The inner maximization problem can be written in the form of a proximal problem
\begin{align}
    \min_{\rho \in \calP} \;
    \int V \, \dd \rho(x) + \tfrac{1}{2\tau} W^2_\epsilon(\rho, {\rho_0})
    \label{eq:proximal-EOT}
\end{align}
If we set $\epsilon = 0 $, we recover the
Wasserstein proximal problem, 
commonly known as the Jordan-Kinderlehrer-Otto (JKO) scheme in the PDE literature,
which
can be viewed as one-step discretization of the continuous time PDE:
\begin{align}
    \partial_t \rho = -\nabla \cdot \left( \nabla \rho \cdot \nabla V \right)
    \label{eq:wasserstein-gradflow-pde}
\end{align}
In this paper, we shift the perspective and view the inner maximization problem as a regularized variational problem that corresponds to a sampling algorithm, without explicitly resorting to gradient ascent dynamics~\eqref{eq:wasserstein-gradflow-pde}.
This change not only simplifies the problem structure but also provides intuition for novel algorithmic design, which we use to effortlessly obtain novel algorithms such as Wasserstein-Fisher-Rao gradient and  Stein variational based DRO algorithms.
\paragraph*{Contributions.}
Importantly, we distinguish between two different notions:
\begin{tightenum}
\item The DRO problem formulation with various ambiguity sets, e.g., Wasserstein DRO~\citep{zhaoDatadrivenRiskaverseStochastic2018,mohajerin2018data,gaoDistributionallyRobustStochastic2016}, KL-DRO~\citep{hu2013kullback,ben-talRobustSolutionsOptimization2013}, Sinkhorn DRO~\citep{wang2021sinkhorn}, Kernel DRO~\citep{zhu2021kernel}, etc.
\item The optimization algorithm and analysis for solving the DRO problem, e.g., WRM~\citep{sinha2020certifyingdistributionalrobustnessprincipled}, reformulation techniques~\citep{mohajerin2018data}, SGD~\citep{levyLargeScaleMethodsDistributionally2020}, etc.
\end{tightenum}
In this paper, We do not invent new DRO problem formulations or ambiguity sets (item 1 above). Instead, we propose a principled mathematical framework for analysis and novel algorithms for solving DRO (item 2 above), based on the theory of gradient flows and PDEs.
As a consequence of our theoretical insights, we have invented novel algorithms such as Stein variational gradient, Wasserstein-Fisher-Rao, rejection sampler-based algorithms, etc., for solving existing DRO problems with great generality and theoretical guarantees.
We provide code implementations of our algorithms in the online repository: \url{https://github.com/ZusenXu/GFS-DRO}.

Furthermore, we see an elegant connection between DRO and variational methods (in the sense of calculus of variations) such as gradient flows and PDE analysis.
This connection allows us to directly apply theoretical results from those domains backed by rigorous mathematical machinery such as metric space gradient flows~\citep{ambrosio2008gradient}, without inventing a new theory or ad-hoc  analysis techniques. 
\vspace{-0.5cm}
\section{Preliminaries}
\subsection{Gradient Systems and Their Gradient Flow Equations}
The Wasserstein gradient flow (WGF) framework of \citep{otto1996double} was introduced into the sampling literature to provide a theoretical foundation; see
\citep{chewi2024statistical,garcia2018continuum}
for recent surveys.
In that framework, one can write a flow equation formally as
\begin{align}
    \dot \rho =
    -  \mathbb G_W(\rho)^{-1}(\rho) \frac{\delta F}{\delta \rho}[ \rho]
    = \nabla \cdot\left(\rho\nabla \frac{\delta F}{\delta \rho}[ \rho]\right)
    \label{eq:wasserstein-gfe}
\end{align}
using
the inverse of the Wasserstein Riemannian metric tensor:
$\mathbb G_W^{-1}(\rho): T^*_\rho \calM \to T_\rho \calM, \xi \mapsto -\nabla \cdot(\rho\nabla \xi),$
where $T_\rho \calM$ is the tangent space of $\Mplus$ at $\rho$ and
$T^*_\rho \calM$ the cotangent space.
With those ingredients, we can formally define the gradient systems that generate gradient flow equations such as the WGF equation~\eqref{eq:wasserstein-gfe}.
We refer to \citep{mielke2023introduction} for more details.
\begin{definition}
    [Gradient system]
    We refer to a tuple $\left( \calM, F, \bbG \right)$
    as a gradient system. 
    It has the gradient structure
    identified by:
\begin{tightenum}
    \item a space $\calM$,
    \item an energy functional $F$,
    \item a dissipation geometry given by either:
        a distance metric defined on $\calM$ or
                a Riemannian metric tensor $\mathbb G$.
\end{tightenum}
\label{def:gradient-system}
\end{definition}

\subsection{Markov Chain Monte Carlo Sampler via Gradient Flow.}
For sampling and inference,
a common choice for the energy functional is the KL divergence,
i.e.,
$F(\rho) = \KL(\rho | \pi)$.
Through elementary calculation, we obtain from \eqref{eq:wasserstein-gfe} the Fokker-Planck equation (FPE)
\begin{align}
    \partial_t \rho =
    \nabla \cdot \left( \rho \nabla \log\frac{\rho}{\pi}  \right)
    =
    { \Delta \rho} - \nabla \cdot \left( \rho \nabla \log \pi  \right)
    .
    \label{eq:fokker-planck}
  \end{align}
When we express the target as $\pi(x) = \frac1Z \exp(-V(x))$
where $Z$ is a normalization constant (partition function), \eqref{eq:fokker-planck} is then 
$\partial_t \rho = { \Delta \rho} + \nabla \cdot \left( \rho \nabla V  \right)$.
Viewed as a dynamic system, the KL divergence energy functional dissipates
along \eqref{eq:fokker-planck} in the steepest descent manner.
Based on Definition~\ref{def:gradient-system}, we say that PDE~\eqref{eq:fokker-planck} has the \emph{gradient structure} that entails the following key ingredients:
\begin{align}
    \begin{cases}
        \textrm{{Space} :}&
        \text{prob. space }
        \mathcal P ,
        \\
        \textrm{Energy functional} :& {F(\cdot):= \KL(\cdot | \pi)},
        \\
        \textrm{Dissipation Geometry} :& { \text{Wasserstein metric}}.
    \end{cases}
    \label{eq:gradient-structure-fkl-wgf}
\end{align}




\subsection{Entropy-regularized Wasserstein DRO.}
Works such as \citep{sinha2020certifyingdistributionalrobustnessprincipled,wang2021sinkhorn} considered DRO problems with rather general loss functions as they are based on general-purpose continuous optimization rather than DRO reformulation techniques for special losses.
A key result from \citep{wang2021sinkhorn} is the characterization of the worst-case distribution $\pi_Y$ that solves problem \eqref{eq:proximal-EOT}, referred to as Sinkhorn DRO therein. It has a closed-form density given by a mixture form:
    $\pi_Y = \E_{x\sim\widehat{\rho_N}} \left[ \rho^*_{Y|X=x} \right]$,
where the conditional density is
\begin{align}
    \rho^*_{Y|X=x} = \frac{1}{Z_x} \exp\left( -\tfrac{2\tau}\epsilon\widetilde V_{x, \tau} \right).
    \label{eq:conditional-sampling-density}
\end{align}
for $\widetilde V_{x, \tau}(y):= V(y)+\tfrac1{2\tau}c(y, x)$, and $Z_x = \int\exp\left( -\tfrac{2\tau}\epsilon\widetilde V_{x, \tau}(y) \right) dy$ is a normalizing constant.


\citep{wang2021sinkhorn} proposed an algorithm based on solving the dual formulation of this DRO problem. 
They developed a specialized two-level sampling procedure to estimate the gradient of the dual objective.
Different from their dual approach, this paper develops a unified gradient flow sampler framework that directly samples from the primal worst-case distribution that solves the inner problem.


\section{A Gradient Flow Framework for Sampling from Worst-case Distributions}

\subsection{Schr\"odinger Bridge Formulation of DRO}
The variational problem \eqref{eq:proximal-EOT}, hence the inner maximization problem of DRO,
is a special case of the Schr\"odinger bridge with one free marginal, i.e., half bridge or one-sided bridge.
\begin{lemma}\label{lem:sb-klform}
    The variational problem \eqref{eq:proximal-EOT} is
    equivalent to the Schr\"odinger half bridge problem
    \begin{align}
        \min _{\Pi}
       \biggl\{
       \int V \dd \Pi + \frac{1}{2\tau }
        \int c(x,y) \dd\Pi + \frac\epsilon{2\tau}\int \log \Pi \dd\Pi
        \, \,
        \bigg |
        \int \Pi \dd y =  \rho_0
        \biggr\}
        .
       \end{align}

    Consequently, it is equivalent to the minimization of the expected KL divergence with respect to the conditional distribution
    \begin{align}\label{eq:sb-klform}
        \min_{\rho_{Y|X}}
    \mathbb{E}_{x\sim\rho_0}
            \KL
            \left(
            \rho_{Y|X=x}(y) 
            \bigg| 
            \dfrac{1}{Z_x}
            \exp \left[
                       - \dfrac{2\tau V(y) +  c(y,x)}{\epsilon} 
            \right]
            \right)
            .
    \end{align}
    The optimal marginal distribution of $Y$ is given by a mixture distribution, for some normalization constant $Z_x$:
    \begin{align}\label{eq:worst-dist}
        \pi_Y \propto \E_{x\sim \rho_0} \left[
            \dfrac{1}{Z_x}\exp
            \left[{-\frac{2\tau V(y) +c(y, x) }{\epsilon}} \right]    
         \right].
    \end{align}
\end{lemma}
This statement gives the overal variaional structure of the DRO problem such as
\eqref{eq:main-ent-dro}.
One may take the initial distribution $\rho_0$ as the empirical distribution $\widehat{\rho_N}$ for the data-driven DRO problem,
resulting in the following proposed gradient flow sampler-based algorithms.
\subsection{Gradient Flow Sampler-based DRO}
In variational problem~\eqref{eq:sb-klform}, 
the KL divergence energy functional (without expectation) can also be written as:
    \begin{align}
        F(\rho) :=
        \frac{\epsilon}{2\tau}
        \KL
            \left(
            \rho\ 
            \bigg| 
            \dfrac{1}{Z_x}
            \exp \left[
                        -\dfrac{2\tau V(y) +  c(y,x)}{\epsilon} 
            \right]
            \right)
            =
        \int V\, \dd \rho
        + \frac{1}{2\tau} W_c^2(\rho, \delta_x) 
         + 
         \frac{\epsilon}{2\tau} \int \rho \log \rho 
         + \const
      \label{eq:conditional-sampling-energy}
    \end{align}
  Note that the formulation on the right-hand side was also observed by \citep{chen2022improved} in studying the proximal sampler.

With those ingredients,
in this paper, we propose to achieve
sampling from the conditional distribution $\rho_{Y\mid X}$ by simulating the \emph{gradient system
$\displaystyle\left(
\calP, F, \mathbb G
\right)$}.
This results in the formal gradient flow equation:
$$
\dot\rho = - \mathbb G^{-1}(\rho)\,  DF(\rho).
$$
where we can freely choose the dissipation geometry $\mathbb G$ for the gradient flow.
Then,
our central methodology is the following perspective of sampling from conditional distributions
connected to the KL-minimization problem~\eqref{eq:sb-klform}.
Considering solving the inner maximization problem~\eqref{eq:main-ent-dro} via the following two-step sampling procedure:
\begin{algorithm}[h!]
\caption{Worst-case Distribution Sampler via Gradient Flows}
\label{alg:sampler}
\begin{algorithmic}[1]
    \State \textbf{Input:} an initial distribution ${\rho}_0$ to sample from (e.g. empirical distribution $\widehat{\rho}_N$ in data-driven DRO) 
    \State Sample $X \sim {\rho}_0$
    \State Sample $Y \sim \rho_{Y\mid X}$
    using a gradient flow
    (e.g. with energy functional given by \eqref{eq:conditional-sampling-energy})
    \State \textbf{Output:} sample $Y$ from the worst-case distribution
\end{algorithmic}
\end{algorithm}


Note that the choice of the dissipation geometry (e.g. Wasserstein) for this flow should not be confused with the entropy-regularized OT ambiguity set of the DRO problem formulation e.g. in \eqref{eq:main-ent-dro}.
Through gradient flow, our goal is to achieve the approximation
$\rho_\infty\approx \rho^*_{Y\mid X=x}$.



Using the gradient flow sampler, we now propose the following general-purpose gradient flow sampler-based DRO framework.
\begin{algorithm}[h!]
\caption{Gradient Flow Sampler-based DRO}
\label{alg:GF-DRO}
\begin{algorithmic}[1]
    \State \textbf{Input:} Initial distribution $\rho_0$, e.g., empirical distribution $\rho_0 = \widehat{\rho_N}$, constraint set $\Theta$, $\tau, \epsilon >0$, stepsize $\eta_t$
    \For{ iteration count $s = 0, \dots$}
        \State Generate a sample from the worst-case distribution $y^s \sim \pi_Y$ by using the gradient flow sampler in Algorithm~\ref{alg:sampler}
        \State DRO step: $\theta^{s+1} \leftarrow \text{Proj}_\Theta (\theta^s-\eta_s \nabla_\theta \ell(\theta^s, y^s))$
    \EndFor
\end{algorithmic}
\end{algorithm}
Other straigforward variants, such as using sample average approximation (SAA) instead of stochastic approximation (SA) above, are possible.
\subsubsection{Wasserstein Gradient Flow Sampler.}
As the first practial outcome of our theoretical insights, we instantiate a Wasserstein gradient flow (WGF) sampler for the entropy-regularized Wasserstein DRO (a.k.a. Sinkhorn DRO) problem using the update rule in \eqref{eq:ula-proximal-general-D-W2}.
\begin{example}
    [Wasserstein GF for SDRO]
    Our goal is to solve the entropy-regularized Wasserstein DRO (a.k.a. Sinkhorn DRO) problem in \eqref{eq:main-ent-dro}.
    We consider the Wasserstein gradient system with the driving functional $J$ and the Wasserstein metric as the dissipation geometry:
    $$\left(
\calP, F, W_2
\right).
$$

In sampling algorithms, this gradient flow is typically
implemented by discretizing the Langevin SDE
$$
dX_t = 
 \tfrac{2\tau}{\epsilon}\nabla\widetilde V_{x, \tau}(X_t) dt + dW_t,
$$
resulting in the following forward Euler discretization known as the unadjusted Langevin algorithm (ULA).
Note that we use a scaled step size.
\begin{lemma}
    The forward Euler discretization of the Wasserstein gradient flow equation of the energy functional $F$ in \eqref{eq:conditional-sampling-energy} is given by the difference equation at step $t$:
    \begin{align}
        X_{t+1} = X_t + \eta_t \nabla \widetilde V_{x, \tau}(X_t) + \sqrt{\eta_t\frac\epsilon\tau} \xi_t
        \label{eq:ula-proximal-general-D-W2}
    \end{align}
    where $\xi_t$ is a standard normal random variable and
    $\eta_t$ is the step size.
\end{lemma}
We note the step size scaling in front of the stochastic variable $\xi^{i,s}_t$ is different from the vanilla ULA update rule.
\end{example}
Adapting the above WGF radient flow dynamics to our GF-DRO framework in Algorithm~\ref{alg:GF-DRO}, we obtain the following discrete-time DRO algorithm summarized in Algorithm~\ref{alg:SDRO-NGD}.
\begin{algorithm}[h!]
\caption{Entropy-regularized Wasserstein DRO via WGF}
\label{alg:SDRO-NGD}
\begin{algorithmic}[1]
  \State \textbf{Input:} Empirical distribution $\widehat{\rho_N}$, constraint set $\Theta$, $\tau, \epsilon >0$, stepsize sequence $\{r_s > 0\}_{s=0}^{S-1}$, inner stepsize $\eta$, inner iterations $T$, number of samples $m$.
  \For{$s = 0, \dots, S-1$}
    \State Sample $x^s \sim \widehat{\rho_N}$
    \State Initialize $y_{0}^{i,s} \leftarrow x^s, i=1,...,m$
    \For{$t = 0, \dots, T-1$}
        \State Sample $\xi^{i,s}_t \sim\mathcal{N}(0,I)$
        \State Update sample using \eqref{eq:ula-proximal-general-D-W2}: $y_{t+1}^{i,s} = y_{t}^{i,s} + \eta\nabla\widetilde V_{x^s, \tau}(y_{t}^{i,s}) + \sqrt{\eta\epsilon/\tau} \xi^{i,s}_t$
    \EndFor
    \State DRO step: $\theta^{s+1} \leftarrow \text{Proj}_\Theta (\theta^s-r_s \sum_{i=1}^m \frac{1}{m}\nabla_\theta \ell(\theta^{s},y_{T}^{i,s}))$
  \EndFor
  \State \textbf{return} $\theta^S$
\end{algorithmic}
\end{algorithm}
\begin{remark}
    [\citep{sinha2020certifyingdistributionalrobustnessprincipled}'s WRM]
    For the penalized Wasserstein DRO problem,
    i.e.
    $\epsilon=0$ in the entropy-regularized Wasserstein DRO problem, 
    the step in \eqref{eq:ula-proximal-general-D-W2} specializes to the update rule:
    \begin{align}
        X_{t+1} = X_t + \eta_t \nabla \widetilde V_{x, \tau}(X_t)
        \label{eq:ula-proximal-general-D-W2-epsilon-0}
    \end{align}
    which coincides with the inner SGD step used in the WRM algorithm by \citep{sinha2020certifyingdistributionalrobustnessprincipled}.
Henece, their WRM algorithm is a special case (with only ODE) of our GF-DRO framework in Algorithm~\ref{alg:GF-DRO}.
\end{remark}
Through the example above, we see the value of this paper's gradient flow perspective: it can easily generalize the WRM algorithm~\citep{sinha2020certifyingdistributionalrobustnessprincipled} from ODE to a novel SDE-based algorithm for the Sinkhorn DRO problem using \eqref{eq:ula-proximal-general-D-W2}.
We see that the gradient flow perspective is not simply a theoretical formulation -- it lets us effortlessly design new algorithms to solve DRO without resorting ad-hoc modifications of other DRO methods.
Moreover, we will later go beyond the standard Wasserstein gradient flow and introduce more advanced gradient flows such as the Wasserstein Fisher-Rao (WFR) and Stein variational gradient (SVG) flows.







\subsubsection{Wasserstein Fisher-Rao Flow Sampler.}
Consider the WFR gradient system of the energy functional $F$, i.e., the triple $\left( \calP, F, \WFR \right)$, where $\WFR$ is the Wasserstein-Fisher-Rao metric, a.k.a. the \emph{Hellinger-Kantorovich} metric restricted to the probability space\footnote{Note that there are many cases of misnomer in the machine learning literature: WFR should technically be defined over the space of positive measures, not the space of probability measures; see \citep{mielkeNotesHellingerDistance2025} for a historical account. The latter corresponds to the spherical Hellinger-Kantorovich metric. Although their gradient flow solutions can be easily related to each other via the mass scaling. See \citep{mielke2025hellinger} for technical details.}.
Specifically, we consider the HK/WFR gradient system associated with reaction-diffusion PDE:
\begin{align}
    \frac{\partial \mu}{\partial t}
    &= \alpha\, \mathrm{div} \left( \mu \nabla \frac{\delta F}{\delta \mu} \right) - \beta\, \mu 
    \left(\frac{\delta F}{\delta \mu} - \int \frac{\delta F}{\delta \mu}\dd \mu \right)
    \label{eq:wfr-gradient-flow}
\end{align}
Discretizing the flow equation, 
we propose a WFR worst-case distribution sampler-based SDRO in Algorithm \ref{alg:SDRO-WFR}.
The main intuition of this sampler, compared with the
WGF version, is the addition of the mass (weights) update of the particles. 

\begin{algorithm}[h!]
\caption{Entropy-regularized Wasserstein DRO via WFR flow}
\label{alg:SDRO-WFR}
\begin{algorithmic}[1]
  \State \textbf{Input:} Empirical distribution $\widehat{\rho_N}$, constraint set $\Theta$, $\tau, \epsilon >0$, stepsize sequence $\{r_s > 0\}_{s=0}^{S-1}$, inner stepsize $\eta$, weight stepsize $\eta_w$, weight threshold $w_{\text{min}}$, inner iterations $T$, number of samples $m$.
\For{$s = 0, \dots, S-1$}
    \State Sample $x^s \sim \widehat{\rho_N}$
    \State Initialize $y_{0}^{i,s} \leftarrow x^s, i=1,...,m$
    \For{$t = 0, \dots, T-1$}
        \State Sample $\xi^{i,s}_t \sim\mathcal{N}(0,I)$
        \State $y_{t+1}^{i,s} = y_{t}^{i,s} + \eta\nabla\widetilde V_{x^s, \tau}(y_{t}^{i,s}) + \sqrt{\eta\epsilon/\tau} \xi^{i,s}_t$
        \State $w_{t+1}^{i,s} = (w_{t}^{i,s})^{1-\epsilon\eta_w/ (2\tau)}e^{\eta_w \widetilde V_{x^t, \tau}(y_{t}^{i,s})}$
        \State Normalize weights: $w_{t+1}^{i,s} = w_{t+1}^{i,s}/\sum_{j=1}^m w_{t+1}^{i,s}$

        
        \State Birth-death sampling: 
            \If {$w_{t+1}^{i,s} < w_{\text{min}}$}
            \State Select $i'$ from $\{1, \dots, m\}$ with probability $w_{t+1}^{i',s}$.
            \State $y_{t+1}^{i,s} = y_{t+1}^{i',s}$, $w_{t+1}^{i,s}= (w_{t+1}^{i,s} + w_{t+1}^{i',s})/2$, $w_{t+1}^{i',s} = w_{t+1}^{i,s}$.
            \EndIf
    \EndFor
    \State DRO step: $\theta^{s+1} \leftarrow \text{Proj}_\Theta (\theta^s-r_s \sum_{i=1}^m w_{T}^{i,s}\nabla_\theta \ell(\theta^{s},y_{T}^{i,s}))$
  \EndFor
  \State \textbf{return} $\theta^S$
        
        
                    
\end{algorithmic}
\end{algorithm}

\subsubsection{Stein Variational Gradient Flow Sampler.}
In an attempt to simulate the \eqref{eq:fokker-planck} using deterministic particle-based methods instead of SDEs,
\citep{liu2016stein}
proposed the Stein variational gradient descent (SVGD) algorithm, which is a 
implementable
discrete-time algorithmic version of \eqref{eq:fokker-planck}.
At time step $t$,
it updates the particle locations via the following gradient descent scheme.
    $$X^i_{t+1} =X^i_t - 
     \eta\cdot 
                    \big(
                    \frac{1}{m}\sum_{j=1}^m
                    \nabla V(X^j_t)k(X^j_t,X^i_t )
                    +
                    \frac{1}{m}\sum_{j=1}^m\nabla_2 k(X^j_t,X^i_t )
                    \big)
                    .
                    $$
Here, $X^i_t$ represents the position of the $i$-th particle at time step $t$, $\eta$ is the step size, $m$ is the total number of particles.
$k(\cdot,\cdot)$ is a positive definite kernel function, and $\nabla_2$ denotes the gradient with respect to the second argument of the kernel.
Different from the typical implementation such as ULA, SVGD offers a deterministic algorithm for sampling with the kernel $k$ modeling the pairwise interaction between particles.
Applying SVGD to the inner problem of DRO, we propose Algorithm~\ref{alg:SDRO-SVG}. See a detailed analysis in Section~\ref{sec:svg}.


\subsubsection{Rejection Sampler.} 
We also propose Algorithm~\ref{alg:SDRO_rgo} based on rejection sampler (RGO). This approach is based on the backward step of the proximal sampler \citep{lee2021structured,chen2022improved,wibisono2025mixing}, as detailed in Section~\ref{sec:rgo}. However, this method requires the loss function to be $L$-smooth and cannot be applied when $\tau$ is relatively large, a limitation we will demonstrate in later experiments.

\section{Analysis of Gradient Flows}
For algorithmic design, one paradigm is to first obtain intuition and insights from principled mathematical analysis before constructing practical numerical algorithms.
While the development in machine learning research does not always follow this pattern, we argue in this paper that, by first understanding the theoretical limits of the gradient flow, we can design novel DRO algorithms that performs as expected theoretically. This is done without inventing new theoretical framework, but by leveraging the existing gradient flow and PDE analysis specialized to the problem.
The discrete-time algorithmic analysis will follow in a later section.



In DRO works such as \citep{sinha2020certifyingdistributionalrobustnessprincipled}
and subsequent variants, 
complexity analysis typically involves controlling the gradient estimation error.
That is, to control
the deviation 
from the true gradient used for DRO.
For illustration, let us consider the
DRO problem \eqref{eq:main-ent-dro}.
Note that
setting $\epsilon = 0$ recovers the
problem of \citep{sinha2020certifyingdistributionalrobustnessprincipled}.
One needs to estimate the gradient w.r.t. $\theta$ for the outer DRO problem, i.e., estimating
\begin{align}
\label{eq:grad-estimate}
\widehat g \approx
\nabla_\theta 
\max _{\rho \in \calP} 
\biggl(
    \int \ell(\theta, z)\dd\rho{(z)}
    -\frac1{2\tau} W_\epsilon^2(\rho, \widehat{\rho_N})
\biggr)
\end{align}
and subsequently, the gradient estimate is used for the outer DRO problem via SGD-like updates, e.g., $\theta^{s+1} \gets \theta^s - \eta \widehat g$.
The bound of $\widehat g$ deviating from the true gradient can be used in the downstream standard optimization error bound for DRO solution.


With our gradient flow sampler-based DRO algorithm,
we aim to generate samples from the worst-case distribution $\pi_Y$, which is the entropic JKO solution
\begin{align}
\label{eq:entropic-jko}
\pi_{Y} \in \argmin_{\rho\in \calP}
\int V \, d\rho
+ \frac{1}{2\tau} W^2_{\epsilon}(\rho, \rho_0)
.
\end{align}
Following the framework of Algorithm~\ref{alg:sampler},
for each $x_i\sim \rho_0$,
we generate samples from the worst-case distribution
by sampling from the conditional distribution, $y_i\sim \rho^*_{Y|X=x_i}$.
Then, for the outer DRO problem, we can use the sample-based gradient estimate $\widehat g = \tfrac1N \sum_{i=1}^N \nabla_\theta \ell(\theta, y_i)$ to update the parameter $\theta$.
The following result uses gradient flow analysis to bound the gradient estimate error.
We note that the result are stated in terms of the ideal continuous-time gradient flow, which is not the mixing time of the discrete-time sampler; we will detail the discrete-time algorithmic analysis in a later section.
The gradient flow analysis provides a key insight and guideline for the design of the gradient flow sampler-based DRO algorithm.



For notational simplicity, we abbreviate $\widetilde V_{x,\tau}$ as $\widetilde V$ in this section and, specifically, for the case of quadratic transport cost $c(y,x) = \|x-y\|^2$, we define the following.
\begin{definition}
\label{as:gf-properties}
\begin{enumerate}
    \item \textbf{Semi-convexity:}
    \begin{align}
        \nabla ^2 \widetilde V \geq \lambda I
        \iff 
        \nabla ^2 V + \frac1{2\tau }I \geq \lambda I
        \label{eq:semi-convexity}
    \end{align}
    \item \textbf{Gradient dominance/PL:}
    There exists a constant $\lambda_{\text{PL}} > 0$ such that
    \begin{align}
        \big|\nabla \widetilde V \big|^2 \geq \lambda_{\text{PL}} \widetilde V,
        \quad \forall x \in \mathcal X
        \label{eq:gradient-dominance-PL}
    \end{align}
    \item \textbf{Generalized PL for functionals of measures along GF}:
    \begin{align}
        \frac{\dd}{\dd t} F(\rho_t) \geq \lambda_{\text{PL}} F(\rho_t),
        \label{eq:generalized-PL}
    \end{align}
\end{enumerate}
\end{definition}

It is immediate that 
$$
\text{semi-convexity with } \lambda >0 \text{ (strong convexity)}
    \implies 
    \text{gradient dominance/PL} 
    \implies 
    \text{generalized PL of }F
    .
$$
We note that
the first condition in Definition~\ref{as:gf-properties}
does not require the function $V$ to be convex. Hence, the DRO loss $\ell(\theta, x)$ does not necessarily need to be concave in variable $x$.
In such cases, we can already obtain the geodesic convexity of $F$ in the Wasserstein space, termed displacement convexity.

Using standard analysis of gradient flow, we obtain the following ideal estimate
\begin{proposition}\label{prop:gradient_oracle_error_control}
    Denote the initial sample $x^i\sim \widehat{\rho_N}$ and samples running the gradient flow sampler for time $t$ are $y^i_t\sim \rho^*_{Y|X=x^i}$.
    Suppose that the loss function $\ell(\theta, x)$
    is $L$-smooth in $x$ and 
    $\widetilde V_{x, \tau}(y) : = -\ell(\theta, y)+\tfrac1{2\tau}c(y, x)$ satisfies the gradient dominance condition in Definition~\ref{as:gf-properties} with $\lambda > 0$.
    Then, in order to generate an $\epsilon$-gradient estimate, i.e.,
    \begin{align*}
    \bigg|\frac1N \sum_{i=1}^N \nabla_\theta \ell(\theta, y^i_t) 
    -
    \nabla_\theta 
    \underbrace{
    \max _{\rho \in \calP} 
    \biggl(
        \int \ell(\theta, z)\dd\rho{(z)}
        -\frac1{2\tau} W_\epsilon^2(\rho, \widehat{\rho_N})
    \biggr)
    }_{\text{DRO objective}}
    \bigg|
    {\leq}
    \epsilon,
    \end{align*}
    the Wasserstein gradient flow sampler needs to run for time at least\;
    $\displaystyle t \gtrsim \calO \left(\frac{1}{\lambda}\log{\frac{L}{\sqrt{\lambda}\epsilon}}\right)$.
\end{proposition}


Results such as Proposition~\ref{prop:gradient_oracle_error_control} are based on the displacement convexity or PL in the Wasserstein space.
The value of our general gradient flow perspective is to let us freely choose the gradient flow geometry, not confined to Wasserstein.
This is not just for theoretical considerations --- the estimate such as
\eqref{eq:gf_time_to_epsilon}
depends crucially on the (PL/convexity) constant $\lambda$, which has been a major limitation of the Langevin type samplers in the literature.
Note that the speed-up is due to
an advantage of the Hellinger or Fisher-Rao geometry when the warm-start initialization is given.
It provides a simple insight for algorithmic design:  the unbalanced WFR gradient flow can be used to speed up the sampling process for DRO.
This insight is later validated in the numerical studies.


As noted earlier, Proposition~\ref{prop:gradient_oracle_error_control} 
is aimed at providing the intuition for the behavior of the gradient flow sampler-based DRO algorithm; it is not the practical complexity estimate based on sampler's mixing time.
As it's based on standard theory of gradient flow, it serves as a guideline for the design of the gradient flow sampler-based DRO algorithm using existing tools from PDE/SDE.
Next, we provide optimization complexity estimate which characterizes the practical sampler-based DRO algorithm.





\section{Optimization Complexity Analysis}
For notational convenience, we define:
$\Phi(\theta): = \max_{\rho\in\calP}
    \int \ell(\theta, z)\dd\rho{(z)} - \tfrac1{2\tau} W^2_{\epsilon}(\rho,\widehat{\rho_N})$.
Then, the DRO problem \eqref{eq:main-ent-dro} can be written succinctly as 
$\min_{\theta\in\Theta} \Phi(\theta)$.
We analyze the computational complexity of our proposed discrete-time algorithms, aiming to find an $\epsilon_{\text{opt}}$-stationary point $\theta^{S}$ at the last iteration $S$,
i.e. $\mathbb{E}[\|\nabla\Phi(\theta^{S})\|^{2}]\le\epsilon_{\text{opt}}^{2}$.

Our analysis uses a standard framework for non-convex stochastic optimization \citep{ghadimi2013stochastic}. The outer loop performs SGD on $\theta$, while the inner loop generates samples to approximate the conditional worst-case distribution. The key challenge is controlling the bias in the stochastic gradient.
For technical reasons, we make the following assumptions.

\begin{assumption}\label{as:totalsmooth}
$\Phi(\theta)$ is $L_{\Phi}$-smooth.
\end{assumption}

\begin{assumption}\label{as:fsmooth}
    The gradient of the loss function $\ell(\theta, z)$ in \eqref{eq:main-ent-dro}, $\nabla_{\theta}\ell(\theta, z)$ is $L_{f}$-Lipschitz with respect to $z$.
\end{assumption}
\begin{assumption}\label{as:gradest}
The stochastic gradient estimator $\widehat g:= \E_{\widehat\rho_{Y|X=x}} [\nabla_{\theta}\ell(\theta, y)]$ has bounded variance $\sigma^2$. The inner-loop sampler (Algorithm~\ref{alg:sampler}) generates a distribution $\widehat\rho_{Y|X=x}$ with a uniformly bounded expected error $\E_{x \sim \widehat{\rho_N}}[W_2(\widehat\rho_{Y|X=x}, \rho^*_{Y|X=x})] \le \delta_{\text{sample}}$.
\end{assumption}

Under these assumptions, we can bound the gradient bias and establish a general convergence result for the outer loop.

\begin{theorem}[Outer Loop Convergence]
\label{thm:outer_loop}
Let Assumptions \ref{as:totalsmooth}, \ref{as:fsmooth}, \ref{as:gradest} hold. With a constant step size $r = O(1/L_\Phi)$, the outer loop requires $S = O(1/\epsilon_{\text{opt}}^2)$ iterations to find an $\epsilon_{\text{opt}}$-stationary point, provided the error $\delta_{\text{sample}}$ in Assumption~\ref{as:gradest} is controlled such that $\delta_{\text{sample}} = O(\epsilon_{\text{opt}}/L_f)$.
\end{theorem}

For the analysis of the ULA-based sampler (Algorithm \ref{alg:SDRO-NGD}), we introduce the following assumption on the geometry of the inner target distribution.

\begin{assumption}\label{as:LSI}
For any fixed $\theta$ and $x$, the conditional distribution $\rho_{Y|X=x}$ is $L_{U}$-smooth and satisfies the $\alpha_U$-log-Sobolev inequality (LSI).
\end{assumption}

Assumption \ref{as:LSI} is motivated by the convergence analysis of ULA \citep{vempala2019rapid, wibisono2025mixing}, which often requires the target distribution to satisfy the $\alpha$-log-Sobolev inequality ($\alpha$-LSI) \citep{gross1975logarithmic}. We recall that a distribution $\rho$ satisfies the $\alpha$-LSI for some $\alpha>0$ if, for any probability distribution $\nu$, the following inequality holds:
\[
\text{FI}(\nu \| \rho) \ge 2\alpha \KL(\nu \| \rho),
\]
where $\KL(\nu \| \rho) := \E_\nu\left[\log\frac{\nu}{\rho}\right]$ is the KL divergence and $\text{FI}(\nu \| \rho) := \E_\nu\left[\left\|\nabla\log\frac{\nu}{\rho}\right\|^2\right]$ is the relative Fisher information. Since the relative Fisher information is the squared gradient norm of the KL divergence in the $W_2$ metric,  the property that $\rho$ satisfies LSI is equivalent to the Polyak-Łojaciewicz (PL) or gradient domination property of the KL divergence objective function. 


It is important to clarify, however, that the PL inequality for the potential function $\widetilde{V}_{x,\tau}$ itself (as in \eqref{eq:gradient-dominance-PL}) is not equivalent to the LSI for the distribution $\rho_{Y|X=x}=\exp\left(-\frac{2\tau}{\epsilon}\widetilde{V}_{x,\tau}\right)$.

Assumption~\ref{as:LSI} allows us to bound the number of inner ULA iterations required to achieve the necessary sampling accuracy. Combining these results yields the total complexity.

\begin{theorem}[Complexity of Algorithm \ref{alg:SDRO-NGD}]\label{thm:ula}
Under Assumptions \ref{as:totalsmooth}, \ref{as:fsmooth}, \ref{as:gradest}, and \ref{as:LSI}, the total complexity for Algorithm 3 to find an $\epsilon_{\text{opt}}$-stationary point is:
\begin{align*}
\text{Complexity} = O\left(\frac{L_{\Phi} L_U^2 L_f^2 d^2}{\alpha_U^3 \epsilon_{\text{opt}}^4}\cdot\log \frac{1}{\epsilon_{\text{opt}}}\right) = \tilde{O}\left(\frac{L_{\Phi} L_U^2 L_f^2 d^2}{\alpha_U^3 \epsilon_{\text{opt}}^4}\right).
\end{align*}
\end{theorem}
The detailed derivations for the gradient bias, the proof of Theorem \ref{thm:outer_loop}, and the complexity analysis are provided in Appendix~\ref{ap:complexity_proofs}.



















\section{Experiments}

We conduct numerical experiments on different tasks to validate our theoretical insights. Our goal is to compare the performance of four different DRO methods-Algorithm~\ref{alg:SDRO-WFR} (WFR) and Algorithm~\ref{alg:SDRO-NGD} (WGF), the dual method for SDRO (Dual)\citep{wang2021sinkhorn}, and WRM \citep{sinha2020certifyingdistributionalrobustnessprincipled}. We also compare Stein Variational Gradient (SVG) (Algorithm~\ref{alg:SDRO-SVG}) method and Restricted Gaussian Oracle (RGO) (Algorithm~\ref{alg:SDRO_rgo}) method with these methods in Section~\ref{sec:experiment3} and Section~\ref{sec:experiment4}. For methodological consistency across all experiments, the Dual method uses Randomized Truncation MLMC\citep{blanchet2015unbiased} as suggested in \citep{wang2021sinkhorn}, and the WRM solves its inner problem using gradient descent. 
All experiments were conducted on a laptop with an NVIDIA RTX A3000 GPU using Python. The code is available at https://github.com/ZusenXu/GFS-DRO 

\subsection{Classification under Data Imbalance}\label{sec:circle-exp}

\begin{figure*}[h!]
\begin{subfigure}[b]{0.32\textwidth}
        \includegraphics[width=\textwidth]{circle_boundary_comparison.pdf}
        \caption{Decision boundaries}
        \label{fig:circle-boundary}
    \end{subfigure}
    \hfill
    \begin{subfigure}[b]{0.32\textwidth}
        \includegraphics[width=\textwidth]{WFR_frame_000_samples.pdf}
        \caption{Perturbation by WFR}
        \label{fig:WFR_circle}
    \end{subfigure}
    \hfill
    \begin{subfigure}[b]{0.32\textwidth}
        \includegraphics[width=\textwidth]{WRM_frame_000_samples.pdf}
        \caption{Perturbation by WRM}
        \label{fig:WRM_circle}
    \end{subfigure}
    \caption{\textbf{Robust Decision Boundaries with Biased Data.} \textbf{(a)} Decision boundaries learned by different methods on the biased circle dataset. The training data are shown as orange (positive class) and blue (negative class) points. The final classification boundaries are shown for the WRM (green), WGF (purple), WFR (red), and Dual (brown) models. All models were trained for 40 epochs. We set the regularization parameter $\tau=2.5$ for all methods and the entropy regularization $\epsilon=0.15$ for methods based on entropy-regularized Wasserstein DRO problem. \textbf{(b)} Samples from the worst-case distribution generated by WFR sampler at the first epoch. 
    \textbf{(c)} Worst-case samples generated by WRM method at the first epoch. WRM can only generate discrete distributions as worst-case distribution while Sinkhorn DRO uses potentially continuous distributions as worst-case distribution. The original data points are shown as circles and the worst-case samples are shown by circles with a shallower color.}
    \label{fig:circle-experiment}
\end{figure*}
This experiment demonstrates the robustness of our proposed methods to data imbalance. The experimental setup is adapted from \citep{sinha2020certifyingdistributionalrobustnessprincipled}. We generate a dataset where features $X \in \mathbb{R}^2$ are drawn from a Gaussian distribution. The labels are assigned based on the rule $Y = \text{sign}(\|X\|_2 - \sqrt{2})$, which creates two classes separated by a circle of radius $\sqrt{2}$. To establish a clear margin, all data points satisfying $\|X\|_2 \in (\sqrt{2}/1.3, 1.3\sqrt{2})$ are excluded. To simulate a biased training distribution, we remove all samples from the first quadrant. This removal introduces a significant imbalance, testing the ability of each algorithm to learn a generalizable decision boundary rather than overfitting to the biased training data. The model is a neural network with a single hidden layer of 4 units.



Figure~\ref{fig:circle-boundary} visualizes the learned decision boundaries. WFR learns a decision boundary that closely matches the true circular boundary, correctly classifying the held-out data in the first quadrant and thus demonstrating robustness to the distributional shift. In contrast, the WRM boundary is overly expansive, misclassifying large regions of the feature space. This indicates a failure to generalize from the biased training set, resulting in a less reliable classifier. The boundary from WGF lies between those of WFR and WRM, highlighting the benefit of the weight flow mechanism in WFR for achieving a more robust solution. The dual method \citep{wang2021sinkhorn} does not perform as well as WFR in this setting. We attribute this to the high sensitivity of the dual method to the choice of hyperparameters ($\epsilon$ and $\lambda$), which can significantly impact the bias of the gradient estimator. 


To further investigate the differing behaviors, Figure~\ref{fig:WFR_circle} and Figure~\ref{fig:WRM_circle} visualize the samples generated by WFR and WRM at the first epoch. We observe that the samples generated by WFR algorithm begin to recover parts of the missing data distribution in the first quadrant. Conversely, the samples generated by WRM fail to do so, providing insights into why it learns a less robust boundary.


\subsection{Two Moon Classification}\label{sec:experiment3}
In this section, we address a binary classification problem using the 'two moons' dataset. The model is a three-layer neural network with ReLU activations and a hidden dimension of 16. The objective is to minimize the cross-entropy loss. The cost function for samples $(x, y)$ and $(x', y')$ is defined as $c((x,y),(x',y'))=\|x-x'\|_2^2 + \infty \cdot \mathbf{1}_{y\neq y'}$.

The training data is generated using \texttt{sklearn.make\_moons} with a noise level of 0.1. To introduce class imbalance, the training set of $N_{\text{train}}=200$ samples consists of 90\% positive ($y=1$) and 10\% negative ($y=0$) samples.

\begin{figure}[h!]
    \centering
    \includegraphics[scale=0.7]{boundary_5_0.01.pdf}
    \caption{Decision boundary comparison for all methods on the two-moon classification task. For all DRO methods, we set $\tau=0.1$, and for SDRO methods, we set $\epsilon=0.01$. In each inner loop, WFR and WGF generate $m=5$ particles.}
    \label{fig:final_boundaries}
\end{figure}


Figure~\ref{fig:final_boundaries} illustrates the decision boundaries learned by various algorithms against the ideal boundary. Notably, the boundary from the WRM method fails to correctly separate the training data. The Dual method, while separating the classes in high-density regions, learns a boundary that deviates significantly from the ground truth. In contrast, WFR, and WGF learn boundaries that more closely approximate the ideal curve. Among them, WFR and WGF achieve the best results, tracking the ideal separator with high fidelity, which indicates superior performance in this setting.


To evaluate the capacity of generating the worst-case distribution of WRM, WGF, WFR, SVG, and RGO, we conducted an experiment on a pre-trained SAA model. This setup isolates the inner-loop optimization process used to find the worst-case distribution. The evolution of the inner objective function, $\mathbb{E}[\widetilde V_{x, \tau}(z)]$, is depicted in Figure~\ref{fig:loss_twomoon}. 


The efficacy of WGF, WFR, and SVG in approximating the worst-case distribution is demonstrated by a significant increase in the objective value. The reweighting mechanism in WFR contributes to its faster convergence compared to WGF. In contrast, SVG's performance is highly dependent on its initialization parameters. An initial standard deviation of 0.1 leads to slower convergence than a standard deviation of 0.2, which achieves a rate similar to WFR. This sensitivity is attributed to the ability of a larger initial standard deviation to propel particles across the decision boundary, thereby facilitating a more thorough exploration of the perturbation space, as illustrated later in Figure~\ref{fig:perturbation_samples}. In contrast, WGF and WFR initiate their optimization from the empirical data distribution, which obviates the need to select initialization parameters.
\begin{figure}[h!]
\centering
    \includegraphics[scale=0.6]{saa_perturb_objective_comparison.pdf}
    \caption{Evolution of $\mathbb{E}[\widetilde V_{x, \tau}(z)]$. We run all methods for 300 steps with a stepsize of 0.01. For RGO (blue), we run a rejection sampling procedure after solving the inner optimization problem. SVG\_0.1 and SVG\_0.2 denote initial distributions with a standard deviation of 0.1 and 0.2, respectively.}
    \label{fig:loss_twomoon}
\end{figure}

\begin{figure*}[h!]
\centering
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{wrm_frame.pdf}
        \caption{Perturbation by WRM}
        \label{fig:wrm_perturb}
    \end{subfigure}
    \hfill
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{rgo_frame.pdf}
        \caption{Perturbation by RGO}
    \end{subfigure}
    \hfill
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{wgf_frame.pdf}
        \caption{Perturbation by WGF}
        \label{fig:wgf_perturb}
    \end{subfigure}
    \hfill
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{wfr_frame.pdf}
        \caption{Perturbation by WFR}
        \label{fig:wfr_perturb}
    \end{subfigure}
    \hfill
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{svg_frame_0.1.pdf}
        \caption{Perturbation by SVG ($\sigma=0.1$)}
        \label{fig:svg0.1}
    \end{subfigure}
    \hfill
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{svg_frame_0.2.pdf}
        \caption{Perturbation by SVG ($\sigma=0.2$)}
        \label{fig:svg0.2}
    \end{subfigure}
    \caption{Visualization of perturbed samples (gray-edged smaller circles) generated from original data (black-edged bigger circles) against the SAA boundary at the final step. For all methods, we use a stepsize of 0.01 and run for 300 iterations. For WFR, the intensity of points visualizes the sample weights. We show perturbations by SVG-DRO with different initializations.}
    \label{fig:perturbation_samples}
\end{figure*}

Conversely, WRM and RGO fail to significantly increase the objective. The non-convex nature of the objective, evidenced by the initial dip and subsequent rise for WGF and WFR, likely presents a challenge for the optimization procedures in WRM and RGO. Additionally, the performance of RGO is hindered by its rejection sampling stage, which is inefficient for non-smooth objectives.



Figure~\ref{fig:perturbation_samples} provides a visual confirmation of these results, displaying the perturbed samples at the final iteration of the inner loop. WRM generates minimal perturbations, with particles remaining in close proximity to the original data points. The perturbations from RGO are qualitatively similar, appearing as slightly noisier versions of the WRM results. In contrast, WGF, WFR, and SVG generate a diverse set of adversarial examples, effectively pushing samples across the decision boundary, including those initially distant from it. Notably, the extent of perturbation for SVG is dependent on the initialization; an initial standard deviation of 0.2 results in more significant perturbations than a standard deviation of 0.1. And the particles generated by SVG for a single data point tend to be highly concentrated.

This phenomenon arises from the interplay between the two forces governing SVGD: a driving force that pushes particles toward regions of higher loss and a repulsive force from the kernel that prevents particle collapse. The optimization dynamics, visualized in Figure~\ref{fig:svg_force}, show that the driving force initially dominates, causing the particles to converge. As the particles draw closer, the repulsive force increases to counteract this convergence. However, since the number of particles is relatively small, the repulsive force only dominates when particles are very close, leading to the observed particle concentration. This result also explains the reason why SVG shows nearly identical performance as WRM in Section~\ref{sec:experiment4}. 

\begin{figure*}[h!]
\centering
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{frame_020.pdf}
        \caption{Iteration 20}
    \end{subfigure}
    \hfill
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{frame_100.pdf}
        \caption{Iteration 100}
    \end{subfigure}
    \hfill
    \begin{subfigure}{0.32\textwidth}
        \includegraphics[width=\textwidth]{frame_288.pdf}
        \caption{Iteration 300}
    \end{subfigure}
    \caption{Evolution of particle positions in the SVG method. The driving force guides initial convergence, while the repulsive force prevents complete collapse. But in this experiment, the repulsive force fails to push particles apart efficiently, leading to the mode collapse. }
    \label{fig:svg_force}
\end{figure*}




\subsection{Adversarial Multi-class Logistic Regression}\label{sec:experiment4}

We evaluate the adversarial robustness of the algorithms on the features extracted from the real-world image dataset CIFAR-10 \citep{krizhevsky2009learning}. The task is multi-class logistic regression, minimizing the negative log-likelihood loss:
$$
h_B(x,y) = -y^T B^T x + \log(\mathbf{1}^T e^{B^T x})
$$
where $B := [w_1, \dots, w_C]$ are the classifier parameters. We solve the DRO problem assuming perturbations affect features $x$ but not labels $y$. The CIFAR-10 images are processed using a ResNet-50 network pre-trained on ImageNet to extract 512-dimensional features. To evaluate robustness, we apply a Projected Gradient Descent (PGD) attack with $l_2$-norm constraints to the test data. The perturbation magnitude, $\Delta$, is normalized by the average $l_2$-norm of the test features, and we vary $\Delta$ from 0 to 0.08. Performance is measured by the misclassification rate. 

\begin{figure*}[h!]
\centering
\begin{subfigure}{0.3\textwidth}
    \includegraphics[width=\textwidth]{rgo_lam=10_eps=0.2.pdf}
    \caption{$\tau=0.05, \epsilon=0.2$}
\end{subfigure}
\hfill
\begin{subfigure}{0.3\textwidth}
    \includegraphics[width=\textwidth]{rgo_lam=10_eps=0.02.pdf}
    \caption{$\tau=0.05, \epsilon=0.02$}
\end{subfigure}
\hfill
\begin{subfigure}{0.3\textwidth}
    \includegraphics[width=\textwidth]{rgo_lam=10_eps=0.002.pdf}
    \caption{$\tau=0.05, \epsilon=0.002$}
\end{subfigure}
\caption{Experiment results of adversarial training on CIFAR-10 features under an $l_2$-norm PGD attack. The plots show the test error (\%) as a function of the normalized perturbation level $\Delta$ for different $\epsilon$. All models are trained for 10 epochs, and for all methods contain an inner loop, they employ a stepsize of 0.01, and run for 100 loops. }
\label{fig:robustness}
\end{figure*}
Figure~\ref{fig:robustness} presents the classification results under adversarial attack. Notably, RGO performs almost identically to the SAA baseline, failing to provide robustness in this task. This poor performance might be attributed to the failure of its rejection sampling mechanism due to the non-smooth loss function setting. In contrast, all other DRO methods demonstrate enhanced robustness. WFR and WGF, in particular, consistently achieve a high degree of robustness across all settings. The performance of the dual method proves sensitive to the choice of $\epsilon$; it performs comparably to WFR and WGF for larger values such as $\epsilon=0.2$, but its robustness degrades as $\epsilon$ decreases, with its error rate approaching that of the SAA baseline when $\epsilon=0.002$. Furthermore, SVG consistently shows nearly the same performance as the WRM method. We hypothesize that this occurs because the number of particles $m$ is relatively small for inner SVGD steps. With a small $m$, the resulting sparse repulsive force is too weak to counteract the drift force, thus failing to push the particles apart and maintain diversity. Thus, SVG practically solves the same Wasserstein DRO problem as WRM in this experiment, rather than an entropy-regularized Wasserstein DRO problem. Methods employing Langevin dynamics (WFR, WGF), conversely, remain robust. Although their formulation also stably converges to that of WRM as $\epsilon\rightarrow 0 $, the introduced stochasticity prohibits the samples from collapsing to a single point even when $m$ is small, thereby preventing the model from simply solving a Wasserstein DRO problem.


\section{Other Works and Discussion}
Recent works have advanced the understanding and application of Distributionally Robust Optimization (DRO) in continuous probability spaces. 
\citep{xu2024flow} proposed FlowDRO, a practical map-based approach that learns a deterministic transport map, parameterized by a sequence of neural networks, to solve the standard Wasserstein DRO problem.
\citep{zhu2024distributionally} provide a foundational theoretical analysis, establishing convergence guarantees for a general iterative min-max framework for discrepancy-based DRO.
Our work is rather orthognal to those efforts in that 
it addresses the DRO problems with a sample-based methodology that offers several aforementioned advantages. 
An important contribution of this paper is allowing optimizers to go beyond the standard (Wasserstein) gradient descent type of algorithm, inspiring novel approaches such SVG and WFR.

In this paper, we have provided a complete framework of (stochastic) optimization algorithms for DRO. While the specific algorithms are designed for the Sinkhorn DRO, our theory treats general DRO by changing the corresponding gradient flow energy function $F$ in \eqref{eq:conditional-sampling-energy}.
Therefore,
many more results in analysis and algorithms are possible by choosing various DRO ambiguity notion and gradient flow geometry.  
















\appendix


\section{Further Background}

\subsection{Dual approach to Entropy-regularized Wasserstein DRO}
\label{sec:dual-review}
The Sinkhorn DRO problem, as formulated by \citep{wang2021sinkhorn}, aims to find a decision that minimizes the worst-case expected loss over an ambiguity set defined by Sinkhorn divergence, which is identical to entropy-regularized Wasserstein distance in this problem. Given a loss function $\ell(\theta, z)$, a nominal distribution $\widehat{\rho_N}$, and a radius $r > 0$, the primal inner problem is:
\begin{equation}
\sup_{\rho\in \calB_{r, \epsilon}(\widehat{\rho_N})} \E_{z \sim \rho}[\ell(\theta, z)],
\label{eq:appendix-primal}
\end{equation}
where the ambiguity set is $\calB_{r, \epsilon}(\widehat{\rho_N}) := \{ \rho \in \calP(\calZ) : W_{\epsilon}^2(\widehat{\rho_N}, \rho) \le r \}$, and the entropy-regularized Wasserstein distance,
\[W^2_{\epsilon}(\rho_1, \rho_2) = \inf_{\gamma\in \Gamma(\rho_1,\rho_2)} \{\E_{(x,y)\sim\gamma}[c(x,y)] + H(\gamma)\}.\]

A dual approach, established by \citep{wang2021sinkhorn}, reformulates the inner maximization problem. Using the strong duality of this problem and substituting \eqref{eq:worst-dist} into this problem, the inner problem can be expressed as a minimization over a dual variable $\tau$:
\begin{equation}
\inf_{\tau \ge 0} \left\{ \frac{r}{2\tau} +\frac{\epsilon}{2\tau} \E_{x\sim\widehat{\rho_N}} \left[ \log \E_{y\sim\rho_{x,\epsilon}} \left[ \exp\left(\frac{2\tau\ell(\theta, y)}{\epsilon}\right) \right] \right] \right\},
\label{eq:appendix-dual}
\end{equation}
where $\rho_{x,\epsilon}(y) = \frac{\exp(\frac{-c(x,y)}{\epsilon})}{Z}$ is a kernel distribution centered at $x$, and $Z$ is a normalizing constant.

To approximate the value and the gradient of the nested expectation in \eqref{eq:appendix-dual}, \citep{wang2021sinkhorn} employ a two-level sampling procedure for a given $\tau$: first, a sample $x$ is drawn from the nominal distribution $\widehat{\rho_N}$; second, a set of samples are drawn from the kernel distribution $\rho_{x,\epsilon}$ to approximate the inner expectation.

However, a major limitation of this dual method is that the nested expectation structure in the objective leads to a biased subgradient estimator when using standard Monte Carlo sampling, especially for small values of $\epsilon$. This bias can impede the convergence of the optimization procedure. This challenge motivates the exploration of alternative approaches, such as the gradient flow methods discussed in this paper, which can directly sample from the primal worst-case distribution to obtain a less biased estimate of the gradient.

\begin{remark}
    Consider the KL-DRO problem:
    \begin{equation}
    \min_\theta \sup_{\rho \ll \rho_0} \left\{ \mathbb{E}_{y \sim \rho}[\ell(\theta, y)] - \frac{1}{2\tau} \KL(\rho \| \rho_0) \right\} \label{eq:dro_objective_modified}
    \end{equation}
    Let us define a reference distribution $\rho_0(y)$ as a kernel density estimator of the empirical distribution $\widehat{\rho_N} = \frac{1}{N}\Sigma_{i=1}^N \delta(x_i)$: $\rho_0(y) = \frac{1}{N}\Sigma _{i=1}^N K_\epsilon(y,x_i)$, where $K_\epsilon(y,x_i)= \alpha_\epsilon \exp(-c(y,x_i)/\epsilon)$ and $\alpha_\epsilon = \E[\exp(-c(y,x_i)/\epsilon)]^{-1}$. With this choice, and by setting $\tau = \tau'/\epsilon$, the solution to the inner maximization problem is given by \citep{hu2013kullback}:
    \begin{align*}
        \rho^{\KL}(y) &= \rho_0(y) \cdot \frac{\exp(2\tau'\ell(\theta, y)/\epsilon)}{\mathbb{E}_{y \sim \rho_0}[\exp(2\tau'\ell(\theta, y)/\epsilon)]} \\
    &= \frac{1}{N}\Sigma _{i=1}^N \alpha_\epsilon \beta \exp((2\tau'\ell(\theta, y)- c(y,x_i))/\epsilon),
    \end{align*}
    where $\beta = \E_{y\sim \rho_0}[\exp(2\tau'\ell(\theta, y)/\epsilon)]^{-1}$. This resulting distribution bears a strong resemblance to the worst-case distribution for entropy-regularized DRO. However, they are not identical, as the normalization constants differ: $\alpha_\epsilon \beta \neq \alpha_x:= \E[\exp((2\tau'\ell(\theta, y)- c(y,x))/\epsilon]^{-1}$. In fact, the term $\alpha_\epsilon \beta \exp((2\tau'\ell(\theta, y)- c(y,x))/\epsilon)$ does not integrate to one and thus does not define a valid conditional probability measure. Actually, $\rho^{\KL}(y)$ can be interpreted as a weighted expectation of conditional distributions:
    \[
    \rho^{\KL} = \E_{x\sim\widehat{\rho_N}}[\frac{\alpha_\epsilon \beta}{\alpha_x} \rho_{Y|X=x}]
    \]
    Consequently, the two-layer sampling procedure requires computing the corresponding weights, sampling from the resulting weighted empirical distribution, and subsequently sampling from the conditional distribution. However, the term $\alpha_x$ is intractable in practice. Therefore, unlike the entropy-regularized DRO case, the KL-DRO problem cannot be directly framed as a two-level sampling task.
\end{remark}

\subsection{Hellinger-Kantorovich a.k.a. Wasserstein-Fisher-Rao Gradient Flows}

Consider the WFR gradient system of the energy functional $F$, i.e., the triple $\left( \calP, F, \WFR \right)$, where $\WFR$ is the Wasserstein-Fisher-Rao metric, a.k.a. the \emph{Hellinger-Kantorovich} metric restricted to the probability space.
Specifically, we consider the HK/WFR gradient system associated with reaction-diffusion PDE:
\begin{align}
    {{\partial_t} \mu}
    &= \alpha\, \mathrm{div} \left( \mu \nabla \frac{\delta F}{\delta \mu} \right) - \beta\, \mu 
    \left(\frac{\delta F}{\delta \mu} - \int \frac{\delta F}{\delta \mu}\dd \mu \right)
    .
\end{align}
The hope of this gradient flow is that the added reaction term on the right-hand side will help the gradient flow to converge more rapidly.
This can be formally seen by checking the time derivative of the energy functional $F$ along the gradient flow \eqref{eq:wfr-gradient-flow}, we have
\begin{align*}
    {\partial_t} F (\rho_t)
    \overset{\text{(EDB)}}{=}
    -\alpha \int \Bigg|\nabla  \log \frac{\rho}{\pi}(x) \Bigg| ^2\dd \rho (x) 
    -\beta \int \Bigg| \log \frac{\rho}{\pi}(x)
    -
    \KL(\rho | \pi)
    \Bigg| ^2\dd \rho (x) 
    \end{align*}
where we used the Energy-Dissipation Balance (EDB, a.k.a. equality) of gradient flows.
We observe the right-hand side is non-positive, hence its dissipation will be faster than the original pure Wasserstein gradient flow.

However, we must note that there is no global PL or strong convexity analogous to the case of the Wasserstein gradient flow, as discussed in the main text.
A ``warm-start'' condition is necessary as it has been shown that global PL/gradient dominance cannot hold for the Fisher-Rao/(spherical)Hellinger gradient flows; see \citep{mielke2025hellinger,carrillo2024fisher} for details.
There exist fine-grained warm-start initialization conditions in the gradient flow literature; see \citep{luBirthdeathDynamicsSampling2023,luAcceleratingLangevinSampling2019,chen2023sampling}.

For example, \citep{luBirthdeathDynamicsSampling2023}
showed that, after a warm-start initialization in the form of a density-ratio lower bound, it is possible to get an improved exponential convergence rate
than the pure Wasserstein gradient flow as characterized in Lemma~\ref{lemma:gradient-dominance-PL-Wasserstein}.
However, it is important to note that
the specific warm-start initialization condition cannot be verified in our DRO problems.
Nonetheless, in practice, we observe significant speedups using the WFR gradient flow as compared to the pure Wasserstein gradient flow.

\section{Further Methods}

\subsection{Using SVGD to Sample from the Worst-Case Distribution}\label{sec:svg}
Stein Variational Gradient Descent (SVGD) is a variational inference algorithm that approximates a target probability density $p(x)$ by iteratively transporting a set of particles $\{ x_i \}_{i=1}^n$ \citep{liu2016stein}. The method is formulated as a functional gradient descent on the Kullback-Leibler (KL) divergence, $\text{KL}(q||p)$, where $q$ represents the empirical distribution of the particles. The particle updates are governed by a velocity field $\phi(x_i)$ that corresponds to the direction of steepest descent. For an empirical measure of $n$ particles, this update is given by:
\begin{equation*}
\boldsymbol{\phi}(x_i) \propto \frac{1}{n} \sum_{j=1}^n \left[ k(x_j, x_i) \nabla_{x_j} \log p(x_j) + \nabla_{x_j} k(x_j, x_i) \right]
\end{equation*}
This update rule can be decomposed into two functional components: (i) a driving force, which is a kernel-smoothed average of the score function $\nabla \log p(x)$ that directs particles towards the modes of the target distribution, and (ii) a repulsive force, which arises from the kernel gradient $\nabla k(\cdot, \cdot)$ and ensures particle diversity to prevent mode collapse \citep{ba2021understanding}.

The iterative application of this update rule forms the basis of the SVGD algorithm. As detailed in \citep{liu2016stein}, its computational complexity is dominated by the pairwise kernel computations, resulting in a cost of $O(m^2d)$ per iteration for $m$ particles in $d$ dimensions. Regarding convergence, theoretical guarantees have been established under certain assumptions. As shown by \citep{salim2022convergence}, for target distributions that satisfy Talagrand's T1 inequality, a finite-iteration complexity bound is derived. A finite-particle convergence analysis is also provided in \citep{shi2023finite}.

By applying SVGD to sample from the worst-case distribution, we arrive at Algorithm \ref{alg:SDRO-SVG}.

\begin{algorithm}[h!]
\caption{Sinkhorn DRO via SVGD}
\label{alg:SDRO-SVG}
\begin{algorithmic}[1]
  \State \textbf{Input:} Empirical distribution $\widehat{\rho_N}$, constraint set $\Theta$, $\tau, \epsilon >0$, stepsize sequence $\{r_s > 0\}_{s=0}^{S-1}$, inner stepsize $\eta$, inner iterations $T$, initial deviation $\sigma$, number of samples $m$.
  \For{$s = 0, \dots, S-1$}
    \State Sample $x^s \sim \widehat{\rho_N}$
    \State Sample $\xi^{s,i}\sim\mathcal{N}(0,\sigma I) ,\, i=1,...,m$
    \State Initialize $y^{s,i}_0 \leftarrow x^s + \xi^{s,i}$
    \For{$t = 0, \dots, T-1$}
        \State $\phi^{s,i}_t = \frac{1}{m}\sum_{j=1}^m\left[k(y^{s,i}_t, y^{s,j}_t)\nabla\widetilde V_{x^s, \tau}(y^{s,j}_t) + \nabla_2 k(y^{s,i}_t, y^{s,j}_t) \right]$
        \State $y^{s,i}_{t+1}= y^{s,i}_t + \phi^{s,i}_t$
    \EndFor
    \State $\theta^{s+1} \leftarrow \text{Proj}_\Theta (\theta^s-r_s \sum_{i=1}^m \frac{1}{m}\nabla_\theta \ell(\theta^{s},y^{s,i}_T))$
  \EndFor
  \State \textbf{return} $\theta^S$
\end{algorithmic}
\end{algorithm}

However, the empirical performance of Algorithm~\ref{alg:SDRO-SVG} is suboptimal. We observe that the convergence is sensitive to the initialization of the particles. Furthermore, when the number of particles $m$ is small, the magnitude of the driving force to dominate the repulsive force, thereby diminishing the variance. In this regime, the SVGD algorithm degenerates into a simple gradient descent method, losing its sampling capabilities. This behavior is demonstrated empirically in Section~\ref{sec:experiment3}.


\subsection{Using RGO to Sample form the Worst-case Distribuion}\label{sec:rgo}

In this section, we explore an alternative approach for sampling from the worst-case distribution based on Proximal Sampler. The \textit{Proximal Sampler} \citep{lee2021structured} was introduced to unbiasedly sample from Gibbs distributions of the form $\nu(x) \propto \exp(-f(x))$. A single iteration of the proximal sampler consists of the following two steps:
\begin{enumerate}
    \item (\textbf{Forward step}): Sample $y_k \mid x_k \sim \nu_\tau^{Y|X}(\cdot \mid x_k) = \mathcal{N}(x_k, \tau I)$. This results in a new iterate $y_k \sim \rho^Y_k$, where $\rho^Y_k = \rho^X_k * \mathcal{N}(0, \tau I)$.
    \item (\textbf{Backward step}): Sample $x_{k+1} \mid y_k \sim \nu_\tau^{X|Y}(\cdot \mid y_k)$. This gives the next iterate $x_{k+1} \sim \rho^X_{k+1}$.
\end{enumerate}


As shown in recent work \citep{chen2022improved}, this forward-backward procedure is equivalent to an entropy-regularized Jordan-Kinderlehrer-Otto (JKO) scheme :
\begin{align}\label{eq:forward_app}
   \rho_k^Y &= \arg\min_{\mu \in P_2(\mathbb{R}^d)} \left\{ \frac{1}{2\tau} W^2_{2\tau}(\rho_{k}^{X}, \mu) \right\}, \\
   \rho_{k+1}^{X} &= \arg\min_{\mu \in P_2(\mathbb{R}^d)} \left\{ \int f \, d\mu + \frac{1}{2\tau} W^2_{2\tau}(\rho_{k}^{Y}, \mu) \right\}, \label{eq:backward_app}
\end{align}
This connection to the entropy-regularized JKO scheme provides a new perspective to our problem. By reformulating our original Lagrangian problem (\eqref{eq:main-ent-dro}) as a minimization and specializing the cost to the squared Euclidean distance, we obtain:
\begin{equation}\label{pb:sinkhornscheme_modified}
    \Phi(\theta)=\min_{\rho\in\calP} \{\E_{y\sim\rho}[-\ell(\theta, y)] + \frac{1}{2\tau} W^2_{\epsilon}(\rho,\widehat{\rho_N})\},
\end{equation}
which precisely matches the form of the problem \eqref{eq:backward_app} solved by backward step. This structural equivalence suggests that methods developed for proximal sampler can be directly applied to solve the entropy-regularized DRO problem. Specifically, we can employ the \textit{Restricted Gaussian Oracle (RGO)} from \citep{lee2021structured} to solve this. The sampling process for a given $x \sim \widehat{\rho_N}$ involves two steps:
\begin{enumerate}\label{alg:rejection-sampling}
    \item Sample $x$ from $\widehat{\rho_N}$, and compute the minimizer \( y^*_{x, \tau} = \arg\min_{y \in \mathbb{R}^d} \{\widetilde V_{x, \tau}(y)\} \).
    \item Repeat until acceptance: draw a sample \( Z \sim \mathcal{N} \left( y^*_{x, \tau}, \frac{\epsilon}{2(1 - L\tau)} I \right) \) and accept it with probability
    \[
    \exp \left(- \widetilde V_{x, \tau}(Z) +\widetilde V_{x, \tau}(y^*_{x, \tau}) + \frac{2(1 - L\tau)}{\epsilon} \| Z - y^*_{x, \tau} \|^2 \right).
    \]
\end{enumerate}


The validity of the RGO sampler requires the loss function $\ell$ to be $L$-smooth, and the penalty parameter $\tau$ must satisfy $\tau < 1/L$. This condition ensures that the auxiliary function $\widetilde V_{x, \tau}(y)$ is $(\frac{2(1-L\tau)}{\epsilon})$-strongly convex, which is crucial for rejection sampling. Applying RGO to the entropy-regularized DRO problem necessitates that the loss function is $L$-smooth and requires solving a convex optimization subproblem in each step. These requirements on the loss function and the computational structure are analogous to the WRM algorithm presented in \citep{sinha2020certifyingdistributionalrobustnessprincipled} for the specific case of a squared Euclidean cost. However, a key difference exists: the WRM algorithm transports each data point to a single worst-case point, whereas in entropy-regularized DRO, each data point induces a full continuous distribution of worst-case scenarios. 
\begin{algorithm}
\caption{Sinkhorn DRO via RGO}
\label{alg:SDRO_rgo}
\begin{algorithmic}[1] %
  \State \textbf{Input:} Empirical distribution $\widehat{\rho_N}$, constraint sets $\Theta$, $\tau, \epsilon > 0$, stepsize sequence $\{r_s > 0\}^{S-1}_{s=0}$
  \For{$s = 0, ..., S-1$} 
    \State Sample $x^s \sim \widehat{\rho_N}$ and find an $\eta$-approximate minimizer $\hat{y}^s$ of $\widetilde{V}_{\tau,x^s}(y)$
    \State Generate samples $\{y^{s,i}\}_{i=1}^m$ via rejection sampling
    \State $\theta^{s+1} = \text{Proj}_\Theta (\theta^s-\frac{r_s}{m}\Sigma^m_{i=1} \nabla_\theta \ell(\theta^{s}, y^{s,i}))$
    \EndFor
  \State \textbf{return} $\theta^S$
\end{algorithmic}
\end{algorithm}


\begin{remark}
Although WRM\citep{sinha2020certifyingdistributionalrobustnessprincipled} assumes an $L$-smooth loss, its algorithm is “parameter-free” – it never needs to know $L$ and works whenever the inner optimization problem admits a closed-form solution. In contrast, Algorithm \ref{alg:SDRO_rgo} requires explicit knowledge of $L$ and cannot be applied if the loss fails to be $L$-smooth. In practice, we may assume $\tau$ is sufficiently small such that $\tau < 1/L$, and $\ell$ is $L$-smooth for all $y$.
\end{remark}


\section{Proofs of Theoretical Results} \label{ap:complexity_proofs}

\subsection{Proof of Lemma \ref{lem:sb-klform}}
\begin{proof}[Proof of Lemma \ref{lem:sb-klform}]
    We rewrite the problem using the definition of the entropy-regularized OT formulation:
    \begin{align}
    \min_{\rho} \min _{\Pi}
    \biggl\{
    \int V \dd \rho + \frac{1}{2\tau }
     \int c(x,y) \dd\Pi + \frac\epsilon{2\tau}\int \log \Pi \dd\Pi
    \biggr\}
    \end{align}
    subject to the margianl constraints $\int \Pi \dd x = \rho$ and $\int \Pi \dd y = \rho_0$.
    Here,
    $\rho_0$ can be set to empirical distribution $\widehat{\rho_N}$ as in the data-driven DRO problem.
    Using the marginal constraints, 
    we can then write the first linear term as
    $\int \int V(y) \dd \Pi(x, y)$ and eliminate the variable $\rho$, i.e., one end is unconstrained -- hence it's the half bridge problem:
    \begin{align}
         \min _{\Pi}
        \biggl\{
        \int V \dd \Pi + \frac{1}{2\tau }
         \int c(x,y) \dd\Pi + \frac\epsilon{2\tau}\int \log \Pi \dd\Pi
         \, \,
         \bigg |
         \int \Pi \dd y =  \rho_0
         \biggr\}
         .
        \end{align}
    After straightforward linear optimization,
    we obtain the optimal solution to the variational problem~\eqref{eq:proximal-EOT} as
    \begin{align}
        \rho^* \propto \E_{x\sim \rho_0} \left[
            \exp
            \left[{-\frac{2\tau V(y) +c(y, x^i_k) }{\epsilon}} \right]    
         \right],
    \end{align}
    which is a mixture of proximal distributions \eqref{eq:conditional-sampling-density}.
    Using elementary manipulations, the variational problem \eqref{eq:proximal-EOT}
    can also be written down as the minimization of the expected KL divergence:
    \begin{align}
        \min_{\rho_{Y|X}}
    \mathbb{E}_{x\sim\rho_0}
    		\KL
            \left(
            \rho_{Y|X=x}(y) 
            \bigg| 
    		\dfrac{1}{Z_x}
            \exp \left[
    					-\dfrac{2\tau V(y) +  c(y,x)}{\epsilon} 
    		\right]
            \right)
            .
    \end{align}
\end{proof}

\subsection{Proof of Proposition~\ref{prop:gradient_oracle_error_control}}
First, we recall some standard results.
\begin{lemma}
    \label{lemma:displacement-convexity}
    Suppose the first condition in Definition~\ref{as:gf-properties} holds for some $\lambda > 0$. Then,
    $F$ is $\lambda$-displacement convex.
\end{lemma}

\begin{lemma}[WGF gradient dominance/PL]
    \label{lemma:gradient-dominance-PL-Wasserstein}
    Let $\rho_t$ be a solution to the Wasserstein gradient flow of the KL divergence energy functional $\KL(\rho | \exp(-\widetilde V_{x, \tau}))$. Then,
\begin{align*}
    \text{$\widetilde V$ is $\lambda$-displacement convex or $\lambda$-PL}
    \implies
    \KL(\rho_t| \exp(-\widetilde V_{x, \tau})) 
    \leq
    \ee^{-2\lambda t} \KL(\rho_0| \exp(-\widetilde V_{x, \tau})).
\end{align*}
\end{lemma}

\begin{lemma}[Talagrand's inequality]
    Under the same assumptions as in Lemma \ref{lemma:gradient-dominance-PL-Wasserstein}, we have
    \label{lemma:talagrand-inequality}
\[
W_2^2(\rho_t, \exp(-\widetilde V_{x, \tau})) \leq
\sqrt{\frac2\lambda \KL(\rho_t | \exp(-\widetilde V_{x, \tau}))} 
\]
\end{lemma}
Recall also that $W_1(\rho_t, \rho^*_{Y|X=x}) \leq W_2(\rho_t, \rho^*_{Y|X=x})$.


\textbf{From conditional to marginal bound}
We have the following relation between the marginal and conditional KL divergences.
Suppose $\rho_X$ is the marginal distribution of $X$.
Recall the relation $\pi_Y = \int_x \rho^* _{Y|X=x} \dd \rho_X(x)$.
Then,
\[
\KL(\rho_Y | \pi_{Y}) \leq \mathbb{E}_{X \sim \rho_X} \left[ \KL\left(\rho_{t, Y|X}  \middle| \rho^*_{Y|X}\right) \right].
\]

Alternatively, the Wasserstein relation gives
\[
W_p(\rho_Y, \pi_Y) \leq \mathbb{E}_{X \sim \rho_X} \left[ W_p(\rho_{Y|X}(\cdot|X), \rho^*_{Y|X}(\cdot|X)) \right].
\]

\textbf{Marginal distribution error control}
Suppose we have the samples
$y^i_t\sim \rho_t$,
Our goal is to use the empirical measure
$\displaystyle\widehat{\rho_Y} = \frac1N \sum_{i=1}^N \delta_{y^i_t}$
to approximate $\pi_Y$.
An elementary manipulation shows that
    \begin{align}
        W_1(\widehat{\rho_Y}, \pi_Y) 
        &\leq \mathbb{E}_{X \sim \rho_X} \left[ W_1(\widehat\rho_{Y|X}, \rho^*_{Y|X}) \right]
        \lesssim \frac{{\ee^{-\lambda t}}}{\sqrt{\lambda}}
        \label{eq:marginal_w1_decay}
    \end{align}

Therefore, to reach an $\epsilon$-solution, we need to run the gradient flow (i.e. simulating the PDE/SDE) for time
\begin{align}
    t \gtrsim \frac{1}{\lambda}\log{\frac1{\sqrt{\lambda}\epsilon}}
    .
    \label{eq:gf_time_to_epsilon}
\end{align}

We will later show the discrete-time gradient descent version of this result.

We emphasize that the continuous-time gradient flow result is important not only for understanding the limit of discrete-time algorithms, but also for serving as a guide for designing new discrete-time algorithms.
Let us get a quick insight from the above estimate:
suppose we discretized algorithm with discretization step size $\eta$ for in total $T$ steps. Ignoring the discretization error, we have $t = T\eta$ and the (ideal) number of iterations needed to reach an $\epsilon$-solution
of the inner maximization problem~\eqref{eq:entropic-jko}
is $T \gtrsim \frac{1}{\eta  \lambda}\log{\frac1{\sqrt{\lambda}\epsilon}}$.

\textbf{DRO gradient oracle error control via gradient flow}
Let us now show how to convert the error estimate in distributions to the error estimate in the gradient oracle, needed for the optimization analysis later.
If $W_1(\widehat{\rho_Y}, \pi_Y) \leq \delta$ for some $\delta > 0$ (as controlled by \eqref{eq:marginal_w1_decay}) and $g$ is $L$-Lipschitz (w.r.t. Euclidean norm).
Then,
taking expectations and using the optimal transport plan
(coupling) $\gamma$ of $(X, Y) \sim (\widehat{\rho_Y}, \pi_Y)$,
we obtain:
\begin{align}
    \label{eq:gradient_oracle_error_control_expectation}
    \left| \int g\, d(\widehat{\rho_Y} - \pi_Y) \right|
=  \left|\int     g(x) - g(y)\dd\gamma(x, y)\right|
\leq L W_1(\widehat{\rho_Y}, \pi_Y)
\leq L\delta.
\end{align}

In the optimization analysis, we can later take $g$ to be the gradient oracle of the DRO loss, $g:=\nabla_\theta \ell(\theta, \cdot)$, to control the gradient oracle error.
Recall that $\widehat{\rho_Y} = \frac1N \sum_{i=1}^N \delta_{y^i_t}$ is the particle approximation obtained by running the gradient flow sampler for $t$ steps, and $\pi_Y$ is the target marginal distribution, which is the worst-case distribution of DRO inner maximization problem~\eqref{eq:entropic-jko}.
Then, 
\eqref{eq:gradient_oracle_error_control_expectation} results in the following
bound on the gradient oracle error:
\begin{multline}
    \label{eq:gradient_oracle_error_control}
\bigg|\frac1N \sum_{i=1}^N \nabla_\theta \ell(\theta, y^i_t) 
-
\nabla_\theta 
\underbrace{
\max _{\rho \in \calP} 
\biggl(
    \int \ell(\theta, z)\dd\rho{(z)}
    -\frac1{2\tau} W_\epsilon^2(\rho, \widehat{\rho_N})
\biggr)
}_{\text{DRO objective}}
\bigg|
\\
=
\bigg|\frac1N \sum_{i=1}^N \nabla_\theta \ell(\theta, y^i_t) 
-
\mathbb{E}_{y \sim \pi_Y}
\nabla_\theta\ell(\theta, y)
\bigg|
\overset{\eqref{eq:gradient_oracle_error_control_expectation}}{\leq}
L\delta.
\end{multline}

Replacing $L\delta$ by $\epsilon$, we obtain the desired bound on the gradient oracle error as in Proposition~\ref{prop:gradient_oracle_error_control}.







\subsection{Proof of Theorem \ref{thm:outer_loop} (Outer Loop Convergence)}
\label{ap:proofth1}
We provide a detailed proof for the convergence of the outer loop, which follows the standard analysis for non-convex Stochastic Gradient Descent with a persistent bias.

\begin{lemma}[Bias of the Stochastic Gradient]
Under Assumptions \ref{as:totalsmooth}, \ref{as:fsmooth}, the bias of the stochastic gradient estimator, defined as $B(\theta) := \E_{} [\widehat g] - \nabla\Phi(\theta)$, is bounded as:
$$\|B(\theta)\| \le L_f \cdot \delta_{\text{sample}}.$$
\end{lemma}
\begin{proof}
By definition, the bias is
\begin{align*}
    B(\theta) &= \E_{x \sim \widehat{\rho_N}}\left[\E_{y \sim \widehat\rho_{Y|X=x}}[\nabla_\theta \ell(\theta, y)]\right] - \E_{x \sim \widehat{\rho_N}}\left[\E_{y \sim \rho^*_{Y|X=x}}[\nabla_\theta \ell(\theta, y)]\right] \\
    &= \E_{x \sim \widehat{\rho_N}}\left[ \E_{y \sim \widehat\rho_{Y|X=x}}[\nabla_\theta \ell(\theta, y)] - \E_{y \sim \rho^*_{Y|X=x}}[\nabla_\theta \ell(\theta, y)] \right].
\end{align*}
Taking norms and using Jensen's inequality:
$$\|B(\theta)\| \le \E_{x \sim \widehat{\rho_N}}\left[ \| \E_{y \sim \widehat\rho_{Y|X=x}}[\nabla_\theta \ell(\theta, y)] - \E_{y \sim \rho^*_{Y|X=x}}[\nabla_\theta \ell(\theta, y)] \| \right].$$
For any two probability measures $\rho_1, \rho_2$ and a function $h$ that is $L_h$-Lipschitz, we know that $\|\E_{\rho_1}[h(y)] - \E_{\rho_2}[h(y)]\| \le L_h W_1(\rho_1, \rho_2)$. Applying this with $h(y) = \nabla_\theta \ell(\theta, y)$, which is $L_f$-Lipschitz in $y$ by Assumption 2, gives:
\begin{align*}
    \|B(\theta)\| &\le \E_{x \sim \widehat{\rho_N}}\left[ L_f \cdot W_1(\widehat\rho_{Y|X=x}, \rho^*_{Y|X=x}) \right] \\
    &\le L_f \cdot \E_{x \sim \widehat{\rho_N}}\left[ W_2(\widehat\rho_{Y|X=x}, \rho^*_{Y|X=x}) \right] \quad (\text{since } W_1 \le W_2) \\
    &\le L_f \cdot \delta_{\text{sample}}.
\end{align*}
\end{proof}

\noindent\textbf{Proof of Theorem \ref{thm:outer_loop}.}
From the $L_\Phi$-smoothness of $\Phi$ (Assumption \ref{as:totalsmooth}), we have the standard descent lemma. For simplicity, we assume $\Theta = \mathbb{R}^d$ and a constant step size $r_s = r$.
\begin{align*}
    \Phi(\theta^{s+1}) &\le \Phi(\theta^s) + \langle \nabla\Phi(\theta^s), \theta^{s+1} - \theta^s \rangle + \frac{L_\Phi}{2}\|\theta^{s+1} - \theta^s\|^2 \\
    &= \Phi(\theta^s) - r \langle \nabla\Phi(\theta^s), \widehat g^s \rangle + \frac{L_\Phi r^2}{2}\|\widehat{g}^s\|^2.
\end{align*}
Taking expectation conditioned on the filtration $\mathcal{F}_s$:
\begin{align*}
    \E[\Phi(\theta^{s+1})|\mathcal{F}_s] &\le \Phi(\theta^s) - r \langle \nabla\Phi(\theta^s), \E[\widehat{g}^s|\mathcal{F}_s] \rangle + \frac{L_\Phi r^2}{2}\E[\|\widehat{g}^s\|^2|\mathcal{F}_s] \\
    &= \Phi(\theta^s) - r \langle \nabla\Phi(\theta^s), \nabla\Phi(\theta^s) + B(\theta^s) \rangle + \frac{L_\Phi r^2}{2} \left( \|\E[\widehat{g}^s|\mathcal{F}_s]\|^2 + \text{Var}(\widehat{g}^s|\mathcal{F}_s) \right).
\end{align*}
Using the variance bound and the bias definition:
\begin{align*}
    \E[\Phi(\theta^{s+1})|\mathcal{F}_s] &\le \Phi(\theta^s) - r\|\nabla\Phi(\theta^s)\|^2 - r\langle \nabla\Phi(\theta^s), B(\theta^s) \rangle + \frac{L_\Phi r^2}{2} \left( \|\nabla\Phi(\theta^s)+B(\theta^s)\|^2 + \sigma^2 \right).
\end{align*}
Applying Young's inequality to the inner product term:
$$-r\langle \nabla\Phi(\theta^s), B(\theta^s) \rangle \le \frac{r}{2}\|\nabla\Phi(\theta^s)\|^2 + \frac{r}{2}\|B(\theta^t)\|^2.$$
Substituting this in and simplifying with $\|a+b\|^2 \le 2\|a\|^2 + 2\|b\|^2$:
$$\E[\Phi(\theta^{s+1})|\mathcal{F}_s] \le \Phi(\theta^s) - \frac{r}{2}\|\nabla\Phi(\theta^s)\|^2 + \frac{r}{2}\|B(\theta^s)\|^2 + L_\Phi r^2 (\|\nabla\Phi(\theta^s)\|^2 + \|B(\theta^s)\|^2) + \frac{L_\Phi r^2 \sigma^2}{2}.$$
Rearranging terms to isolate $\|\nabla\Phi(\theta^s)\|^2$:
$$r\left(\frac{1}{2} - L_\Phi r\right)\|\nabla\Phi(\theta^s)\|^2 \le \Phi(\theta^s) - \E[\Phi(\theta^{s+1})|\mathcal{F}_s] + \left(\frac{r}{2} + L_\Phi r^2\right)\|B(\theta^s)\|^2 + \frac{L_\Phi r^2 \sigma^2}{2}.$$
Choosing a step size $r \le \frac{1}{4L_\Phi}$, we have $(\frac{1}{2} - L_\Phi r) \ge \frac{1}{4}$ and $(\frac{r}{2} + L_\Phi r^2) \le r$. This gives:
$$\frac{r}{4}\|\nabla\Phi(\theta^s)\|^2 \le \Phi(\theta^s) - \E[\Phi(\theta^{s+1})|\mathcal{F}_s] + r\|B(\theta^s)\|^2 + \frac{L_\Phi r^2 \sigma^2}{2}.$$
Taking total expectation and using Lemma 1, $\|B(\theta^s)\|^2 \le (L_f \delta_{\text{sample}})^2$:
$$\frac{r}{4}\E[\|\nabla\Phi(\theta^s)\|^2] \le \E[\Phi(\theta^s)] - \E[\Phi(\theta^{s+1})] + r (L_f \delta_{\text{sample}})^2 + \frac{L_\Phi r^2 \sigma^2}{2}.$$
Summing from $s=0$ to $S-1$ yields a telescoping sum:
$$\frac{r}{4}\sum_{s=0}^{S-1}\E[\|\nabla\Phi(\theta^s)\|^2] \le \Phi(\theta^0) - \E[\Phi(\theta^S)] + Sr(L_f \delta_{\text{sample}})^2 + \frac{S L_\Phi r^2 \sigma^2}{2}.$$
Dividing by $\frac{r S}{4}$ and using $\E[\Phi(\theta^S)] \ge \Phi_{\inf}$:
$$\frac{1}{S}\sum_{s=0}^{S-1}\E[\|\nabla\Phi(\theta^s)|^2] \le \frac{4(\Phi(\theta^0) - \Phi_{\inf})}{r S} + 4(L_f \delta_{\text{sample}})^2 + 2L_\Phi r \sigma^2. $$




\subsection{Derivation of ULA Sampler Complexity (for Theorem~\ref{thm:ula})}
\label{ap:ula_derivation}
The complexity of the ULA sampler depends on the geometric properties of the inner target distribution $\rho_{Y|X=x}$.
Under Assumption \ref{as:LSI}, standard results on ULA convergence \citep{vempala2019rapid} state that after $T$ steps, the KL-divergence to the target is bounded. To achieve a final KL-divergence of $\delta_{KL}$, the number of iterations required is $T = O\left(\frac{L_U^2 d}{\alpha_U^2 \delta_{KL}}\log \frac{1}{\delta_{KL}}\right) = \tilde{O}\left(\frac{L_U^2 d}{\alpha_U^2 \delta_{KL}}\right)$.

Our goal is to connect the required outer-loop sampling accuracy $\delta_{\text{sample}}$ (in $W_2$ distance) to the required inner-loop KL-divergence accuracy $\delta_{KL}$. Under the LSI assumption, Talagrand's inequality gives the relationship:
$$ W_2^2(\widehat \rho_{Y|X=x} , \rho_{Y|X=x}) \le \frac{2}{\alpha_U} \KL(\widehat \rho_{Y|X=x} || \rho^*_{Y|X=x}) $$
From Theorem \ref{thm:outer_loop}, we require $W_2(\widehat \rho_{Y|X=x}, \rho_{Y|X=x}) \le \delta_{\text{sample}} = O(\epsilon_{\text{opt}}/L_f)$. This implies we need to achieve a KL-divergence of:
$$ \delta_{KL} \le \frac{\alpha_U}{2} W_2^2(\widehat \rho^*_{Y|X=x}, \rho_{Y|X=x}) = O\left(\alpha_U \frac{\epsilon_{\text{opt}}^2}{L_f^2}\right). $$
Substituting this required $\delta_{KL}$ into the ULA iteration complexity gives the number of inner loop steps:
$$ T_{ULA} = \tilde{O}\left(\frac{L_U^2 d}{\alpha_U^2 \cdot \alpha_U \frac{\epsilon_{\text{opt}}^2}{L_f^2}}\right) = \tilde{O}\left(\frac{L_U^2 L_f^2 d}{\alpha_U^3 \epsilon_{\text{opt}}^2}\right). $$
The total complexity is the product of outer iterations $S = O(1/\epsilon_{\text{opt}}^2)$, inner iterations $T_{ULA}$, and the cost per inner gradient step $C_{\nabla_z} = O(d)$.
\begin{align*}
\text{Complexity} &= S \times T \times C_{\nabla_{z}} \\
&= O\left(\frac{1}{\epsilon_{\text{opt}}^{2}}\right) \times \tilde{O}\left(\frac{L_U^2 L_f^2 d}{\alpha_U^3 \epsilon_{\text{opt}}^2}\right) \times O(d) \\
&= \tilde{O}\left(\frac{L_U^2 L_f^2 d^2}{\alpha_U^3 \epsilon_{\text{opt}}^4}\right).
\end{align*}
Note that we absorb constants like $L_{\Phi}$ into the $\tilde{O}(\cdot)$ notation. This completes the derivation for Theorem \ref{thm:ula}.

\subsection{Complexity of RGO}
We analyze the computational complexity of RGO method under specific regularity conditions on the loss function.

\begin{theorem}[Complexity of RGO Sampling]
\label{thm:rgo_complexity}
Let the loss function $\ell: \mathbb{R}^d \to \mathbb{R}$ be $L$-smooth. For parameter $\tau = \frac{1}{Ld}$ and a target distribution $\rho^*_{Y|X=x}$, though the rejection sampling procedure descrided in Section \ref{sec:rgo}, we can draw a sample from a distribution $\tilde{\rho}_{Y|X=x}$ such that the Kullback-Leibler (KL) divergence satisfies $\KL(\tilde{\rho}_{Y|X=x} || \rho^*_{Y|X=x}) < \delta$ with the iteration-complexity bound: 
\[O(\log\left(\frac{1}{\epsilon\delta}\right)) \]
\end{theorem}

\begin{proof}
Under the assumption that $\ell$ is $L$-smooth and $\tau < \frac{1}{L}$, $\widetilde V_{x, \tau}(y)$, is $\frac{2-2\tau L}{\epsilon}$-strongly convex and $\frac{2+2\tau L}{\epsilon}$-smooth.

Let the minimizer of $\widetilde V_{x, \tau}(y)$ as $y^*_{x, \tau}$. To find the minimum of $\widetilde V_{x, \tau}(y)$, we can apply gradient descent. The number of iterations required to achieve a $\delta'$-approximate optimal solution $\tilde{y}$ is determined by the condition number $\kappa = \frac{1+\tau L}{1-\tau L}$. Specifically, the iteration complexity is:
$$ O\left(\kappa \log\left(\frac{1}{\delta'}\right)\right) = O\left(\frac{1+\tau L}{1-\tau L}\log\left(\frac{1}{\delta'}\right)\right) $$


\textbf{2. Distributional Proximity.}
Let the distribution obtained via rejection sampling with a $\delta'$-approximate optimal solution for the proposal be denoted by $\tilde{\rho}_{Y|X=x}$. The procedure defines the proposal distributions $p(y)$ (using the true minimizer $y^*_{x,\tau}$) and $\tilde{p}(y)$ (using the approximate minimizer $\tilde{y}$) as Gaussian distributions:
\[ p(y) = \mathcal{N} \left( y \mid y^*_{x,\tau}, \frac{\epsilon}{2 - 2\tau L} I \right), \quad \tilde{p}(y) = \mathcal{N} \left( y \mid \tilde{y}, \frac{\epsilon}{2 - 2\tau L} I \right) \]
The corresponding acceptance probabilities $a(y)$ and $\tilde{a}(y)$ are given by:
\begin{align*}
    a(y) &= \exp \left( - \widetilde V_{x, \tau}(y) + \widetilde V_{x, \tau}(y^*_{x,\tau}) + \frac{1 - \tau  L }{\epsilon} \| y - y^*_{x,\tau} \|^2 \right) \\
    \tilde{a}(y) &= \min\left\{\exp \left( - \widetilde V_{x, \tau}(y) + \widetilde V_{x, \tau}(\tilde{y}) + \frac{1 - \tau  L }{\epsilon} \| y - \tilde{y} \|^2 \right), 1\right\}
\end{align*}
The resulting KL divergence from the distribution obtained with the true minimizer, $\rho_{\theta,x}$, is
\begin{align}
    \KL(\tilde{\rho}_{Y|X=x}||\rho_{Y|X=x}) &= \int \frac{\tilde{a}(y)\tilde{p}(y)}{\int \tilde{a}(y)\tilde{p}(y)} \log \left(\frac{\frac{\tilde{a}(y)\tilde{p}(y)}{\int \tilde{a}(y)\tilde{p}(y)}}{\frac{a(y)p(y)}{\int a(y)p(y)}}\right) dy\\
    &= \int \frac{\tilde{a}(y)\tilde{p}(y)}{\int \tilde{a}(y)\tilde{p}(y)} \left(\log \left(\frac{\tilde{a}(y)\tilde{p}(y)}{a(y)p(y)}\right) + \log\left(\frac{\int a(y)p(y)}{\int \tilde{a}(y)\tilde{p}(y) }\right) \right) dy
    \\
    &\leq\int \frac{\tilde{a}(y)\tilde{p}(y)}{\int \tilde{a}(y)\tilde{p}(y)}\cdot 2 \left(\widetilde V_{x, \tau}(\tilde{y})-\widetilde V_{x, \tau}(y^*_{x, \tau}) \right) dy
    \\
    &\leq\int \frac{\tilde{a}(y)\tilde{p}(y)}{\int \tilde{a}(y)\tilde{p}(y)}\cdot \frac{4+4\tau L}{\epsilon} \delta'  dy
    = \frac{4+4\tau L}{\epsilon}\delta' 
\end{align}

To ensure the final KL divergence is less than $\delta$, we must set the optimization accuracy to $\delta' < \frac{\epsilon}{4+4\tau L}\delta$. This implies an optimization complexity of $O\left(\frac{1+\tau L}{1-\tau L}\log\left(\frac{4+4\tau L}{\epsilon\delta}\right)\right)$.

\textbf{3. Rejection Sampling Efficiency.}
The expected number of iterations until acceptance is at most $\left(\frac{1+\tau L}{1-\tau L}\right)^{d/2}$\citep{chewi2022query}. Note that if we choose $\tau = \frac{1}{Ld}$, then the
expected number of iterations until acceptance is at most $\left(\frac{1+\tau L}{1-\tau L}\right)^{d/2}  = \left(\frac{1+1/d}{1-1/d}\right)^{d/2} = O(1)$. 
As $d \to \infty$, this expression converges to $e$. Thus, for this choice of $\tau$, the expected number of rejection sampling trials is $O(1)$.

\textbf{4. Total Complexity.}
The total complexity is the product of the optimization complexity and the expected number of rejection sampling iterations. With $\tau = \frac{1}{Ld}$, the condition number becomes $\kappa = \frac{d+1}{d-1}$.
More precisely, the complexity is:
\[
O\left(\left(\frac{d+1}{d-1}\right)^{\frac{d}{2} + 1} \times\log\left(\frac{4d+4}{d\epsilon\delta}\right)  \right) = O(\log\left(\frac{1}{\epsilon\delta}\right))
\]
This completes the proof.
\end{proof}

\begin{remark}[Practical Limitations]
The theoretical complexity presented in Theorem \ref{thm:rgo_complexity} is highly compelling when compared to alternative sampling methodologies. However, its applicability is constrained by stringent underlying assumptions.
\begin{enumerate}
    \item \textbf{Smoothness Requirement:} The analysis presupposes that the loss function $\ell$ is $L$-smooth. In many practical applications, loss functions are non-smooth or exhibit a very large smoothness constant $L$.
    \item \textbf{Parameter Dependency:} The optimal setting for the parameter, $\tau = \frac{1}{Ld}$, is inversely proportional to both the smoothness constant $L$ and the dimension $d$. If $L$ is large, the resulting $\tau$ will be small. This makes the worst-case distribution very close to the original distribution, thereby limiting the robustness conferred by the method.
\end{enumerate}
Consequently, while the theoretical result is intriguing, the method's performance in empirical settings is often suboptimal due to the difficulty in satisfying these idealized conditions.
\end{remark}


\section{Additional Experiments}
\subsection{Classification under Data Imbalance}



To further validate the findings in \ref{sec:circle-exp}, we conducted an experiment with an increased training set of 2000 samples while keeping all other settings unchanged. With a larger dataset, the decision boundaries of all methods are expected to converge toward the true circular boundary. In this high-data regime, WGF, Dual, and WRM learn similar boundaries that are nearly circular but still exhibit noticeable deviations. WFR, however, learns a tight and circular boundary, demonstrating its consistent superior performance (see Figure~\ref{fig:circle_2000}).
\begin{figure}[h]
    \centering
    \includegraphics[width=0.45\linewidth]{2000.pdf}
    \caption{Decision boundaries on the Biased Circle dataset with 2000 training samples. WFR still learns a more accurate circular boundary compared to other methods.}
    \label{fig:circle_2000}
\end{figure}













\subsection{Uncertain Least Square}
We evaluate our methods on a distributionally robust least squares problem, following the setup from \citep{zhu2021kernel}, which was adapted from \citep{el1997robust}. The objective is to find a parameter vector $\theta \in \mathbb{R}^{10}$ that minimizes the loss $f_{\theta}(\xi) = \|A(\xi)\theta - b\|_2^2$, where the system matrix $A(\xi)$ is subject to uncertainty. The matrix $A(\xi) = A_0 + \xi A_1 \in \mathbb{R}^{10 \times 10}$ is an affine function of an uncertain scalar parameter $\xi \in [-1, 1]$. The matrices $A_0, A_1$ and the vector $b$ are fixed, with entries drawn independently from a standard normal distribution $\mathcal{N}(0, 1)$.

\begin{figure}[h!]
    \centering
    \includegraphics[scale=0.6]{uncertain_least_squares_results.pdf}
    \caption{Test loss as a function of the perturbation level $\Delta$ for the uncertain least squares problem. All methods were trained for 10 epochs with $\lambda=0.1$. For methods with an inner loop, the step size was $0.0001$ for $3000$ iterations. SDRO-based methods used $\epsilon=0.1$. WGF and WFR used $m=8$ particles.}
    \label{fig:uls_results}
\end{figure}
The training set comprises $N=10$ samples $\{\xi_i\}_{i=1}^{N}$ drawn uniformly from $[-0.5, 0.5]$. To evaluate robustness, test samples are drawn from a shifted distribution, specifically uniform on $[-0.5(1+\Delta), 0.5(1+\Delta)]$. We vary the shift magnitude $\Delta$ from $0$ to $10$, where a larger $\Delta$ signifies a greater departure from the training distribution. 



Figure~\ref{fig:uls_results} shows the test loss as a function of the perturbation level $\Delta$. All distributionally robust methods maintain a significantly lower test loss than the empirical risk minimization baseline, demonstrating their robustness to distributional shifts. At low perturbation levels, all robust methods perform comparably. As $\Delta$ increases, the performance of the dual method degrades more rapidly than the others. WFR shows a marginal improvement over WGF and WRM. This limited advantage is attributable to the one-dimensional nature of the inner problem, which constrains worst-case samples to be either -1 or 1, thus minimizing the differences between the generated adversarial distributions. 






\bibliographystyle{apalike}
\bibliography{Bibliography}






\end{document}