Journal Article

Wasserstein-Robust Training for One-Hidden-Layer ReLU Networks with Distributional Guarantees

Wang, Yang, Liu, Xiaona, Zou, Difan

2026

Overview

Neural networks are vulnerable to adversarial perturbations, which motivates training procedures with formal robustness guarantees. In this paper, we study one-hidden-layer ReLU networks from the perspective of Wasserstein distributional robustness. Leveraging the network structure, we derive an upper bound for the intractable robust surrogate in the form of a tractable regularized empirical risk objective whose regularizer is computed through a low-rank optimization problem based on Burer–Monteiro factorization. This reformulation yields a distributional robustness certificate on the worst-case expected loss over a Wasserstein ball. The upper bound construction and distributional certificate are developed for the shallow fixed-output multiclass formulation, while the optimization analysis focuses on a binary specialization with margin loss and exact linear separability. We also analyze a modified stochastic gradient descent scheme for the resulting regularized problem in this binary linearly separable setting, and we establish a corresponding generalization bound. The experiments validate the proposed surrogate and training procedure on binary MNIST and CIFAR-10 tasks, and we added a 10-class MNIST experiment to further check the multiclass trainability of the surrogate.

• Publisher: MDPI AG
• Subject: Analysis / Classification / Geometry / Guarantees / Neural networks / Optimization / Propagation / Rankings / Robustness / Specialization / Upper bounds