Brenier Meets Adversarial Training: Optimal Transport Geometry for Robust Learning

Alireza Abdollahpoorrostam; Ehsan Sharifian; Buse Şen; Marco Cuturi; Daniel Kuhn

cs.LGSep 25, 2026

Brenier Meets Adversarial Training: Optimal Transport Geometry for Robust Learning

Authors: Alireza Abdollahpoorrostam, Ehsan Sharifian, Buse Şen, Marco Cuturi, Daniel Kuhn

Organizations: EPFL · Apple

Abstract

Distributionally robust optimization (DRO) provides a principled framework for learning under distribution shift, but its practical use is hindered by the difficulty of evaluating worst-case risks for nonconvex loss functions. We study a penalized DRO formulation in which the adversary may choose any distribution but incurs a Wasserstein penalty for deviating from the empirical distribution. We show that the adversary's problem can be reformulated as an optimization problem over transport maps that push empirical samples to adversarial ones, and we prove that optimal maps are cyclically monotone. We also show that standard adversarial training---based on per-sample local optimization---violates cyclical monotonicity and wastes transport costs unless the adversary is severely restricted. We propose two remedies. First, we introduce multi-start particle ascent, which alternates parallel gradient ascent with reassignment to enforce cyclical monotonicity across samples. Second, we parameterize adversarial maps as gradients of input-convex neural networks, which guarantees cyclical monotonicity by construction. Experiments on robust regression, image classification, and robust control show that our methods consistently outperform standard adversarial training and state-of-the-art baselines, achieving improved robustness and better generalization under distribution shift.

Figures & tables

Figure 1: PA and MPA trajectories over the contours of $f(\theta,z)-\lambda\|z-\widehat{z}_{i}\|_{2}^{2}$ for $i=1$ (left) and $i=2$ (right).

Figure 2: Test error vs. perturbation level

\Delta

on multi-class logistic regression.

Method	CIFAR-10				CIFAR-10.1	CIFAR-10.2	CIFAR-10-C (Common Corruption)						Runtime
Method	Clean	PGD	AA	Monge Gap	Clean	Clean	1	2	3	4	5	Avg	(s)
ERM	93.12	14.76	00.20	—	87.12	83.14	91.50	87.95	83.84	77.89	68.41	81.92	—
PA	90.08	52.41	48.86	8.01e-4	80.27	76.16	87.11	84.35	80.21	74.62	69.74	79.20	1320
RO	90.05	53.37	50.26	8.06e-4	80.28	76.81	87.15	85.06	81.75	75.35	70.12	79.89	659
WFR	91.03	53.28	50.37	5.63e-3	81.24	77.09	88.60	85.34	82.93	77.56	71.09	81.10	6396
SDRO	90.00	53.12	49.61	1.37e-3	80.20	76.48	87.12	84.82	80.33	74.70	70.09	79.41	6241

Table 1: Accuracy (%) under adversarial attacks, natural domain shifts and image corruptions. The left section reports results for the original CIFAR-10 test set (Clean, PGD, and AA), along with the Monge gap. The center columns evaluate accuracy on the unseen domains of CIFAR-10.1 and CIFAR-10.2. The right section (CIFAR-10-C) details performance across five increasing levels of corruption severity, followed by the mean accuracy across all five levels and the runtime. Bold values indicate the best results, and underlined values denote the second best.

Figure 3: Mean episode length for

\lambda{=}1

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Loss landscape of ICNN-DRO on CIFAR-10. Panels (a) and (b) correspond to test sample #4731 (Deer), while panels (c) and (d) detail test sample #5118 (Bird). For each sample, the left panel (a, c) displays the 2D decision regions. The dashed circle marks the $\ell_{2}$ budget $\varepsilon=0.5$ ; orange and blue indicate correct and incorrect classification, respectively; and point markers specify the exact coordinates of the initial clean image (black circle) and the final adversarial image (black triangle). The right panel (b, d) plots the corresponding 3D cross-entropy surface where the red curve traces the loss along the adversarial direction.

(a) $\lambda{=}0.1$
Method	Original	Easier Environments			Harder Environments
Method	Original	Light	Long	Soft Gravity	Heavy	Short	Strong Gravity
ERM	$186.4\pm 1.9$	$188.7\pm 1.9$	$191.4\pm 1.2$	$217.3\pm 1.9$	$178.6\pm 1.5$	$155.6\pm 0.8$	$157.6\pm 1.1$
PA	$382.9\pm 7.4$	$395.1\pm 7.4$	$266.9\pm 4.1$	$543.3\pm 7.1$	$289.0\pm 5.8$	$205.1\pm 1.7$	$207.0\pm 2.8$
RO	$380.1\pm 5.1$	$398.3\pm 5.5$	$270.5\pm 3.4$	$556.6\pm 5.3$	$295.1\pm 2.1$	$209.0\pm 1.6$	$212.4\pm 1.2$
WFR	$507.3\pm 7.6$	$532.9\pm 8.0$	$434.5\pm 5.8$	$756.8\pm 7.2$	$387.0\pm 4.6$	$313.4\pm 2.5$	$323.9\pm 3.0$

