Channel-Directed Gradients for Optimization of Convolutional Neural Networks

1 INTRODUCTION . Stochastic gradient descent ( SGD ) is currently the dominant algorithm for optimizing large-scale convolutional neural networks ( CNNs ) ( LeCun et al . ( 1998 ) ; Simonyan & Zisserman ( 2014 ) ; He et al . ( 2016b ) ) . Although there has been large activity in optimization methods seeking to improve performance , SGD still dominates in terms of its generalization ability . Despite SGD ’ s dominance , there is still often a gap between training and real-world test accuracy performance , which motivates research in improved optimization methods . In this paper , we derive new optimization methods that are simple modifications of SGD . The methods implicitly induce correlation in the output direction of parameter tensors in CNNs . This is based on the empirical observation that parameter tensors in trained networks typically exhibit correlation over output channel dimension ( see Figure 1 ) . We thus explore encoding correlation by constructing smooth gradients in the output direction , which we show improves generalization accuracy . This is done by introducing new Riemmanian metrics on the parameter tensors , which changes the underlying geometry of the space of tensors , and reformulating the gradient with respect to those metrics . Our contributions are as follows . First , we formulate output channel-directed Riemannian metrics ( a re-weighted version of the standard L2 metric and another that is a Sobolev metric ) over the space of parameter tensors . This encodes channel-directed correlation in the gradient optimization without changing the loss . Second , we compute Riemannian gradients with respect to the metrics showing linear complexity ( in the number of parameters ) over standard gradient computation , and thus derive new optimization methods for CNN training . Finally , we apply the methodology to training CNNs and show the empirical advantage in generalization accuracy , especially with small batch sizes , over standard optimizers ( SGD , Adam ) on numerous applications ( image classification , semantic segmentation , generative adversarial networks ) with simple modification of existing optimizers . 1.1 RELATED WORK . We discuss related work in deep network optimization ; for a detailed survey , see Bottou et al . ( 2018 ) . SGD , e.g. , Bottou ( 2012 ) , samples a batch of data to tractably estimate the gradient of the loss function . As the stochastic gradient is a noisy version of the gradient , learning rates must follow a decay schedule in order to converge . Many methods have been formulated to choose learning rate over epochs and components of the gradient , including adaptive learning rates ( e.g. , Duchi et al . ( 2011 ) ; Zeiler ( 2012 ) ; Kingma & Ba ( 2014 ) ; Bengio ( 2015 ) ; Loshchilov & Hutter ( 2017 ) ; Luo et al . ( 2019 ) ) . For instance , Adam Kingma & Ba ( 2014 ) adaptively adjusts the learning rate so that parameters that have changed infrequently based on historical gradients are updated more quickly than parameters that have changed frequently . Another way to interpret such methods is that they change the underlying metric on the space on which the loss function is defined to an iso-tropically scaled version of the L2 metric given by a simple diagonal matrix ; we change the metrics an-isotropically . We show that our method can be used in conjunction with such methods by simply using the stochastic gradient computed with our metrics to boost performance . As the stochastic gradient is computed based on sampling , different runs of the algorithm can result in different local optima . To reduce the variance , several methods have been been formulated , e.g. , Defazio et al . ( 2014 ) ; Johnson & Zhang ( 2013 ) . We are not motivated by variance reduction , rather , inducing correlation in the parameter tensor to improve generalization . However , as our method smooths the gradient , our experiments show reduced variance with our metrics compared to SGD . Another method motivated by variance reduction is Osher et al . ( 2018 ) ( see applications Wang et al . ( 2019 ) ; Liang et al . ( 2020 ) ; Wang et al . ( 2020 ) ) , where the stochastic gradient is pre-multiplied with an inverse Laplacian smoothing matrix . For CNNs , the gradient with respect to parameters is rasterized in row or column order of network filters before smoothing . Our work is inspired by Osher et al . ( 2018 ) , though we are motivated by correlation in the parameter tensor . Osher et al . ( 2018 ) can be interpreted as using the gradient of the loss with respect to a Sobolev metric . One insight over Osher et al . ( 2018 ) is that keeping the structure of the parameter tensor and defining the Sobolev metric with respect to the output-channel direction boosts accuracy , while other directions do not . Secondly , we introduce a re-weightedH0 metric that preferentially treats the output-channel direction , and can be implemented with a line of Pytorch code , has linear ( in parameter size ) complexity , and performs comparably ( in many cases ) to our channel-directed Sobolev metric , boosting accuracy of SGD . Third , our Sobolev gradient , a variant of the ordinary one , has linear complexity rather than quasi-linear ( not requiring FFT as Osher et al . ( 2018 ) ) . Sobolev gradients have been used in computer vision Sundaramoorthi et al . ( 2007 ) ; Charpiat et al . ( 2007 ) for their coarse-to-fine evolution Sundaramoorthi et al . ( 2008 ) ; we adapt that formulation to CNNs . We formulate Sobolev gradients by considering the space of parameter tensors as a Riemannian manifold , and choosing the Sobolev metric on the tangent space . By choosing a metric , gradients intrinsic to the manifold can be computed and gradient flows are guaranteed to decrease loss . Other Riemannian metrics have been used for optimization in neural networks , e.g. , Amari ( 1998 ) ; MarceauCaron & Ollivier ( 2016 ) ; Hoffman et al . ( 2013 ) ; Gunasekar et al . ( 2020 ) and tangentially relate to our work . These works are based on Amari ’ s Amari ( 1998 ) information geometry on probability measures , and the metric considered is the Fisher information metric . The motivation for these methods is re-parametrization invariance of optimization , whereas our motivation is imposing correlation in the parameter space . Other works Gunasekar et al . ( 2020 ) use the Hessian metric ( in the convex case ) , but these metrics are data-dependent and the gradient is challenging to compute , requiring ( a large ) inverse matrix computation . 2 CHANNEL-DIRECTED GRADIENTS . We now present the theory to define channel-directed gradients . To do this , we formulate new metrics on the space of tensors , and then derive analytic formulas for channel-directed gradients in terms of the standard L2 gradient . As we show , our channel-directed gradients effectively smooth the components of the L2 gradient across the output direction of the parameter tensors of the CNN , which induces correlation in that direction in the gradient and thus also the parameter tensor . Another interpretation is we are changing the geometry of the loss landscape ( without changing the loss ) to a more smooth one by changing the metric of the space on which the loss is defined . 2.1 BACKGROUND ON RIEMANNIAN GRADIENTS . We present the definition of gradient on a Riemannian manifold , and show the dependence of the gradient on the chosen metric on the manifold ( see Carmo ( 1992 ) ; Abraham et al . ( 2012 ) for more details ) . A manifold X is a space that is locally linear around each point X ∈ X ; this linear space is the tangent space , denoted TXX . A Riemannian manifold has a smoothly varying positive definite bilinear form 〈· , ·〉 ( called the metric ) on the tangent space . This metric allows one to define the notion of lengths of curves on the space , in addition to other operations , including gradients . Definition 1 ( Gradient of a Function ) Let X be a Riemannian manifold , and f : X → R be a function . The directional derivative of f at X ∈ X along a direction k ∈ TXX is defined as df ( X ) · k = ddε f ( X + εk ) |ε=0 . The gradient of f at X ∈ X is the vector , ∇f ( X ) ∈ TXX , that satisfies the relation df ( X ) · k = 〈∇f ( X ) , k〉 , for all k ∈ TXX . ( 1 ) Note that “ the ” gradient will depend on the choice of the metric on the manifold . We note that any such gradient will decrease the the function f by moving infinitesimally in the tangent space in the direction of negative the gradient as df ( X ) · k = −‖∇f ( X ) ‖2 < 0 when k = −∇f ( X ) , where ‖ · ‖ is the norm induced from the metric . The gradient flow , defined by the differential equation Ẋt = −∇f ( Xt ) , will converge to a local minimum . In our application of this theory to CNN optimization , f will be the loss function , and X will be the space of parameter tensors . A consequence of this definition is that the gradient is the direction ( up to a scale factor ) in the tangent space that optimizes the following problem : arg max k∈TXX\ { 0 } |df ( X ) · k| ‖k‖ . ( 2 ) Thus , the gradient can be regarded as the most efficient direction as it maximizes the ratio of the change in energy by perturbing in a direction k over the cost ( defined by the metric ) of k. Thus , by constructing the metric to have small costs for perturbations ( directions ) that we prefer for gradients , the gradient flow will move in these preferential directions while minimizng the function , and thus land in favorable local minima . 2.2 CHANNEL-DIRECTED METRICS . In existing deep network gradient-based optimization schemes , the underlying metric on the loss function is assumed to be the standard Euclidean L2 metric . We will consider a re-weighted version of the L2 metric and a Sobolev metric that favor correlation in the output channel direction of the gradient and thus the parameter tensors . To formulate the methodology , we start from a continuum formulation , where we treat weight tensors in the continuum , formulate the metrics in the continuum and then in the next sub-section derive the gradients with respect to these metrics . Finally , we discretize gradient flows in the implementation to derive iterative schemes . Let X : O × I × H × W → R denote a parameter tensor of a layer of a convolutional neural network . Here O = [ 0 , O ] denotes indices to the output channel dimension of the tensor , I = [ 0 , I ] Figure 2 : Visualization of kernels applied to the H0 gradient under different metrics for λ = 1 . This illustrates the smoothing effect of the metrics . In computation , linear cost formulas are applied to compute the gradients not using the convolution interpretation . denote the indices to the input channel , and H = [ 0 , H ] , W = [ 0 , W ] denote indices to the height and width dimension of the spatial filters of the tensor . The metric is defined on the tangent space to the space of such X . An element of the tangent space will have the same form of the tensor , i.e. , k : O × I ×H×W → R. The L2 ( called H0 from now on ) metric is defined as 〈k1 , k2〉 H0 = ∫ O , I , H , W k1 ( o , i , h , w ) · k2 ( o , i , h , w ) do didhdw , ( 3 ) where k1 , k2 are in the tangent space of tensors . We now define a re-weighted version of H0 that favors tangent vectors that have global smoothness in the direction of the O dimension : 〈k1 , k2〉H0λ = ∫ I , H , W k̄1 ( i , h , w ) · k̄2 ( i , h , w ) didhdw + λ O 〈 k1 − k̄1 , k2 − k̄2 〉 H0 , ( 4 ) where λ > 0 is a hyper-parameter , and k̄ is the average value in the output channel direction , i.e. , k̄ ( i , h , w ) = 1 O ∫ O k ( o , i , h , w ) do . ( 5 ) The metric in ( 4 ) splits the tangent vector into global translations in the output channel direction and its orthogonal complement , i.e. , the deformation . The weight λ is used to control the weighting between the translation and deformation components , i.e. , larger values of λ means that deformations more heavily influence the norm of the perturbation . As shown in the next sub-section that means gradients with respect to this metric have higher weighted channel-directed translations than deformations . Next , we introduce a channel-directed version of a Sobolev metric , defined as follows : 〈k1 , k2〉H̃1 = ∫ I , H , W k̄1 ( i , h , w ) · k̄2 ( i , h , w ) di dhdw + λO 〈 ∂k1 ∂o , ∂k2 ∂o 〉 H0 , ( 6 ) where ∂∂o indicates the partial derivative with respect to the the output channel direction . The partial derivative in the o-direction implies that tensor perturbations that are smooth along the o-direction are close with respect to these metrics , which will imply that the corresponding gradients will exhibit smoothness in this direction , i.e. , convolution filters that are nearby in the output direction will exhibit correlation . The metric is a weighted combination of the H0 metric of the derivative in the output direction , and the H0 metric of the output-directed translation . Note that the traditional Sobolev metric uses the H0 metric of the perturbation rather than the translation . Our choice is motivated by computational efficiency of the corresponding gradient , to be discussed below . The scale factors of O in the expressions above are so that the metric is scale invariant with respect to different sizes of output channels . The part of the metric with the partial derivative component implies that tensors that differ in the output channel direction by a non-smooth perturbation are far away in distance . Tensors that differ by just a channel-directed translation are close .

Summary: This paper follows a very active line of research on gradient-based optimization methods for convolutional neural networks. In particular, this work proposes a method that can be applied as an extension to popular optimization methods such as SGD and Adam. The main idea is to introduce a modification of the tensor space in order to encourage correlations in the output-channel dimension. In particular, they propose two  channel re-weighting strategies: H^0 and H^1.  The theoretical foundation behind the paper seems to be sound and the overall writing is clear. Results indicate that, indeed, the proposed modifications help to improve the performance of the baselines: SGD and Adam. This is shown using several datasets. Furthermore, the proposed method does not produce a significant overhead with respect to the regular operation of the baselines.