Grokking, Visualized
Chapters
  1. The whole picture, first
  2. Counting on a clock
  3. The task as a dataset
  4. Vectors and matrices
  5. From scores to probabilities
  6. The model, step by step
  7. Simplifying the model
  8. How the model learns
  9. Grokking
  10. Angles and waves
  11. The Fourier basis
  12. Measuring fit and sparsity
  13. The Fourier multiplication algorithm
  14. Evidence 1: everything is periodic
  15. Evidence 2: neurons and heads
  16. Evidence 3: reading off the answer
  17. Evidence 4: interventions
  18. Progress measures and the three phases
  19. Watch it form
  20. Gradients and weight decay
  21. Is it robust? Seeds and controls
  22. The whole picture, again

Chapter 11 of 0–21

Measuring fit and sparsity

Variance explained, least squares, SVD/PCA and the Gini coefficient.

Sources for this chapter

To claim “the network computes X” we must say how well X describes it. This chapter collects the yardsticks the paper and ARENA use.

Mean, variance and how much an approximation explains

The mean of a list is its average. The variance is the average squared distance from the mean — how spread out the numbers are.

Worked example

For y = (1, 2, 3, 6): mean = 12/4 = 3; variance = ((1−3)² + (2−3)² + (3−3)² + (6−3)²)/4 = (4 + 1 + 0 + 9)/4 = 3.5.

If we approximate data y by some prediction ŷ, the residual is what is left, y − ŷ. The fraction of variance explained compares the leftover with the original spread:

\text{FVE} = 1 - \frac{\sum (y - \hat y)^2}{\sum (y - \bar y)^2}.

FVE = 1 means the approximation is perfect; FVE = 0 means it is no better than always predicting the mean.

Least squares

Least squares chooses the free coefficients of an approximation to make the sum of squared residuals as small as possible. The paper fits the logits this way (chapter 15), ARENA fits neuron activations this way (chapter 14).

Worked example

Fitting a line ŷ = α + βx to y = (1, 2, 3, 6) at x = (0, 1, 2, 3). Least squares gives α = 0.6, β = 1.6, predictions (0.6, 2.2, 3.8, 5.4), residuals (0.4, −0.2, −0.8, 0.6). Sum of squared residuals 1.2, total sum of squares around the mean 14, so FVE = 1 − 1.2/14 ≈ 0.914: a line explains 91 % of the variance.

How many directions does a matrix use? SVD and PCA

The singular value decomposition writes any matrix as rotation × stretching × rotation, M = U\,\mathrm{diag}(s_1, s_2, \ldots)\,V^\top [29]. The stretch factors s_1 \ge s_2 \ge \ldots \ge 0 are the singular values; the number of non-zero ones is the rank (chapter 3). Many tiny singular values mean the matrix effectively uses only a few directions.

Worked example

\begin{pmatrix}3&0\\0&1\end{pmatrix} has singular values 3 and 1. \begin{pmatrix}1&2\\2&4\end{pmatrix} has singular values 5 and 0 — its second row is twice its first, so it only ever outputs multiples of one direction (rank 1).

Principal component analysis applies the same idea to a cloud of points: it finds the directions along which the points vary most, so we can draw a 128-dimensional cloud faithfully in 2D [30]. ARENA uses the SVD of the embedding matrix to see how many directions it uses during training [3]:

Singular values of the embedding

The 113 singular values of W_E (largest first) at the selected step.

Loading…

How to read it 
<div class="fig-read"><p>A <a class="term" data-term="svd" href="../glossary.html#svd">singular value decomposition (SVD)</a> shows how many independent directions the embedding really uses.</p>

What to expect: Starts as a smooth random-matrix profile; ends with two large singular values per key frequency standing out, plus (on the paper's run) two smaller ones from frequency 31, which only the embedding uses.

Learn step by step → · Source: [3] section 3, “Examine the SVD of W_E” · Computed by checkpoint_metrics.py (code)

The same singular values at every checkpoint of training (largest at the bottom): by the end, ten of them stand well above the rest (3.5–4.2 on the paper's run at step 40,001), followed by two smaller ones (2.04 and 2.02); all others are at most 0.41.

The embedding's singular values over training

All 113 singular values of W_E at every checkpoint (largest at the bottom).

Loading…

How to read it 
<p class="fig-expect"><strong>What to expect:</strong> Late in training two large singular values per key frequency (cos and sin; ten on the paper's run) stand above the rest, plus, on the paper's run, two smaller ones from frequency 31, which only the embedding uses.</p>
<p class="fig-meta">
  <a href="../chapters/11-measuring-fit.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> section 3, “Examine the SVD of W_E”
  · Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Embedding geometry (PCA)

The 113 number embeddings in their two main principal component analysis (PCA) directions.

Loading…

How to read it 
<p class="fig-meta">
  <a href="../chapters/11-measuring-fit.html">Learn step by step →</a>
  
  · Computed by <code>geometry.pca</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

How concentrated is a list? The Gini coefficient

To measure sparsity — “are only a few Fourier components large?” — the paper uses the Gini coefficient of the list of component sizes [1], citing Hurley and Rickard, who show it satisfies all six intuitive requirements of a sparsity measure they list [15]. With the N values sorted from smallest to largest, c_{(1)} \le \ldots \le c_{(N)}:

G(c) = 1 - 2\sum_{k=1}^{N}\frac{c_{(k)}}{\lVert c\rVert_1}\cdot\frac{N - k + \tfrac12}{N}, \qquad \lVert c\rVert_1 = \sum_k |c_k|.

The paper does not print the formula; we use Hurley and Rickard's definition, which gives final values consistent with the paper's Fig. 7 (chapter 17).

Worked example

Gini of four values. Equal values (1, 1, 1, 1): G = 0. All mass in one (0, 0, 0, 4): G = 1 − 1/4 = 0.75 (the maximum for N = 4 is 1 − 1/N). Steadily increasing (1, 2, 3, 4): G = 0.25. Scaling all values by the same factor never changes G.

Press ▶ next to “units moved onto c1” to move the eight units onto one number, one at a time: G climbs from 0 to its maximum 1 − 1/8 = 0.875.

Loading…