Feature vs. Lazy Training

Stefano Spigler

Matthieu Wyart

Mario Geiger

Arthur Jacot

https://arxiv.org/abs/1906.08034

Two different regimes in the dynamics of neural networks

feature training

the network learns features

lazy training

no features are learned

\(f(w) \approx f(w_0) + \nabla f(w_0) \cdot dw\)

\(f(w,x)\)

characteristic training time separating the two regimes

training procedure to force each regim

How does the network perform in the infinite width limit?

\(n\to\infty\)

There exist two limits in the literature

\(n\) neurons per layer

Overparametrization

  • perform well
  • theoretically tractable

The 2 limits can be understood in the context of the central limit thm

\(\displaystyle Y = \frac{1}{\color{red}\sqrt{n}} \sum_{i=1}^n X_i \quad \)

As \({\color{red}n} \to \infty, \quad Y \longrightarrow \) Gaussian

\(\langle Y \rangle={\color{red}\sqrt{n}}\langle X_i \rangle\)   

\(\langle Y^2\rangle - \langle Y\rangle^2 = \langle X_i^2 \rangle - \langle X_i\rangle^2\)

\(\langle Y \rangle=\langle X_i \rangle\)   

\(\langle Y^2\rangle - \langle Y\rangle^2 = {\color{red}\frac{1}{n}} (\langle X_i^2 \rangle - \langle X_i\rangle^2)\)

\(\displaystyle Y = \frac{1}{\color{red}n} \sum_{i=1}^n X_i\)

Central Limit Theorem:

\(X_i\)

(i.i.d.)

(Law of large numbers)

regime 1: kernel limit

\(x\)

\(f(w,x)\)

\(W^1_{ij}\)

\(\displaystyle z^{\ell+1}_i = {\frac1{\color{red}\sqrt{n}}} \sum_{j=1}^n W^\ell_{ij} \ \phi(z^\ell_j)\)

At initialization

\(W_{ij} \longleftarrow \mathcal{N}(0, 1)\)

\(w\) = all the weights

\(z^1_j\)

\(z^2_i\)

With this scaling, small correlations adds up significantly => the weights will change only a little during training

[1995 R. M. Neal]

Independent terms in the sum, CLT

=> when \({\color{red}n} \to\infty\) output is a Gaussian process

space of functions \(\mathbb{R}^d \to \mathbb{R}\)

(dimension \(\infty\))

regime 1: kernel limit

\(f(w_0)\)

\(f(w_0) + \sum_\mu c_\mu \Theta(w_0, x_\mu)\)

kernel space

(dimension \(m\))

tangent space

(dimension \(N\))

\(f(w_0) + \nabla f(w_0) \cdot dw\)

\(d\) input dimension

\(N\) number of parameters

\(m\) size of the trainset

\(f(w)\)

network manifold

(dimension \(N\))

[2018 Jacot et al.]

[2018 Du et al.]

[2019 Lee et al.]

The NTK is independent of the initialization
and is constant through learning
=> the network behaves like a kernel method
 

The weights barely change
\(\|dw\|\sim \mathcal{O}(1)\) and \(dW_{ij} \sim \mathcal{O}({\color{red}1/n})\)   (\(\mathcal{O}({\color{red}1/\sqrt{n}})\) at the extremity)
 

The internal activations barely change

\( dz \sim \mathcal{O}({\color{red}1/\sqrt{n}})\)
=> no feature training

for \({\color{red}n}\to\infty\)

neural tengant kernel

\(\Theta(w,x_1,x_2) = \nabla_w f(w, x_1) \cdot \nabla_w f(w, x_2)\)

regime 1: kernel limit

regime 2: mean field limit

\(\displaystyle f(w,x) = \frac{1}{\color{red}n} \sum_{i=1}^n W_i \; \phi \! \left(\frac{1}{\color{red}\sqrt{n}} \sum_{j=1}^n W_{ij} x_j \right)\)

\(W_{ij}\)

\(x\)

\(f(w,x)\)

\(W_{i}\)

was \(\frac{1}{\color{red}\sqrt{n}}\) in the kernel limit

studied theoretically for 1 hidden layer

Another limit !

for \({\color{red}n}\longrightarrow \infty\)

\(\displaystyle f(w,x) = \frac{1}{\color{red}n} \sum_{i=1}^n W_i \; \phi \! \left(\frac{1}{\color{red}\sqrt{n}} \sum_{j=1}^n W_{ij} x_j \right)\)

\(\frac{1}{\color{red}n}\) instead of \(\frac{1}{\color{red}\sqrt{n}}\) implies

  • no output fluctuations at initialization
     
  • we can replace the sum by an integral

\(\displaystyle f(w,x) \approx f(\rho,x) = \int d\rho(W,\vec W) \; W \; \phi \! \left(\frac{1}{\sqrt{n}} \sum_{j=1}^n W_j x_j \right)\)

where \(\rho\) is the density of neuron's weights

regime 2: mean field limit

