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

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.