# Improving Model Generalization by On-manifold Adversarial Augmentation in the Frequency Domain

## TL;DR

This paper introduces AdvWavAug, a wavelet-based on-manifold adversarial data augmentation method that improves deep neural network generalization, especially on out-of-distribution data, achieving state-of-the-art results on ImageNet and its variants.

## Contribution

Proposes a novel wavelet domain augmentation technique for on-manifold adversarial examples, enhancing model robustness and out-of-distribution generalization.

## Key findings

- Improves OOD generalization across multiple models and datasets
- Achieves state-of-the-art results on transformer-based models
- Demonstrates effectiveness of wavelet-based augmentation in adversarial training

## Abstract

Deep neural networks (DNNs) may suffer from significantly degenerated performance when the training and test data are of different underlying distributions. Despite the importance of model generalization to out-of-distribution (OOD) data, the accuracy of state-of-the-art (SOTA) models on OOD data can plummet. Recent work has demonstrated that regular or off-manifold adversarial examples, as a special case of data augmentation, can be used to improve OOD generalization. Inspired by this, we theoretically prove that on-manifold adversarial examples can better benefit OOD generalization. Nevertheless, it is nontrivial to generate on-manifold adversarial examples because the real manifold is generally complex. To address this issue, we proposed a novel method of Augmenting data with Adversarial examples via a Wavelet module (AdvWavAug), an on-manifold adversarial data augmentation technique that is simple to implement. In particular, we project a benign image into a wavelet domain. With the assistance of the sparsity characteristic of wavelet transformation, we can modify an image on the estimated data manifold. We conduct adversarial augmentation based on AdvProp training framework. Extensive experiments on different models and different datasets, including ImageNet and its distorted versions, demonstrate that our method can improve model generalization, especially on OOD data. By integrating AdvWavAug into the training process, we have achieved SOTA results on some recent transformer-based models.