$$

Suppose a product $w_1\cdots w_d$ is trained to approach a target value $s$, with gradient descent. How fast will it approach, and how does this depend on the degree $d$?

Setup

More concretely, let’s define a squared loss

$$\mathcal{L}:=\frac{1}{2}(w_1\cdots w_d-s{)}^2,$$

and make each “weight” $w_i$ evolve over time according to its gradient

$$\frac{\mathrm{d}{w}_i}{\mathrm{d}t}=-\frac{\partial \mathcal{L}}{\partial w_i}=(s-w_1\cdots w_d)\prod _{j\ne i}{w}_j.$$

This corresponds to what would happen if we were somehow training a linear network of depth $d$ and width $1$ on a 1D linear regression task where the ground truth is $f(x)=sx$. But more generally, as shown in Deep linear networks, this also roughly describes the dynamics with which linear networks of any width learn the singular modes of the linear relationship between inputs and outputs.

Degree $d=2$

When the degree is $2$, we have just

$$\{\begin{array}{rl}\frac{\mathrm{d}{w}_1}{\mathrm{d}t}& =(s-w_1{w}_2)w_2\\ \frac{\mathrm{d}{w}_2}{\mathrm{d}t}& =(s-w_1{w}_2)w_1.\end{array}$$

Phase 1: growth

Let’s first think about the early stage, when $w_1{w}_2\le s/2$, where we have

$$\{\begin{array}{rl}\frac{\mathrm{d}{w}_1}{\mathrm{d}t}& =\mathrm{\Theta }(s{w}_2)\\ \frac{\mathrm{d}{w}_2}{\mathrm{d}t}& =\mathrm{\Theta }(s{w}_1).\end{array}$$

Intuitively, we should expect $w_1$ and $w_2$ to approach each other (at least in relative terms): indeed, suppose that initially $w_1\gg w_2$, then that would cause $w_2$ to grow much faster than $w_1$ until it mostly catches up to it. We describe this process of catching up in more details in the section How fast do weights equalize? below, but for now let’s just assume that $w_1$ and $w_2$ have caught up, which means there is now just a single variable $w=w_1=w_2$. The early stage corresponds to $w^2\le s/2$, and

$$\frac{\mathrm{d}w}{\mathrm{d}t}=\mathrm{\Theta }(sw)⟹\frac{\mathrm{d}\left(w^2\right)/\mathrm{d}t}{w^2}=\mathrm{\Theta }(s)$$

The relative derivative is constant, so this solves to an exponential $w^2(t)=e^{\mathrm{\Theta }(st)}{w}^2(0)$, which means that it would take time $\mathrm{\Theta }\left(\frac{\log (s/ϵ)}{s}\right)$ to go from a small value $w^2=ϵ$ to the halfway point $w^2=s/2$.

Phase 2: approach

In the later stage of learning, when $w^2\ge s/2$, we have

$$\frac{\mathrm{d}\left(w^2\right)}{\mathrm{d}t}=2(s-w^2)w^2,$$

so

$$\frac{\mathrm{d}(s-w^2)/\mathrm{d}t}{s-w^2}=-2w^2=-\mathrm{\Theta }(s),$$

which solves to $s-w^2(t)=e^{-\mathrm{\Theta }(s)t}(s-w^2(0))$, which means that it would take time $\mathrm{\Theta }\left(\frac{\log (1/ϵ)}{s}\right)$ to go from the halfway point $w^2=s/2$ to $w^2=s(1-ϵ)$.

Learning rate

#to-write

• derive the maximum learning rate
• adapt the conclusion in term of training steps
• but think about how this would change things with width >1 when there’s several different objectives with different singular values $s$: make the point that larger singular values are generally favored (which is not the case when $d=1$)
• (note that in Deep linear networks we didn’t worry about the learning rate because it was all relative between different singular values; on the other hand you can’t make claims about (absolute) acceleration without worrying about those learning rates)

General case

In general, similar dynamics will force the weights $w_1,\dots ,w_d$ to approach each other in relative terms,¹ and we would be left with one variable $w$ with

$$\frac{\mathrm{d}w}{\mathrm{d}t}=(s-w^d)w^{d-1}⟹\frac{\mathrm{d}\left(w^d\right)}{\mathrm{d}t}=d(s-w^d)w^{2d-2}.$$

Phase 1: growth

When $w^d\le s/2$, treating $d$ as a constant, we have

$$\frac{\mathrm{d}\left(w^d\right)/\mathrm{d}t}{w^d}=\mathrm{\Theta }\left(s{w}^{d-2}\right)=\mathrm{\Theta }\left(s{\left(w^d\right)}^{1-2/d}\right).$$

Depending on the sign of the exponent $1-2/d$ in this relative derivative, $w^d$ will experience different types of growth:

• If $d<2$, then $w^d$ will grow polynomially. In particular, the only relevant case is $d=1$, and gives $w(t)=w(0)+\mathrm{\Theta }(s)$. So going from a small value $w\approx 0$ to the halfway point $w=s/2$ takes $\mathrm{\Theta }(1)$ time.
• If $d=2$, as seen above, $w^2$ grows exponentially, and going from $ϵ$ to $s/2$ takes $\mathrm{\Theta }\left(\frac{\log (s/ϵ)}{s}\right)$ time.
• If $d>2$, then $w^d$ will grow hyperbolically as $w^d(t)=\mathrm{\Theta }{\left(\frac{1}{s(t-t_0)}\right)}^{1+\frac{2}{d-2}}$, and going from a small value $w^d=ϵ$ to the halfway point $w^d=s/2$ takes $\mathrm{\Theta }\left(\frac{{ϵ}^{-(1-2/d)}}{s}\right)$ time.

In particular, if $w^d$ is starting from a fairly low value, degrees $d>2$ will take a long time to take off, and in general the trade-off depends on $s$.

Phase 2: approach

When $w^d\ge s/2$, treating $d$ as a constant, we have

$$\frac{\mathrm{d}(s-w^d)/\mathrm{d}t}{s-w^d}=\mathrm{\Theta }\left(w^{2d-2}\right)=\mathrm{\Theta }\left({\left(w^d\right)}^{2-2/d}\right)=-\mathrm{\Theta }\left(s^{2-2/d}\right),$$

so in any case the approach is exponential, with

$$s-w^d(t)=e^{-\mathrm{\Theta }\left(s^{2-2/d}t\right)}(s-w^d(0)),$$

which means that going from the halfway point $w^d=s/2$ to $w^d=(1-ϵ)s$ takes

$$\mathrm{\Theta }\left(\frac{\log (1/ϵ)}{s^{2-2/d}}\right)$$

time. In this phase, bigger degree is always better.

Learning rate

#to-write same as above

Nonlinearities

Now suppose that we’re trying to approach $s$ with a monomial like $w_1^2{w}_2$, where some of the weights are squared, and the loss is still the square

$$\mathcal{L}:=\frac{1}{2}{(w_1^2{w}_2-s)}^2.$$

This is roughly the situation that would arise if you’re trying to approximate function $f(x)=s{x}^2$ using a depth-2 network with a quadratic nonlinearity $\varphi (h):=h^2$ on its hidden layer.

Then the dynamics are

$$\{\begin{array}{rl}\frac{\mathrm{d}{w}_1}{\mathrm{d}t}& =2(s-w_1^2{w}_2)w_1{w}_2\\ \frac{\mathrm{d}{w}_2}{\mathrm{d}t}& =(s-w_1^2{w}_2)w_1^2.\end{array}$$

Again, in the growth phase where $w_1^2{w}_2\le s$, we have

$$\{\begin{array}{rl}\frac{\mathrm{d}{w}_1}{\mathrm{d}t}& =\mathrm{\Theta }(s{w}_1{w}_2)\\ \frac{\mathrm{d}{w}_2}{\mathrm{d}t}& =\mathrm{\Theta }(s{w}_1^2),\end{array}$$

so again, $w_1$ and $w_2$ will tend to equalize, because if e.g. $w_1\gg w_2$, then $w_2$ will grow much faster. So we can approximate the dynamics by

$$\frac{\mathrm{d}w}{\mathrm{d}t}=\mathrm{\Theta }\left((s-w^3)w^2\right).$$

Not too surprisingly, this is exactly the same dynamics that we got for a the problem of approximating $s$ by $w_1{w}_2{w}_3$. So, depending on how a nonlinearity acts, it seems like it could make a neural network act like it has more depth than it actually has!

How fast do weights equalize?

Let’s study how fast the ratio $w_1/w_2$ approaches $1$ when weights are subject to the dynamics

$$\{\begin{array}{rl}\frac{\mathrm{d}{w}_1}{\mathrm{d}t}& =(s-w_1{w}_2)w_2\\ \frac{\mathrm{d}{w}_2}{\mathrm{d}t}& =(s-w_1{w}_2)w_1.\end{array}$$

For this, it’s convenient to look at the relative derivative of both the ratio $w_1/w_2$ itself, but also compare it with the relative derivative of their product $w_1{w}_2$, since it will allow us to compare the timescales at which equalizing vs learning happen.

We have

$$\begin{array}{rl}\frac{\mathrm{d}(w_1{w}_2)/\mathrm{d}t}{w_1{w}_2}& =\frac{\mathrm{d}{w}_1/\mathrm{d}t}{w_1}+\frac{\mathrm{d}{w}_2/\mathrm{d}t}{w_2}=(s-w_1{w}_2)(\frac{w_2}{w_1}+\frac{w_1}{w_2})\\ \frac{\mathrm{d}(w_1/w_2)/\mathrm{d}t}{w_1/w_2}& =\frac{\mathrm{d}{w}_1/\mathrm{d}t}{w_1}-\frac{\mathrm{d}{w}_2/\mathrm{d}t}{w_2}=(s-w_1{w}_2)(\frac{w_2}{w_1}-\frac{w_1}{w_2}),\end{array}$$

which at the least confirms that if $w_1>w_2$ then $w_1/w_2$ is decreasing over time, and vice versa. Let’s assume we start out with $w_1\gg w_2>0$, then we have

$$\begin{array}{rl}\frac{\mathrm{d}(w_1{w}_2)/\mathrm{d}t}{w_1{w}_2}& \approx (s-w_1{w}_2)\frac{w_1}{w_2}\\ \frac{\mathrm{d}(w_1/w_2)/\mathrm{d}t}{w_1/w_2}& \approx -(s-w_1{w}_2)\frac{w_1}{w_2}.\end{array}$$

That is, $w_1{w}_2$ increases at the same relative rate as $w_1/w_2$ decreases, which means (to a first approximation), we can expect that $w_1$ and $w_2$ will roughly equalize as long as we start out with

$$\underset{\text{how much }{w}_1{w}_2\text{ needs to grow}}{\underbrace{\frac{s}{w_1{w}_2}}}\ge \underset{\text{how much }{w}_1/w_2\text{ needs to decrease}}{\underbrace{\frac{w_1}{w_2}}}⟺w_1\le \sqrt{s}.$$

1. or at least, it seems like their absolute values would ↩