DRO-Augment Framework: Robustness by Synergizing Wasserstein Distributionally Robust Optimization and Data Augmentation

TL;DR

This paper introduces DRO-Augment, a framework combining Wasserstein Distributionally Robust Optimization with data augmentation to significantly enhance neural network robustness against various corruptions and adversarial attacks, while maintaining accuracy.

Contribution

The paper presents a novel integration of W-DRO with data augmentation, providing improved robustness and new theoretical generalization bounds for neural networks.

Findings

Outperforms existing augmentation methods under severe perturbations

Maintains accuracy on clean datasets across multiple benchmarks

Establishes new theoretical generalization error bounds

Abstract

In many real-world applications, ensuring the robustness and stability of deep neural networks (DNNs) is crucial, particularly for image classification tasks that encounter various input perturbations. While data augmentation techniques have been widely adopted to enhance the resilience of a trained model against such perturbations, there remains significant room for improvement in robustness against corrupted data and adversarial attacks simultaneously. To address this challenge, we introduce DRO-Augment, a novel framework that integrates Wasserstein Distributionally Robust Optimization (W-DRO) with various data augmentation strategies to improve the robustness of the models significantly across a broad spectrum of corruptions. Our method outperforms existing augmentation methods under severe data perturbations and adversarial attack scenarios while maintaining the accuracy on the clean datasets on a range of benchmark datasets, including but not limited to CIFAR-10-C, CIFAR-100-C, MNIST, and Fashion-MNIST. On the theoretical side, we establish novel generalization error bounds for neural networks trained using a computationally efficient, variation-regularized loss function closely related to the W-DRO problem.