These are my (still under construction) notes on some of the key results in Kumar et al.’s “Grokking as the Transition from Lazy to Rich Dynamics.”
What is Grokking
Grokking was the observation made by Power et al. who trained small transformers on modular arithmetic and related tasks, of a large gap between training accuracy and test accuracy reaching 100%:
Although training accuracy reaches 100% early, it takes 3 orders of magnitude more iterations before the network “grok”s the problem, and test accuracy starts to rise. Why does this occur?
Lazy Learning
The solutions suggested by Kumar and colleagues is that grokking reflects the transition from a “lazy” learning regime, in which the network tries to fit the training data without moving too far from its initialization, to a “rich” regime, in which the network finally moves outs of the neighbourhood of its initialization, and learns task-relevant features, which generalize to the test set.
The basic idea is that we can always view learning by gradient descent locally as a regression problem, whose kernel determines directions that are easy (large eigenvalue) and hard (small eigenvalue) to learn. The lazy regime corresponds to learning along the easy directions around the initial point.
Let’s get concrete. We are given a dataset consisting of $D$-dimensional inputs $\xx_i$ and corresponding real-valued targets, $y(\xx_i)$: $$\text{Data} = \{(\xx_1, y(\xx_1)), (\xx_2, y(\xx_2)), \dots, \}.$$
We have a neural network, with weights $\ww$, that will map each $\xx_i$ to an output $f(\xx_i, \ww)$. We want to adjust its weights so that this output is as close as possible to the corresponding target values in the training set.
We can measure how well a given set of weights $\ww$ solve the problem by adding up the squared error on each of the training examples. This gives us the loss $$ L(\ww) = {1\over 2N} \sum_{i=1}^N (y(\xx_i) – f(\ww, \xx_i))^2.$$
We do this using gradient descent, so $$ \dot \ww \propto -\nabla_\ww L = {1 \over N} \sum_i (y(\xx_i) – f(\ww, \xx_i)) \nabla_\ww f(\ww, \xx_i).$$
The Simplest Network: Linear Regression
Before going further with our nonlinear neural network, we can consider a much simpler network: one with a single, linear layer, where the outputs are a simple linear combination of the inputs $$f(\ww, \xx_i) = \xx_i^T \ww.$$
Our loss function is then simply \begin{equation} L(\ww) = {1\over 2N} \sum_{i=1}^N (y_i – \xx_i^T \ww)^2. \label{eq:linloss} \end{equation}
We can express this more compactly if we first stack our inputs and target outputs into a matrix and vector, respectively,$$ \XX \triangleq [\xx_1, \xx_2, \dots, \xx_N]^T, \quad \yy \triangleq [y_1, y_2, \dots, y_N]^T.$$
Our loss is then $$ L(\ww) = {1 \over 2 N} (\yy – \XX \ww)^T (\yy – \XX \ww) = {1 \over 2N} \| \yy – \XX \ww\|_2^2.$$
Gradient descent on this loss gives the weight dynamics $$ \dot \ww = -\nabla_\ww L = {1 \over N} \XX^T (\yy – \XX \ww).$$
Output dynamics
The weights are a high-dimensional object. What we’re ultimately interested in is how our output changes, so let’s study that instead.
Call the output of our network $\aa$, that is $$ a_i \triangleq f(\ww, \xx_i) = \xx_i^T \ww, \quad \aa \triangleq \XX \ww.$$ Then its dynamics are $$\dot \aa = \XX \dot \ww = {1 \over N} \XX \XX^T (\yy – \XX \ww) = {1\over N} \XX \XX^T (\yy – \aa).$$
This expression says two important things.
First, the $(\yy – \aa)$ term says that our outputs $\aa$ are moving towards the targets $\yy$. In fact, if we didn’t have the $\XX \XX^T$ term (more on that below), the output dynamics would be $$ \dot \aa ={1 \over N} (\yy – \aa).$$ We could then consider output output individually, $$ \dot a_i = {1 \over N} (y_i – a_i).$$ The solution would then be exponential decay of all the outputs, at the same rate, towards their targets: $$ a_i(t) = a_i(0) e^{-t/N} + y_i (1 – e^{-t/N}).$$
The Kernel
Our output dynamics are made a little more interesting by the presence of $\XX \XX^T$. This is the matrix, stores the similarity of each of the inputs with eachother, $$ \XX \XX^T = \begin{bmatrix} \xx_1^T \\ \xx_2^T \\\vdots\\ \xx_N^T \end{bmatrix} \begin{bmatrix} \xx_1, \xx_2, \dots, \xx_N \end{bmatrix}, \quad [\XX \XX^T ]_{ij} = \xx_i^T \xx_j.$$ We can see this as in instance of a more general kernel that measures the similarity of any two inputs, not just those in our dataset, the same way: $$ k(\xx, \xx’) = \xx^T \xx’.$$ We’ll therefore call $\XX \XX^T$ the kernel matrix, and give it the symbol $$\KK \triangleq {1 \over N} \XX \XX^T.$$ Our dynamics are therefore $$ \dot \aa = \KK (\yy – \aa).$$
Spectral Bias
How does the kernel affect the output dynamics? We can see that by applying an eigendecomposition to our kernel: $$ \KK = \UU \SS \UU^T,$$ where the columns of $\UU$ are the eigenvectors of $\KK$, an the diagonal matrix $\SS$ has the eigenvalues. Then $$ \dot \aa = \UU \SS \UU^T (\yy – \aa).$$
If we switch coordinates to the eigenvectors of the kernel, $$\tilde \aa \triangleq \UU^T \aa, \quad \tilde \yy \triangleq \UU^T \yy,$$ the dynamics again decouple, $$ \dot{ \wt{\aa}} = \SS (\wt{\yy} -\wt{\aa}),$$ or elementwise, \begin{equation} \dot{\wt{a_i}} = S_i (\wt{y_i} – \wt{a_i}).\label{eq:daidt}\end{equation}
The solutions are again exponential decay, $$ \wt{a_i}(t) = \wt{a_i}(0) e^{- S_i t} + \wt{y_i} (1 – e^{-S_i t}).$$ But notice the important difference: the decays are at different rates, determined by the kernel eigenvalues $S_i$, with earlier elements decaying faster than later ones. Since the output is constructed by combining these terms, \begin{align*} \aa(t) &= \sum_{i=1}^N \uu_i \wt{a_i}(t)\\ &= \sum_{i=1}^N \uu_i (\wt{a_i}(0) e^{- S_i t} + \wt{y_i}(1 – e^{-S_i t})) \\ &= \sum_{i=1}^N \uu_i (\uu_i^T \aa (0) e^{- S_i t} + \uu_i^T \yy (1 – e^{-S_i t})) \end{align*} we see that the components of the target $\yy$ that align with principal eigenvectors are learned faster than those that align with the rest. In other words, learning is biased towards target functions that align with the kernel.
In other words, equation \eqref{eq:daidt} tells us two things
- The predictions along the principal eigenvectors of the kernel, those for which $S_i$ is large, are learned faster than those along other directions. Lazy learning cooresponds to learning along these directions.
- We want those directions to lineup with $y$. It doesn’t matter if we have fast decay along some directions, if our targets don’t lie along those directions.
The second point motivates using centered kernel alignment, CKA, as a measure of how well the kernel can capture the targets: $$\text{CKA} = { \yy^T \KK \yy \over \|\KK\| \|\yy\|^2 } = {\sum_{i=1}^P S_i \tilde{y}_i^2 \over \sqrt{\sum_i S_i^2} \left( \sum_i \tilde y_i^2 \right)}. $$ The numerator is a projection of $\yy$ along the eigenvectors of $\KK$, weighted by the eigenvalues, literally measuring alignment. The denominator normalizes the value to lie between 0 and 1.
Back to the nonlinear case
How does the linear case relate to our nonlinear neural network? The main idea is to consider infinitesimal weight changes around the initial weights $\ww_0$. For such tiny weight changes, the behaviour of our network is approximately linear, and we can bring in the linear regression insights we developed above.
Since we’re assume the weight update from $\ww_0$ to $\ww$ is infinitesimal, we can write \begin{align*} f(\xx, \ww) &= f(\xx, \ww_0) + \nabla_\ww f(\xx,\ww_0) (\ww – \ww_0)\\ &= f(\xx, \ww_0) – \nabla_\ww f(\xx, \ww_0) \ww_0 + \nabla_\ww f(\xx,\ww_0) \ww. \end{align*}
The first two terms don’t depend on $\ww$, so we can think of them as a constant offset (relative to $\ww$): $$ b(\xx, \ww_0) \triangleq f(\xx, \ww_0) – \nabla_\ww f(\xx, \ww_0) \ww_0.$$
The final term is linear in $\ww$, with the coefficients being determined by the gradient of $f$. That defines a feature map, mapping each input $\xx$ into the gradient, $$ \xx \mapsto \nabla_\ww f(\xx, \ww_0) \triangleq \phi(\xx).$$
So infinitesimal changes around an initial value are affine (linear + constant offset) in $\ww$: $$ f(\xx, \ww) = b(\xx, \ww_0) + \phi(\xx)^T \ww.$$
We can then write our loss function for these infinitesimal changes as \begin{align*} L(\ww) &= {1 \over 2 N} \sum_{i=1}^N (y(\xx_i) – f(\xx_i, \ww))^2\\ &= {1 \over 2N } \sum_{i=1}^N (y(\xx_i) – b(\xx, \ww_0) – \phi(\xx_i)^T \ww)^2. \end{align*}
To avoid introducing new symbols, we can absorb the constants $\bb(\xx_i, \ww_0)$ into the corresponding targets $y(\xx_i)$, so $$ y(\xx_i) \leftarrow y(\xx_i) – \bb(\xx_i, \ww_0).$$
Our loss then becomes \begin{align*} L(\ww) &= {1 \over 2 N} \sum_{i=1}^N (y(\xx_i) – \phi(\xx_i)^T \ww)^2. \end{align*}
This is just like the linear regression loss in equation \eqref{eq:linloss}, just with the raw inputs $\xx_i$ replaced by the features $\phi(\xx_i)$. So our full, nonlinear, problem, for infinitesimal weight changes, is linear regression.
Everything carries forward just as in the linear regression case, just with the raw inputs replaced by the corresponding features:
The weight dynamics become $$ \dot{\ww} = {1 \over N} \sum_i (y(\xx_i) – \phi(\xx_i)^T \ww) \phi(\xx_i).$$
If we stack the features into a matrix as $$ \bPhi \triangleq [\phi(\xx_1), \phi(\xx_2),\dots,\phi(\xx_N)],$$ then we can write the weight update as $$ \dot \ww = \bPhi (\yy – \bPhi^T \ww).$$
We again shift to studying not the weights, but the (linear) predictions, $\aa = \bPhi^T \ww$. These change as $$ \dot \aa = \bPhi^T \bPhi \yy – \bPhi^T \bPhi \aa.$$ Just as before, we can think of these dynamics in terms of a kernel,$$ \KK \triangleq \bPhi^T \bPhi, \quad K(\xx, \xx’) = \phi(\xx)^T \phi(\xx).$$ This is the kernel of our regression problem, the neural tangent kernel.
In terms of this kernel, the prediction dynamics are just like the linear case, $$ \dot \aa = \KK \yy – \KK \aa.$$
We can eigendecompose the kernel as $$ \KK = \UU \SS^2 \UU^T.$$ If we switch coordinates to these eigenvectors, so $\tilde{\aa} = \UU^T \aa,$ then $$ \dot{\tilde{a_i}} = S_i^2 \tilde{y_i} – S_i^2 \tilde{a}_i,$$ and we see that the dynamics have spectral bias.
Polynomial Regression
To demonstrate how lazy vs rich learning explains grokking, the authors study a simple polynomial regression problem. Instead of using a transformer, they study a two layer network with fixed output weights. So their network output is $$ f(\ww, \xx) = {\alpha \over N} \sum_{i=1}^N g(\ww_i^T \xx),$$ where $\xx \in \RR^D$, and $g$ is a polynomial activation function.
Previous work by others has shown that when the activation function is an order $k$ polynomial, the functions can be learned using as many examples $P$ as the dimensionality of the input $D$. The authors’ previous work has shown that lazy learning requires $D^k$ examples, where $k$ is the oder of the nonlinearity. This gap suggested that grokking was possible.
To determine the lazy and rich directions, the authors comptued the NTK for the regression problem, using the quadratic activation function $$g(h) = h + {\veps h^2 \over 2}.$$ To compute the NTK we need the partial derivatives of the output $f$ relative to the weights. These split in $N$ independent blocks, one for each hidden unit. The component from block $i$ is $$ {\partial f \over \partial \ww_i} = {\alpha \over N} g'(\ww_i^T \xx) \xx = {\alpha \over N} (1 + \veps \ww_i^T\xx) \xx.$$
The kernel measures how these features overlap for different inputs, \begin{align*} K(\xx, \xx’) &= \phi(\xx)^T \phi(\xx’)\\ &= \sum_i {\partial f (\xx, \ww_0) \over \partial \ww_i}^T {\partial f (\xx’, \ww_0) \over \partial \ww_i} \\ &= {\alpha^2 \over N^2} \sum_i (1 + \veps \ww_i^T \xx)(1 + \veps \ww_i^T \xx’) \xx^T \xx’ \\ &={\alpha^2 \over N^2} \left[ N \xx^T \xx’ + \veps \sum \ww_i^T (\xx + \xx’) \xx^T \xx’ + \veps^2 (\xx^T \xx’) \xx^T \sum_i \ww_i \ww_i^T \xx’ \right]\\ &= {\alpha^2 \over N} \left[ \xx^T \xx’ + \veps \overline \ww (\xx + \xx’) \xx^T \xx’ + \veps^2 \xx^T \MM \xx (\xx^T \xx’)\right],\end{align*} where $$ \overline \ww \triangleq {1 \over N} \sum \ww_i, \quad \MM \triangleq {1 \over N} \sum \ww_i \ww_i^T.$$
At their initialization $\overline \ww = \bzero$ and $\MM = \II$. So, ignoring the constant of proportionality, \begin{equation} K(\xx, \xx’) = \xx^T \xx’ + \veps^2 (\xx^T \xx’)^2.\label{eq:ker}\end{equation}
Now for a given dataset, $K$ will have a specific value and will give us some specific eigenvectors that will be easy or hard to learn. We’ll get more insight if we consider these eigenvectors in the limit of infinite data.
Mercer Eigenvalue Problem
To do that, let’s write the eigenvector equation in the finite data case: $$ {1 \over P} \sum_\nu K(\xx_\mu, \xx_\nu) r^{(a)}_\nu = \lambda^{(a)} r^{(a)}_\mu.$$
Now, let’s think of $r^a_\nu$ as some function $\psi$ evaluated at $\xx_\nu$. So, $$ {1 \over P} \sum_\nu K(\xx_\mu, \xx_\nu) \psi(\xx_\nu) = \lambda^{(a)} \psi(\xx_\mu).$$
We can then think of this as an empirical estimate for points $\xx_\nu$ drawn from some distribution $P(\xx_\mu).$ Taking the limit of infinite data points, our sum above tends to the integral $$ \int K(\xx’, \xx)\, \psi(\xx)\,P(\xx) d\xx = \lambda^{(a)} \psi(\xx’).$$
This is the Mercer Eigenvalue Problem. The advantage of thinking in these infinite data terms is that we can focus on the kernel itself, rather than the sampling noise involved in any particular finite sample.
The authors now look for the eigenfunctions and corresponding eigenvalues of the kernel in equation \eqref{eq:ker}.
The Kernel Eigenfunctions
The first function they consider is $\psi(\xx) = x_i$. Our integral above becomes \begin{align*} \langle K(\xx, \xx’) x_i \rangle &= \langle (\xx^T \xx’)x_i + \veps^2 (\xx^T \xx’)^2 \rangle \\ &= \left\langle \sum_j (x_j x_j’)x_i + \veps^2 \sum_j \sum_k x_j x_k x_{j’} x_{k’} x_i \right\rangle. \end{align*}
Since $\xx \sim \mathcal{N}(0, 1/D),$ the elements of $\xx$ are independent, and we only need to consider even powers. So, the first term becomes $$ \left \langle \sum_j (x_j x_j’) x_i \right\rangle = \langle x_i^2 x_i’ \rangle = \langle x_i^2 \rangle x_i’= {1 \over D} x_i’.$$
In the second term, we only ever get even powers of one element accompanied by an odd power of another, so that term averages to zero. We therefore arrive at $$ \langle K(\xx, \xx’) x_i \rangle = {1 \over D} x_i, $$ showing that our kernel has $D$ eigenfunctions $x_i$ with eigenvalue $D^{-1}$.
To be continued…
Leave a Reply