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: The paper proposes a modification of the gradients of objective function wrt model parameters by changing the underling metric on the manifold of model parameters.  It proposes two metrics, a reweighed L2 metric, and a specific Sobolev metric, and presents how to compute modified gradients in a compute efficient manner.  Motivated by empirical observation that parameters of trained CNN tensors tend to be correlated along output dimension, the paper focuses on metrics that promote such correlation.