Scaling Laws and Superposition

Pavan Katta

Update (July 2024): A formal version of this work is now on arXiv.

Abstract

Using results from scaling laws, this short note argues that the following two statements cannot be simultaneously true:

1. Superposition hypothesis where sparse features are linearly represented across a layer is a complete theory of feature representation.
2. Features are universal, meaning two models trained on the same data and achieving equal performance will learn identical features.

Introduction

Scaling laws for Language Models gives us a relation for a model's macroscopic properties such as cross entropy loss $L$, Amount of Data $D$ used and Number of non-embedding parameters $N$ in the model.

L(N,D) = \left[ \left( \frac{N_c}{N} \right)^{\frac{\alpha_N}{\alpha_D}} + \left( \frac{D_c}{D} \right)^{\alpha_D} \right]^{\alpha_D}

where $N_c$, $D_c$, $\alpha_N$, and $\alpha_D$ are constants for a given task such as Language modeling

The scaling laws are not mere empirical observations and can be seen as a predictive laws on limits of language model performance. During training of GPT-4, OpenAI was able to predict the final loss of GPT-4 early in the training process using scaling laws with high accuracy.

An important detail is that the relation is expressed in terms of the number of parameters. It's natural to think of a model's computational capacity in terms of parameters, as they are the fundamental independent variables that the model can tune during learning. The amount of computation that a model performs in FLOPs for each input is also estimated to be $2N$ .

Let's compare this with Interpretability, where the representation of a feature is defined in terms of neurons or groups of neuronsIntentionally left out defining a feature, as there's no universally accepted formal definition. Refer to Neel Nanda's explainer for a good review.. At first glance, it might seem unnecessary to distinguish between computational capacity and feature representational capacity, as parameters are connections between neurons after all. However, we can change the number of neurons in a model while keeping the number of parameters constant. Kaplan et al. found that Transformer performance depends very weakly on the shape parameters $n_{layer}$ (number of layers), $n_{heads}$ (number of attention heads), and $d_{ff}$ (feedforward layer dimension) when we hold the total non-embedding parameter count $N$ fixed . The paper reports that the aspect ratio (the ratio of number of neurons per layer to the number of layers) can vary by more than an order of magnitude, with performance changing by less than 1%.Refer to Sec 3.1 and Figure 5 in the paper

In this paper, we assume the above to be true and consider the number of parameters to be the true limiting factor, and we can achieve similar model performance for a range of aspect ratios. We then apply this as a postulate to the superposition hypothesis, our current best and successful theory of feature representation, and explore the implications.

The superposition hypothesis states that models can pack more features than the number of neurons they have. There will be interference between the features as they can't be represented orthogonally, but when the features are sparse enough, the benefit of representing a feature outweighs the cost of interference. Concretely, given a layer of activations of $m$ neurons, we can decompose it linearly into activations of $n$ features, where $n$ > $m$, as:

activation_{layer} = x_{f_1} W_{f_1} + x_{f_2} W_{f_2} + \cdots + x_{f_n} W_{f_n}

where $activation_{layer}$ and $W_{f_i}$ are vectors of size $m$, and $x_{f_i}$ represents the magnitude of activation of the $i$-th feature. Sparsity means that for a given input, only a small fraction of features are active, which means $x_{f_i}$ is non-zero for only a few values of $i$.

Case study on changing Aspect Ratio

Let's consider two models, Model A and Model B, having the same macroscopic properties. Both have an equal number of non-embedding parameters, are trained on the same dataset, and achieve similar loss according to scaling laws. However, their shape parameters differ. Using the same notation as Kaplan et al., let's denote the number of layers as $n_{layer}$, and number of neurons per layer as $d_{model}$Usually also referred as dimension of the model. Model B has twice the number of neurons per layer compared to A. As the number of parameters is approximated by $d_{model}^2n_{layer}$, The number of non-embedding parameters is equal to $12d_{model}^2n_{layer}$. For simplicity, we can ignore the constant factor Model B must have $\frac{1}{4}$ the number of layers to maintain the same number of parameters as Model A. This means Model B has 8 times the aspect ratio($\frac{d_{model}}{n_{layer}}$) of A which falls under the reported range in Kaplan et al.

