1️⃣ Periodicity & Fourier basis

utils.lines(
    W_attn,
    labels=[f"head {hi}" for hi in range(4)],
    xaxis="Input token",
    yaxis="Contribution to attn score",
    title="Contribution to attention score (pre-softmax) for each head",
)

All this periodicity might make us think that the vocabulary basis isn't the most natural one to be operating in. The question is - what is the appropriate basis? ## Fourier Transforms > TL;DR: > > * We can define a Fourier basis of cosine and sine waves with period dividing $p$ (i.e. frequency that's a multiple of $2 \pi / p$). > * We can apply a change of basis of the vocab space into the Fourier basis, and periodic functions are sparse in the Fourier basis. > * For activations that are a function of just one input we use the 1D Fourier transform; for activations that are a function of both inputs we use the 2D Fourier transform. A natural way to understand what's going on is using Fourier transforms. This represents any function as a sum of sine and cosine waves. Everything here is discrete, which means our functions are just $p$ or $p^2$ dimensional vectors, and the Fourier transform is just a change of basis. All functions have *some* representation in the Fourier basis, but "this function looks periodic" can be operationalised as "this function is sparse in the Fourier basis". Note that we are applying a change of basis to $\mathbb{R}^p$, corresponding to the vocabulary space of one-hot encoded input vectors. (We are abusing notation by pretending that `=` is not in our vocabulary, so $d_\mathrm{vocab} = p$, allowing us to take Fourier transforms over the input space.) ### 1D Fourier Basis We define the 1D Fourier basis as a list of sine and cosine waves. **We begin with the constant wave (and then add cosine and sine waves of different frequencies.** The waves need to have period dividing $p$, so they have frequencies that are integer multiples of $\omega_1 = 2 \pi / p $. We'll use the shorthand notation that $\vec{\textbf{x}} = (0, 1, ..., (p-1))$, and so $\cos (\omega_k \vec{\textbf{x}})$ actually refers to the following vector in $\mathbb{R}^p$: (after being scaled to unit norm), where $\omega_k = 2 \pi k / p$. We will also denote $F$ as the $p \times p$ matrix where each **row** is one such wave: Again, we've omitted the normalization constant, but you should assume each row is a basis vector with norm 1. This means the constant term $\vec{\textbf{1}}$ is scaled by $\sqrt{\frac{1}{p}}$, and the rest by $\sqrt{\frac{2}{p}}$. Note also that the waves (esp. at high frequencies) look jagged, not smooth. This is because we discretise the inputs to just be integers, rather than all reals. ### Exercise - create the 1D Fourier basis Complete the function below. Don't worry about computational efficiency; using a for loop is fine.

Solution

