Distribution of gradients across dimensions for neural networks

Question

Does the distribution of gradients for neural networks known to follow a particular distribution?

That is, suppose I've a model with $N$ parameters. Then, the (stochastic) gradient at some point is a $N$ dimensional vector $(x_1, x_2, \dots, x_N)$.

Can we say anything regarding how $x$ values are distributed over this vector?

Answer 1

The exact answer is going to depend greatly on the type of network, the inputs, how it's trained....

For a simple way to see this:

• If we're at a (local) optimum, the full gradient (across the entire training dataset) will be zero. In the interpolating regime common to modern neural networks, the individual gradient for each training point may even be exactly zero; depending on the loss function and how much you've trained / etc, it might instead be mean zero but approximately normal, etc.

• At initialization, when we're very far from a solution, the gradient may be extremely similar for different datapoints (and very far from zero).

In some particular limits (network becoming infinitely wide, initialized near zero, trained via SGD, no batch normalization, ...), the "neural tangent kernel" regime has an answer here: activations are distributed according to a particular Gaussian process, and their derivatives are too. (When using square loss, this means the final result corresponds to kernel ridge regression / Gaussian process regression with a particular kernel.) See e.g.

• Appendix 2 of Jacot et al. (2018), Neural Tangent Kernel: Convergence and Generalization in Neural Networks;
• Lee et al. (2019), Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent;
• Appendix D of Arora et al. (2019), On Exact Computation with an Infinitely Wide Neural Net.