Appendix

Table 2: Mean episode length

\pm

standard error over

1{,}000

deterministic rollouts across different

\lambda

. Panel (a) reports

\lambda=0.1

; panel (b) reports

\lambda=0.5

and panel (c) reports

\lambda=1

. The leftmost column corresponds to the original training configuration, the middle group to easier physics variants (Light, Long, Soft Gravity), and the right group to harder variants (Heavy, Short, Strong Gravity). Bold entries indicate the best result per column.

Figure 5: Test loss vs. perturbation level

\Delta

on uncertain least squares.

	Accuracy (%)			Transfer
Model	Clean $\uparrow$	Adv. $\uparrow$	Drop $\downarrow$	Ovlp.	% src	Lbl mis.%
R1-ICNN (src)	$92.38$	$74.29$	$18.09$	$2{,}336$	—	—
ViT1	$89.37$	$49.39$	$39.98$	$1{,}903$	$81.46$	$56.96$
ViT2	$78.22$	$53.54$	$24.68$	$1{,}696$	$72.60$	$54.13$
ViT3	$82.32$	$54.38$	$27.95$	$1{,}763$	$75.47$	$55.30$
ViT4	$92.42$	$57.29$	$35.12$	$1{,}839$	$78.72$	$55.41$

Appendix

Table 3: Cross-architecture transfer of the ICNN adversary trained against R1. Drop : Clean-Adv. gap; Ovlp : counts examples fooled by both the source and the target; % src err : percentage of source-model failures that transfer to the target; Lbl mis.% : percentage of overlap examples on which source and target predict different incorrect labels.

		Regularization strength $\lambda$
	Method	$2$	$5$	$10$	$15$	$30$
	ERM	$0$	$0$	$0$	$0$	$0$
Baselines	PA	$1.24\text{\times}{10}^{-2}$	$1.15\text{\times}{10}^{-2}$	$3.01\text{\times}{10}^{-3}$	$7.93\text{\times}{10}^{-4}$	$7.97\text{\times}{10}^{-4}$
	WFR	$1.13\text{\times}{10}^{-2}$	$1.02\text{\times}{10}^{-2}$	$1.73\text{\times}{10}^{-3}$	$5.68\text{\times}{10}^{-4}$	$5.58\text{\times}{10}^{-4}$
	SDRO	$1.83\text{\times}{10}^{-0}$	$1.04\text{\times}{10}^{-0}$	$4.52\text{\times}{10}^{-3}$	$1.27\text{\times}{10}^{-3}$	$1.29\text{\times}{10}^{-3}$
	NN-DRO	$1.37\text{\times}{10}^{-2}$	$1.21\text{\times}{10}^{-2}$	$3.54\text{\times}{10}^{-3}$	$3.28\text{\times}{10}^{-3}$	$3.16\text{\times}{10}^{-3}$

Appendix

Table 4: Monge gap on CIFAR10.

		Regularization strength $\lambda$
	Method	$0.1$	$0.5$	$1$
	ERM	$0$	$0$	$0$
Baselines	PA	$4.87\text{\times}{10}^{-2}$	$2.01\text{\times}{10}^{-2}$	$1.99\text{\times}{10}^{-2}$
	WFR	$1.87\text{\times}{10}^{-2}$	$1.41\text{\times}{10}^{-2}$	$1.13\text{\times}{10}^{-2}$
	SDRO	$2.18\text{\times}{10}^{-2}$	$2.03\text{\times}{10}^{-2}$	$1.93\text{\times}{10}^{-3}$
	NN-DRO	$2.17\text{\times}{10}^{-2}$	$2.07\text{\times}{10}^{-2}$	$1.98\text{\times}{10}^{-3}$

Appendix

Table 5: Monge gap on robust control (CartPole).

Complex backbone [1024, 512, 512, 256, 128, 64]
Method	$r_{\mathrm{out}}$	Total params	CIFAR-10		Wall-clock Runtime (s)
Method	$r_{\mathrm{out}}$	Total params	Clean	PGD ε=0.5	Wall-clock Runtime (s)
ICNN-DRO-Q	8	31,693,314	91.37	57.36	10854
ICNN-DRO-Q	32	31,767,042	91.30	57.41	10908
ICNN-DRO-Q	128	32,061,954	91.21	57.58	11325
ICNN-DRO	8	8,665,602	91.45	57.28	6439