\(\displaystyle f(\rho,x) = \int d\rho(W,\vec W) \; W \; \phi \! \left(\frac{1}{\sqrt{n}} \sum_{j=1}^n W_j x_j \right)\)

\(\rho\) follows a differential equation

[2018 S. Mei et al], [2018 Rotskoff and Vanden-Eijnden], [2018 Chizat and Bach]

In this limit, the internal activation do change

\(\langle Y \rangle=\langle X_i \rangle\) => feature training

regime 2: mean field limit

What is the difference between the two limits

  • which limit describe better finite \(n\) networks?
  • are there corresponding regimes for finite \(n\)

kernel limit              and              mean field limit

\(\frac{1}{\color{red}\sqrt{n}}\)

\(\frac{1}{\color{red}n}\)

frozen

change

internal activations \(z^\ell\)

\(\displaystyle{\color{red}\alpha} f(w,x) = \frac{{\color{red}\alpha}}{\sqrt{n}} \sum_{i=1}^n W_i \ \phi(z_i)\)

[2019 Chizat and Bach]

  • if \(\alpha\) is fixed constant and \(n\to\infty\) then => kernel limit
  • if \({\color{red}\alpha} \sim \frac{1}{\sqrt{n}}\) and \(n\to\infty\) then => mean field limit

\(z_i\)

\(W_i\)

\(f(w,x)\)

we use the scaling factor \(\alpha\) to investigate

the transition between the two regimes

\( \alpha \cdot ( f(w,x) - f(w_0, x) ) \)

linearize the network with \(f - f_0\)

We would like that for any finite \(n\), in the limit \(\alpha \to \infty\), the network behaves linearly

This \(\alpha^2\) is here to converge in a time that does not scale with \(\alpha\) in the limit \(\alpha \to \infty\)

\(\displaystyle \mathcal{L}(w) = \frac{1}{\alpha^2 |\mathcal{D}|} \sum_{(x,y)\in\ \mathcal{D}} \ell\left( {\color{red}\alpha (f(w,x) - f(w_0,x))}, y \right) \)

loss function

\(\displaystyle f(w,x) =\\ \frac{1}{{\color{red}\sqrt{n}}} \sum_{i=1}^n W^3_i \ \phi(z^3_i)\)

\(W_{ij}(t=0) \longleftarrow \mathcal{N}(0, 1)\)

\(\displaystyle z^{\ell+1}_i = {\frac1{\color{red}\sqrt{n}}} \sum_{j=1}^n W^\ell_{ij} \ \phi(z^\ell_j)\)

\(z^3_i\)

\(W^3_i\)

\(z^2_i\)

\(z^1_i\)

\(W^0_{ij}\)

\(W^1_{ij}\)

\(W^2_{ij}\)

\(x_i\)

\(\dot w = -\nabla_w \mathcal{L}(w)\)

Implemented with a dynamical adaptation of the time step \(dt\) such that,

\(10^{-4} < \frac{\|\nabla \mathcal{L}(t_{i+1})- \nabla \mathcal{L}(t_i)\|^2}{\|\nabla \mathcal{L}(t_{i+1})\|\cdot\|\nabla \mathcal{L}(t_i)\|} < 10^{-2}\)

(works well only with full batch and smooth loss)

continuous dynamics

continuous dynamics

\(10^{-4} < \frac{\|\nabla \mathcal{L}(t_{i+1})- \nabla \mathcal{L}(t_i)\|^2}{\|\nabla \mathcal{L}(t_{i+1})\|\cdot\|\nabla \mathcal{L}(t_i)\|} < 10^{-2}\)

momentum dynamics

\(\dot v = -\frac{1}{\tau}(v + \nabla \mathcal{L})\)

\(\dot w = v\)

there is a plateau for large values of \(\alpha\)

MNIST 10k parity, FC L=3, softplus, gradient flow with momentum

lazy regime

MNIST 10k parity, FC L=3, softplus, gradient flow with momentum

ensemble average

the ensemble average converge with \(n \to \infty\)

no overlap

\(\alpha\)

the ensemble average

\(\displaystyle \bar f(x) = \int f(w(w_0),x) \;d\mu(w_0)\)

MNIST 10k parity, FC L=3, softplus, gradient flow with momentum

plot in function of \(\sqrt{n} \alpha\) overlap the lines

overlap !

\(\sqrt{n}\alpha\)

\(\alpha\)

\(\sqrt{n}\alpha\)

\(\alpha\)

\(\frac{\|\Theta - \Theta_0 \|}{\|\Theta_0\|}\)

the kernel evolution displays two regimes

MNIST 10k parity, FC L=3, softplus, gradient flow with momentum

the phase space is split in two by \(\alpha^*\) who decays with \(\sqrt{n}\)

\(n\)

\(\alpha\)

feature training

kernel limit

mean field limit

\(\alpha^* \sim \frac{1}{\sqrt{n}}\)

lazy training

same for other datasets: the trends depends on the dataset

MNIST 10k

reduced to 2 classes

10PCA MNIST 10k

reduced to 2 classes

FC L=3, softplus, gradient flow with momentum