def make_fourier_basis(p: int) -> tuple[Tensor, list[str]]:
    """
    Returns a pair fourier_basis, fourier_basis_names, where fourier_basis is
    a (p, p) tensor whose rows are Fourier components and fourier_basis_names
    is a list of length p containing the names of the Fourier components (e.g.
    ["const", "cos 1", "sin 1", ...]). You may assume that p is odd.
    """
    # Define a grid for the Fourier basis vecs (we'll normalize them all at the end)
    # Note, the first vector is just the constant wave
    fourier_basis = t.ones(p, p)
    fourier_basis_names = ["Const"]
    for i in range(1, p // 2 + 1):
        # Define each of the cos and sin terms
        fourier_basis[2  i - 1] = t.cos(2  t.pi  t.arange(p)  i / p)
        fourier_basis[2  i] = t.sin(2  t.pi  t.arange(p)  i / p)
        fourier_basis_names.extend([f"cos {i}", f"sin {i}"])
    # Normalize vectors, and return them
    fourier_basis /= fourier_basis.norm(dim=1, keepdim=True)
    return fourier_basis.to(device), fourier_basis_names

Once you've done this (and passed the tests), you can run the cell below to visualise your Fourier components.

Click to see the expected output

*Note - from this point onwards, the `fourier_basis` and `fourier_basis_names` variables are global, so you'll be using them in other functions. We won't be changing the value of `p`; this is also global.* Now, you can prove the fourier basis is orthonormal by showing that the inner product of any two vectors is one if they are the same vector, and zero otherwise. Run the following cell to see for yourself:

Click to see the expected output

Now that we've shown the Fourier transform is indeed an orthonormal basis, we can write any $p$-dimensional vector in terms of this basis. The **1D Fourier transform** is just the transformation taking the components of a vector in the standard basis to its components in the Fourier basis (in other words we project the vector along each of the Fourier basis vectors). ### Exercise - 1D Fourier transform You should now write a function to compute the Fourier transform of a vector. Remember that the **rows** of `fourier_basis` are the Fourier basis vectors.

Solution

def fft1d(x: Tensor) -> Tensor:
    """
    Returns the 1D Fourier transform of x, which can be a vector or a batch of vectors.
    x.shape = (..., p)
    """
    return x @ fourier_basis.T

Note - if x was a vector, then returning fourier_basis @ x would be perfectly fine. But if x is a batch of vectors, then we want to make sure the multiplication happens along the last dimension of x.

We can demonstrate this transformation on an example function which looks periodic. The key intuition is that **'function looks periodic in the original basis'** implies **'function is sparse in the Fourier basis'**. Note that functions over the integers $[0, p-1]$ are equivalent to vectors in $\mathbb{R}^p$, since we can associate any such function $f$ with the vector:

Click to see the expected output

You should observe a jagged but approximately periodic function in the first plot, and a very sparse function in the second plot (with only three non-zero coefficients). ### 2D Fourier Basis **All of the above ideas can be naturally extended to a 2D Fourier basis on $\mathbb{R}^{p \times p}$, ie $p \times p$ images. Each term in the 2D Fourier basis is the outer product $v w^T$ of two terms $v, w$ in the 1D Fourier basis.** Thus, our 2D Fourier basis contains (up to a scaling factor): * a constant term $\vec{\textbf{1}}$, * linear terms of the form $\,\cos(\omega_k \vec{\textbf{x}}),\,\sin(\omega_k \vec{\textbf{x}}),\,\cos(\omega_k \vec{\textbf{y}})$, and $\sin(\omega_k \vec{\textbf{y}})$, * and quadratic terms of the form: Although we can think of these as vectors of length $p^2$, it makes much more sense to think of them as matrices of size $(p, p)$. > Notation - $\cos(\omega_i \vec{\textbf{x}})\cos(\omega_j \vec{\textbf{y}})$ should be understood as the $(p, p)$-size matrix constructed from the outer product of 1D vectors $\cos (\omega_i \vec{\textbf{x}})$ and $\cos (\omega_j \vec{\textbf{y}})$. In other words, the $(x, y)$-th element of this matrix is $\cos(\omega_i x) \cos(\omega_j y)$. ### Exercise - create the 2D Fourier basis Complete the following function. Note that (unlike the function we wrote for the 1D Fourier basis) this only returns a single basis term, rather than the entire basis.

Solution

def fourier_2d_basis_term(i: int, j: int) -> Float[Tensor, "p p"]:
    """
    Returns the 2D Fourier basis term corresponding to the outer product of the i-th component of
    the 1D Fourier basis in the x direction and the j-th component of the 1D Fourier basis in
    the y direction.
    Returns a 2D tensor of length (p, p).
    """
    return fourier_basis[i][:, None]  fourier_basis[j][None, :]

Note, indexing with None is one of many ways to write this function. A few others are:

torch.outer Using einsum: torch.einsum('i,j->ij', ...) Using the unsqueeze method to add dummy dimensions to the vectors before multiplying them together.

Once you've defined this function, you can visualize the 2D Fourier basis by running the following code. Verify that they do indeed look periodic.

Click to see the expected output

What benefit do we get from thinking about $(p, p)$ images? Well, the batch dimension of all our data is of size $p^2$, since we're dealing with every possible value of inputs `x` and `y`. So we might think of reshaping this batch dimension to $(p, p)$, then applying a 2D Fourier transform to it. Let's implement this transform now! ### Exercise - Implementing the 2D Fourier Transform This exercise should be pretty familiar, since you've already done this in 1D.

Solution

def fft2d(tensor: Tensor) -> Tensor:
    """
    Retuns the components of tensor in the 2D Fourier basis.
    Asumes that the input has shape (p, p, ...), where the
    last dimensions (if present) are the batch dims.
    Output has the same shape as the input.
    """
    # fourier_basis[i] is the i-th basis vector, which we want to multiply along
    return einops.einsum(tensor, fourier_basis, fourier_basis, "px py ..., i px, j py -> i j ...")

While working with the 1D Fourier transform, we defined simple periodic functions which were linear combinations of the Fourier basis vectors, then showed that they were sparse when we expressed them in terms of the Fourier basis. That's exactly what we'll do here, but with functions of 2 inputs rather than 1. Below is some code to plot a simple 2D periodic function (which is a linear combination of 2D Fourier basis terms). Note that we call our matrix `example_fn`, because we're thinking of it as a function of its two inputs (in the x and y directions).

Click to see the expected output

Code to show this function is sparse in the 2D Fourier basis (you'll have to zoom in to see the non-zero coefficients):

Click to see the expected output

You can run this code, and check that the non-zero components exactly match the basis terms we used to construct the function. ## Analysing our model with Fourier Transforms So far, we've made two observations: * Many of our model's activations appear periodic * Periodic functions appear sparse in the Fourier basis So let's take the obvious next step, and apply a 2D Fourier transformation to our activations! Remember that the batch dimension of our activations is $p^2$, which can be rearranged into $(p, p)$, with these two dimensions representing the `x` and `y` inputs to our modular arithmetic equation. These are the dimensions over which we'll take our Fourier transform. ### Plotting activations in the Fourier basis Recall our previous code, to plot the heatmap for the attention scores token 2 pays to token 0, for each head: In that plot, the x and y axes represented the different values of inputs `x` and `y` in the modular arithmetic equation. The code below takes the 2D Fourier transform of the attention matrix, and plots the heatmap for the attention scores token 2 pays to token 0, for each head, in the Fourier basis:

Click to see the expected output

You should find that the result is extremely sparse - there will only be a few cells (mostly on the zeroth rows or columns, i.e. corresponding to the constant or linear terms) which aren't zero. This suggests that we're on the right track using the 2D Fourier basis! Now, we'll do ths same for the neuron activations. Recall our previous code: We'll do the exact same here, and plot the activations in the Fourier basis:

Click to see the expected output

### Exercise - spot patterns in the activations Increase `top_k` from 3 to a larger number, and look at different neurons. What do you notice about the patterns in the activations?

Answer (what you should see)

Again, you should see sparsity, although this time some quadratic terms should be visible too.

Beyond this, there are 2 distinct patterns worth commenting on:

Each neuron has the same pattern of non-zero terms: for some value of $k$, the non-zero terms are the constant term plus all four linear and four quadratic terms involving just frequencies $k$ (i.e. terms like $\cos(\omega_k \vec{\textbf{x}})$, $\sin(\omega_k \vec{\textbf{y}})$, $\cos(\omega_k \vec{\textbf{x}})\sin(\omega_k \vec{\textbf{y}})$, etc). There are only a handful of different values of $k$ across all the neurons, so many of them end up having very similar-looking activation patterns.

#### Aside: Change of basis on the batch dimension A change of basis on the batch dimension is a pretty weird thing to do, and it's worth thinking carefully about what happens here (note that this is a significant deviation to the prior Transformer Circuits work, and only really makes sense here because this is such a toy problem that we can enter the entire universe as one batch). There are *four* operations that are not linear with respect to the batch dimension. As above, the attention softmax, ReLU and final softmax. But also the elementwise multiplication with the attention pattern. In particular, ReLU becomes super weird - it goes from an elementwise operation to the operation 'rotate by the inverse of the Fourier basis, apply ReLU elementwise in the *new* basis, rotate back to the Fourier basis'. ### Plotting effective weights in the Fourier basis As well as plotting our activations, we can also look at the weight matrices directly. *Note - this section isn't essential to understanding the rest of the notebook, so feel free to skip it if you're short on time.* We'll now adjust our previous code which plotted the `W_neur` matrix in the standard basis, so it's visualized in the Fourier basis instead. In the plot below, here, each line shows the activations for some neuron as a function of the input token `x` (if the `=` token at position 2 only paid attention to a token `x`), with the neuron index determined by the slider value. If this seems a bit confusing, you can use the dropdown below to remind yourself of the functional form of this transformer, and the role of $W_{neur}$.

Functional Form

$$ f(t)=\operatorname{ReLU}\Bigg(\sum_h{\bigg(\alpha^h t_0\;+\;\left(1\;-\;\alpha^h\right) t_1 \bigg)^T}{W_{neur}^h}\Bigg) \; W_{logit} $$

From this, we can see clearly the role of $W_{neur}^h$. It is a matrix of shape $(d_{vocab}, d_{mlp})$, and its rows are the vectors we take a weighted average of to get our MLP activations (pre-ReLU).

The code below makes the same plot, but while the previous one was in the standard basis (with the x-axis representing the input token), in this plot the x-axis is the component of the input token in each Fourier basis direction. Note that we've provided you with the helper function `fft1d_given_dim`, which performs the 1D Fourier transform over a given dimension. This is necessary for `W_neur`, since it has shape `(n_heads, d_vocab, d_mlp)`, and we want to transform over the `d_vocab` dimension.

Click to see the expected output

Note that each line plot generally has $\sin k$ and $\cos k$ terms non-zero, rather than having one but not the other. Lastly, we'll do the same with `W_attn`:

Click to see the expected output

You may have noticed that the handful of non-zero frequencies in both these last two line charts exactly match the important frequencies we read off the attention patterns! ## Recap of section Let's review what we've learned in this section. We found that: > - The simple architecture of our 1-layer model heavily constrains the functional form of any learned solutions. > - In particular, we can define a handful of matrices which fully describe the model's behaviour (after making some simplifying assumptions). > - Many of our model's internal activations appear periodic in the inputs `x`, `y` (e.g. the attention patterns and neuron activations). > - The natural way to represent a periodic function is in the Fourier basis. Periodic functions appear sparse in this basis. > - This suggests our model might only be using a handful of frequencies (i.e. projecting the inputs onto a few different Fourier basis vectors), and discarding the rest. > - We confirmed this hypothesis by looking at: > - The model's activations (i.e. attention patterns and neuron activations) > - The model's effective weight matrices (i.e. $W_{attn}$ and $W_{neur}$) > - Both these observations confirmed that we have sparsity in the Fourier basis. Furthermore, the same small handful of frequencies seemed to be appearing in all cases.