Channel-Directed Gradients for Optimization of Convolutional Neural Networks

We briefly survey related work in deep network optimization; a detailed survey is Bottou et al. (2018). Stochastic gradient descent (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 recent work on 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 $\mathbb{L}^{2}$ 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 we are motivated by inducing regularity in optimization to improve generalization. However, as our method effectively smooths the gradient, our empirical experiments do indicate reduced variance with our metrics compared to SGD.

Another method motivated by variance reduction is the recent work of Osher et al. (2018), 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, which still lowers variance. Our work is inspired by Osher et al. (2018), though we are motivated by incorporating structured regularity of the parameter tensor inherently in the optimization. Osher et al. (2018) can be interpreted as using the gradient of the loss with respect to a Sobolev metric. Our major insight over Osher et al. (2018) is that keeping the multi-dimensional structure of the parameter tensor (rather than rasterizing) and preferentially defining the Sobolev metric with respect to the output-channel direction boosts generalization accuracy, while other directions appear to have no boost. Secondly, we introduce a re-weighted $\mathbb{L}^{2}$ metric that preferentially treats the output-channel direction, and can be implemented with one line of Pytorch code, is linear (in parameter size) complexity, and achieves similar results (in many cases) to our channel-directed Sobolev metric, boosting generalization of SGD and Adam. Third, our channel-directed Sobolev gradient can be implemented in linear cost rather than quasi-linear (not requiring FFT to compute). Sobolev gradients have been used in computer vision Sundaramoorthi et al. (2007); Charpiat et al. (2007) for their coarse-to-fine evolution properties Sundaramoorthi et al. (2008) and we adapt that formulation to channel-directed metrics for CNNs.

We formulate Sobolev gradients by considering the space of parameter tensors as a Riemannian manifold, and choosing the metric (i.e., inner product) on the tangent space to be a Sobolev metric. By choosing a metric, gradients intrinsic to the manifold can be computed and gradient flows decrease loss. Other Riemannian metrics have been used for optimization in machine learning, e.g., Amari (1998); Marceau-Caron & Ollivier (2016); Hoffman et al. (2013); Gunasekar et al. (2018; 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 regularity directly in the parameter space. Most of these methods relate to density estimation as the metric is on probability measures. Gunasekar et al. (2018) notes that even vanilla gradient descent has certain implicit bias relating to the underlying metric on the space. In Gunasekar et al. (2020) the Hessian metric (in the convex case) is analyzed and relates to mirror descent. These metrics are data-dependent and the gradient is challenging to compute, requiring (a large) inverse matrix computation. Moreover, they do not exploit the channel dimension regularity, the main purpose of our work.

From the definition, we 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 $\,\mathrm{d}f(X)\cdot k=-\|\nabla f(X)\|^{2}<0$ when $k=-\nabla f(X)$, where $\|\cdot\|$ is the norm induced by $\left<\cdot,\cdot\right>$. The gradient flow, defined by the differential equation $\dot{X}_{t}=-\nabla f(X_{t})$, will converge to a local minimum. In our application of this theory to CNN optimization, $f$ will be the loss function, and $\mathcal{X}$ will be the space of parameter tensors. In this case, as the tensor is multi-dimensional, the gradient flow will be a partial differential equation.

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:

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 fucntion, and thus land in more favorable local minima. In particular, we construct metrics that favor gradients with output channel direction regularity and this does induce regularity in the final tensor.

As gradients of a function depend on the metric structure on the underlying space, we re-define the metric on the underlying space so that tensors that differ smoothly along the output channel direction have small distance. In existing deep network gradient-based optimization schemes, the underlying metric on the loss function is assumed to be the standard Euclidean $\mathbb{L}^{2}$ metric. We will consider a re-weighted version of the $\mathbb{L}^{2}$ metric and Sobolev metrics that favor regularity in the output channel direction of 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:\mathcal{O}\times\mathcal{I}\times\mathcal{H}\times\mathcal{W}\to\mathbb{R}$ denote a parameter tensor of a deep network (from a layer of a convolutional network). Here $\mathcal{O}=[0,O]$ denotes indices to the output channel dimension of the tensor, $\mathcal{I}=[0,I]$ denote the indices to the input channel, and $\mathcal{H}=[0,H],\mathcal{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:\mathcal{O}\times\mathcal{I}\times\mathcal{H}\times\mathcal{W}\to\mathbb{R}$. The $\mathbb{L}^{2}$ (called $H^{0}$ from now on) metric is defined as

where $k_{1},k_{2}$ are in the tangent space of tensors. We now define a re-weighted version of $H^{0}$ that favors tangent vectors that have global smoothness in the direction of the $\mathcal{O}$ dimension:

where $\lambda>0$ is a hyper-parameter, and $\bar{k}$ is the average value in the output channel direction, i.e.,

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 $\lambda$ is used to control the weighting between the translation and deformation components, i.e., larger values of $\lambda$ 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 channel-directed versions of the Sobolev metric, defined as follows:

where $\frac{\partial{}}{\partial{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 first metric is analogous to the usual Sobolev metric being a weighted combination of the $H^{0}$ metric and the $H^{0}$ metric of the derivative, except only considering the derivative with respect to one direction. The second metric is similar to the first except that we use the $H^{0}$ metric of the channel-directed average rather than that of the perturbation itself. As we will see, both have similar properties, but the latter is computationally less costly to compute. 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. In the latter metric, tensors that differ by just a channel-directed translation are close. Compared with the re-weighted $H^{0}$ metric (4), the latter Sobolev metric promotes smooth perturbations beyond global translations.

We have presented only channel-directed metrics that preferentially treat the output channel dimension of the tensor as our empirical experiments demonstrate that promoting regularity in other directions is detrimental to optimization performance.

We now compute gradients with respect to the metrics defined in the previous sub-section in terms of the $H^{0}$ gradient so that simple processing of the existing gradient can be done with no other modification of existing optimization code. To compute the relation between the channel-directed gradients and the usual $H^{0}$ gradient, we note the relation between the directional derivative of a loss function, the gradient and the metric:

where $dL(X)\cdot k=\lim_{\varepsilon\to 0}\frac{L(X+\varepsilon k)-L(X)}{\varepsilon}$ is the directional derivative in the direction of the perturbation $k$. Note that the expression above indicates the directional derivative is equal to the inner product (metric) between the gradient with respect to the metric and the perturbation. This holds for any metric. With this relation, we may compute the channel-directed gradients in terms of the $H^{0}$ gradient. Derivations are in the Supplementary materials. Letting $f=\nabla_{H^{0}}L(X)$, we have

where the last two expressions are second order ordinary differential equations (ODE), whose solution we discuss next. Notice that the re-weighted $H^{0}$ gradient (9) simply re-weights the channel-directed translation component and the deformation component of the $H^{0}$ gradient differently, i.e., as $\lambda$ gets larger, the channel-directed translation becomes more prominent.

In obtaining the ODE expressions for the Sobolev gradients above, we have assumed periodic boundary conditions in the $\mathcal{O}$ dimension¹¹1Since ordering of filters along the channel direction in a CNN has no particular significance, choosing periodic or non-periodic boundary conditions is arbitrary. Periodic conditions induces smoothness between the starting and ending filters in the $o$-dimension. The periodic condition enables a simpler computational solution.. In this case, the Sobolev gradients can be interpreted as the circular convolution of the $H^{0}$ gradient with convolution kernels given as

for each of the $H^{1}$ and $\tilde{H}^{1}$ metrics, respectively, where $o$ above is scaled by $O$ to be between 0 and 1, and the circular convolution is given by

Note that the re-weighted $H^{0}$ solution also has an interpretation of convolution with respect to a smoothing kernel. Figure 2 shows plots of various kernels for the parameter $\lambda$ chosen in experiments. For each $o$, the Sobolev or re-weighted $H^{0}$ is a local average whose weights die far away from $o$. Thus, we can see that the effect of the metrics is to induce smoothness of the gradient along the output channel direction.

The second version of the Sobolev gradient need not use the convolution formula for computation, as one can just integrate the ODE twice (after noting that the channel-directed averages for both the $H^{0}$ and Sobolev gradient are the same). This saves one from having to compute the convolution directly, and hence a reduction in computational cost from quadratic (or quasi-linear with an FFT) to linear in $O$ given the $H^{0}$ gradient. The second version of the Sobolev gradient can be computed as

where $g=\nabla_{\tilde{H}^{1}}L(X)$ is the second version of the Sobolev gradient and $f=\nabla_{H^{0}}L(X)$, which are just three integrals that can be computed in linear complexity with respect to $O$. The gradient flows under these metrics are given by

where $t$ denotes the artificial time variable, $\dot{X}$ is the time derivative of the parameter tensor, and $\nabla$ denotes the gradient with respect to the desired metric. Under any metric, this reduces the loss.

We describe some properties of the resulting gradient flows according to the metrics defined in the previous sections.

Coarse-to-Fine Evolution and Removal of Some Local Minima: In Sundaramoorthi et al. (2008), it is shown that gradient flows with respect to Sobolev metrics evolve in a coarse-to-fine fashion, deforming according to coarse-scale perturbations before moving to finer scale perturbations. This can avoid being trapped in local minima due to fine-scale structures. It is also shown that when we change the metric on the space $\mathcal{X}$, the loss landscape changes and some local minima with respect to $H^{0}$ may change to other critical points with respect to Sobolev, and numerically some local minima may cease to exist. That is, local minima due to fine-scale structures can be removed by switching to a Sobolev metric. As local minima that are from wide, flat minima generalize well, the removal of local minima due to localized fine structures may encourage convergence at wider, flat minima and hence generalize better than ordinary SGD.

Regularity of the Weight Tensor: By the convolution formulas above, we can see that the Sobolev gradients are a smoothing of the usual $H^{0}$ gradient. Noting that the gradient flow (16) integrates the (smooth) gradients over time, the final tensor will also be smooth in the output direction provided the initialization is smooth. In practice, in applications with deep networks, one typically initializes the weight tensor to be random noise, so the final tensor may exhibit some randomness, but the final tensor is sum of a smooth component and the randomness, and so it exhibits regularity, e.g., strong correlation across nearby output-channel components in the weight tensor, as we verify in experiments. Further, experiments (see Section 4) indicate that final results with respect to our metrics are less dependent on initialization than SGD, which suggests that the initial randomness may die out.