Robustmix: Improving Robustness by Regularizing the Frequency Bias of Deep Nets

Deep neural networks have achieved state-of-the-art accuracy across a range of benchmark tasks such as image segmentation (Ren et al., 2015) and speech recognition (Hannun et al., 2014). These successes have led to the widespread adoption of neural networks in many real-life applications. However, while these networks perform well on curated benchmark datasets, their performance can suffer greatly in the presence of small data corruptions (Szegedy et al., 2014; Goodfellow et al., 2014; Moosavi-Dezfooli et al., 2017; Athalye et al., 2018; Hendrycks & Dietterich, 2018). This poses significant challenges to the application of deep networks.

Hendrycks & Dietterich (2018) show that the accuracy of a standard model on ImageNet can drop from 76% to 20% when evaluated on images corrupted with small visual transformations. This shows modern networks are not robust to certain small shifts in data distribution. That is a concern because such shifts are common in many real-life applications. Secondly, Szegedy et al. (2014) show the existence of adversarial perturbations, which are imperceptible to humans but have a disproportionate effect on the predictions of a network. This raises significant concerns about the safety of using deep networks in critical applications such as self-driving cars (Sitawarin et al., 2018).

These problems have led to numerous proposals to improve the robustness of deep networks. Some of these methods, such as those proposed by Hendrycks et al. (2019), require a priori knowledge of the visual transformations in the test domain. Others, such as Geirhos et al. (2018) use a deep network to generate transformations which comes with significant computation cost.

This paper proposes a new data augmentation technique to improve the robustness of deep networks by regularizing frequency bias. This new regularization technique is based on Mixup. It has many advantages compared to related robustness regularizers: (1) it does not require knowledge of a large set of priori transformations, (2) it is inexpensive, and (3) it doesn’t have many hyper-parameters. The key idea is to bias the network to rely more on lower spatial frequencies to make predictions.

We demonstrate on ImageNet-C that this method works well with recent advances and reaches a state-of-the-art mCE of 44.8, 85.0 clean accuracy with EfficientNet-B8 and RandAugment(Cubuk et al., 2019). This is an improvement of 16 mCE compared to the baseline EfficientNet-B8 and matches ViT-L/16 (Dosovitskiy et al., 2020), which is trained on $300\times$ more data. Our implementation of the method with DCT transform adds negligible overhead in our experiments. We find that Robustmiximproves accuracy on Stylized-ImageNet by up to 15 points, and we show that it can increase adversarial robustness.

The proposed approach can be seen as a generalization of Mixup (Zhang et al., 2018), a data augmentation method that regularizes models by training them on linear interpolations of two input examples and their respective labels. These new examples are generated as follows

with mix being the linear interpolation function

where $\lambda\sim\mathrm{Beta}(\alpha,\alpha)$, $\alpha$ is the Mixup coefficient hyper-parameter.

Zhang et al. (2018) show that Mixup improves the accuracy of networks and can also improve the robustness of the network. In the past years, several versions of Mixup were proposed with application in Fairness (Chuang & Mroueh, 2021), 3D reconstruction (Cheng et al., 2022), semi-supervised learning (Beckham et al., 2019), as well as robustness ((Mai et al., 2021; Yun et al., 2019; Faramarzi et al., 2020; Kim et al., 2020; Verma et al., 2019)). The novel version we propose here is frequency-based and does not include additional learnable parameters.

Augmix (Hendrycks et al., 2019) is a data augmentation technique to improve robustness by training on a mix of known image transformations. It adds little computational overhead but requires knowledge of a diverse set of domain-specific transformations. Hendrycks et al. (2019) mixes a set of 9 different augmentations to reach $68.4$ mCE on ImageNet. In contrast, the proposed method does not rely on specific image augmentations but on the more general principle that natural images are a kind of signal where most of the energy is concentrated in the lower frequencies.

The idea of frequency filtering is popular in Deep learning frameworks and has numerous applications, including unsupervised domain adaptation (Yang & Soatto (2020)) and adversarial perturbation attacks (Guo et al. (2018); Li et al. (2021)). Unlike the latter papers, which focus on measuring the accuracy of a model after an adversarial attack, we focus on common (noise) corruptions by measuring mCE as a robustness assessment.

