Geometry of the Loss Landscape in Invariant Deep Linear Neural Networks

Thanks for the response and various clarifications, I hope they'll help future readers. Regarding some of the points, starting from minor:

3 - it might be useful giving this intuition explicitly, it's helpful.

4 - word-searching the document I can't find "tangent". I don't think this is an issue but FYI.

7 - excellent. Note that proposition 5, line 831, has "group Lie group"

8 - I believe you mean a unique PSD root, no? https://mathoverflow.net/questions/266286/is-the-square-root-of-a-matrix-unique would appreciate clarifying this in the main text.

9 - I see, recommend clarifying this.

10 - Fascinating, thanks for clarifying.

11 - then why mention them? Doesn't matter much either way as you've now clarified it.

12 - likely not exactly the same but I believe [1] for example has that decomposition. It isn't a main result of 4.3 and hence I don't think detracts from it, but perhaps worth mentioning.

13 - this is an interesting conjecture, it's good you added it to the main text.

And regarding more major, conceptual ones:

2 - believe my issue here is mostly the discussion part, apologies for not making it clearer. Specifically - the different kinds of solutions for each setting are found, but their implications are not deeply discussed. The solutions themselves are interesting but it is not said why they are interesting, or what insight we get from them. This is related to what I wrote under Decision, where currently the paper's novelty is presented via its results but insufficiently interpreted and spoon-fed to the reader. Even if the results are interesting only to someone working on theoretical geometric deep learning, ideally any deep learning researcher can read the paper and understand why people in that community find it interesting.

Clarity - the first paragraph of clarity issues mentioned in the weaknesses still for the most part hold. Small changes here and there have helped alleviate them but insufficiently. For example, the conclusion section states the results but doesn't say what's interesting about them, how they may carry over to real networks, etc. The related work section still goes over existing works without the reader always understanding why they are mentioned, or explicitly how the current work answers gaps/limitations they have - see [2] for a good example of a related work section which gives good context.

Rank constraint - An earlier version of this comment was posted before/synchronously to your reply regarding the rank constraint. I think that explanation is extremely elucidating and implore you to add it to the paper.

If I didn't comment on something it means it's a good change or I didn't notice it.

[1] Gideoni, Yonatan. "Implicitly Learned Invariance and Equivariance in Linear Regression."

[2] Finn, Chelsea, Pieter Abbeel, and Sergey Levine. "Model-agnostic meta-learning for fast adaptation of deep networks." International conference on machine learning. PMLR, 2017.

Thank you for your very detailed review. We have updated the manuscript in response to your review. Answers to some of the main questions you asked are provided below:

1. Line 41 - “have shown” how? Feels like a citation’s needed.

Citations have been added.

2. Lines 47-49 - where in the paper are the solutions studied/referred to? Eg. when are the regularised solutions invariant? This is implicitly shown but not discussed.

The solutions are studied in Section 3 (Section 3.1 for hard-wiring, Section 3.2 for regularization, Section 3.3 for data augmentation). The regularized solutions are not invariant when the penalty coefficient $\lambda$ is finite.

3. Lines 74-76 - how is this intuitive/obvious? It’s important to give that intuition. To play devil’s advocate, if it’s obvious then it’s even moreso important to clarify why this is studied and what new insight was achieved if any.

We stated that this might seem intuitive to some readers, but that we certainly do not consider it to be obvious. We have rephrased to avoid confusion. To explain what that intuition might be: augmenting the data across the group orbits does not bring new information other than enforcing the group symmetry. Therefore, one might expect that data augmentation should have a similar effect as constraining (hard-wiring) the model.

4. Lines 145-146 - why is the tangent space relevant here? I didn’t understand where it was heavily used throughout the paper.

It is used in proving Proposition 7. Since it is not heavily used, we have removed this line, and added it in Proposition 7.

5. Lines 197 - shouldn’t det-variety be defined? It’s not a well known term, and if it’s not important enough to be defined it should be delegated to the appendix or not used.

We added a reference in the main text. A brief description is included in accessible language. We mentioned this to emphasize how this set is different from a linear space. To explain the terminology: a matrix has rank at most $r$ if all of its sub-matrices of size $(r+1)\times(r+1)$ have determinant zero. Thus, the set of matrices of rank at most $r$ is the set of matrices whose entries satisfy a list of polynomial equations (this makes it a variety) and these equations can be written as determinants (making it a determinantal variety).

6. Line 211 - what’s r in this context?