The total number of neurons in a model is calculated by multiplying the number of neurons per layer by the number of layers. As a result, Model B has half the total number of neurons compared to Model A.

Now, let's apply the superposition hypothesis, which states that features can be linearly represented in each layer. Since both models achieve equal loss on the same dataset, it's reasonable to assume that they have learned the same features. Let's denote the total number of features learned by both models as $F$.

The above three paragraphs are summarized in the table below:

	Model A	Model B
Total Parameters	$d_{model}^2n_{layer}$	$d_{model}^2n_{layer}$
Neurons per Layer	$d_{model}$	$2d_{model}$
Number of Layers	$n_{layer}$	$\frac{n_{layer}}{4}$
Total Number of Neurons	$d_{model}n_{layer}$	$\frac{d_{model}n_{layer}}{2}$
Total Number of Features Learned	$F$	$F$
Number of Features per Layer	$\frac{F}{n_{layer}}$	$\frac{4F}{n_{layer}}$
Features per Neuron	$\frac{F}{d_{model}n_{layer}}$	$\frac{2F}{d_{model}n_{layer}}$

The average number of features per neuron is calculated by dividing the number of features per layer by the number of neurons per layer. In Model B, this value is twice as high as in Model A, which means that Model B is effectively compressing twice as many features per neuron, in other words, there's a higher degree of superposition. However, superposition comes with a cost of interference between features, and a higher degree of superposition requires more sparsity.

Elhage et al. show that, using lower bounds of compressed sensing , if we want to recover $n$ features compressed in $m$ neurons (where $n$ > $m$), the bound is $m=\Omega(-n(1-S)\log(1-S))$, where $1-S$ is the sparsity of the features. For example, if a feature is non-zero only 1 in 100 times, then $1-S$ equals 0.01. We can define the degree of superposition as $\frac{n}{m} = \frac{1}{(1-S)\log(1-S)}$, which is a function of sparsity, inline with our theoretical understanding.

So Model B, with higher degree of superposition, should have sparser features compared to Model A. But, sparsity of a feature is a property of the data itself, and the same feature can't be sparser in Model B if both models are trained on the same data. This might suggest that they are not the same features, which breaks our initial assumption of two models learning the same features. So either our starting assumption of feature representation through superposition or feature universality needs revision. In the next section, we discuss how we might modify our assumptions.

Discussion

To recap, we started with the postulate that model performance is invariant over a wide range of aspect ratios and arrived at the inconsistency between superposition and feature universality. Though we framed the argument through the lens of superposition, the core issue is that the model's computational capacity is a function of parameters where as model's representational capacity is a function of total neurons.

A useful, though non-rigorous analogy, is to visualise a solid cylinder of radius $d_{model}$ and height $n_{layer}$. The volume (parameters) of the cylinder can be thought of as computational capacity whereas features are represented on the surface (neurons). We can change the aspect ratio of the cylinder while keeping the volume constant by stretching or squashing it. This changes the surface area accordingly. Though this analogy doesn't include sparsity, it captures the essentials of the argument in a simple way.

Coming to solutions, I do not have one that's consistent with scaling laws, superposition hypothesis and feature universality, but will speculate on what a possible one might look like.

Schemes of Compression Alternative to Superposition: A crude and simple way to convert the total number of features into a function of parameters is to add a square term to compressed sensing bounds so it becomes $n = m^2.f(1-S)$ . But this would require a completely new compression scheme compared to superposition. Methods such as Dictionary learning which disentangle features assuming superposition hypothesis have been successful for extracting interpretable features. So it's not ideal to ignore it, representation schemes whose first-order approximation looks like superposition might be more viable.

This isn't to say there's nothing we can improve on in the superposition hypothesis. Although dictionary learning features in Bricken et al. are much more mono-semantic than individual neurons, the lower activation levels in these features still look quite polysemantic.

Cross Layer Superposition: Previously, we used to look for features in a single neuron, now we extended it to a group of neurons in a layer. A natural progression is to look for features localising to neurons across multiple layers. But Model B from the above section, has half the number of neurons as A and the same inconsistencies would arise if features grow linearly on number of neurons. Number of features represented across two or more layers by cross layer superposition should grow superlinearly if Model B were to compensate for fewer neurons and still have the same representational capacity.