Enhancing Adversarial Defense by k-Winners-Take-All

1 INTRODUCTION . In the tremendous success of deep learning techniques , there is a grain of salt . It has become wellknown that deep neural networks can be easily fooled by adversarial examples ( Szegedy et al. , 2014 ) . Those deliberately crafted input samples can mislead the networks to produce an output drastically different from what we expect . In many important applications , from face recognition authorization to autonomous cars , this vulnerability gives rise to serious security concerns ( Barreno et al. , 2010 ; 2006 ; Sharif et al. , 2016 ; Thys et al. , 2019 ) . Attacking the network is straightforward . Provided a labeled data item ( x , y ) , the attacker finds a perturbation x′ perceptually similar to x but misleading enough to cause the network to output a label different from y . By far , the most effective way of finding such a perturbation ( or adversarial example ) is by exploiting the gradient information of the network with respect to its input : the gradient indicates how to perturb x to trigger the maximal change of y . The defense , however , is challenging . Recent studies showed that adversarial examples always exist if one tends to pursue a high classification accuracy—adversarial robustness seems at odds with the accuracy ( Tsipras et al. , 2018 ; Shafahi et al. , 2019a ; Su et al. , 2018a ; Weng et al. , 2018 ; Zhang et al. , 2019 ) . This intrinsic difficulty of eliminating adversarial examples suggests an alternative path : can we design a network whose adversarial examples are evasive rather than eliminated ? Indeed , along with this thought is a series of works using obfuscated gradients as a defense mechanism ( Athalye et al. , 2018 ) . Those methods hide the network ’ s gradient information by artificially discretizing the input ( Buckman et al. , 2018 ; Lin et al. , 2019 ) or introducing certain randomness to the input ( Xie et al. , 2018a ; Guo et al. , 2018 ) or the network structure ( Dhillon et al. , 2018 ; Cohen et al. , 2019 ) ( see more discussion in Sec . 1.1 ) . Yet , the hidden gradient in those methods can still be approximated , and as recently pointed out by Athalye et al . ( 2018 ) , those methods remain vulnerable . Technical contribution I ) Rather than obfuscating the network ’ s gradient , we make the gradient undefined . This is achieved by a simple change to the standard neural network structure : we advocate the use of aC0 discontinuous activation function , namely the k-Winners-Take-All ( k-WTA ) activation , to replace the popular activation functions such as rectified linear units ( ReLU ) . This is the only change we propose to a deep neural network . All other components ( such as BatchNorm , convolution , and pooling ) as well as the training methods remain unaltered . With no significant overhead , k-WTA activation can be readily used in nearly all existing networks and training methods . k-WTA activation takes as input the entire output of a layer , retains its k largest values and deactivates all others to zero . As we will show in this paper , even an infinitesimal perturbation to the input may cause a complete change to the network neurons ’ activation pattern , thereby resulting in a large jump in the network ’ s output . This means that , mathematically , if we use f ( x ; w ) to denote a k-WTA network taking an input x and parameterized by weights w , then the gradient∇xf ( x ; w ) at certain x is undefined—f ( x ; w ) is C0 discontinuous . Technical contribution II ) More intriguing than the mere replacement of the activation function is why k-WTA helps improve the adversarial robustness . We offer our theoretical reasoning of its behavior from two perspectives . On the one hand , we show that the discontinuities of f ( x ; w ) is densely distributed in the space of x . Dense enough such that a tiny perturbation from any x almost always comes across some discontinuities , where the gradients are undefined and thus the attacker ’ s search of adversarial examples becomes blinded ( see Figure 1 ) . On the other hand , a paradox seemingly exists . The discontinuities in the activation function also renders f ( x ; w ) discontinuous with respect to the network weights w ( at certain w values ) . But training the network relies on the presumption that the gradient with respect to the weights is almost always available . Interestingly , we show that , under k-WTA activation , the discontinuities of f ( x ; w ) is rather sparse in the space of w , intuitively because the dimension of w ( in parameter space ) is much larger than the dimension of x ( in data space ) . Thus , the network can be trained successfully . Summary of results . We conducted extensive experiments on multiple datasets under different network architectures , including ResNet ( He et al. , 2016 ) , DenseNet ( Huang et al. , 2017 ) , and Wide ResNet ( Zagoruyko & Komodakis , 2016 ) , that are optimized by regular training as well as various adversarial training methods ( Madry et al. , 2017 ; Zhang et al. , 2019 ; Shafahi et al. , 2019b ) . In all these setups , we compare the robustness performance of using the proposed k-WTA activation with commonly used ReLU activation under several white-box attacks , including PGD ( Kurakin et al. , 2016 ) , Deepfool ( Moosavi-Dezfooli et al. , 2016 ) , C & W ( Carlini & Wagner , 2017 ) , MIM ( Dong et al. , 2018 ) , and others . In all tests , k-WTA networks outperform ReLU networks . The use of k-WTA activation is motivated for defending against gradient-based adversarial attacks . Our experiments suggest that the robustness improvement gained by simply switching to k-WTA activation is universal , not tied to specific network architectures or training methods . To promote reproducible research , we will release our implementation of k-WTA networks , along with our experiment code , configuration files and pre-trained models1 . 1.1 RELATED WORK : OBFUSCATED GRADIENTS AS A DEFENSE MECHANISM . Before delving into k-WTA details , we review prior adversarial defense methods that share the same philosophy with our method and highlight our advantages . For a review of other attack and defense methods , we refer to Appendix A . Methods aiming for concealing the gradient information from the attacker has been termed as obfuscated gradients ( Athalye et al. , 2018 ) or gradient masking ( Papernot et al. , 2017 ; Tramèr et al. , 2017 ) techniques . One type of such methods is by exploiting randomness , either randomly transforming the input before feeding it to the network ( Xie et al. , 2018a ; Guo et al. , 2018 ) or introducing stochastic layers in the network ( Dhillon et al. , 2018 ) . However , the gradient information in these methods can be estimated by taking the average over multiple trials ( Athalye et al. , 2018 ; 2017 ) . As a result , these methods are vulnerable . Another type of obfuscated gradient methods relies on the so-called shattered gradient ( Athalye et al. , 2018 ) , which aims to make the network gradients nonexistent or incorrect to the attacker , by purposely discretizing the input ( Buckman et al. , 2018 ; Ma et al. , 2018 ) or artificially raising numerical instability for gradient evaluation ( Song et al. , 2018 ; Samangouei et al. , 2018 ) . Unfortunately , these methods are also vulnerable . As shown by Athalye et al . ( 2018 ) , they can be compromised by backward pass 1https : //github.com/a554b554/kWTA-Activation differentiable approximation ( BPDA ) . Suppose fi ( x ) is a non-differentiable component of a network expressed by f = f1 ◦ f2 ◦ · · · ◦ fn . The gradient ∇xf can be estimated as long as one can find a smooth delegate g that approximates well fi ( i.e. , g ( x ) ≈ fi ( x ) ) . In stark contrast to all those methods , a slight change of the k-WTA activation pattern in an earlier layer of a network can cause a radical reorganization of activation patterns in later layers ( as shown in Sec . 3 ) . Thereby , k-WTA activation not just obfuscates the network ’ s gradients but destroys them at certain input samples , introducing discontinuities densely distributed in the input data space . We are not aware of any possible smooth approximation of a k-WTA network to launch BPDA attacks . 2 k-WINNERS-TAKE-ALL ACTIVATION The debut of the Winner-Takes-All ( WTA ) activation on the stage of neural networks dates back to 1980s , when Grossberg ( 1982 ) introduced shunting short-term memory equations in on-center off-surround networks and showed the ability to identify the largest of N real numbers . Later , Majani et al . ( 1989 ) generalized the WTA network to identify the K largest of N real numbers , and they termed the network as the K-Winners-Take-All ( KWTA ) network . These early WTA-type activation functions output only boolean values , mainly motivated by the properties of biological neural circuits . In particular , Maass ( 2000a ; b ) has proved that any boolean function can be computed by a single KWTA unit . Yet , the boolean nature of these activation functions differs starkly from the modern activation functions , including the one we use . 2.1 DEEP NEURAL NETWORKS ACTIVATED BY k-WINNERS-TAKE-ALL We propose to use k-Winners-Take-All ( k-WTA ) activation , a natural generalization of the boolean KWTA2 ( Majani et al. , 1989 ) . k-WTA retains the k largest values of an N × 1 input vector and sets all others to be zero before feeding the vector to the next network layer , namely , φk ( y ) j = { yj , yj ∈ { k largest elements of y } , 0 , Otherwise . ( 1 ) Here φk : RN → RN is the k-WTA function ( parameterized by an integer k ) , y ∈ RN is the input to the activation , and φk ( y ) j denote the j-the element of the output φk ( y ) ( see the rightmost subfigure of Figure 2 ) . Note that if y has multiple elements that are equally k-th largest , we break the tie by retaining the element with smaller indices until the k slots are taken . When using k-WTA activation , we need to choose k. Yet it makes no sense to fix k throughout all layers of the neural network , because these layers often have different output dimensions ; a small k to one layer ’ s dimension can be relatively large to the other . Instead of specifying k , we introduce 2In this paper , we use k-WTA to refer our activation function , while using KWTA to refer the original boolean version by Majani et al . ( 1989 ) . a parameter γ ∈ ( 0 , 1 ) called sparsity ratio . If a layer has an output dimension N , then its k-WTA activation has k = bγ ·Nc . Even though the sparsity ratio can be set differently for different layers , in practice we found no clear gain from introducing such a variation . Therefore , we use a fixed γ—the only additional hyperparameter needed for the neural network . In convolutional neural networks ( CNN ) , the output of a layer is a C ×H ×W tensor . C denotes the number of output channels ; H and W indicate the feature resolution . While there are multiple choices of applying k-WTA on the tensor—for example , one can apply k-WTA individually to each channel—empirically we found that the most effective ( and conceptually the simplest ) way is to treat the tensor as a long C ·H ·W × 1 vector input to the k-WTA activation . Using k-WTA in this way is also backed by our theoretical understanding ( see Sec . 3 ) . The runtime cost of computing a k-WTA activation is asymptotically O ( N ) , because finding k largest values in a list is asymptotically equivalent to finding the k-th largest value , which has an O ( N ) complexity ( Cormen et al. , 2009 ) . This cost is comparable to ReLU ’ s O ( N ) cost on a N -length vector . Thus , replacing ReLU with k-WTA introduces no significant overhead . Remark : other WTA-type activations . Relevant to k-WTA is the local Winner-Take-All ( LWTA ) activation ( Srivastava et al. , 2013 ; 2014 ) , which divides each layer ’ s output values into local groups and applies WTA to each group individually . LWTA is similar to max-pooling ( Riesenhuber & Poggio , 1999 ) for dividing the layer output and choosing group maximums . But unlike ReLU and max-pooling being C0 continuous , LWTA and our k-WTA are both discontinuous with respect to the input . The differences among ReLU , max-pooling , LWTA , and k-WTA are illusrated in Figure 2 . LWTA is motivated toward preventing catastrophic forgetting ( McCloskey & Cohen , 1989 ) , whereas our use of k-WTA is for defending against adversarial threat . Both are discontinuous . But it remains unclear what the LWTA ’ s discontinuity properties are and how its discontinuities affect the network training . Our theoretical analysis ( Sec . 3 ) , in contrast , sheds some light on these fundamental questions about k-WTA , rationalizing its ability for improving adversarial robustness . Indeed , our experiments confirm that k-WTA outperforms LWTA in terms of robustness ( see Appendix D.3 ) . WTA-type activation , albeit originated decades ago and widely studied in computational neuroscience ( Douglas et al. , 1989 ; Douglas & Martin , 2004 ) , remains elusive in modern neural networks . Perhaps this is because it has not demonstrated a considerable improvement to the network ’ s standard test accuracy , though it can offer an accuracy comparable to ReLU ( Srivastava et al. , 2013 ) . Our analysis and proposed use of k-WTA and its enabled improvement on adversarial defense may suggest a renaissance of studying WTA . 2.2 TRAINING k-WTA NETWORKS k-WTA networks require no special treatment in training . Any optimization algorithm ( such as stochastic gradient descent ) for training ReLU networks can be directly used to train k-WTA networks . Our experiments have found that when the sparsity ratio γ is relatively small ( ≤ 0.2 ) , the network training converges slowly . This is not a surprise . A smaller γ activates fewer neurons , effectively reducing more of the layer width and in turn the network size , and the stripped “ subnetwork ” is much less expressive ( Srivastava et al. , 2013 ) . Since different training examples activate different subnetworks , collectively they make the training harder . Nevertheless , we prefer a smaller γ . As we will discuss in the next section , a smaller γ usually leads to better robustness against finding adversarial examples . Therefore , to ease the training ( when γ is small ) , we propose to use an iterative fine-tuning approach . Suppose the target sparsity ratio is γ1 . We first train the network with a larger sparsity ratio γ0 using the standard training process . Then , we iteratively fine tune the network . In each iteration , we reduce its sparsity ratio by a small δ and train the network for two epochs . The iteration repeats until γ0 is reduced to γ1 . This incremental process introduces little training overhead , because the cost of each fine tuning is negligible in comparison to training from scratch toward γ0 . We also note that this process is optional . In practice we use it only when γ < 0.2 . We show more experiments on the efficacy of the incremental training in Appendix D.2 . 3 UNDERSTANDING k-WTA DISCONTINUITY We now present our theoretical understanding of k-WTA ’ s discontinuity behavior in the context of deep neural networks , revealing some implication toward the network ’ s adversarial robustness . Activation pattern . To understand k-WTA ’ s discontinuity , consider one layer outputting values x , passed through a k-WTA activation , and followed by the next layer whose linear weight matrix is W ( see adjacent figure ) . Then , the value fed into the next activation can be expressed as Wφk ( x ) , where φk ( · ) is the k-WTA function defined in ( 1 ) . Suppose the vector x has a length l. We define the k-WTA ’ s activation pattern under the input x as A ( x ) : = { i ∈ [ l ] | xi is one of the k largest values in x } ⊆ [ l ] . ( 2 ) Here ( and throughout this paper ) , we use [ l ] to denote the integer set { 1 , 2 , ... , l } . Discontinuity . The activation pattern A ( x ) is a key notion for analyzing k-WTA ’ s discontinuity behavior . Even an infinitesimal perturbation of x may change A ( x ) : some element i is removed from A ( x ) while another element j is added in . Corresponding to this change , in the evaluation of Wφk ( x ) , the contribution of W ’ s column vector Wi vanishes while another column Wj suddenly takes effect . It is this abrupt change that renders the result of Wφk ( x ) C0 discontinuous . Such a discontinuity jump can be arbitrarily large , because the column vectors Wi and Wj can be of any difference . Once W is determined , the discontinuity jump then depends on the value of xi and xj . As explained in Appendix B , when the discontinuity occurs , xi and xj have about the same value , depending on the choice of the sparsity ratio γ ( recall Sec . 2.1 ) —the smaller the γ is , the larger the jump will be . This relationship suggests that a smaller γ will make the search of adversarial examples harder . Indeed , this is confirmed through our experiments ( see Appendix D.6 ) . Piecewise linearity . Now , consider an n-layer k-WTA network , which can be expressed as f ( x ) = W ( 1 ) · φk ( W ( 2 ) · φk ( · · ·φk ( W ( n ) x + b ( n ) ) ) + b ( 2 ) ) + b ( 1 ) , where W ( i ) and b ( i ) are the i-th layer ’ s weight matrix and bias vector , respectively . If the activation patterns of all layers are fixed , then f ( x ) is a linear function . When the activation pattern changes , f ( x ) switches from one linear function to another linear function . Over the entire space of x , f ( x ) is piecewise linear . The specific activation patterns of all layers define a specific linear piece of the function , or a linear region ( following the notion introduced by Montufar et al . ( 2014 ) ) . Conventional ReLU ( or hard tanh ) networks also represent piecewise linear functions and their linear regions are joined together at their boundaries , whereas in k-WTA networks the linear regions are disconnected ( see Figure 1 ) . Linear region density . Next , we gain some insight on the distribution of those linear regions . This is of our interest because if the linear regions are densely distributed , a small ∆x perturbation from any data point x will likely cross the boundary of the linear region where x locates . Whenever boundary crossing occurs , the gradient becomes undefined ( see Figure 3-a ) . For the purpose of analysis , consider an input x passing through a layer followed by a k-WTA activation ( see adjacent figure ) . The output from the activation is φk ( Wx + b ) . We would like to understand , when x is changed into x′ , how likely the activation pattern of φk will change . First , notice that if x′ and x satisfy x′ = c · x with some c > 0 , the activation pattern remains unchanged . Therefore , we introduce a notation d ( x , x′ ) that measures the “ perpendicular ” distance between x and x′ , one that satisfies x′ = c · ( x + d ( x , x′ ) x⊥ ) for some scalar c , where x⊥ is a unit vector perpendicular to x and on the plane spanned by x and x′ . With this notion , and if the elements in W is initialized by sampling from N ( 0 , 1l ) and b is initialized as zero , we find the following property : Theorem 1 ( Dense discontinuities ) . Given any input x ∈ Rm and some β , and ∀x′ ∈ Rm such that d2 ( x , x′ ) ‖x‖22 ≥ β , if the following condition l ≥ Ω ( ( m γ · 1 β ) · log ( m γ · 1 β ) ) is satisfied , then with a probability at least 1− ·2−m , we have A ( Wx + b ) 6= A ( Wx′ + b ) . Here l is the width of the layer , and γ is again the sparsity ratio in k-WTA . This theorem informs us that the larger the layer width l is , the smaller β—and thus the smaller perpendicular perturbation distance d ( x , x′ ) —is needed to trigger a change of the activation pattern , that is , as the layer width increases , the piecewise linear regions become finer ( see Appendix C for proof and more discussion ) . This property also echos a similar trend in ReLU networks , as pointed out by Raghu et al . ( 2017 ) . Why is the k-WTA network trainable ? While k-WTA networks are highly discontinuous as revealed by Theorem 1 and our experiments ( Figure 3-a ) , in practice we experience no difficulty on training these networks . Our next theorem sheds some light on the reason behind the training success . Theorem 2 . Consider N data points x1 , x2 , · · · , xN ∈ Rm . Suppose ∀i 6= j , xi‖xi‖2 6= xj ‖xj‖2 . If l is sufficiently large , then with a high probability , we have ∀i 6= j , A ( Wxi + b ) ∩ A ( Wxj + b ) = ∅ . This theorem is more formally stated in Theorem 10 in Appendix C together with a proof there . Intuitively speaking , it states that if the network is sufficiently wide , then for any i 6= j , activation pattern of input data xi is almost separated from that of xj . Thus , the weights for predicting xi ’ s and xj ’ s labels can be optimized almost independently , without changing their individual activation patterns . In practice , the activation patterns of xi and xj are not fully separated but weakly correlated . During the optimization , the activation pattern of a data point xi may change , but the chance is relatively low—a similar behavior has also been found in ReLU networks ( Li & Liang , 2018 ; Du et al. , 2018 ; Allen-Zhu et al. , 2019a ; b ; Song & Yang , 2019 ) . Further , notice that the training loss takes a summation over all training data points . This means a weight update would change only a small set of activation patterns ( since the chance of having the pattern changed is low ) ; the discontinuous change on the loss value , after taking the summation , will be negligible ( see Figure 3-c ) . Thus , the discontinuities in k-WTA is not harmful to network training .

Summary: This paper addresses the important question of improving the robustness of deep neural networks against adversarial attacks. The authors propose a surprisingly simple measure to improve adversarial robustness, namely replacing typical activation functions such as ReLU with a k-winners-take-all (k-WTA) functions, whereby the k largest values of the input vector are copied, and all other elements of the output vector are set to zero. Since the size of input and output maps varies drastically within networks, the authors instead use a sparsity ratio \gamma that calculates k as a function of input size. k-WTA networks can be trained without special treatment, but for low \gamma values the authors propose a training schedule, whereby \gamma is slowly reduced, then re-training takes place, until the desired value of \gamma is reached. The presented effect is backed up by extensive theoretical investigations that relate the increased robustness to the dense introduction of discontinuities, which makes gradient-based adversarial attacks harder. A small change in an input signal can change the identity of the "winning" inputs, and thus in a sub-sequent matrix multiplication make use of other rows or columns, thus allowing arbitrarily large effects due to small input variations. Empirical evaluations in CIFAR and SVHN for a variety of attacks and defense mechanisms demonstrate the desired effects, and illustrate the loss landscapes due to using k-WTA.