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: The authors propose using k-winner take all (k-WTA) activation functions to prevent white box adversarial attacks. A k-WTA activation functions outputs the k highest activations in a layer while setting all other activations to zero. The reasoning given by the authors is that k-WTA activation functions have many discontinuities with respect to the input space. This makes it more difficult for attacks to use gradient information. The authors note that networks with k-WTA activation functions are still easy to train because, for a given input, the sub-network that is activated becomes more stable as training progresses. Therefore, it is not as discontinuous in the parameter space.