Certifying Some Distributional Robustness with Principled Adversarial Training 1: SGD on the Wasserstein Penalty Surrogate Finds an Approximate Stationary Point at Rate 4σ√(L_φΔ_F/T)Research Paper
Motivation
Neural networks trained by empirical risk minimization are sensitive to small, deliberately chosen perturbations of their inputs ("adversarial examples"). A standard defence is adversarial training: at each step, perturb the training example to increase the loss, then take a gradient step at the perturbed point. Sinha, Namkoong, Volpi and Duchi (arXiv:1710.10571, ICLR 2018) recast adversarial training as distributionally robust optimization over a Wasserstein ball around the data distribution, and replaced the hard ball constraint by a Lagrangian penalty. The resulting procedure perturbs each example by solving a penalized inner maximization, and the paper shows that, when the penalty is large enough relative to the smoothness of the loss, this inner problem is strongly concave and therefore tractable.
This mission formalizes the paper's computational guarantee, Theorem 2: stochastic gradient descent on the penalized robust objective, with an inexact adversary, reaches an approximate stationary point at the rate of smooth nonconvex stochastic optimization. The analysis follows Ghadimi and Lan's argument for nonconvex SGD (arXiv:1309.5549), adapted to a gradient oracle that is biased by the inexact inner maximization.
Setting
Parameters live in and data in a set ; both spaces carry the Euclidean norm . Data are drawn from a probability distribution on .
A transportation cost measures how expensive it is to move a data point to . Assumption A: is continuous, , and for each the map is -strongly convex. The canonical example is .
A loss has partial gradients and . Assumption B: these gradients are Lipschitz in each argument,
For a penalty , the robust surrogate and the robust objective are
By Wasserstein duality, equals the worst-case expected loss over all distributions , penalized by times the transport distance from to .
Algorithm 1 (with and constant stepsize ): from , for , draw independently, find an -approximate maximizer , meaning , and set .
Formalization targets
Goal: Theorem 2 (p. 6)
Let , , , and assume for all . With and ,
Milestones
- (4), p. 4: for , is -strongly concave.
- Lemma 1, p. 5: the supremum of a -strongly concave family satisfying Assumption B is differentiable, with , a -Lipschitz maximizer and a -Lipschitz gradient.
- §2.1, p. 6: has -Lipschitz gradient .
- (31) and (32), p. 40: the descent inequality for an -smooth objective, and its form with the bias separated, for .
- The bias bound, p. 41: .
- The one-step bound in conditional expectation and the averaged bound, p. 41:
for every constant stepsize . The goal is this bound at the balancing stepsize.
Significance
The result. Theorem 2 certifies that the robust training procedure is no harder, in iteration count, than ordinary nonconvex SGD: the stationarity measure decays as , and the inexactness of the adversary enters only as an additive term proportional to , independent of . This is what licenses running a few steps of gradient ascent for the inner problem rather than solving it exactly. It also explains the role of the penalty: the guarantee requires , so smoothness of the loss in the data (smooth activations rather than ReLUs) is what makes the robust problem tractable.
Formalizing it. The result and its proof are published; nothing here is open. To our knowledge none of it is machine-checked. A formal proof would add a verified Danskin-type theorem for strongly concave suprema with explicit Lipschitz constants (Lemma 1), which is reusable well beyond this paper, and a verified nonconvex SGD analysis with a biased oracle, including the measure-theoretic bookkeeping (independence of the fresh sample from the current iterate, differentiation under the integral) that the paper's proof compresses into one line.
Difficulty
The deterministic inequalities (31), (32) and the bias bound are short. The weight sits elsewhere. First, Lemma 1 requires differentiability of a supremum over a possibly unbounded set, with the maximizer depending on ; the paper invokes a Danskin-type theorem that is not in Mathlib, and the cost is not assumed differentiable, so the first-order optimality conditions for the inner problem must be derived from strong convexity alone. Second, the expectation step uses , which needs both differentiation under the integral sign for and the independence of from along the run. Third, finiteness of every expectation along the run is not assumed and must be propagated from step to step.
Formalization scope
Lean representation: and are EuclideanSpace ℝ (Fin d) and EuclideanSpace ℝ (Fin m); is a set in the latter. The loss is a function on all of ; its partial gradients are required at points of and the Lipschitz conditions on only. Strong convexity is Mathlib's StrongConvexOn, whose modulus convention matches the paper. Suprema are sSup over the image of . is a Bochner integral, and is assumed -integrable for every . The run is driven by (Measure.pi). The expectations in the conclusion and in the variance hypothesis are lower Lebesgue integrals in ; the -term of Theorem 2 is moved to the right-hand side.
Hypotheses made explicit that the paper leaves implicit: (the paper writes but divides by ); nonempty, convex and closed; Lipschitz constants nonnegative; , , , ; bounded below (so is a real number); integrability of ; ; joint measurability of the oracle ; and, in (32) and the later milestones, , which the paper assumes "in the rest of the proof". Lemma 1 is specialized to the Euclidean norm, the case Theorem 2 uses. The cost is not assumed differentiable.
A trivializing formalization is ruled out: the run uses an oracle tied to by the -gap and to by membership, is the supremum of the penalized loss over , is its -integral, and the hypotheses are satisfiable (for instance , , , a point mass, ). Expectations as lower integrals prevent a non-integrable from being read as .
Infrastructure a complete development needs: the descent lemma for functions with Lipschitz gradient; first-order characterizations of strong concavity on a convex set; a Danskin/envelope theorem for strongly concave suprema; differentiation under the integral for ; and conditional-expectation bookkeeping for Markov-type recursions driven by i.i.d. samples on a finite product space. The first three are reusable in any minimax or bilevel analysis. Contributions to any milestone, or to these supporting facts as separate lemmas, are welcome.
Selected references
- A. Sinha, H. Namkoong, R. Volpi, J. Duchi, Certifying Some Distributional Robustness with Principled Adversarial Training, ICLR 2018; arXiv:1710.10571v5, 2020. https://arxiv.org/abs/1710.10571
- S. Ghadimi, G. Lan, Stochastic First- and Zeroth-Order Methods for Nonconvex Stochastic Programming, SIAM Journal on Optimization 23(4), 2013. https://doi.org/10.1137/120880811
- J. F. Bonnans, A. Shapiro, Perturbation Analysis of Optimization Problems, Springer, 2000. https://doi.org/10.1007/978-1-4612-1394-9
- R. Gao, A. Kleywegt, Distributionally Robust Stochastic Optimization with Wasserstein Distance, Mathematics of Operations Research 48(2), 2023; arXiv:1604.02199. https://arxiv.org/abs/1604.02199