Interpolation for Robust Learning: Data Augmentation on Wasserstein Geodesics
It addresses robustness in machine learning models, particularly for image classification tasks, by offering a novel data augmentation approach that can be combined with existing methods.
The paper tackles model robustness by interpolating training data distributions using Wasserstein geodesics, achieving improvements such as a 7.7% increase in certifiable robustness on CIFAR-10 and 16.8% on empirical robustness on CIFAR-100.
We propose to study and promote the robustness of a model as per its performance through the interpolation of training data distributions. Specifically, (1) we augment the data by finding the worst-case Wasserstein barycenter on the geodesic connecting subpopulation distributions of different categories. (2) We regularize the model for smoother performance on the continuous geodesic path connecting subpopulation distributions. (3) Additionally, we provide a theoretical guarantee of robustness improvement and investigate how the geodesic location and the sample size contribute, respectively. Experimental validations of the proposed strategy on \textit{four} datasets, including CIFAR-100 and ImageNet, establish the efficacy of our method, e.g., our method improves the baselines' certifiable robustness on CIFAR10 up to $7.7\%$, with $16.8\%$ on empirical robustness on CIFAR-100. Our work provides a new perspective of model robustness through the lens of Wasserstein geodesic-based interpolation with a practical off-the-shelf strategy that can be combined with existing robust training methods.