FLOAT:  FAST LEARNABLE ONCE-FOR-ALL ADVERSARIAL TRAINING FOR TUNABLE TRADE-OFF BETWEEN ACCURACY AND ROBUSTNESS

1 INTRODUCTION . As artifical intelligence ( AI ) -enabled applications proliferate , the robustness concerns of deep neural networks ( DNNs ) has grown in importance Zhang et al . ( 2019 ) . These concerns are largely associated with the vulnerability of machine learning ( ML ) classifiers to adversarial examples , i.e. , images that have well-crafted barely-visible perturbations from corresponding clean images and have the ability to fool classifier models into making wrong predictions Carlini & Wagner ( 2017 ) ; Goodfellow et al . ( 2014 ) . With the growing usage of DNNs in safety-critical and sensitive applications including autonomous-driving Bojarski et al . ( 2016 ) and medical image analysis Han et al . ( 2021 ) , it therefore becomes crucial that they have high accuracy on both adversarially-perturbed and clean images . To improve the model performance of DNNs against such adversarial attacks , various defense mechanisms have been proposed including hiding gradients Tramèr et al . ( 2017 ) , adding noise to parameters He et al . ( 2019 ) , and detection of adversaries Meng & Chen ( 2017 ) . In particular , adversarial training Madry et al . ( 2017 ) has proven to be a consistently effective approach . These defenses , however , come at various costs . Firstly , most of these methods suffer from increased training time due to the additional back-propagation overhead to generate perturbed images . Secondly , adversarial defenses sometimes cause a significant drop in clean-image accuracy Tsipras et al . ( 2018 ) , highlighting an accuracy-robustness trade-off that has been explored both theoretically and experimentally Sun et al . ( 2019 ) ; Tsipras et al . ( 2018 ) ; Schmidt et al . ( 2018 ) . Moreover , the defenses rely on several hyperparameters whose settings force the model to work at a specific point along this trade-off . This is disadvantageous in applications in which the desired trade-off depends on context Wang et al . ( 2020 ) . A naive solution to this problem is to use multiple networks trained with different priorities between clean and adversarial images . However , this increases both training time and inference memory . Alternatively , recent work has proposed training a once-for-all adversarial network ( OAT ) that supports conditional learning Wang et al . ( 2020 ) , enabling the network to adjust to different input distributions . In particular , after each batch-normalization ( BN ) layer , they add a feature-wise linear modulation ( FiLM ) module Perez et al . ( 2018 ) whose weights are controlled by a parameter λ . For inference , the user sets λ to enable an in-situ trade-off between accuracy and robustness . The disadvantage with this approach is that the added FiLM modules increase the parameter count , training time , and network latency , limiting applicability in resource-constrained , real-time applications . Our contributions . In view of the concerns above , we present fast learnable once-for-all adversarial training ( FLOAT ) . In FLOAT , we train a model using a novel mechanism wherein each weight tensor of the model is transformed by conditionally adding a noise tensor based on a binary parameter λ . This yields state-of-the-art ( SOTA ) test accuracy for clean and adversarial images by in-situ setting λ = 0.0 and 1.0 , respectively . For inference , we further show that model robustness can be correlated to the strength of the noise-tensor scaling factor . This motivates a simple yet effective noise re-scaling approach controlled by an added floating-point parameter that can help the user to have a practical accuracy-robustness trade-off . Because FLOAT does not require additional layers to perform conditioning , it incurs no increase in latency and causes only a negligible increase in parameter count compared to the baseline models . Moreover , compared to OAT , FLOAT training is up to 1.43× faster , as it does not explicitly require learning intermediate fine-grained conditioning using different λs . Moreover , we show how FLOAT can be extended to slimmable networks Yu et al . ( 2018 ) to enable an in-situ inference trade-off across three dimensions , namely , accuracy , robustness , and complexity . To evaluate the merits of FLOAT , we conduct extensive experiments on CIFAR-10 , CIFAR-100 , Tiny-ImageNet , SVHN , and STL10 datasets with ResNet34 ( on both CIFAR and Tiny-ImageNet datasets ) , WRN16-8 , WRN40-2 , respectively . As shown in Fig . 1 , compared to existing methods , FLOAT models can provide improved accuracy of up to ∼6.5 % , and ∼14.5 % , on clean and perturbed images , respectively . The remainder of this paper is organized as follows . Section 2 presents the necessary preliminaries and prior work . We present our approach in Section 3 and analyze experimental results in Section 4 . Finally , the paper concludes in Section 5 with discussions of broader impact in Section 6 . 2 PRELIMINARIES . 2.1 NOTATION . Consider a model Φ with L layers parameterized by Θ that learns a function fΦ ( . ) . For a classification task on dataset X with distribution D , the model parameters Θ are learned by minimizing the empirical risk ( ERM ) as follows minimize Θ L ( fΦ ( x , Θ ; t ) ) , ( 1 ) where t is the ground-truth class label , x is the vectorized input drawn from X , and L is the crossentropy loss function . 2.2 ROBUST MODEL TRAINING . Several forms of adversarial training ( AT ) have been proposed to improve robustness Madry et al . ( 2017 ) ; Samangouei et al . ( 2018 ) ; Buckman et al . ( 2018 ) . They all use both clean and adversariallyperturbed images to train a model . Projected gradient descent ( PGD ) attack , recognized as one of the strongest L∞ adversarial example generation algorithm Madry et al . ( 2017 ) , is typically used to create adversarial images during training . The perturbed image for a PGD-k attack with k as the number of steps is given by , x̂k = ProjP ( x ) ( x̂ k−1 + σ × sign ( ∇xL ( fΦ ( x̂k−1 , Θ ; t ) ) ) ( 2 ) Here , the scalar corresponds to the perturbation constraint that determines the severity of the perturbation . Proj projects the updated adversarial sample onto the projection space P ( x ) which is the -L∞ neighbourhood of the benign sample x1 . σ is the attack step-size . For PGD-AT , the model parameters are then learned by the following ERM minimize Θ [ ( 1− λ ) L ( fΦ ( x , Θ ; t ) ) ︸ ︷︷ ︸ LC +λL ( fΦ ( x̂ , Θ ; t ) ) ︸ ︷︷ ︸ LA ] , ( 3 ) where LC and LA corresponds to the clean and adversarial image classification loss components , respectively , weighted by the scalar λ . Hence , for a fixed λ and adversarial strength , the model learns a fixed tradeoff between accuracy and robustness . For example , an AT with λ value of 1 will allow the model to completely focus on perturbed images , resulting in a significant drop in clean-image classification accuracy . Another strategy to improve model robustness is through the addition of noise to the model weight tensors . For example , He et al . ( 2019 ) introduced the idea of noisy weight tensors with a learnable noise scaling factor to improve model robustness against gradient-based attacks . However , this strategy also incurs a significant drop in clean accuracy . 2.3 CONDITIONAL LEARNING . Conditional learning involves training a model with multiple computational paths that can be selectively enabled during inference Wang et al . ( 2018 ) . For example , Teerapittayanon et al . ( 2016 ) ; Huang et al . ( 2017 ) ; Kaya et al . ( 2019 ) enhanced a DNN model with multiple early exit branches at different architectural depths to allow early predictions of various inputs . Yu et al . ( 2018 ) introduced switchable BNs that enable the network to adjust the channel widths dynamically , providing a in-situ efficient trade-off between complexity and accuracy . Recently , Bulat & Tzimiropoulos ( 2021 ) used switchable BNs to support runtime bit-width selection of a mixed-precision network . Another conditional learning approach used feature transformation to modulate intermediate DNN features Huang & Belongie ( 2017 ) ; Yang et al . ( 2019 ) ; De Vries et al . ( 2017 ) ; Wang et al . ( 2020 ) . In particular , Wang et al . ( 2020 ) used FiLM Perez et al . ( 2018 ) to adaptively perform a channel-wise affine transformation after each BN stage that is controlled by the hyperparameter λ of Equation 3 . Each FiLM module is composed of two fully-connected ( FC ) layers with leaky ReLU activation functions and dimensions that match the output feature-map channel size . Despite requiring a relatively small number of additional FLOPs , however , the FiLM module can significantly increase the number of model parameters and associated memory access cost Horowitz ( 2014 ) . Moreover , the increased number of layers can significantly increase inference latency Singh et al . ( 2019 ) , potentially prohibiting its use in real-time applications . 1Note that the generated x̂ are clipped to a valid range which , for our experiments , is [ 0 , 1 ] . 3 PROPOSED APPROACH . 3.1 FLOAT . In this section , we disclose our fast learnable once-for-all adversarial training ( FLOAT ) . We first focus on defining the conditions for a model being trained on either clean or adversarial images , as the two boundary conditions for training . We use a binary conditioning parameter λ , that during training , forces the model to focus on either of these two conditions . To formalize our approach , consider a L-layer DNN parameterized by Θ and let θl ∈ Rkl×kl×cli×clo represent the layer l weight tensor , where clo and c l i represent the number of filters and channels per filter , respectively , and kl represents the kernel height/width . We transform each parameter of θl , by adding a noise tensor ηl ∈ Rkl×kl×cli×clo scaled by a parameter αl and conditioned by λ , as follows , θ̂l = θl + λ · αl · ηl ; ηl∼N ( 0 , ( σl ) 2 ) . ( 4 ) Note that the standard deviation ( σl ) of the noise matches that of its weight tensor . For λ = 0 and 1 , we obtain the original weight tensor or a noisy variant , respectively . As illustrated in Algorithm 1 , we train our models by partitioning an image batch B into two equal sub-batches B1 and B2 , one with clean ( IFMC ) images and the other with perturbed variants ( IFMA ) ( lines 4 and 6 in Algorithm 1 ) . We use the PGD-7 attack to generate perturbations on the image batch B2 . As illustrated in Fig . 2 ( b ) , the original and noisy weight tensors are convolved only with clean and perturbed variants , respectively . Similar to He et al . ( 2019 ) , we also train the noise scaling factor αl ( line 9 ) . As exemplified in Fig . 3 , the post-convolution feature maps for clean and adversarial inputs can differ significantly in their respective mean/variances Xie et al . ( 2020 ) ; Xie & Yuille ( 2019 ) . Therefore , the use of a single BN to learn both distributions may limit model performance Wang et al . ( 2020 ) . To solve this problem , we extend the λ-conditioning to choose between two BNs , BNC and BNA , dedicated for IFMC and IFMA , respectively . This approach differs from previous efforts in several ways . Earlier research performed noise-injection via regularization Bietti et al . ( 2018 ) ; Lecuyer et al . ( 2019 ) and perturbed weight tensors He et al . ( 2019 ) to boost model robustness at the cost of a significant accuracy drop on clean images . In contrast , we use noise tensors to transform a shared weight tensor yielding a model that can be configured in-situ to provide SOTA accuracy on either clean or perturbed images . Our approach is similar to λ-conditioning used by Wang et al . ( 2020 ) . However , instead of transforming activations using added FiLMbased layers trained with multiple values of λ Wang et al . ( 2020 ) , we transform weight tensors using added noise conditioned by binary λ . Compared to Wang et al . ( 2020 ) , we thus require models with significantly fewer parameters and training scenarios , yielding faster training ( up to 1.43× ) .

Summary: This paper tackles the problem of in-situ trade-off between model robustness and accuracy, which is of great importance in real-world applications. The proposed method, FLOAT, makes improvements on the previous state-of-the-art method OAT (NeurIPS'20) by using a more efficient model conditioning method. Experimental results show that FLOAT outperforms OAT in terms of accuracy, robustness and model size. The FLOPs of FLOAT, OAT, and the PGDAT baseline are almost the same.