The Gauss-Newton (GN) matrix plays an important role in machine learning, most evident in its use as a preconditioning matrix for a wide family of popular adaptive methods to speed up optimization. Besides, it can also provide key insights into the optimization landscape of neural networks. In the context of deep neural networks, understanding the GN matrix involves studying the interaction between different weight matrices as well as the dependencies introduced by the data, thus rendering its analysis challenging. In this work, we take a first step towards theoretically characterizing the conditioning of the GN matrix in neural networks. We establish tight bounds on the condition number of the GN in deep linear networks of arbitrary depth and width, which we also extend to two-layer ReLU networks. We expand the analysis to further architectural components, such as residual connections and convolutional layers. Finally, we empirically validate the bounds and uncover valuable insights into the influence of the analyzed architectural components.
Paper
Similar papers
Peer review
Summary
The authors derived upper bounds for the condition number of the outer product of the jacobian of the neural network output (the Gauss Newton matrix) in the case of deep linear networks, and non-linear networks with a single hidden layer with piecewise-linear activation functions. They empirically evaluate these bounds using MNIST (with the exception of one plot in the main text that used CIFAR-10) and find their bounds reflect the general trend of the empirical condition number as depth (or in some cases width) are varied. The work is motivated by highlighting the importance of Hessian information in the optimisation of neural networks.
Strengths
- The empirical evaluations show that the bounds are actually quite informative, following the general behaviour of the condition number and in some cases being quite tight as well. They point out that if they replace the convex combination in their bound(s) with a maximum, the bound becomes much looser and does not show any behaviour following trends. - The understanding of neural network loss landscapes and how best to optimise them is still in its relative infancy, so works like this can be valuable. - Extensive proofs and extra plots are provided in the appendices
Weaknesses
- ~~The paper generally uses the same dataset (MNIST) for almost all the plots, so it's hard to tell if these good behaviour of their bounds hold on other, potentially harder to optimise, problems such as CIFAR-10.~~ **It was pointed out in the rebuttal that CIFAR-10 was included in the appendix.** - The precise motivation of studying the Gauss Newton is slightly weak. It's not clear to me that these results actually tell us anything substantial about optimising the neural network. - The authors were not able to extend their bound to the deep non-linear network case.
Questions
- To be clear, on line 82 where we refer to "outer gradient product of the loss function", this is actually the outer gradient product of the network function and has no dependence on the loss function, right? - Suggestion: In the equation below line 182 (which should probably be numbered) some extra brackets and the use of \cdots rather than \dots would make this expression clearer to read.
Rating
8
Confidence
2
Soundness
4
Presentation
3
Contribution
3
Limitations
The authors briefly discuss some future directions their work could be taken in, but do not seriously critique the work they present here.
Summary
This paper examines the condition number of the Gauss-Newton matrix [1] in neural networks. It shows that normalization techniques, such as Batch Normalization [2], initial normalization, skip connections [3], and appropriate layer dimensions, reduce the condition number and therefore enhance the training stability. The objective is to provide insights into the training of deep neural networks by characterizing a new, tight upper bound for the condition number of the Gauss-Newton matrix. The paper primarily focuses on linear neural networks and shallow non-linear neural networks. [1] SCHRAUDOLPH, Nicol N. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 2002, vol. 14, no 7, p. 1723-1738. [2] IOFFE, Sergey et SZEGEDY, Christian. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In : International conference on machine learning. pmlr, 2015. p. 448-456. [3] HE, Kaiming, ZHANG, Xiangyu, REN, Shaoqing, et al. Deep residual learning for image recognition. In : Proceedings of the IEEE conference on computer vision and pattern recognition. 2016. p. 770-778.
Strengths
This paper provides new theoretical insight into the training dynamics of Deep Neural Networks (DNNs) by examining the conditioning number of the Gauss-Newton Matrix. The experimental results are robust and shed light on often obscure aspects of DNN optimization, such as the impacts of normalization, layer size, residual connections, and ReLU activation function. Meanwhile, the way the authors interpret the experimental results seems reasonable. Importantly, the paper also clearly outlines its limitations in the conclusion.
Weaknesses
The paper is difficult to read and follow due to several typos and unclear ideas. The main contribution is not significant and heavily relies on Singh's work [4]. The authors discuss the K-FAC method [5], which uses the EFIM and not the Hessian, even though it can be an approximation in certain cases where the model's likelihood is in the exponential family [6]. The paper only addresses the Gauss-Newton matrix and not the Generalized Gauss-Newton matrix [6], which could be more relevant, especially when using a cross-entropy loss function instead of MSE. Additionally, the references are outdated (e.g., line 49 and most of the related work), and there are typographical errors, such as "generalized" being repeated on line 140 and "collapse" being misspelled on line 124. There is also a lack of references related to the Gauss-Newton matrix, Kronecker properties PyTorch library, etc... [4] SINGH, Sidak Pal, BACHMANN, Gregor, et HOFMANN, Thomas. Analytic insights into structure and rank of neural network hessian maps. Advances in Neural Information Processing Systems, 2021, vol. 34, p. 23914-23927. [5] MARTENS, James et GROSSE, Roger. Optimizing neural networks with kronecker-factored approximate curvature. In : International conference on machine learning. PMLR, 2015. p. 2408-2417. [6] MARTENS, James. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 2020, vol. 21, no 146, p. 1-76.
Questions
1. Could you clarify the statement on line 31: “imagine an entire set of neurons being dead”? What implications does this scenario have for your study? 2. How good is the approximation of the Gauss Newton matrix $\mathbf{G}_O$ without considering the term $\mathbf{H}_F$? 3. Can you also conduct experiments on more challenging datasets like ImageNet [7] or TinyImageNet [8]? 4. Can you develop more about the time and space complexity of the computation of the condition number? 5. Can you explain and rephrase the lines 719 and 720 with the associated Figures 19 and 21? 6. Why did you compute the condition number on CPU rather than on GPUs (line 725)? 7. Can you also explore the relation between te condition number and the batch size? [7] DENG, Jia, DONG, Wei, SOCHER, Richard, et al. Imagenet: A large-scale hierarchical image database. In : 2009 IEEE conference on computer vision and pattern recognition. Ieee, 2009. p. 248-255. [8] Le, Ya and Xuan S. Yang. “Tiny ImageNet Visual Recognition Challenge.” (2015).
Rating
6
Confidence
4
Soundness
2
Presentation
3
Contribution
3
Limitations
The limitations of the work are well addressed by the authors. I do not believe this work has any particular negative social impact.
Answers to Questions 3-7
**Question 3**: > Can you develop more about the time and space complexity of the computation of the condition number? **Answer**: Thank you for this question. In order to compute the condition number we need to compute the eigenspectrum of the GN matrix $\mathbf{G}_O$, which has dimension $p \times p$, where $p$ is the number of parameters or of the matrix $\hat{\mathbf{G}}_O$, which has dimensions $kd \times kd$, where $d$ and $k$ are the input, respectively output dimension of the network. The time complexity of calculating the eigenvalue decomposition has a cubic complexity. That is, in order to compute the condition number, we have a computational complexity of $\mathcal{O}(\min(p, kd)^3)$. **Question 4**: > Can you also conduct experiments on more challenging datasets like ImageNet [7] or TinyImageNet [8]? **Answer**: Thank you for this suggestion. - **Only marginal gain in insight is expected from experiments on ImageNet.** We would like to argue that we expect only a marginal gain in insights from additional experiments on ImageNet compared to experiments on Cifar-10, which we have already conducted. In particular, note that the complexity of a given dataset appears only through the input covariance matrix $\boldsymbol{\Sigma}$, which is a separate factor in the upper bound. Furthermore, as we have mentioned in Remark R1, the effect of the conditioning of the input data on the conditioning of the GN spectra can be largely reduced by whitening or normalizing the input data. Thus, the effect that different datasets have on the condition number of the GN matrix can be largely alleviated through preprocessing, which is common practice. - **Scaling up experiments is expensive and challenging.** An additional aspect that we want to emphasize is the additional effort to scale up the computation of the condition number, which is challenging and not straightforward, given the time and space complexity elaborated in the previous question. Given the above explanations, we believe that experiments on ImageNet are less relevant to our work and currently also out of scope. **Question 5**: > Can you explain and rephrase the lines 719 and 720 with the associated Figures 19 and 21? **Answer**: Thank you for this question. As one can see from Eq. (6), the condition number for each term in the sum of the first upper bound in Eq.(6) improves when $\beta$ increases. This can be seen by the fact that the ratio will be dominated by $\beta$ and will go to 1 for $\beta \to \infty$. This is also what we observe empirically in Figure 19 an 21, where the condition number is smaller for $\beta = 1$ compared to the other two settings, where $\beta = 1/L < 1/\sqrt{L} < 1$ for deeper networks with $L > 1$. **Question 6**: > Why did you compute the condition number on CPU rather than on GPUs (line 725)? **Answer**: Thank you for this question. - In the case of linear and residual networks (with no activation function), where the GN matrix can be expressed analytically through Eq. (2), we did not find a significant time advantage of running the code on GPU compared to CPU. - In the other case, where the GN matrix had to be computed numerically through automatic differentiation (for instance the experiments on pruning weights at initialization), we ran into memory problems with the GPUs. The way we resolved this was to build up the GN matrix through backpropagation on the GPU and move it to CPU to finally compute the condition number. Despite the slow down by moving the GN matrix between GPU and CPU, this still led to a speed up in time. **Question 7**: > Can you also explore the relation between the condition number and the batch size? **Answer**: Thank you for this interesting question! - Our work does not explicitly consider training dynamics, so generally speaking, our bounds consider $n$ to be the total number of sample points. However, one could also interpret $n$ as the batch size during training, in which case the condition number could be interpreted as the condition number of the "local loss landscape" of a single mini-batch. Note that $n$ only appears implicitly in our theoretical upper bounds through the condition number of the empirical input covariance matrix $\boldsymbol{\Sigma}$. We kindly request the reviewer to let us know if there are any remaining questions they may have. If they find that their queries have been sufficiently addressed, we would greatly appreciate it if they could reconsider their evaluation of our paper.
Official Comment by Reviewer 6bdM
Thank you for the authors' comprehensive rebuttal and the additional analyses provided in response to the feedback from myself and the other reviewers. The authors have well addressed my concerns, which has significantly improved my understanding of the paper. Based on this, I will raise my score to 6 (Weak Accept).
Summary
This paper characterizes the conditioning of Gauss-Newton (GN) matrix. The contribution of this paper is clear and straightforward: for deep linear networks, it establishes a bound on the condition number of GN matrix, which is further extended to 2-layer ReLU networks. These bounds could be useful in certain scenarios. Numerical experiments are conducted to support the theoretical claims.
Strengths
1. This paper is clearly written and well-organized. The motivation of conducting the proposed research is clearly demonstrated in the introduction, i.e., why it is interesting to study the condition number of GN matrix for certain types of simplified deep neural networks. The main results and empirical results are also clearly presented. 2. Numerical experiments verify that the bounds are tight under conditions imposed by the authors.
Weaknesses
1. My major concern is about the implications of the derived bounds. Specifically, I note that training/learning such as gradient descent learning dynamics, which is crucial in practice, is not involved in deriving the bound. Thus it is hard to see the implication of these derived bounds since they are the same for both before and after training of the model parameters, i.e., it provides no information about the effects of training and fails to characterize the properties of solutions of deep learning models in a specific task. This might even lead the bounds to be meaningless in certain scenarios. For example, in Lemma 1, let $k = 1$, then according to [1], after enough iterations of gradient descent, the largest singular value $\sigma_{\max} \to \infty$ (as there is no formal definition of $\sigma_{\max}$ in Lemma 1 I assume that my understanding of its definition is correct) while $\sigma_{\min} \to 0$ for both $W$ and $V$, thus the bound blows up and becomes meaningless. 2. I find the introduction part of this paper a bit verbose, e.g., it spends about 4 pages (almost half of the main body) to discuss existing results and related works before presenting details of main results (starting from Lemma 1). I think it would be better to emphasize more about the technical contributions of the current work, which is not clear to me since it seems that many important steps have already been solved by previous works, e.g., Eq. (2), (3), and Lemma 3. **Reference** [1] Ji and Telgarsky. Gradient descent aligns the layers of deep linear networks.
Questions
Could the authors give some implications of the proposed bounds and how we can better use it in practice?
Rating
5
Confidence
4
Soundness
3
Presentation
3
Contribution
2
Limitations
The authors addressed their limitations in Section 7. In addition, in my view, the bounds fail to capture the distinctions between the solutions of deep learning models and random parameters such as the initialization, which limits its significance.
Summary
This paper is dedicated to the theoretical characterization of the condition number of the Gauss-Newton (GN) matrix in neural networks. By studying deep linear networks and two-layer nonlinear networks, the authors establish tight bounds on the GN matrix's condition number and extend this analysis to architectures incorporating residual connections and convolutional layers. The methodology is rigorous, and the experimental validation is thorough, making significant contributions to understanding optimization processes in deep learning.
Strengths
1. This paper deeply studies the properties of Gauss-Newton matrices, especially in terms of condition numbers in deep linear networks and two-layer two-layer nonlinear networks (Leaky ReLU activation), and provides rigorous theoretical derivation and proof. 2. This paper experimentally shows that the width and depth of the network, when the parameters are initialized with a certain distribution, have a strong correlation with the condition number.
Weaknesses
1.The discussion in this paper on the relationship between the Gauss-Newton matrix condition number and the convergence rate of network optimization could be richer. 2.Convergence Rate Analysis in Figure 17: The network's optimal solution and the corresponding minimum loss differ under various settings, making it difficult to analyze the convergence rate from the loss changes. For instance, at epoch 300, the network with a width of 15 may have nearly converged, while the network with a width of 200 still shows a significant downward trend. Thus, it is hard to draw conclusions about the convergence speed. 3.The caption of Figure 9 does not match the subfigures. 4.Figures 12-16 are not referenced or analyzed in the paper. 5.This paper lacks exploration of more general networks and recent network structures.
Questions
1. Can this paper be extended to study the effect of network initialization on convergence rate? 2. Is it more appropriate to adjust the experiment on the effect of condition number and convergence rate in this paper to study different initializations or study the depth of the current network? They are more likely to keep the minimum loss close to each other, and make more use of the analysis of convergence rate.
Rating
6
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
As stated in the weaknesses.
Answers to Question 1 and 2
**Question 1**: > Can this paper be extended to study the effect of network initialization on convergence rate? **Answer**: This is a very interesting question. As the condition number is a very local property, it is in general hard to connect the conditioning at network initialization to a global convergence rate. However, we would like to argue below that an ill-conditioned network initialization will still affect the rate of convergence for gradient descent (GD) in the initial phase of training. For this we will present a modified analysis of GD for strongly convex functions, where we use local constants $\mu(k)$ and $L(k)$ instead of the global smoothness and Lipschitz constant, respectively. Let us denote the Lipschitz constant by $L$ and the smoothness constant by $\mu$. Furthermore, let the step size be such that $\eta_k \leq \frac{1}{L}$. Then by the definition of gradient descent, we have: \begin{align*} ||\boldsymbol{\theta}\_{k+1} - \boldsymbol{\theta}^* ||^2 &= ||\boldsymbol{\theta}_k - \boldsymbol{\theta}^* - \eta_k \nabla f(\boldsymbol{\theta}_k) ||^2 \\\\ &= ||\boldsymbol{\theta}_k - \boldsymbol{\theta}^* ||^2 -2 \eta_k \nabla f \left( \boldsymbol{\theta}_k \right)^\top \left( \boldsymbol{\theta}_k - \boldsymbol{\theta}^* \right) + \eta_k^2 ||\nabla f \left( \boldsymbol{\theta}_k \right) ||^2 \\\\ & \stackrel{\text{Strong convexity}}{\leq} (1-\eta_k \mu) ||\boldsymbol{\theta}_k - \boldsymbol{\theta}^* ||^2 - 2\eta_k (f(\boldsymbol{\theta}_k) - f(\boldsymbol{\theta}^*)) + \eta_k^2 ||\nabla f \left( \boldsymbol{\theta}_k \right) ||^2 \\\\ & \stackrel{\text{Smoothness}}{\leq} (1-\eta_k \mu) ||\boldsymbol{\theta}_k - \boldsymbol{\theta}^* ||^2 - 2\eta_k (f(\boldsymbol{\theta}_k) - f(\boldsymbol{\theta}^*)) + 2 \eta_k^2 L (f(\boldsymbol{\theta}_k) - f(\boldsymbol{\theta}^*)) \\\\ &= (1-\eta_k \mu) ||\boldsymbol{\theta}_k - \boldsymbol{\theta}^* ||^2 - 2 \eta_k (1 - \eta_k L) (f(\boldsymbol{\theta}_k) - f(\boldsymbol{\theta}^*)) \end{align*} Since we assumed that $\eta_k \leq \frac{1}{L}$, the last term is negative. Therefore: \begin{equation} ||\boldsymbol{\theta}_{k+1} - \boldsymbol{\theta}^* ||^2 \leq (1-\eta_k \mu) ||\boldsymbol{\theta}_k - \boldsymbol{\theta}^* ||^2 \end{equation} So by recursively applying Eq. (1) and replacing $\mu$ by the local smoothness constants $\mu(k)$: \begin{equation} ||\boldsymbol{\theta}_k - \boldsymbol{\theta}^* ||^2 \leq \prod\_{i=0}^{k-1} (1-\eta_i \mu(i)) ||\boldsymbol{\theta}_0 - \boldsymbol{\theta}^* ||^2 \end{equation} One can clearly see the effect of $\mu(0)$ in the bound, which is even more dominant when $\mu(k)$ changes slowly. Of course, the effect of $\mu(0)$ attenuates over time, and that's why we are talking about a local effect. However, one should keep in mind that overparametrization leads the parameter to stay closer to initialization (at least in the NTK regime). We are happy to add a detailed discussion in the final version of the paper. **Question 2**: > Is it more appropriate to adjust the experiment on the effect of condition number and convergence rate in this paper to study different initializations or study the depth of the current network? They are more likely to keep the minimum loss close to each other, and make more use of the analysis of convergence rate. **Answer**: We are not sure what the Reviewer is asking precisely and kindly ask the Reviewer to rephrase the question, so we can dedicate an answer to it during the Author-Reviewer discussion phase. We kindly request the reviewer to let us know if there are any remaining questions they may have. If they find that their queries have been sufficiently addressed, we would greatly appreciate it if they could reconsider their evaluationof our paper.
Reply to author rebuttal
**Response to rebuttal of weakness 1** My concern still remains. Fig. 9 is not a practical setting: the model for Fig. 9 is a 3-layer linear model and cannot directly perform classification for data that is not linearly separable. The first plot of Fig. 9 only reveals that the loss converges, then 1. if the model fits the data perfectly, then the data is linearly separable, which contradicts the second point of the author response. 2. if the model does not fit the data perfectly, then the first point of the author response is far from sufficient as the training dynamics is actually not a successful one. --- Overall, my point lies in that the proposed methods overlook many aspects of practical settings, therefore the range that the proposed methods can be applied to is rather limited.
Thank you for your comprehensive response. Yes, I see now that you have also run the main experiments on CIFAR-10, so this does address my concern. I wonder whether you could fit both MNIST and CIFAR-10 in figure 5, perhaps side by side? I believe the camera-ready version allows an extra page (please double check) so do consider this in case of acceptance. I am convinced by your explanation of the motivation. On second inspection, it does appear that this is alluded to in the introduction, but perhaps some more explicit references to work like KFAC and Sophia might make the motivation clearer, as you said in your rebuttal. I also agree that it is fair to leave the deep non-linear case to future work. Thank you for clarifying my points of confusion. Please make modifications to wording if you believe it will help readability. Overall, I think my initial score of 6 (weak accept) was slightly harsh. After realising that the experiments were also run on CIFAR-10, I am happy to raise my score to an 8 (strong accept) as I believe this is good work that has been carried out to an excellent standard.
Decision
Accept (poster)