EMNIST 10k

reduced to 2 classes

Fashion MNIST 10k

reduced to 2 classes

FC L=3, softplus, gradient flow with momentum

same for other datasets: the trends depends on the dataset

CIFAR10 10k

reduced to 2 classes

 

CNN SGD ADAM

CNN: the tendency is inverted

\(\sqrt{n}\alpha\)

how does the learning curves depends on \(n\) and \(\alpha\)

MNIST 10k parity, L=3, softplus, gradient flow with momentum

\(\sqrt{n} \alpha\)

overlap !

same time in lazy

lazy

MNIST 10k parity, L=3, softplus, gradient flow with momentum

there is a characteristic time in the learning curves

\(\sqrt{n} \alpha\)

overlap !

characteristic time \(t_1\)

lazy

\(t_1\) characterise the curvature of the network manifold

\(t_1\)

\(t\)

\(f(w)\)

network manifold

(dimension \(N\))

\(d\) input dimension

\(N\) number of parameters

\(m\) size of the trainset

tangent space

(dimension \(N\))

\(f(w_0) + \nabla f(w_0) \cdot dw\)

\(f(w_0) + \sum_\mu c_\mu \Theta(w_0, x_\mu)\)

kernel space

(dimension \(m\))

\(f(w_0)\)

space of functions \(\mathbb{R}^d \to \mathbb{R}\)

(dimension \(\infty\))

\(t_1\) is the time you need to drive to realize the earth is curved

\(v\)

\(R\)

\(t_1 \sim R/v\)

the rate of change of \(W\) determines when we leave the tangent space, aka \(t_1 \sim \sqrt{n}\alpha\)

\(x\)

\(f(w,x)\)

\(W_{ij}\)

\(\displaystyle z^{\ell+1}_i = \frac1{\sqrt{n}} \sum_{j=1}^n W^\ell_{ij} \ \phi(z^\ell_j)\)

\(z_j\)

\(z_i\)

\(W_{ij}\) and \(z_i\) are initialized \(\sim 1\)

\(\dot W_{ij}\) and \(\dot z_i\) at initialization \(\Rightarrow t_1\)

\(\dot W^\ell = \mathcal{O}(\frac{1}{\alpha n})\)

\(\dot W^0 = \mathcal{O}(\frac{1}{\alpha \sqrt{n}})\)

\(\dot z = \mathcal{O}(\frac{1}{\alpha \sqrt{n}})\)

actually the correct scaling

Upper bound: consider weight of last layer

\(\Rightarrow t_1 \sim \sqrt{n} \alpha \)

\(\dot W^L= -\frac{\partial\mathcal{L}}{\partial W^L} = \mathcal{O}\left(\frac{1}{\alpha \sqrt{n}}\right)\)

convergence time is of order 1 in lazy regime 

  • \(\displaystyle \mathcal{L}(w) = \frac{1}{\alpha^2 |\mathcal{D}|} \sum_{(x,y)\in\ \mathcal{D}} \ell\left( {\color{red}\alpha (f(w,x) - f(w_0,x))}, y \right) \)
     
  • \(\dot f(w, x) = \nabla_w f(w,x) \cdot \dot w\)
  • \(= - \nabla_w f(w,x) \cdot \nabla_w \mathcal{L} \)
  • \(\sim - \nabla_w f \cdot (\frac{\partial}{\partial f} \mathcal{L} \; \nabla_w f)\)
  • \(\sim \frac{1}{\alpha} \Theta\)
     
  • "\( \alpha \dot f t = 1 \)" \(\Rightarrow\) "\( t_{lazy} = \|\Theta\|^{-1} \)"

  \(\Longrightarrow\)  \({\color{red}\alpha^* \sim \frac{1}{\sqrt{n}}}\)

when the dynamics stops before \(t_1\) we are in the lazy regime

\(t_1 \sim \sqrt{n} \alpha \)

\(t_{lazy} \sim 1 \)

(time to converge in the lazy regime)

(time to exit the tangent space)

\(\alpha^* \sim 1/\sqrt{n}\)

then for large \(n\)

\(\Rightarrow\)

\( \alpha^* f(w_0, x) \ll 1 \)

\(\Rightarrow\)

\( \alpha^* (f(w,x) - f(w_0, x)) \approx \alpha^* f(w,x) \)

for large \(n\) our conclusions should holds without this trick

linearize the network with \(f - f_0\) was not necessary

arxiv.org/abs/1906.08034

github.com/mariogeiger/feature_lazy

\(n\)

\(\alpha\)

lazy training

stop before \(t_1\)

feature training

go beyond \(t_1\)

kernel limit

mean field limit

\(\alpha^* \sim \frac{1}{\sqrt{n}}\)

time to leave the tangent space \(t_1 \sim \sqrt{n}\alpha\)
 => time for learning features

[SPML] feature-lazy

By Mario Geiger

[SPML] feature-lazy

Presentation of the paper https://arxiv.org/abs/1906.08034

  • 156
Loading comments...

More from Mario Geiger