A Features' Perspective on Neural Scaling Laws

January 15, 2025

Notes of a talk I recently gave on feature superposition (a microscopic phenomenon) and neural scaling laws (a macroscopic phenomenon).

Guiding Question For My Research

Can feature superposition (a microscopic phenomenon) explain neural scaling laws (a macroscopic phenomenon)?

Specific Questions For This Talk

Consider Transformer models of fixed depth but different widths.

Do different models store the same features in different fashion?
OR: Do larger models store more features, giving them higher capability?
In either case, can we see its effect on scaling laws?

Outline

Background
- Neural Scaling Laws
- Feature Superposition
This work
- Feature Importance and Universality
- Relationship with Scaling Laws

Background

Neural Scaling Laws

Neural scaling laws describe how model performance improves with scale (Kaplan et al., 2020; Hoffmann et al., 2022, and others).

Loss scales with compute, dataset size, and parameter count

Definition: Neural Network Features

Activations can be decomposed into (overcomplete) bases.
For token $j$, we can write activations as a sum over features $i$.
Each feature $\mathbf{d}_i$ has a specific interpretation.
$f_i(x) \ge 0$ represents feature activations.
For a given token, only a few features are active.

\[ \mathbf{x}^{j} \approx \mathbf{b} + \sum_i f_i(\mathbf{x}^{j})\mathbf{d}_i \]

Example: The Golden-Gate-Bridge feature of Claude 3 Sonnet (Templeton et al., 2024). Activating this feature to 10x its value changes the model behavior.

Text activating the Golden Gate Bridge feature

Claude output before and after amplifying the Golden Gate Bridge feature

Definition: Superposition

Neural networks store more features than the number of available dimensions.
Hence, some features interfere with others.
Intuitively, larger models perform better as they have more “capacity”: they can store more features without interference (Elhage et al., 2022).

Interference between features sharing a two-dimensional space

This talk: Make this more precise.

How are features learned?

Reconstruct activations using autoencoders and let the decomposition be sparse — Sparse Autoencoders (SAEs) (Bricken et al., 2023; Cunningham et al., 2023).

\[ \mathcal{L} = \frac{1}{|X|} \sum_{x \in X} \left\lVert \mathbf{x} - \hat{\mathbf{x}} \right\rVert_2^2 + \lambda \left\lVert \mathbf{f} \right\rVert_1 \]

\[ f_i(x) = \operatorname{ReLU}\!\left(W_e(\mathbf{x} - \mathbf{b}_d) + \mathbf{b}_e\right)_i \]

New Work

Feature Importance

Can we define a notion of importance of a feature for features of a real model?

Important features must be more universal across models of different widths.
Important features may be learned early in training. (Not answered today, but I have some observations.)
Hence, scaling laws could be studied from the perspective of feature importance.

Proposal (Definition): Feature Importance

\[ I_i = \max_{x \in X} f_i(x) \]

Let the importance of feature $i$ be its maximum activation value over a large dataset.

Experiments: Transformer Models

Parameter counts and training and validation losses for four model widths

Train four 8-layer models with varying embedding dimensions: 128, 256, 512, 768
Train SAEs on the residual stream entering the seventh Transformer block, all with 24,576 latents.
Reconstructed Losses with the SAEs spliced on.

Training loss versus non-embedding parameter count

Original and SAE-reconstructed losses versus non-embedding parameter count

Are important features more universal?

Measure of universality: activation similarity. Roughly, how much do the activations of two features correlate?

More mathematically:

Assign to each feature, a vector of length $\lvert X\rvert$ of its activations.
Compute Pearson Correlation between features of one model with another.

Results:

Relative feature importance versus maximum activation similarity — Relative Feature Importance (x-axis) vs Maximum Activation Similarity (y-axis)

Fitted Equation: $y = 1 - e^{-b(x-c)}$
$r = 0.7002$

Pairwise feature importance and activation similarity across four model widths

Table of Correlation Coefficients:

width \ width	128	256	512	768
128	-	0.7735	0.7751	0.7640
256	0.6446	-	0.7026	0.7352
512	0.7136	0.7152	-	0.7216
768	0.6975	0.7096	0.7002	-

New Insight #1

Important features tend to be more universal amongst models of fixed depth and various widths.

Dependence of Features Importance on Scaling Laws

Training loss versus model size for different fractions of retained features

Larger models perform better for any fixed number of features.
Reason: less interference.
In fact, models’ differences increases as you pack more features.
From the same number of features, larger models gain more capabilities.

Scaling exponent versus the fraction of retained features

But do they also store more features?

Reconstructed loss versus the fraction of retained features for each model width

Yes! Larger models also pack more features!

New Insight #2

Larger models store more features. But even for a smaller number of features, they extract more performance.

Answer to Our Questions

Consider Transformer models of fixed depth but different widths.

It’s a combination of both!

Neural Scaling Laws get contribution from both factors:

Larger models extract more from a similar set of features.
Larger models also store more features.

Quantifying the contribution from both is an interesting problem.

Some Limitations

We ignored any inductive biases SAEs bring to the family of features learned.
We ignored part of original loss that is not reconstructed by the SAEs.
We studied only 4 models of varying widths.

References

Jared Kaplan, et al. Scaling Laws for Neural Language Models. arXiv:2001.08361 (2020). arXiv ↩
Jordan Hoffmann, et al. Training Compute-Optimal Large Language Models. arXiv:2203.15556 (2022). arXiv ↩
Trenton Bricken, et al. Towards Monosemanticity: Decomposing Language Models With Dictionary Learning. Transformer Circuits (2023). Article ↩
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse Autoencoders Find Highly Interpretable Features in Language Models. arXiv:2309.08600 (2023). arXiv ↩
Nelson Elhage, et al. Toy Models of Superposition. Transformer Circuits (2022). Article ↩
Adly Templeton, et al. Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet. Transformer Circuits (2024). Article ↩