Zhang (2019) uses low pass filters directly inside the model to improve the frequency response of the network.Wang et al. (2019) uses a differentiable neural network to extract textual information from images without modelling the lower frequency. Our method also uses low-pass filtering but does not entirely remove high-frequency features. Additionally, we only use frequency filtering during training; therefore, no computational overhead is incurred during evaluation.

In this section, we introduce a novel extension of Mixup called Robustmix that increases robustness by regularizing the network to focus more on the low-frequency features in the signal.

Motivation Wang et al. (2020) suggest that convolutional networks trade robustness for accuracy in their use of high-frequency image features. Such features can be perturbed in ways that change the model’s prediction, even though humans cannot perceive the change. This can lead models to make puzzling mistakes, such as with adversarial examples. We aim to increase robustness while retaining accuracy by regularizing how the model uses high-frequency information.

Robustmix We propose to regularize the model’s sensitivity to each frequency band by extending Mixup’s linear interpolations with a new type of band interpolation. The key insight is that we can condition the sensitivity to each band using images that mix the frequency bands of two different images. Suppose that we mix the lower-frequency band of an image of a boathouse with the high-frequency band of an image of a dog. We can encourage sensitivity to the lower band by training the model to predict "dog" for this mixed image. However, this approach is too simplistic because it completely disregards the impact of the image in the high band. Indeed, the ablation study in section 4.4 shows that it is insufficient.

Instead, we interpolate the label of such mixed images according to an estimate of the importance of each frequency band. We propose using the relative amount of energy in each band to estimate the importance. Thus the sensitivity of the model to high-frequency features will be proportional to their energy contribution in natural images. As shown in Figure 2, most of the spectral energy in natural images is concentrated in the lower end of the spectrum. This should limit the ability of high-frequency perturbations to change the prediction unilaterally.

Furthermore, we use linear interpolations of images like in mixup within each band instead of raw images. This closely reflects the more common case where the features in the bands are merely corrupted instead of entirely swapped. It also has the benefit of encouraging linearity inside the same frequency band.

Specifically, the mixing formula for Robustmix is given by

where $\lambda_{L},\lambda_{H}\sim\mathrm{Beta}(\alpha,\alpha)$, $\alpha$ is the Mixup coefficient hyper-parameter, and $\texttt{Low}(\cdot,c),\texttt{High}(\cdot,c)$ are a low pass and high pass filter respectively with a uniformly sampled cutoff frequency $c\in[0,1]$. And $\lambda_{c}$ is the coefficient that determines how much weight is given to the lower frequency band. It is given by the relative amount of energy in the lower frequency band for natural images

This coefficient can be efficiently computed on a mini-batch of examples.

Implementation Computational overhead is an important consideration for data augmentation techniques since training deep networks is computationally intensive, and practitioners have limited computational budgets. We note that many popular techniques such as Mixup (Zhang et al., 2018) add little overhead.

The frequency separation is implemented using a Discrete Cosine Transform (DCT) to avoid the complex multiplication required by a Discrete Fourier Transform. We directly multiply the images with the 224x224 DCT matrix because the spatial dimensions are relatively small, and (non-complex) matrix multiplication is well-optimized on modern accelerators. A batch of images is transformed into frequency space and the low and high-pass filtered images must be transformed back to image space. Additionally, we must apply the DCT transform over the x and y dimensions separately. Thus, 6 DCT matrix multiplications are required, resulting in $0.2$ GFLOPs per image. In contrast, just the forward pass of ResNet50 requires $3.87$ GFLOPs (Hasanpour et al., 2016).

In our implementation of Robustmix, we reorder commutative operations (low pass and mixing) to compute the DCT only once per minibatch. The pseudocode is provided in Algorithm 1, where $\mathrm{reverse}$ is a function that reverses the rows of its input matrix.

In this paper, we have introduced a new method to improve robustness called Robustmix, which regularizes models to focus more on lower spatial frequencies to make predictions. We have shown that this method yields improved robustness on a range of benchmarks, including ImageNet-C and Stylized ImageNet. In particular, this approach attains an mCE of 44.8 on ImageNet-C with EfficientNet-B8, which is competitive with models trained on $300\times$ more data.

Our method offers a promising new research direction for robustness with several open challenges. We have used a standard DCT-based low-pass filter on images and an L2 energy metric to determine the contribution of each label. This leaves many alternatives to be explored, such as different data modalities like audio, more advanced frequency separation techniques like Wavelets, and alternative contribution metrics for mixing labels.