Modern neural network architectures often generalize well despite containing many more parameters than the size of the training dataset. This paper explores the generalization capabilities of neural networks trained via gradient descent. We develop a data-dependent optimization and generalization theory which leverages the low-rank structure of the Jacobian matrix associated with the network. Our results help demystify why training and generalization is easier on clean and structured datasets and harder on noisy and unstructured datasets as well as how the network size affects the evolution of the train and test errors during training. Specifically, we use a control knob to split the Jacobian spectum into "information" and "nuisance" spaces associated with the large and small singular values. We show that over the information space learning is fast and one can quickly train a model with zero training loss that can also generalize well. Over the nuisance space training is slower and early stopping can help with generalization at the expense of some bias. We also show that the overall generalization capability of the network is controlled by how well the labels are aligned with the information space. A key feature of our results is that even constant width neural nets can provably generalize for sufficiently nice datasets. We conduct various numerical experiments on deep networks that corroborate our theoretical findings and demonstrate that: (i) the Jacobian of typical neural networks exhibit low-rank structure with a few large singular values and many small ones leading to a low-dimensional information space, (ii) over the information space learning is fast and most of the labels falls on this space, and (iii) label noise falls on the nuisance space and impedes optimization/generalization.
The authors use the empirically supported assumption of the low-rank nature of the neural network Jacobian to provide new elements of data-dependent optimization and generalization theory. By modelling the data as a low-rank object itself, they analytically study the evolution of the train and test errors. The paper divides the space of weights and biases into the “information” and “nuisance” subspaces, spanned by the top largest and the remaining singular vectors of the Jacobian respectively. They use this division and its alignment with the low-rank structure of the data to talk about convergence speed. Finally, they provide numerical experiments to back their claims.
I enjoyed the paper, however, there were many points where I was unclear on the precise nature of the assumptions used / the strength of the results.
Disclaimer: I didn’t manage to read through the proofs in the appendix and cannot therefore vouch for its correctness.
-- Point 1 --
Leveraging the data structure
I am unclear on how exactly you were modelling the structure of the data. From your proofs, it seems that you have been dealing with the matrix X comprising the concatenated flatted vectors of the raw input features (e.g. pixels) of the input data [x1,x2,...,x_datasetsize]^T. In particular, the only place where I see data explicitly enter is in Definition 2.1, where you look at the X X^T and fi’(w X) fi’(w X)^T.
If the data is linearly separable in the raw input space on its own, then I see that the matrix X X^T will be low-rank (related to the number of classes). I also see your point about the connection of the y to the relevant (semantic) clusters. The same argument could by applied to fi’(w X) fi’(w X)^T -- provided that the features produced are again linearly separable, we will observe this object to have a low (number of classes - 1) rank. 
What is unclear to me is whether these assumptions are warranted. I understand that some simple datasets, e.g. MNIST, are essentially linearly separable in the raw pixels, and therefore the X X^T indeed is low rank. However, I doubt anything like that is true for big datasets, such as ImageNet. For deep networks that are used on these big datasets, such a modelling assumptions would likely not be true. I wonder how this relates to your results, since the low-rank nature of the Hessian is observed even for those networks, which is in turn related to the low rank nature of the Jacobian.
-- Point 2 --
Data implicitly present in the Jacobian tensor.
I wonder how you modelled the Jacobian tensor that you started using on page 3. Since the Jacobian -- the derivative of the output logits with respect to the weights, has to be evaluated at a particular input X, the assumptions you make on the data are in turn having an effect on the Jacobian, and vice versa.
I am unclear on exactly what assumptions you make about the object, and whether you are actually saying that its low rank structure comes from the data, is empirical observed and therefore assumed, or due to the network regardless of the data.
I recently saw a new arXiv submission that seems to be looking into this on real networks: https://arxiv.org/pdf/1910.05929.pdf Their model explicitly assumes that logit gradients cluster in the weight space in a particular way.
-- Point 3 --
Square loss vs softmax
You are using the square loss |f(X) - y|^2 throughout your work. Many of the empirical low-rank observations of the Hessian (related to the JJ^T) are performed on real networks with the cross-entropy loss. While the Hessian with the square loss is of the form JJ^T + terms, the softmax in the cross-entropy loss introduces an additional cross-term (let us call it P for now), which in turn makes it JP(PJ)^T + terms. Do you know how this relates to your results?
More generally, does the square loss you use make the results significantly different from what we would get for a softmax?
-- Point 4 --
Neural Tangent Kernel (NTK) -- assumptions
Under the NTK assumption, you still need to model the derivatives of each logit with respect to each weight on each input in order to obtain the Jacobian matrix and in turn JJ^T. I am therefore very confused by “Based on our simulations the M-NTK indeed haslow-rank structure with a few large eigenvalues and many smaller ones. “ on page 5. What assumptions exactly do you use in your model?
-- Point 5 --
Neural Tangent Kernel (NTK) -- validity
By assuming the NTK holding, do you limit the validity of your results? I think it is believed that NTK might not, generally speaking, be enough to capture the complexity of DNNs, and therefore assuming it might limit the range of applicability of any results derived assuming it.
-- Conclusion --
In general, my points of confusion often stem from being unsure as to what parts of the argument were assumed and based on what empirical / theoretical evidence, and what parts were generically true. While the paper seems interesting, I am not sure what its novel contribution is and how broad the claims made actually are in their applicability.
Appendix: I was not able to judge the proofs in the appendix.