Fantastic Four: Differentiable and Efficient Bounds on Singular Values of Convolution Layers

1 INTRODUCTION . Bounding singular values of different layers of a neural network is a way to control the complexity of the model and has been used in different problems including robustness , generalization , optimization , generative modeling , etc . In particular , the spectral norm ( the maximum singular value ) of a layer bounds the factor by which the norm of the signal increases or decreases during both forward and backward propagation within that layer . If all singular values are all close to one , then the gradients neither explode nor vanish ( Hochreiter , 1991 ; Hochreiter et al. , 2001 ; Klambauer et al. , 2017 ; Xiao et al. , 2018 ) . Spectral norm regularizations/bounds have been used in improving the generalization ( Bartlett et al. , 2017 ; Long & Sedghi , 2020 ) , in training deep generative models ( Arjovsky et al. , 2017 ; Gulrajani et al. , 2017 ; Tolstikhin et al. , 2018 ; Miyato et al. , 2018 ; Hoogeboom et al. , 2020 ) and in robustifying models against adversarial attacks ( Singla & Feizi , 2020 ; Szegedy et al. , 2014 ; Peck et al. , 2017 ; Zhang et al. , 2018 ; Anil et al. , 2018 ; Hein & Andriushchenko , 2017 ; Cisse et al. , 2017 ) . These applications have resulted in multiple works to regularize neural networks by penalizing the spectral norm of the network layers ( Drucker & Le Cun , 1992 ; Yoshida & Miyato , 2017 ; Miyato et al. , 2018 ; 2017 ; Sedghi et al. , 2019 ; Singla & Feizi , 2020 ) . For a fully connected layer with weights W and bias b , the lipschitz constant is given by the spectral norm of the weight matrix i.e , ∥W∥2 , which can be computed efficiently using the power iteration method ( Golub & Van Loan , 1996 ) . In particular , if the matrix W is of size p × q , the computational complexity of power iteration ( assuming convergence in constant number of steps ) is O ( pq ) . Convolution layers ( Lecun et al. , 1998 ) are one of the key components of modern neural networks , particularly in computer vision ( Krizhevsky et al. , 2012 ) . Consider a convolution filter L of size cout × cin × h ×w where cout , cin , h and w denote the number of output channels , input channels , height and width of the filter respectively ; and a square input sample of size cin ×n×n where n is its height and width . A naive representation of the Jacobian of this layer will result in a matrix of size n2cout × n 2cin . For a typical convolution layer with the filter size 64 × 3 × 7 × 7 and an ImageNet sized input 3 × 224 × 224 ( Krizhevsky et al. , 2012 ) , the corresponding jacobian matrix has a very large size : 802816 × 150528 . This makes an explicit computation of the jacobian infeasible . Ryu et al . ( 2019 ) provide a way to compute the spectral norm of convolution layers using convolution and transposed convolution operations in power iteration , thereby avoiding this explicit computation . This leads to an improved running time especially when the number of input/output channels is small ( Table 1 ) . However , in addition to the running time , there is an additional difficulty in the approach proposed in Ryu et al . ( 2019 ) ( and other existing approaches described later ) regarding the computation of the spectral norm gradient ( often used as a regularization during the training ) . The gradient of the largest singular value with respect to the jacobian can be naively computed by taking the outer product of corresponding singular vectors . However , due to the special structure of the convolution operation , the jacobian will be a sparse matrix with repeated elements ( see Appendix Section D for details ) . The naive computation of the gradient will result in non-zero gradient values at elements that should be in fact zeros throughout training and also will assign different gradient values at elements that should always be identical . These issues make the gradient computation of the spectral norm with respect to the convolution filter weights using the technique of Ryu et al . ( 2019 ) difficult . Recently , Sedghi et al . ( 2019 ) provided a principled approach for exactly computing the singular values of convolution layers . They construct n2 matrices each of size cout × cin by taking the Fourier transform of the convolution filter ( details in Appendix Section B ) . The set of singular values of the jacobian equals the union of singular values of these n2 matrices . However , this method can have high computational complexity since it requires SVD of n2 matrices . Although this method can be adapted to compute the spectral norm of n2 matrices using power iteration ( in parallel with a GPU implementation ) instead of full SVD , the intrinsic computational complexity ( discussed in Table 2 ) can make it difficult to use this approach for very deep networks and large input sizes especially when computational resources are limited . Moreover , computing the gradient of the spectral norm using this method is not straightforward since each of these n2 matrices contain complex numbers . Thus , Sedghi et al . ( 2019 ) suggests to clip the singular values if they are above a certain threshold to bound the spectral norm of the layer . In order to reduce the training overhead , they clip the singular values only after every 100 iterations . The resulting method reduces the training overhead but is still costly for large input sizes and very deep networks . We report the running time of this method in Table 1 and its training time for one epoch ( using 1 GPU implementation ) in Table 4c . Because of the aforementioned issues , efficient methods to control the spectral norm of convolution layers have resorted to heuristics ( Yoshida & Miyato , 2017 ; Miyato et al. , 2018 ; Gouk et al. , 2018 ) . Typically , these methods reshape the convolution filter of dimensions cout × cin × h ×w to construct a matrix of dimensions cout × hwcin , and use the spectral norm of this matrix as an estimate of the spectral norm of the convolution layer . To regularize during training , they use the outer product of the corresponding singular vectors as the gradient of the largest singular value with respect to the reshaped matrix . Since the weights do not change significantly during each training step , they use only one iteration of power method during each step to update the singular values and vectors ( using the singular vectors computed in the previous step ) . These methods result in negligible overhead during the training . However , due to lack of theoretical justifications ( which we resolve in this work ) , they are not guaranteed to work for all different shapes and weights of the convolution filter . Previous studies have observed under estimation of the spectral norm using these heuristics ( Jiang et al. , 2019 ) . On one hand , there are computationally efficient but heuristic ways of computing and bounding the spectral norm of convolutional layers ( Miyato et al. , 2017 ; 2018 ) . On the other hand , the exact computation of the spectral norm of convolutional layers proposed by Sedghi et al . ( 2019 ) ; Ryu et al . ( 2019 ) can be expensive for commonly used architectures especially with large inputs such as ImageNet samples . Moreover , the difficulty in computing the gradient of the spectral norm with respect to the jacobian under these methods make their use as regularization during the training process challenging . In this paper , we resolve these issues by deriving a differentiable and efficient upper bound on the spectral norm of convolutional layers . Our bound is provable and not based on heuristics . Our computational complexity is similar to that of heuristics ( Miyato et al. , 2017 ; 2018 ) allowing our bound to be used as a regularizer for efficiently training deep convolutional networks . In this way , our proposed approach combines the benefits of the speed of the heuristics and the theoretical rigor of Sedghi et al . ( 2019 ) . Table 2 summarizes the differences between previous works and our approach . In Table 1 , we empirically observe that our bound can be computed in a time significantly faster than Sedghi et al . ( 2019 ) ; Ryu et al . ( 2019 ) , while providing a guaranteed upper bound on the spectral norm . Moreover , we empirically observe that our upper bound and the exact value are close to each other ( Section 3.1 ) . Below , we briefly explain our main result . Consider a convolution filter L of dimensions cout × cin × h ×w and input of size cin × n × n. The corresponding jacobian matrix J is of size n2cout × n2cin . We show that the largest singular value of the jacobian ( i.e . ∥J∥2 ) is bounded as : ∥J∥2 ≤ √ hwmin ( ∥R∥2 , ∥S∥2 , ∥T∥2 , ∥U∥2 ) , where R , S , T and U are matrices of sizes hcout×wcin , wcout×hcin , cout×hwcin and hwcout×cin respectively , and can be computed by appropriately reshaping the filter L ( details in Section 3 ) . This upper bound is independent of the input width and height ( n ) . Formal results are stated in Theorem 1 and proved in the appendix . Remarkably , ∥T∥2 is the heuristic suggested by Miyato et al . ( 2018 ) . To the best of our knowledge , this is the first work that derives a provable bound on the spectral norm of convolution filter as a constant factor ( dependant on filter sizes , but not filter weights ) times the heuristic of Miyato et al . ( 2018 ) . In Tables 1 and 3 , we show that the other 3 bounds ( using ∥R∥2 , ∥S∥2 , ∥U∥2 ) can be significantly smaller than √ hw∥T∥2 for different convolution filters . Thus , we take the minimum of these 4 quantities to bound the spectral norm of a convolution filter . In Section 4 , we show that our bound can be used to improve the generalization and robustness properties of neural networks . Specifically , we show that using our bound as a regularizer during training , we can achieve improvement in accuracy on par with exact method ( Sedghi et al. , 2019 ) while being significantly faster to train ( Table 4 ) . We also achieve significantly higher robustness certificates against adversarial attacks than CNN-Cert ( Boopathy et al. , 2018 ) on a single layer CNN ( Table 5 ) . These results demonstrate potentials for practical uses of our results . Code is available at the github repository : https : //github.com/singlasahil14/fantastic-four . 2 NOTATION . For a vector v , we use vj to denote the element in the jth position of the vector . We use Aj , ∶ and A∶ , k to denote the jth row and kth column of the matrix A respectively . We assume both Aj , ∶ , A∶ , k to be column vectors ( thus Aj , ∶ is the transpose of jth row of A ) . Aj , k denotes the element in jth row and kth column of A . The same rules can be directly extended to higher order tensors . For a matrix A ∈ Rq×r and a tensor B ∈ Rp×q×r , vec ( A ) denotes the vector constructed by stacking the rows of A and vec ( B ) denotes the vector constructed by stacking the vectors vec ( Bj , ∶ , ∶ ) , j ∈ [ p−1 ] : vec ( A ) T = [ AT0 , ∶ , A T 1 , ∶ , ⋯ , A T q−1 , ∶ ] vec ( B ) T = [ BT0 , ∶ , ∶ , B T 1 , ∶ , ∶ , ⋯ , B T p−1 , ∶ , ∶ ] We use the following notation for a convolutional neural network . L denotes the number of layers and φ is the activation function . For an input x , we use z ( I ) ( x ) ∈RNI and a ( I ) ( x ) ∈RNI to denote the raw ( before applying φ ) and activated ( after applying φ ) neurons in the Ith hidden layer respectively . a ( 0 ) denotes the input image x . To simplify notation and when no confusion arises , we make the dependency of z ( I ) and a ( I ) to x implicit . φ ′ ( z ( I ) ) and φ ′′ ( z ( I ) ) denotes the elementwise first and second derivatives of φ at z ( 1 ) . W ( I ) denotes the weights for the Ith layer i.e W ( I ) will be a tensor for a convolution layer and a matrix for a fully connected layer . J ( I ) denotes the jacobian matrix of vec ( z ( 1 ) ) with respect to the input x. θ denotes the neural network parameters . fθ ( x ) denotes the softmax probabilities output by the network for an input x . For an input x and label y , the cross entropy loss is denoted by ` ( fθ ( x ) , y ) .

Summary: This paper propose to study the Lipschitz constant of convolutional layers and to give an easy to compute and differentiable upper bound. The upper bound is composed of 4 different bounds, based on tensor unfolding of the Jacobian, and taking the min of these 4 values. This upper bound is then used to train networks with spectral norm regularization and compared with network trained with singular value clipping from `Sedghi et al. (2019)`. The proposed bound gives similar performances with much cheaper computational time.