Optimization on Submanifolds of Convolution Kernels in CNNs
Mete Ozay, Takayuki Okatani
Introduction
Over the last decade, convolutional neural networks (CNNs) have been utilized to perform various tasks such as image classification and scene analysis . The unprecedented performance of CNNs on these tasks has been attributed, in practice, to design of large-scale datasets, and modeling of more complex architectures with larger number of layers and parameters .
While performing optimization using stochastic gradient descent (SGD) algorithms with backpropagation (BP) in CNNs, we observe that norm of gradients may exponentially increase or decrease . Exploding and vanishing gradients trigger several open problems such as convergence of SGD and its robustness to reparametrization of convolution kernels, and internal covariate shift. In order to cope with these open problems, various normalization methods have been proposed on kernels and/or gradients at initialization, and/or at each update of SGD using different orthogonality constraints . However, various kernel normalization methods may affect geometry of search spaces, dissimilarly. More precisely, level sets of classification loss functions, critical points residing in level sets, and convergence properties of SGD may be different for different kernel normalization methods. In order to assure convergence of SGD algorithms to solutions at single minimum, we need to identify search spaces and compute steps according to the geometry of kernel spacesIn this work, we refer to convolution kernels used in CNNs by kernels. in CNNs. In addition, various orthogonality constraints are nonconvex and nonlinear. Therefore, embedding constraints into cost functions may lead to many local minimizers .
In this work, we address the aforementioned problems by identifying kernel spaces as topological smooth manifolds under a geometric optimization framework for training of CNNs. We pose the kernel estimation problem in CNNs as optimization on embedded and/or immersed submanifolds of kernels which are described according to different geometric properties of the kernels, such as orthonormal rectangular or orthogonal square kernels. Thereby, we can define constraints of optimization problems of CNNs in search spaces of SGD algorithms, instead of embedding the constraints into cost functions of the problems . To this end, we first provide theoretical analyses and results to explore geometry of kernels in CNNs. Then, we employ our theoretical results for image classification. In our framework, we first construct kernel submanifolds at each layer of a CNN such that a kernel resides as a point on a kernel submanifold. Then, we employ a SGD algorithm for optimization on kernel submanifolds using BP in CNNs.
Related Work and Summary of Contributions
Popular kernel normalization methods have been implemented using reparametrizations , and additional constraints, such as orthogonality , in order to preserve unit norm property of the kernels for forward propagation , at initialization , or at each epoch of SGD . Unit norm kernels were used for symmetry invariant optimization at the first and second layers of a network in . One of the challenges of these approaches, besides the lack of theoretical understandings mentioned above, is the employment of the reparametrization and rescaling methods at new layers before/after convolution layers, resulting in an increase of complexity of the network structure by aggregation of the new layers. In addition, statistical properties of data need to be recorded during training, and testing. Therefore, kernel normalization methods may increase computational overhead of CNNs for both training and testing.
Removal of scale and translation from kernels by normalization can be interpreted as imposition of a geometric structure such that the kernels lie on the sphere . In our approach, embedded kernel submanifolds can be described using the sphere, oblique and/or the Stiefel manifold. Additional constraints can also be imposed using immersed submanifolds such as rotation groups. Thus, our approach can be considered as generalization of the aforementioned approaches such that we can employ our methods to model different submanifolds according to various constraints, such as orthonormal kernels. Thereby, we employ geometry of kernels to identify the constraints on the optimization problem of CNNs. Moreover, gradient descent of natural gradient (NG) methods can be cast as an approximation to SGD for submanifolds which are equipped with Riemannian structure and employed in our framework . In this aspect, our proposed methods can be considered as generalization of NG methods . Our contributions can be summarized as follows:
Analysis of geometry and smooth structures of kernel submanifolds: One of our main motivations for employment of kernel submanifolds is to assure existence of singular minimum of a loss function of CNNs in the search space of a SGD. For this purpose, the loss function is defined as a smooth map , where is a kernel manifold, and is a space of loss values such as a set of classification errors. Thereby, we can formulate the relationship between level sets of and submanifolds of . We also analyze the conditions under which level sets are submanifolds that contain critical points in Section 3.
A SGD algorithm for optimization on kernel submanifolds in CNNs, and analysis of convergence properties: By making use of our theoretical results, we propose a SGD algorithm for optimization on kernel submanifolds for training of CNNs in Section 4. More precisely, our theoretical results first enable us to employ various smooth manifolds with different metrics to describe submanifolds. For computational efficiency, we then employ Riemannian manifolds to perform steps of SGD methods on submanifolds. We compute steps of SGD according to smooth structures of submanifolds, such as metrics and differential maps, defined on submanifolds, and their topological properties, such as compactness. Moreover, in our proposed SGD algorithm, we can employ momentum and Euclidean gradient decay for optimization on submanifolds extending the methods proposed in . In Section 4.1, we provide two theorems to analyze the convergence of the proposed SGD algorithm. We provide a discussion on computational complexity of the proposed algorithm in Section 4.2.
To the best of our knowledge, this is the first comprehensive work which employs stochastic optimization methods on embedded and immersed submanifolds of kernels in CNNs with convergence properties.
The paper is organized as follows. Our proposed mathematical framework is introduced in Section 3. We provide the proposed SGD algorithm and the convergence properties in Section 4. In Section 5, we examine the proposed algorithm, methods and theoretical results for different manifolds using several benchmark datasets in comparison with state-of-the-art methods. Section 6 concludes the paper. Proofs of the theorems and implementation details are given in the supplemental material (sup. mat.).
Geometry of Kernel Submanifolds
where is an image for , and . The channel of the data matrix is convolved with the kernel to obtain the feature map by We ignore the bias terms in the notation for the sake of simplicity.. Given a batch of samples , we denote a value of a classification loss function for a kernel by , and the loss function of kernels utilized in the CNN by . If we assume that consists of a single sample , then, an expected loss or cost function of the CNN is computed by
The expected loss for is denoted by . For a finite set of samples , is approximated by an empirical loss , where is the size of (similarly, denotes the empirical loss for ). Then, feature representations are learned using SGD by solving
Restriction on the search space can be imposed into (3) by defining constraints expressed as a function of the variables , such as normalization . If the search space is equipped with a manifold structure, then constrained optimization problem of CNNs can be converted into that of unconstrained optimization. We explore the geometric relationship between the loss (2) and kernels, by first defining
where are called slice coordinates of for constants . Next, we define embedded submanifolds.
Definition 3.1 describes a relationship between partitioning of into level sets that contain critical kernels using (4) (see Fig. 1), and into embedded kernel submanifolds that satisfy (5) (see Fig. 2). In other words, we consider an approach to construct embedded kernel submanifolds which correspond to level sets of loss functions on manifolds that contain critical kernels. However, not every level set corresponds to an embedded kernel submanifold. In the next theorem, we introduce a condition to associate a level set to an embedded kernel submanifold.
If a loss {\color[rgb]{0,0,1}\mathcal{L}}:\mathcal{M}\to{\color[rgb]{0,0,1}\mathcal{N}} is a smooth map with constant rank (i.e. the rank of the Jacobian matrix of {\color[rgb]{0,0,1}\mathcal{L}}), then each level set of {\color[rgb]{0,0,1}\mathcal{L}} is an embedded kernel submanifold {\color[rgb]{1,0,0}\hat{\mathcal{M}}}\subseteq\mathcal{M} (see Fig. 1).
In addition, not all embedded kernel submanifolds can be expressed as level sets of a loss that contain critical kernels even if the loss is a smooth submersion. However, the next proposition shows that every embedded kernel submanifold can be locally expressed as a level set.
In order to perform optimization on kernel submanifolds whose inclusion maps are injective submersions, such as rotation groups, we need to consider a more general notion of submanifolds characterized by immersed kernel submanifolds (see the sup. mat. for details). A Lie subgroup of a Lie group , e.g. the rotation group, is endowed with a topology and smooth structure making it into a Lie group and an immersed kernel submanifold of . Consequently, embedded kernel submanifolds, which are also subgroups of , are automatically Lie subgroups. Therefore, our framework enables us to employ both embedded and immersed kernel submanifolds in SGD. In practice, this property is required by the SGD algorithms in order to train CNNs using various normalized kernels including orthogonal square kernels with determinant 1 (i.e. members of the rotation group) with assurance of convergence. Next, we use our framework to explore geometry of space of normalized kernels.
Analysis of the geometry of submanifolds of normalized kernels is an open and understudied problem. This is a crucial problem since gradient steps should be identified by the geometry of submanifolds as explained in the previous sections. In the next theorem, we explore this conjecture using concrete examples for different normalization methods.
This theorem shows that different kernel normalization methods, even if they are implemented in an intuitively similar manner, imply different geometric properties. For instance, if the kernels are normalized to the unit Frobenius norm, then they reside on the dimensional sphere. Besides, if the kernels are first orthonormalized as suggested in (iii), then each column of the kernel also resides on the sphere . However, the kernel resides on the Stiefel manifold. In addition, if the shape of is a square, then kernels reside on . If each column is first normalized with unit norm, then resides on a kernel submanifold isometric to .
In SGD, we need to pay attention to restriction of gradients and steps employed on kernel submanifolds to assure convergence to a solution. The first reason is that kernels should reside in locally compact sets to assure existence of critical kernels (see Theorem 3.2 and Proposition 3.3). Second, while performing SGD steps, gradients of kernels should be bounded in the compact sets, which can be achived by employing mappings from Euclidean gradients obtained using BP to submanifold gradients residing on tangent spaces of kernel submanifolds. However, these two requirements are not considered in the state-of-the-art normalization methods, resulting in exploding and vanishing gradients. Note that compact sets and mappings of gradients are computed according to manifold and smooth structures of submanifolds. Therefore, we need to employ appropriate mappings of kernels and gradients in SGD while performing steps. In order to perform SGD for different kernel submanifolds assuring almost sure convergence to a solution, we suggest a SGD algorithm considering an optimization approach on Riemannian manifolds in the next section.
A SGD Algorithm for Optimization on Kernel Submanifolds in CNNs
Various optimization algorithms have been developed to solve optimization problems on matrix manifolds . However, development of SGD algorithms on kernel submanifolds, and analysis of their convergence properties in CNNs have not been addressed yet. In our framework, we perform optimization on kernel submanifolds at each convolution layer of an -layer CNN. An algorithmic description of our proposed methods is given in Algorithm 1:
Initialization: We first define a KS , for each convolution layer whose members are , where , , , .
For each epoch , and for each , following steps are performed (see Fig. 3);
- Line 4: The Euclidean gradient is computed and obtained using backpropagation (see Fig. 3).
- Line 5: Momentum and Euclidean gradient decay methods are employed on the Euclidean gradient using \mu_{t}:=q\Big{(}{\rm grad}_{E}\;\mathcal{L}(\omega_{l}^{t}),\mu_{t},\Theta\Big{)} (see Fig. 3). We can employ state-of-the-art acceleration methods modularly in this step. For example, momentum can be employed with the Euclidean gradient decay using
where is the parameter employed on the momentum variable . We consider as the decay parameter for the Euclidean gradient. The reason is that and affect the step performed in the ambient Euclidean space while the learning rate (LR) is employed on the submanifold gradient. A detailed discussion of methods that are used to compute is given in the sup. mat.
- Line 6: The moved vector is projected to the tangent space at , to compute the submanifold gradient , where is a projection operator defined according to the geometry of (see Fig. 3), and is used to bound the norm of the gradient.
- Line 7: The learning rate is updated by , where is a function that controls the convergence rate . We choose which satisfies the following as suggested in ;
- Line 8: A vector is computed using , where is a function that defines the next step on the tangent space at (see Fig. 3). In this work, we employed to move the solution in a gradient descent direction with step size .
- Line 9: Compute the next iterate using , where is a mapping from onto (see Fig. 3). We employ exponential maps and/or retractions for implementation of (see the supp. mat. for details). Therefore, this step enables us also to keep on a locally compact subset of the kernel submanifold , .
In , a procedure was developed to analyze convergence properties of SGD methods for a particular class of manifolds following the proof methods suggested in . In this work, we extend and employ their results to train CNNs using different kernel submanifolds. We first consider a collection of kernels computed at the layer of a CNN as a stochastic process. Then, the expected value of a submanifold gradient of a classification loss can be computed by In practice, we receive a batch of samples at each epoch. Assuming that each batch contains a single sample, denotes an average gradient computed by .. In the following theorems, we provide convergence properties of Algorithm 1 for two cases where we use i) exponential maps, and ii) retractions for at the step of the algorithm.
i) Exponential maps for : An exponential map is used to map a vector to a kernel along a geodesic curve on which goes through in the direction of .
Suppose that the following conditions are satisfied;
(1) Condition for maps onto kernel submanifolds: is a connected compact Riemannian kernel submanifold.
(2) Condition for kernels: There exists a compact set such that , . The minimal distance between conjugate kernels and denoted by is a geodesic that satisfies .
(3) Condition for gradients: The gradient is bounded on , such that , , and .
(4) Condition for the classification loss function: We use a three times continuously differentiable function .
Then, the loss function and the gradient converges almost surely (a.s.) by , where is a minimum, and .
ii) Retractions for : In the experimental analysis, we used retractions to compute numerical approximations to exponential maps onto kernel submanifolds . Thus, we next provide the convergence properties for the case where is implemented using retractions.
Suppose that the conditions (1)-(4) of Theorem 4.1 are satisfied, and that is a twice continuously differentiable retraction. Then, we have and .
2 Computational complexity of Algorithm 1
Compared to traditional SGD algorithms , the computational complexity of Algorithm 1 is dominated by computation of the maps and at line 6 and 9, depending on the structure of the kernel submanifold used in the algorithm at the layer. Concisely, the computational complexity of is determined by computation of different norms that identify the submanifolds. For instance, for the sphere, we use . Thereby, for an kernel, the complexity is bounded by . Similarly, the computational complexity of depends on the submanifold structure. For example, the exponential maps on the sphere and oblique manifold can be computed using functions of and functions, while that on the Stiefel manifold is a function of matrix exponential. For computation of matrix exponential, various numerical approximations with complexity were proposed for different approximation order . However, unit norm matrix normalization is used for computation of retractions on the sphere and the oblique manifold. Moreover, QR decomposition of matrices is computed with for retractions on the Stiefel manifold. In addition, the computation time of maps can be reduced using parallel computation methods. For instance, a rotation method was suggested to compute QR using processors in unit time in . Therefore, computation of retractions is computationally less complex compared to the exponential maps. Since the complexity analysis of these maps is beyond the scope of this work, and they provide the same convergence properties for our proposed algorithm, we used the retractions in the experiments. Detailed formulations of maps, and implementation details are given in the sup. mat.
Experimental Analyses and Results
The proposed framework and the algorithm can be employed to train CNNs using different manifolds. We train state-of-the-art CNNs using our algorithm on benchmark datasets by optimization on three kernel submanifolds, namely the sphere, the oblique and the Stiefel manifold. In order to perform a fair performance comparison, we used the same code and hyperparameters provided by the authors. Implementation details are given in the sup. mat.
At the lower layers (), if we use kernels with and , then we perform additional regularization on spatially distributed patterns within a neighborhood determined by and . For instance, square shape kernels of the Stiefel manifold construct the orthogonal group by Proposition 3.1. Then, the orthogonality constraints imposed by these kernels enable us to learn representations of shape variation caused by both shape deformation and viewpoint changes . Moreover, translation and scaling variability is removed from kernels of the sphere. This property has been used for statistical shape analysis to learn representations of unit length curves, shape primitives , and deformable shapes . Therefore, these shape representations can be learned by training CNNs using kernels on the sphere. Moreover, the constraints determined by the oblique manifold induce oblique rotation. Thereby, features which are mutually independent, and the orthogonal transformations that minimize the dependence between features, can be learned using the kernels on the oblique manifold.
Since a detailed analysis of each of these properties is beyond the scope of this work, we explore them experimentally by training different CNNs with the aforementioned manifolds, and analyzing their performance.
In a recent work , kernels are reparameterized by a fixed norm that is initialized by the inverse of standard deviation of pre-activations. Thereby, their proposed method constructs a space of kernels identified by the sphere with radius . In this aspect, their proposed method can be considered as a realization of our proposed algorithm for the sphere. In other words, we can perform optimization on other kernel manifolds such as the oblique and the Stiefel manifold in addition to the sphere. Moreover, each kernel can reside in a different manifold endowed with a different geometry. For instance, we can perform optimization on different kernels that reside on the spheres with different radii. Therefore, our proposed methods enable us to have a better control on the geometry of kernel spaces for training of CNNs compared to their method .
We examine this property by training their proposed CNN architecture , which is a variation of All-CNN architecture and denoted by SK, using our proposed methods for the sphere (Sp), the oblique manifold (Ob) and the Stiefel manifold (St). The results given in Table I are obtained by training CNNs on the Cifar-10 dataset without using data augmentation (DA). In Table I, we observe that our methods that use the sphere (SK (Sphere)) outperform the methods proposed in (SK (WN)). This observation supports our claim for the benefit of employment of kernels belonging to spaces with various manifold structures, e.g. the sphere with varying radii (which are determined by statistical properties of the data and the gradients of the loss). We also observe that we further boost the performance using the oblique and the Stiefel manifolds. This result also propounds conjectures and results provided in the previous works regarding regularization and invariance properties of models learned using different manifolds.
We aim to learn representations robust to statistical variance and mean of pre-activations by normalizing them using batch normalization (BN). If features input to a layer are i.i.d. with zero mean and unit variance, then we can equivalently obtain these robust representations using normalized kernels at that layer . This property is explored in by proving that approximation error to covariance of pre-activations is upper bounded by a function of kernel norms. Therefore, normalized kernels that reside in the sphere are used to train CNNs in . Following this property, BN is employed by removing just mean for the features obtained using normalized kernels, and this method is called mean-only BN (MOBN) . The results given in Table I show that we can further boost the performance using MOBN for different manifolds. We should also notice that, MOBN boosts the performance only if normalized kernels are used for training, such that the error of SK (MOBN) is 8.52% while the error of SK (BN) is 8.05%.
2 Results for Training Large-scale CNNs
We first employ our methods for training of residual networks with constant depth (RCD) and stochastic depth (RSD) consisting of 110 layers . In order to explore how the proposed methods enable us to learn invariance properties as discussed above, we also analyze the results for Cifar and Imagenet datasets that are augmented using standard DA methods (details are given in the sup. mat.). The results given in Table II, show that the performance boost is larger for datasets prepared w/o DA compared to the augmented datasets. In addition, we can further boost the performance even for augmented datasets, since data augmentation is conducted using large scale transformations, while the kernels computed at different layers can learn the invariants at different resolutions.
Since Res with less number of layers do not perform as well as SK and NiN on Cifar dataset, we also provide the results for the Cifar-10 with DA in Table III. The results show that our methods can boost the performance of the baseline Res. However, for a smaller Res (Res-20), the kernels of the sphere may outperform the kernels of the oblique as also observed in Table II. Moreover, we observe that the amount of performance boost (for St) decreases from 0.78% to 0.35% as the number of layers increases to 44 in Table III. On the other hand, for St, we obtain 0.65% and 2.11% boost for Cifar 10 and 100 with DA, and 0.72% and 4.98% boost for Cifar 10 and 100 without DA, using Res consisting of 110 layers equipped with pre-activations (RCD) in Table II. Therefore, the amount of boost also depends on the number of classes and use of the augmentation methods.
We also observe that performance boosts more for Cifar-100 compared to the results obtained for Cifar-10. This result suggests that we can learn feature representations of diverse patterns observed in large number of classes using kernels belonging to the manifolds. In order to scrutinize this observation, we provide the results for training of residual networks (Res) using the Imagenet dataset in Table IV. The results given in Table IV complement the previous observations such that we have 2.45%, 1.72% and 1.50% performance boost (for St) using Res-18, Res-34 and Res-50, respectively. We also provide the performance of a recent method proposed for optimization on a probabilistic manifold of network parameters, called PRONG . In other words, PRONG implements an approximation for natural gradient descent. An interesting result is that Res-18 with 18 convolution layers, which was trained using our methods with manifolds, outperforms Inception (22 conv. layers) which was trained using PRONG (see Table IV).
Conclusion
We proposed a mathematical framework to explore and make use of geometric properties of spaces of convolution kernels in CNNs. More precisely, we suggested several mathematical methods and tools to describe and utilize kernel spaces by particular topological smooth manifolds, namely embedded and immersed kernel submanifolds. Following our theoretical results, we proposed a SGD algorithm for optimization on kernel submanifolds to train CNNs with assurance of convergence to a solution at single minimum of loss.
We employed our algorithm to train the CNNs using benchmark datasets. We observed that our methods boost the performance of the CNNs for various datasets prepared with and without using data augmentation methods. We believe that our results will guide researchers to develop geometry-aware training algorithms that employ powerful regularization methods and take advantage of invariance properties of kernels. In the feature work, we plan to apply our framework for other tasks such as segmentation, detection, pose estimation, action recognition and video analysis. Moreover, we will use our algorithm to train other deep networks such as auto-encoders and recurrent neural networks.