Appendix

Table 6: Adversary capacity versus robustness. Best results per column are in bold, and second-best results are underlined. PGD cells are shaded green for configurations satisfying our robustness criterion (PGD accuracy

\geq 56\%

), with darker green indicating stronger robustness, and red otherwise, with darker red indicating a larger shortfall. Runtime cells are shaded green only for configurations satisfying this robustness criterion, with darker green indicating lower computational cost. A runtime marked with

\dagger

corresponds to a configuration with PGD accuracy below

56\%

. Although such configurations may be substantially faster, we do not regard them as successful accelerations of ICNN training because they fail our minimum robustness criterion. The objective of this ablation is specifically to reduce the approximately

10^{4}

s wall-clock training time of the full-width quadratic ICNN while retaining at least

56\%

PGD accuracy.

Method	$K$	CIFAR-10			CIFAR-10.1	CIFAR-10.2	CIFAR-10-C (Common Corruption)
Method	$K$	Clean	PGD ε=0.5	AA ε=0.5	Clean	Clean	1	2	3	4	5	Avg
ICNN-DRO (fixed $\lambda=30$ )	50	91.20	59.83	58.81	82.57	79.98	89.28	86.85	84.10	80.00	74.06	82.85
ICNN-DRO (fixed $\lambda=30$ )	20	91.41	58.11	57.71	82.56	80.12	89.20	86.79	84.01	79.92	73.95	82.77
ICNN-DRO + $\lambda$ -annealing	20	91.48	59.94	58.92	82.60	80.10	89.31	86.93	84.15	80.04	74.12	82.91

Appendix

Table 7: Penalty annealing versus fixed-penalty adversarial training on CIFAR-10. The

K=20

fixed-

\lambda

baseline is reproduced from Table 1 for comparison. Best results per column are in bold, and the second best are underlined.

Distribution Shift Robustness

Distribution shift robustness studies how to keep machine-learning predictions reliable when deployment data differ from training data.

See all 280 papers

Statistical Guarantees for Distributionally Robust Optimization with Optimal Transport and OT-Regularized Divergences
Mar 29, 2026Jeremiah Birrell, Xiaoxi ShenDistributionally Robust Optimization Wasserstein Distributionally Robust Optimization
Department of Mathematics Texas State University San Marcos, TX, USA
Wasserstein Robust Classification Error: Surrogate Loss Misalignment and Distributional Attacks
Oct 11, 2025Bach C. Le, Tung V. Dao, Binh T. Nguyen +1Robust Classification Adversarial Robustness
College of Engineering & Computer Science, VinUniversity, Hanoi, Vietnam
Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spaces
Dec 29, 2024Linglingzhi Zhu, Yunqin Zhu, Yao XieDistributionally Robust Optimization Wasserstein Distributionally Robust Optimization
1H. Milton Stewart School of Industrial and Systems Engineering Georgia Institute of Technology

Mar 29, 2026J/K move · Enter open · S save

Statistical Guarantees for Distributionally Robust Optimization with Optimal Transport and OT-Regularized Divergences

Jeremiah Birrell, Xiaoxi Shen

Department of Mathematics Texas State University San Marcos, TX, USA

Distributionally Robust Optimization Wasserstein Distributionally Robust Optimization Statistical Learning Theory Adversarial Training Optimal Transport

Open paper AskPDF arXiv

We study finite-sample statistical performance guarantees for distributionally robust optimization (DRO) with optimal transport (OT) and OT-regularized divergence model neighborhoods. Specifically, we derive concentration inequalities for supervised learning via DRO-based adversarial training, as commonly employed to enhance the adversarial robustness of machine learning models. Our results apply to a wide range of OT cost functions, beyond the

p

-Wasserstein case studied by previous authors. In particular, our results are the first to: 1) cover soft-constraint norm-ball OT cost functions; soft-constraint costs have been shown empirically to enhance robustness when used in adversarial training, 2) apply to the combination of adversarial sample generation and adversarial reweighting that is induced by using OT-regularized

f

-divergence model neighborhoods; the added reweighting mechanism has also been shown empirically to further improve performance. In addition, even in the

p

-Wasserstein case, our bounds exhibit better behavior as a function of the DRO neighborhood size than previous results when applied to the adversarial setting.

Abstract

Distribution Shift Robustness

Explore similar work

Statistical Guarantees for Distributionally Robust Optimization with Optimal Transport and OT-Regularized Divergences

Wasserstein Robust Classification Error: Surrogate Loss Misalignment and Distributional Attacks

Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spaces