Consider a fully-connected linear network where the narrowest layer has width $r$. The function space is the set of end-to-end matrices, which are $n \times m$ matrices of rank at most $r$. The manuscript has been updated to clarify this.

7. Lines 234-236, Finitely generated groups:

We presented our results for the case of cyclic groups or finitely generated groups. However, extending our results to continuous groups (such as $SO(n)$ and $SE(n)$) is straightforward, since Lie theory provides a way of analyzing continuous groups in terms of their infinitesimal generators. To explain this point we have updated the manuscript to include Remark 1 and an Appendix A.1 with details about the extension of our results to Lie groups.

8. Lines 243-245 - do all roots work equally well? Assuming this corresponds to several classes of solutions no?

It is known that a positive definite matrix has a unique square root.

9. Line 303-305, Continuity of regularization path:

Yes, due to the rank constraints, it is not clear whether the regularization path is continuous. We are able to show that it is indeed a continuous path.

10. Lines 373-377 does spurious here mean suboptimal? Recommend clarifying.

Pure critical points are critical points that arise from the geometry of the function space, while spurious ones result solely from the parametrization map and do not correspond to critical points in the function space.

11. Line 395-396 why change Adam’s betas?

We haven't changed betas of Adam. These are the default values in PyTorch.

12. Lines 486-496 - this is a nice decomposition but I believe I saw it in other works, are you aware of it appearing elsewhere in the literature?

It is a relatively intuitive decomposition that other people may also have come up with. We are not aware of other work using exactly the same setup as here.

13. Line 496-497 - why is the double descent interesting? Do you mean the orange line in figure 3.a?

All the lines in 3(a) have different degrees of "double descent", especially the blue (data augmentation) and the orange ones. Our conjecture is that the loss may also be decomposed into two parts, one controlling the error of invariance, and the other one controlling the error from the target. Therefore, the gradient of the weights during training can be decomposed into two directions as well, and their differences may result into this phenomenon. This can help us better understand the training dynamics of those models, which eventually could shed light on methods to accelerate training.

We appreciate the reviewer’s comments and feedback.

1. Regarding the limitation to linear networks:

We consider a model of rank constrained linear maps. While this is indeed a relatively simple model, one of its interesting properties is that it has a nonlinear function space. Part of our motivation to study this model is that it could inform us about how invariance constraints might interact with a nonlinear function space. Networks with nonlinear activation functions are typically also nonlinear. We think that rank bounded linear functions are an interesting and natural place to start, but agree that studying other types of constraints will be important towards obtaining a more complete picture of general invariant neural networks. Linear convolutional networks might be an interesting model to consider next in conjunction with invariance, which in contrast to fully connected linear networks may have spurious local minima and richer types of constraints in function space.

2. Regarding the limitation to invariance in the linear context:

It is interesting to consider more general types of invariances. Most group invariances people use in practice are linear, such as rotation and translation. In representation theory, a representation is defined as a group homomorphism between $\mathcal{G}$ and $GL(V)$, i.e., $\rho: \mathcal{G} \rightarrow GL(V)$. Thus, when people set up the framework for invariant or equivariant models, they implicitly assume that the transformations are linear.

3. Regarding other architectures:

We agree that linear networks are a relatively simple model. One of the interesting aspects of our work is that we consider function spaces that are not required to be linear. This means that some phenomena and conclusions might be transferrable to other architectures. We suggest that data-augmented model may have similar performance to hard-wired ones, though at the cost of data and computing efficiency. Regularized models do not require more data and should have close performance to the hard-wired model, but may induce more critical points than the other two methods. Hard-wired model should have the best performance, though one might need to design the invariant architecture carefully before feeding the data to the model. These propositions of course will necessitate further investigation, and we hope our work might serve as an inspiration or starting point.

Dear Reviewer,

Thank you for your time and feedback on our paper. We have provided a detailed response addressing your concerns. Since then, we have not seen any follow-up comments from you. If there are additional questions or concerns that remain unresolved, we would be happy to address them further. Your feedback is valuable to us, and we sincerely hope to help clarify any lingering doubts.

Thank you to the authors for their response. However, I would like to maintain my current score, as focusing on linear networks substantially limits the contributions of this paper.

We appreciate the reviewer’s insightful comments and feedback.

1. Regarding the limitation of cyclic groups or finitely generated groups: We presented our results for the case of cyclic groups or finitely generated groups. However, this is not a serious limitation in the following sense. Extending our results to continuous groups (such as $SO(n)$ and $SE(n)$) is straightforward, since Lie theory provides a way of analyzing continuous groups in terms of their infinitesimal generators. To explain this point, we have updated the manuscript to include Remark 1 and an Appendix A.1 with details about the extension of our results to Lie groups.

