A Theoretical Study of the Jacobian Matrix in Deep Neural Networks

1. It is not clear whether Approximation 1 is an existing result that has been proved or an assumption. It is presented like an assumption, but in section 4 (also slightly mentioned in the abstract), it is claimed as a recent breakthrough from Random Matrix Theory. I hope the authors can clarify: what is exactly the “recent breakthrough”? how does Approximation 1 relate to this result.

2. This is one of my major concerns. The paper assumed (in Approximation 2) that the weights matrices are i.i.d. across layers. I don’t think this is true. Think about the simple case – a two layer linear network: $f= v^T u x$, where $u$ and $v$ are the first and second layer weights, and $x$ is the input. During training, $u$ and $v$ are clearly correlated, as $v^T u$ must change its value. This should also be true with non-linear activation. The authors probably mistakenly think that: the learn rate is of small order resulting in a small change in each weight, hence keeping independent weight matrices. However, that is not enough for independence.

3. In the discussion after Theorem 3, the paper claims that the required step size is “mild” and is generally satisfied in practice. I could not agree on this point. The theorem is in the infinite-width limit, i.e., $n\to\infty$, hence one can not ignore the factor $log(n)$, which should be very big resulting in an infinitesimal step size. To me, Theorem 3 is a trivial result, as the step size is so tiny that the training does not really go far away from the initialization.

4. The paper claims “the stagnation phase spans the whole training procedure” without justification. As I mentioned above, the so-called stagnation phase is probably due to the small learning rate so that the training has not gone far yet. It is totally unclear whether the stagnation phase lasts until convergence.

5. The descent phase is only shown on a very simple model – one-dimensional linear model. It seems the analysis does not extend to more complicated models, hence, it is unclear whether this phase happen for other models, e.g., neural networks.

Moreover, this descent phenomenon on linear models has been shown in a prior work [1]. Hence, I don’t see much novelty on this point.

[1] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. “The large learning rate phase of deep learning: the catapult mechanism”.

1. The assumptions are strong and unrealistic. In Assumption 2, the authors assume that up to time $t$, the entries of the $D$ matrix follow i.i.d Bernoulli distribution, matrix $D$ are independent from the weight matrix $W$, and the weight matrix are independent across layers. On theoretical side, I can not see why these assumptions hold for even a toy model. On the practical side, although the authors conduct experiments on synthetic datasets to validate this, I am not convinced that these assumptions hold for the general case. Therefore, I can not tell whether the claimed Jacobian analysis beyond initialization comes from the overassumptions.
2. The classifications of Jacobian behavior into three phases is not well validated. For the first phase, the authors only provide theoretical justifications, but contradict the experiment results in figure 2, since the plots do not show that the Jacobian norm can stays approximately constant in the early phase of training. For the second and the third phase, the authors mainly provide empirical evidences but without a satisfactory theoretical justifications. I encourage the authors to conduct more extensive experiments to validate the claimed three stage behavior.
3. The presentations are not clear. For example, the contents are Section 4 and Section 5 are cut apart and I cannot see their internal connection upon first reading. I would recommend moving the stability theorem to the main text, explain the theorem before using it for Section 4 and Section 5. Furthermore, there are many typos that can cause misunderstanding of the main results. For example in Theorem 3. Assumptions 1 be assumption 2, if I understand it correctly. Overall, I recommend the authors to polish the writing and make a clear presentation of the results.

Assumption 2 appears to be quite stringent, as it presumes certain conditions throughout the entire training process. While there is empirical support provided, it lacks a solid theoretical basis, and the given intuitive reasoning doesn't resonate with me. This assumption seems crucial for this study.

Regarding Theorem 3, I'm not entirely certain, but my interpretation is that it essentially says: "In the initial stages of training, we are in the kernel regime with a stagnant feature, leading to a static norm of the Jacobian." Is the challenge in proving this mainly technical? That's how I currently see it. Furthermore, the example in Theorem 4 seems overly simplistic. Perhaps analyzing it using multi-dimensional input in deep linear networks would be more illuminating? Lastly, the assertion that gradient-based training inherently pushes the network towards smaller gradients, causing the Jacobian norm to decrease, doesn't sit right with me. In the end, even with a zero gradient, the Jacobian norm remains non-zero.

1. Evolution of the Jacobian: to me, the results are not surprising. If gradient descent can optimize the network, then it is clear the Jacobian would not blow up. That said, similar results for the 3 phases can be found here: https://arxiv.org/pdf/2106.16225.pdf (Figure 8, Lemma 8). While the paper above is on the Hessian rank, arguably, the effects are similar. But anyway, I like the fact that the authors are being more precise and that they provide bounds.

2. Figure 2: the discussion around these results, I think, is misleading. I am not able to spot a stagnation phase in any curve. What do you mean? I think here the interesting thing is the gap between Jacobians for methods that reach the same accuracy. However, I think one can draw a link to sharpness and say that the higher learning rate converges to flatter regions (for the edge of stability), and therefore, the Hessian (related to the Jacobian) will be smaller. I do not think your results can completely capture this phenomenon, and I would definitely not link this to feature learning and muP.

3. pruning: all makes sense, and again, I think this is very interesting. I think what is missing though is a trainability test: can you train your pruned networks after the jacobian has been stabilized?

Hope you can address some of the points above so I can raise my score!

• The paper is written fairly well, but I think that Sections 3 and 4 do not communicate well the intended information/contributions.
• The main result of the paper is that if the Jacobian is stable at initialization, then for small enough learning rates it remains almost constant during training, mainly focusing on the infinite-width regime where lazy-training typically applies. Instead, for the non-infinite widths the Jacobian seems to change. I think it is not clear what is the actual goal of the paper.