On Implications of Scaling Laws on Feature Superposition
Abstract
Using results from scaling laws, this theoretical 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.
1 Introduction
Scaling laws for language models give us a relation for a modelβs macroscopic properties such as cross entropy loss , Amount of Data used and Number of non-embedding parameters in the model (Kaplan etΒ al., 2020).
(1) |
where , , , and are constants for a given task such as language modeling.
The scaling laws are not mere empirical observations and can be seen as 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 (OpenAI, 2023).
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 (Kaplan etΒ al., 2020).
Letβs compare this with Interpretability, where the representation of a feature is defined in terms of neurons or groups of neurons. 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 kee** the number of parameters constant. Kaplan et al. found that Transformer performance depends very weakly on the shape parameters (number of layers), (number of attention heads), and (feedforward layer dimension) when we hold the total non-embedding parameter count fixed (Kaplan etΒ al., 2020). 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
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 (Elhage etΒ al., 2022). 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 neurons, we can decompose it linearly into activations of features, where , as:
(2) |
where and are vectors of size , and represents the magnitude of activation of the -th feature. Sparsity means that for a given input, only a small fraction of features are active, which means is non-zero for only a few values of .
2 Case study on changing Aspect Ratio
Model A | Model B | |
---|---|---|
Total Parameters | ||
Neurons per Layer | ||
Number of Layers | ||
Total Number of Neurons | ||
Total Number of Features Learned | ||
Number of Features per Layer | ||
Features per Neuron |
Letβs consider two Transformer 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 (Kaplan etΒ al., 2020). However, their shape parameters differ. Using the same notation as Kaplan et al., letβs denote the number of layers as , and number of neurons per layer as . Model B has twice the number of neurons per layer compared to A. As the number of parameters is approximated by , Model B must have the number of layers to maintain the same number of parameters as Model A. This means Model B has 8 times the aspect ratio () 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 (Olah etΒ al., 2020). Letβs denote the total number of features learned by both models as .
The above three paragraphs are summarized in Table 1.
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.(Elhage etΒ al., 2022) show that, using lower bounds of compressed sensing (Ba etΒ al., 2010), if we want to recover features compressed in neurons (where ), the bound is
(3) |
where is the sparsity of the features. For example, if a feature is non-zero only 1 in 100 times, then equals 0.01. We can define the degree of superposition as
(4) |
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.
3 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 whereas the modelβs representational capacity is a function of total neurons.
A useful, though non-rigorous analogy, is to visualize a solid cylinder of radius and height . 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 kee** 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.
![Refer to caption](extracted/5703347/visualization.png)
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.
3.1 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 . 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 (Bricken etΒ al., 2023). 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.
3.2 Cross Layer Superposition
Previously, we used to look for features in a single neuron (Radford etΒ al., 2017), now we extended it to a group of neurons in a layer. A natural progression is to look for features localizing 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 the 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.
Acknowledgements
Iβm thankful to Jeffrey Wu and Tom McGrath for their helpful feedback on this topic. Thanks to Vinay Bantupalli for providing feedback on the draft. Earlier version of this work was supported by an Open Philanthropy research grant.
References
- Ba etΒ al. (2010) KhanhΒ Do Ba, Piotr Indyk, Eric Price, and DavidΒ P Woodruff. Lower bounds for sparse recovery. In Proceedings of the twenty-first annual ACM-SIAM symposium on Discrete Algorithms, pages 1190β1197, 2010.
- Bricken etΒ al. (2023) Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, JosiahΒ E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023.
- Elhage etΒ al. (2022) Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of superposition. Transformer Circuits Thread, 2022.
- Kaplan etΒ al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, TomΒ B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
- Olah etΒ al. (2020) Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 10.23915/distill.00024.001.
- OpenAI (2023) OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
- Radford etΒ al. (2017) Alec Radford, Rafal Jozefowicz, and Ilya Sutskever. Learning to generate reviews and discovering sentiment. arXiv preprint arXiv:1704.01444, 2017.