2. Regarding the assumptions in Remark 2:

It is true that the noise can be structured in real data. In our assumption, we only need that the noises for different data entries are independent and identically distributed. While simplistic, this is a relatively common assumption and is relatively mild in practice. We added an appendix where we empirically computed the spectrum of target matrices in MNIST and verified that they satisfy the assumptions.

3. Regarding citations:

Thank you for providing these relevant references. We have updated our manuscript to include them.

4. Regarding figure 3 and figure 5 (originally figure 4):

In Figure 3(a), we are plotting the Frobenius norm of the non-invariant component of the end-to-end matrix, i.e., $||W^{\perp}||_F$, for data augmentation and regularization trained with mean squared loss. For data augmentation, it actually converges to zero if we allow more epochs. For regularization, since the penalty coefficient $\lambda$ is finite, the critical points are actually not invariant. Therefore, in those cases $||W^{\perp}||_F$ does not converge to zero. However, the larger the regularization strength, the smaller the Frobenius norm of the non-invariant component at convergence, in line with the theory. In Figure 5(a), we again are plotting $||W^{\perp}||_F$ for data augmentation and regularization trained on the same dataset, but with cross entropy loss. It is observed that, for larger $\lambda$, the dynamics of $||W^{\perp}||_F$ resemble those when trained with MSE. On the other hand, for small $\lambda$, $||W^{\perp}||_F$ may increase at first, and then decrease. For data augmentation, we observe that $||W^{\perp}||_F$ actually decreases after increasing if we allow more epochs. Our theoretical results only support the scenario for mean squared loss. Thus, when trained with cross entropy, we cannot say whether all the critical points are invariant or not. Further research needs to be done to investigate this.

We appreciate the reviewer’s insightful comments and feedback.

1. Regarding your comment about the settings being limited to vector spaces:

Kindly observe that the model we are considering consists of functions that are linear over the input variable but that the function space itself is not a vector space. Our function space consists of linear maps with bounded rank. In general this is a so-called determinantal variety and is not a convex set. It is a vector space only in the special case where the bound on the rank is trivial (i.e., equal to either the number of rows or columns). Due to the non-convexity of the function space, the effect of the constraints imposed by the invariances is not obvious. We regard this as one of the most interesting parts of our work. We consider this setting precisely to gain a better understanding of function spaces that are nonlinear subsets of vector spaces, which is also the case in networks with nonlinear activation functions.

2. Regarding your comment about the settings being limited to cyclic and finitely generated groups:

We presented our results for the case of cyclic groups or finitely generated groups. However, we contend that this is not a serious limitation. Extending our results to continuous groups (such as $SO(n)$ and $SE(n)$) is straightforward, since Lie theory provides a way of analyzing continuous groups in terms of their infinitesimal generators. To explain this, we have updated the manuscript to include Remark 1 and an Appendix A.1 with details about the extension of our results to Lie groups.

3. Regarding the comparison with Nordenfors et al., (2024) [1]:

Regarding Theorem 3.7 in the work of Nordenfors et al., 2024 [1]:

1. That result shows that the data augmented model and the hard-wired model have the same stationary points within the set of representable equivariant maps $\mathcal{E}$. As they state in their Remark 3.8.2, the set $\mathcal{E}$ is assumed to be a vector space. This means that the result considers only a limited class of architectures. In our work, in contrast, the function space is not required to be a vector space. Our function space is a vector space only in the special case where the rank constraint is trivial. The fact that our function space is not required to be a linear space is in our opinion one of the most interesting aspects of our work.
2. Furthermore, as stated in Nordenfors' Remark 3.8.1, their theorem only applies to points in $\mathcal{E}$. This means that their result does not say anything about stationary points of the data augmented model that are not in $\mathcal{E}$. In contrast, our results describe all critical points for a non-linear rank-constrained function space and we show that all of them are indeed invariant. Finally, we also compare the regularized model with the other two, whereas Nordenfors et al. do not consider regularization in their work.

4. Regarding Section 4.3:

The experiments in Section 4.3 show the training dynamics for linear networks under the settings that we consider. Since the parametrized function $\hat{W}$ is linear, we can decompose it orthogonally into two parts, a non-invariant part $\hat{W}^{\perp}$ and an invariant part $\hat{W}-\hat{W}^{\perp}$. We are plotting the Frobenius norm of both components during training. This is to monitor whether the non-invariant part converges to zero and at what rate for the case of data augmentation and regularization. The experiments show that as the regularization strength increases, the result of training will tend to have a smaller non-invariant component. This is in line with our theory in Theorem 2. Our theoretical results focus on the static loss landscape rather than training dynamics. The experiments section complements this experimentally and hints at interesting phenomena that merit further investigation. We have updated the manuscript for better illustration.

