{'fraction_non_alphanumeric': 0.02820000677070991, 'fraction_numerical': 0.009952943566132908, 'mean_word_length': 5.114676050507142, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
AugNet: Dynamic Test-Time Augmentation via Differentiable Functions

Distribution shifts, which often occur in the real world, degrade the accuracy of deep learning systems, and thus improving robustness is essential for practical applications. To improve robustness, we study an image enhancement method that generates recognition-friendly images without retraining the recognition model. We propose a novel image enhancement method, AugNet, which is based on differentiable data augmentation techniques and generates a blended image from many augmented images to improve the recognition accuracy under distribution shifts. In addition to standard data augmentations, AugNet can also incorporate deep neural network-based image transformation, which further improves the robustness. Because AugNet is composed of differentiable functions, AugNet can be directly trained with the classification loss of the recognition model. AugNet is evaluated on widely used image recognition datasets using various classification models, including Vision Transformer and MLP-Mixer. AugNet improves the robustness with almost no reduction in classification accuracy for clean images, which is a better result than the existing methods. Furthermore, we show that interpretation of distribution shifts using AugNet and retraining based on that interpretation can greatly improve robustness.


Introduction
With the development of deep learning, the field of visual recognition has made great progress, and services using deep learning are becoming more practical. Many deep learning models are trained and validated on images from the same distribution. However, images in the real world are subject to distribution shifts due to a variety of factors such as weather conditions, sensor noise, blurring, and compression artifacts. Usually, deep learning models do not take these distribution shifts into account, causing reduction in the accuracy in the presence of such artifacts. This is fatal for safety-critical applications, such as automated driving, where the environment changes frequently.
There are two main approaches to solving this problem: training robust recognition models and image enhancement to improve recognition. The former approach studies data augmentation or training algorithms such that the recognition model has high robustness. The latter approach uses test-time augmentation or deep neural networks to transform the distorted images into recognition-friendly images that are easily recognized by the recognition model. These approaches can be used in combination to achieve higher robustness. In this paper, we focus on an image enhancement approach. It is highly practical because it is used before inference by the pretrained recognition model without retraining the recognition model. Deep neural networks based image enhancement [37] excessively transforms images to remove corruption, significantly improving robustness but reducing accuracy for clean images. Test-time augmentation based image enhancement [14], which dynamically selects the best one from several data augmentations for an input image, improves robustness while maintaining clean accuracy. However, its augmentation search space is very limited, so the improvement in robustness is smaller than by deep neural network-based image enhancement.
To remove this limitation and further improve the robustness while maintaining clean accuracy, we propose a novel image enhancement method, which we call AugNet. AugNet uses differentiable data augmentation techniques [6,20,31] and image blending to dynamically generate the recognition-friendly image for an input image from a huge augmentation space. Image transformation using deep neural networks can be incorporated into AugNet by considering it as a kind of data augmentation. An overview diagram of AugNet is shown in Fig. 1. AugNet takes an image as input and outputs the magnitude parameters and blend weights for predefined data augmentation. Each augmentation is performed based on the magnitude parameters. These augmented images are linearly combined with blend weights to generate the output image. The classification model takes AugNet's output images as input and performs inference. Because AugNet is composed of differentiable functions, AugNet can be directly trained with the task loss function of the recognition model. As a result, AugNet transforms distorted images to improve recognition Figure 1. AugNet is used before inference by the classification network to generate recognition-friendly images. First, AugNet outputs the magnitude parameters and blend weights. Next, predefined data augmentations are performed using the magnitude parameters. Finally, the output image is generated by linearly combining the augmented images with the blend weights. The generated images are input to the classification network. In this paper, we refer to these image transformation models that are used before the classification network as enhancement models.
accuracy. In addition, we introduce a practical training and evaluation scenario, the blind setting, which does not assume the type of test-time distribution shifts and was not conducted in existing image enhancement studies.
AugNet is evaluated on two datasets and a variety of classification models, including state-of-the-art classification models such as Vision Transformer [41] and MLP-Mixer [40]. AugNet improves the accuracy for distorted images while maintaining the classification accuracy for clean images, which is a better result than existing methods.
We further show that AugNet enables us to estimate effective training-time data augmentations for distribution shifts. AugNet weights effective test-time augmentations for distribution shifts. By using the inverse operations of the highly weighted data augmentations during training, we make classification models robust for distribution shifts.
The main contributions of this paper are as follows.
• We propose AugNet, a novel image enhancement method based on differentiable data augmentation techniques and image blending. Using it before the pretrained classification model improves robustness.
• We show that AugNet improves robustness compared to existing methods with almost no reduction in clean accuracy.
• We use AugNet to estimate the effective data augmentation during training for distribution shift, and show that retraining the classification model with the estimated data augmentations significantly improve robustness.

Robustness of Deep Neural Networks
Deep neural networks are known to be vulnerable to image distortions, which often occur in the real world. These image distortions have been studied in several works. [3] showed that defects such as noise and blur in real-world sensors cause performance degradation in image recognition networks. [27] investigated that haze reduces the accuracy of image classification and experimented that the dehaze method for human visibility is not effective in improving image classification performance. [28] empirically studied real-world image degradation problems for nine kinds of degraded images-hazy, motion-blurred, fish-eye, underwater, low resolution, salt-and-peppered, white Gaussian noise, Gaussian-blurred, and out-of-focus.
Various datasets have been proposed to evaluate the robustness of deep neural networks. [5] made a styletransformed dataset to demonstrate that deep neural networks recognize objects by their texture rather than their shape. [9] proposed 19 algorithmically generated corruption datasets in five levels, belonging to four categories: noise, blur, weather, and digital. They showed that deep neural networks trained on a clean dataset decrease accuracy on these corruption datasets. In addition, they introduced the dataset of naturally occurring adversarial examples in the real world [11], and the dataset containing distribution shifts that occur in the real world, such as image style, image blurriness, geographic location, camera operation, and more [8].
In this paper, we evaluate our method on these datasets.

Training Robust Classification Models
Many works have been studied to improve the robustness of deep neural networks to naturally occurring image distortions. One way to improve the robustness is the data augmentation. [4,10,44] proposed a method to improve robustness by using a mix of multiple data augmentations. [32,42] showed that adding noise during training improves the robustness to noise and other types of distortions. [23] investigated that the similarity between data augmentation and test-time image corruption is strongly correlated with the performance.
Training algorithms to improve robustness have also been studied. [46] proposed Advprop, a method to improve adversarial training, which improves not only robustness but also clean accuracy. [21] proposed an algorithm to train model by paying complementary attention to the shape and texture of objects. Since these methods are used when training recognition models, they can be used in combination with image enhancement methods, which assume that the recognition models are pretrained. Our experiments show that combining the recognition models pretrained by these methods with image enhancement methods further improves the robustness.

Test-time Augmentation
It is possible to improve recognition accuracy by using data augmentation not only during training-time but also during test-time. Various test-time augmentation methods have been proposed, using simple geometric transformations such as flip and crop [2,7,17,29,38,47], mixup [18,26], and augmentation in the embedding space [1]. [15,34] experimentally analyzed why and when test-time augmentation works, and gave theoretical guarantees. These studies use only specific data augmentations for all test set images. Dynamically selecting the best data augmentation for each input image is expected to improve the accuracy. [24] proposed a greedy search-based test-time augmentation method to find the augmentation policy on the test set, but it is not optimal for each input image. [14] proposed a module called Loss Predictor to predict the classification loss of augmented images, allowing dynamic selection of the best one for the input image from the 12 data augmentations. However, despite the fact that the augmentation space is actually infinite, these methods have significant limitations on the augmentation space due to computational cost and the improvement in robustness is small. Our method eliminates this limitation and achieves a better robustness than existing methods by generating the best augmented image from a large augmentation space. In addition, our method is highly interpretable than existing methods because it handles a large number of data augmentations.

Image Transformation
Several studies have been conducted to improve the recognition accuracy by image transformation using deep neural networks. [35] proposed convolutional neural networkbased enhancement filters that enhance image-specific details to improve recognition accuracy. [19,22,25,37] used a deep neural network model to transform degraded images into recognition-friendly images. [39] showed that not only image enhancement but also resizing at the same time improves the recognition accuracy. These methods use deep neural networks to transform images, but tend to perform excessive transformations to remove distortions. This reduces the accuracy of clean images and makes image transformation difficult to interpret. By considering these image enhancement methods as a kind of data augmentation, we treat them as part of AugNet.

Key Ideas
Many data augmentations have continuous magnitude parameters (e.g., rotation has a continuous magnitude parameter from 0 to 360 degrees). In addition, there are a large number of combinations of two or more data augmentations (e.g., augmentation is performed in the order of rotation, contrast, and sharpness). These facts make the augmentation space infinite. To select the best combination of data augmentations and their magnitude parameters for the test image from this infinite space, we introduce two key ideas. The first is to find the local optimal magnitude parameter by optimization. Since most data augmentations are differentiable with respect to the magnitude parameter, the local optimal magnitude parameter is obtained by optimization algorithms, such as gradient descent. This idea is inspired by recent differentiable data augmentation techniques [6,20,31]. The second is to represent the combination of augmentation with image blending, which is the process of affine combination of multiple images. For example, if the weights are [1, 0, 0], only the 1st augmented image is selected, and if the weights are [0.5, 0, 0.5], the 1st and 3rd augmented images are combined. Note that image blending does not provide the performing order of the augmentation. We call these ideas "Blending" and "Optimization", respectively. Based on these key ideas, we propose a novel image enhancement method, AugNet.

Overview of AugNet
The overview diagram is shown in Fig. 1. First, we define the differentiable data augmentations used in AugNet and their range of magnitude parameters M . Some data augmentations, such as rotation, have a magnitude parameter. On the other hand, other data augmentations do not have a magnitude parameter, such as auto-contrast. Details of the data augmentations used in AugNet are shown in Section 3.3.
Any neural network can be used as the backbone of AugNet by modifying its output layer. AugNet takes an image as input and outputs the magnitude parameters m 1 , · · · , m K and blend weights w 1 , · · · , w K . K is the number of data augmentations. The magnitude parameter is mapped to a value in the range (0, 1) or (−1, 1) using an activation function f act k such as sigmoid function or hyperbolic tangent function. It is then multiplied by the predefined magnitude range M . The blend weights are converted to weights that sum to 1 by the softmax function. We show m and w in the following equation.
Using the magnitude parameters and blend weights, AugNet generates a recognition-friendly image. K data augmentations are performed with magnitude parameters m, then AugNet generates a blended image by linearly combining these augmented images with w. The output of AugNet is obtained by the following equation.
Where x is the input image, O k is the k-th augmentation operation, x aug k is the k-th augmented image, and x out is the output image.

Augmentation Space
The data augmentations and their f act k and M used by AugNet are shown in Tab. 1. We use only differentiable data augmentations.
Placing low-pass and high-pass filters within the deep neural network architecture [12,30,43,49] or using them as training-time augmentation [13,36,48] improve robustness. Inspired by these studies, we try to improve the robustness by using low-pass and high-pass filters as test-time augmentation. These filters have a filter size parameter which cannot be obtained via gradient descent. Therefore, we discretize this parameter and prepare multiple filters. In our experiments, we make 19 low-pass and 19 high-pass filters, each with a filter size from 0.05 to 0.95 in 0.05 increments.
URIE [37] uses a deep neural network to generate recognition-friendly images that improve robustness. Such an image transformation model can be applied to AugNet by considering it as a kind of augmentation. In this case, the image transformation model is trained simultaneously with AugNet.
AugNet handles many data augmentations, but is also computationally expensive. By not executing data augmentation with small blend weights, the number of data augmentation executions are significantly reduced while maintaining accuracy. We describe the details in the Appendix.

Loss Function
AugNet is trained by using it before inference of a pretrained classification model. When training AugNet, the

Training and Evaluation Settings
In this paper, we use two settings, blind and non-blind settings, for training and evaluation. In the non-blind setting adopted by existing studies [14,37], the corruption dataset [9] is used. This dataset consists of 19 types of corruption with 5 severity levels generated by the algorithm. 15 corruptions are used for training as Seen corruptions 1 , and 4 corruptions are used for testing as Unseen corruptions 2 . Seen and Unseen corruptions consist of the same four categories: noise, blur, weather, and digital, so this is a setting in which we know in advance the type of corruption that occurs during the test time. Therefore, we introduce a new training and evaluation setting, the blind setting, which has no assumption on test-time distribution shifts. In the blind setting, no corrupted images were used and only data augmentations that improve the robustness of classification models are used for training, such as AugMix [10]. The Seen and Unseen corruption datasets are used for only testing.

Experiments
In this section, we describe the experimental setup, experimental results comparing AugNet to other image enhancement methods, and detailed experiments with 1 Gaussian noise, Shot noise, Impulse noise, Defocus blur, Glass blur, Motion blur, Zoom blur, Snow, Frost, Fog, Brightness, Contrast, Elastic transform, Pixelate, JPEG. 2 Speckle noise, Gaussian blur, Spatter, Saturation.
AugNet. We also show that robustness can be significantly improved by using AugNet to estimate the most effective training-time data augmentation for distribution shift, which is then used for retraining. Some detailed experiments are described in the Appendix.

Experimental Setup
We evaluate our method on two widely used image recognition datasets, CUB [45] and ImageNet [33]. In experiments on CUB, we run all experiments three times and report the average results. We use Loss Predictor [14] and URIE [37] as comparison methods. The backbone network of AugNet and Loss Predictor is prepared by fine-tuning a ResNet18 [7] pretrained on ImageNet. Classification models are prepared in advance by standard training or AugMix training. For training in the non-blind setting, corruption is randomly selected for each mini-batch to train enhancement models. The corruptions consist of 15 Seen corruptions in five levels and clean (no corruption). For testing, we use 4 Unseen corruptions. In the blind setting, the enhancement models are trained using AugMix for CUB and AugMix+DeepAugment [8] for ImageNet. For testing, we use 19 corruptions (Seen+Unseen). When evaluating on the corruptions dataset, we use the average accuracy as evaluation metric. In addition to evaluating the robustness of these distribution shift datasets, we also evaluate our results on the standard clean dataset. The baseline does not use the enhancement model and the difference from the baseline is reported in the results.

Classification on CUB
CUB is an image dataset of 200 class bird species. It consists of 5994 training images and 5794 test images. We used ResNet50 [7], Mixer-B16 [40] and DeiT-base [41] as the classification model. For optimization, we employ Adam [16] with learning rate 0.001 for URIE, 0.001 for AugNet and Loss Predictor, and decay the learning rate every 10 epochs by 2. These models are trained in 60 epochs.
We show the results in Tab. 2. In the non-blind setting, URIE improves Unseen accuracy but reduces the clean accuracy by a maximum of about 3.3 points. Loss Predictor does not decrease the clean accuracy, but the improvement in Unseen accuracy is small. AugNet outperforms the comparison method in both clean accuracy and Unseen accuracy. AugNet greatly improves the Unseen accuracy with almost no decrease in clean accuracy. Our experiments show that image enhancement methods are effective for state-of-the-art classification models such as DeiT and MLP-Mixer. It is also possible to further increase the robustness of classification models trained by AugMix. In the blind setting, similar to the non-blind setting, URIE tends to reduce the clean accuracy by a maximum of about 1.0 points. Loss Predictor improves both clean accuracy and robustness a little. AugNet has less degradation in clean accuracy than URIE, and it is also more robust. In particular, AugNet significantly improves the robustness of Mixer-B16 compared to existing methods. Our experiments show that AugMix-trained enhancement models slightly improve the robustness of AugMix-trained classification models.

Classification on ImageNet
ImageNet is the most widely used 1000-class image recognition dataset. It is composed of 1.3 million training images and 50,000 test images. We used ResNet50 and DeiT as the classification model. For optimization, we employ Adam with learning rate 0.001 for URIE, 0.0001 for AugNet and Loss Predictor, and decay the learning rate every 8 epochs by 10. These models are trained in 30 epochs.
We show the results in Tab. 3. In the non-blind setting, URIE improves Unseen accuracy but reduces the clean accuracy by a maximum of about 1.6 points. Loss Predictor does not decrease the ResNet50 clean accuracy, but the improvement in Unseen accuracy is small. In DeiT, Loss Predictor showed a decrease in both clean and Unseen accuracy. AugNet slightly reduces the clean accuracy, but the improvement in Unseen accuracy is higher than the comparison method. In the blind setting, URIE and Loss Predictor experiments showed almost the same trend results as in the non-blind setting. AugNet improves robustness over URIE without degrading the clean accuracy.

Effects of Key Ideas
We conduct an ablation study of our key ideas by using ResNet50 on the CUB dataset and comparing it to Loss Predictor. Loss Predictor selects one best augmentation for a test image by learning the classification losses of 12 augmented images 3 as a label. One of our key ideas, "Blending", improves upon Loss Predictor by using not just one image but a combination of many augmented images. Image blending is lightweight because it does not require the computation of classification losses, which has the advantage of being easily extendable in augmentation space by adding data augmentations. We also measure the effect of "Optimization" by using Rotate, {0.8, 1.2} Zoom, Saturate, Auto-contrast, and Sharpness with the same magnitude parameters as in Tab. 1. Identity is not used because it is included in some data augmentations (e.g., a rotation of 0 degrees).
The results are shown in Tab. 4. "Blending" is much more robust than the Loss Predictor. This result means that multiple-image blending is better than choosing one best image. In the non-blind setting, "Optimization" shows a worse Unseen accuracy than "Blending". This is because the magnitude parameters used by "Blending" are tuned manually for the non-blind setting. On the other hand, in the blind setting, "Optimization" achieved better robustness than "Blending". This is because "Optimization" automatically find the local optimal magnitude parameters using gradient descent.

Effects of Augmentation Space
In this paper, the 14 data augmentations shown in Tab. 1 are used for AugNet. We investigate the contribution of each data augmentation to classification accuracy in the nonblind setting using ResNet50 on the CUB dataset. Fig. 2 shows the accuracy when each augmentation is excluded from AugNet. All data augmentations, except Equalize contribute to improving robustness with almost no reduction in clean accuracy. In particular, the contribution of URIE to improving robustness is very significant, but URIE tends to degrade clean accuracy. AugNet blends URIE with other  Figure 2. Impact of individual augmentation on classification accuracy. The y-axis is the difference between excluding one augmentation and using all data augmentations. A-Contrast, LPFs, and HPFs denote auto-contrast, low-pass filters, and high-pass filters, respectively. data augmentations to maintain the clean accuracy.

Detailed Experiments
In this section, we experiment to understand the behavior of AugNet. In all experiments, we use ResNet50 as the classification model on the CUB dataset.

Effects of Magnitude Range
We investigate the effect of different magnitude ranges on classification accuracy. In the experiments in the blind setting here, we also experimented with the stylized dataset [5] in addition to the corruption dataset. The range in Tab. 1 multiplied by 1.5 is referred to "Large", and the range multiplied by 0.5 is referred to "Small".
The results are shown in the Tab. 5. "Small" showed a better result for corruption dataset and "Large" showed a better result for stylized dataset. The clean accuracy of "Large" is a little lower than the others. We have observed that a large magnitude range is effective for large distribution shifts such as stylized datasets, but large image transformations degrade clean accuracy. "Small" showed more robust than the standard range defined in Tab. 1. This may be caused by the fact that "Small" has a smaller search space and the optimization is more stable.

Visualization of AugNet Output
We visualize the output of AugNet in the non-blind setting. We show augmented images and output image by AugNet for level 5 Unseen corruptions in Fig. 3. The bottom two rows of augmented images are almost black, except for the

Interpretation of distribution shift by AugNet
In this section we describe a different use of AugNet, the interpretation of distribution shifts. For distributionshifted datasets, AugNet weights effective test-time augmentations that improve accuracy. The distribution shift is expected to consist of inverse operations of the highly weighted data augmentations. For example, we expect the distribution shift for which low-pass filters are effective in test-time to be the corruption of the high-frequency domain by high-pass filters. Therefore, we consider that retraining the recognition model using inverse operations of highly weighted data augmentations improve the accuracy for the distribution-shifted dataset.
In the blind setting, we split the level 5 Unseen corruption test set into a validation set and a test set in a 2:8 ratio. Fig. 4 shows a visualization of the magnitude parameters and the blend weights, which are the outputs of AugNet for the validation set. The weights of the low-pass filters are higher for high-frequency corruptions such as Speckle Noise and Spatter, and the saturation weight is high for Saturate. These results are consistent with the intuitive interpretation of humans. Low-pass filters, gamma, contrast, and saturation are commonly high for all corruptions. Many of the magnitude parameters have high variance. This is because when the blend weight is close to 0, any value of the magnitude parameter has almost no effect on the output image.
Next, from the test-time augmentation that was effective for the validation set, we estimate the training-time augmentation that is most effective for the target corruption. We choose the 4 data augmentations used to retrain the classification model: high-pass filters, gamma, contrast, and saturation, which are the inverse operations of 4 data augmentations that were commonly effective for all corruptions in the validation set. We experiment by varying the data augmentation space used by the AugMix. The results of the Unseen test set are shown in Tab. 6. "Normal" is the normal Aug-Mix and uses 9 data augmentations 4 . "All" is the normal 9 data augmentations plus saturation, contrast, brightness and sharpness, and "Estimated" is the normal 9 data augmentations plus saturation, contrast, high-pass filters and gamma. The classification model trained with estimated data augmentations achieved the best accuracy for all corruptions except Saturate. On average, it was approximately 10 points more accurate than "Normal", and in particular, approximately 20 points more accurate than "Normal" for Speckle Noise.

Conclusions
In this paper, we propose AugNet, a novel image enhancement method based on differentiable data augmentation techniques and image blending. AugNet uses a gradient descent algorithm to find magnitude parameters and blend weights from a huge augmentation space. This augmentation space includes deep neural networks such as URIE as well as standard data augmentation. AugNet improves robustness by using it before inference of the classification model without retraining the classification model. Existing image enhancement studies are evaluated under the assumption that they know the type of test-time distribution shifts. We introduce a practical training and evaluation scenario that does not assume the type of test-time distribution shifts. Our experiments, which included state-of-theart classification models, showed that AugNet outperforms existing methods and improves robustness with almost no reduction in clean accuracy. We showed that image enhancement methods further improve the robustness of such robust classification models rained by AugMix and modern classification models such as ViT and MLP-Mixer are highly robust. Furthermore, we showed that interpretation of distribution shifts using AugNet and retraining based on that interpretation greatly improve robustness. In this paper, we experimented only with the image classification task. We will apply and evaluate AugNet to other tasks such as object detection and segmentation. The search for the best architecture for AugNet and the reduction of computational cost are also future works.