5. Useful insights for practitioners:

Though our theoretical results are for linear networks with non-convex function space, it is natural to conjecture that some phenomena might carry over to other overparametrized models with non-convex function space. We suggest that the data-augmented model may have similar performance to hard-wired ones, but will incur higher data and computing costs. The regularized model does not require more data and should have a performance close to the hard-wired model, but it may induce more critical points than the other two methods. The hard-wired model should have the best performance, though one might need to design the invariant architecture carefully before feeding the data to the model.

References:

[1] Oskar Nordenfors, Fredrik Ohlsson, and Axel Flinth. Optimization dynamics of equivariant and augmented neural networks, 2024.

We appreciate the reviewer’s comments and feedback.

1. Regarding novelty and differences from previous literatures:

As indicated in the introduction, Baldi & Hornik (1989) [1] studied two-layer linear networks with no constraints on the weight matrices other than their format and trained with the square loss. They show that this optimization problem has single minimum up to a trivial symmetry and all other critical points are saddles. Notice that this general structure fundamentally results from the properties of low-rank matrix approximation with the square loss and corresponding results from the 1930s. That work does not investigate learning invariances in the data. In contrast, we aim to study how data augmentation, regularization in function space, and imposing constraints based on invariances can affect the optimization problem and the loss landscape. We specifically consider a function space subject to rank constraints, which is a determinantal variety and not a vector space, and for which it is not clear how it interplays with the constraints imposed by invariance. We are not aware of previous works studying this setting.

2. Regarding Proposition 2 and Proposition 7(originally Theorem 4):

1. For Proposition 2: We obtain this by applying manifold regularization results from Zhang & Zhao (2013) [3], Theorem 1. The difference is that the symmetric positive semidefinite matrix that characterizes the low-dimensional structure of the manifold behind the weight matrix is $\tilde{G}\tilde{G}^T$ in our context.
2. For Proposition 7: We had included a version of Trager's Theorem 28 for clarity of presentation. We have now moved it to the appendix. Observe that this result is a version of results that have appeared in diverse works on low rank matrix approximation following from Mirski's characterization. A minor difference is that we included the description of the critical points for the case of low rank targets, which was not included by Trager et al. (2020) [4].

3. Regarding your question about finite data:

In our paper, we don't have any assumptions related to infinite data. All the results hold for a finite amount of data in general position. We are able to prove that the critical points in function space for data augmentation and hard-wired invariant model are the same, implying that the generalization performance is be the same for both methods if global optima are achieved.

4. Regarding the comparison with Levin et al. (2024) [2]:

We are familiar with the work of Levin et al. (2024) [2], which studies the effect that the parametrization map has on an optimization landscape and specifically establishes conditions for when a critical point in the parameter space corresponds to a critical point in the function space. In contrast, we aim to characterize the global optima in a rank-constrained function space for the optimization problem in data augmentation, regularization, and hard-wired models. Discussing the effect of parametrization in general is not a main focus of our paper. Nonetheless, we can show that there are no local minima other than global minima in the parameter space when the model is parametrized as the product of linear weight matrices.

References:

[1] P. Baldi and K. Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2(1):53–58, 1989.

[2] Eitan Levin, Joe Kileel, and Nicolas Boumal. The effect of smooth parametrizations on nonconvex optimization landscapes. Math. Program., March 2024. ISSN 0025-5610, 1436-4646. doi: 10.1007/s10107-024-02058-3.

[3] Zhenyue Zhang and Keke Zhao. Low-rank matrix approximation with manifold regularization. IEEE Transactions on Pattern Analysis and Machine Intelligence, 35(7):1717–1729, 2013. doi: 10.1109/TPAMI.2012.274.

[4] Matthew Trager, Kathlen Kohn, and Joan Bruna. Pure and spurious critical points: a geometric study of linear networks. In International Conference on Learning Representations, 2020.

Dear Reviewer,

Thank you for your time and feedback on our paper. We have provided a detailed response addressing your concerns. Since then, we have not seen any follow-up comments from you. If there are additional questions or concerns that remain unresolved, we would be happy to address them further. Your feedback is valuable to us, and we sincerely hope to help clarify any lingering doubts.