Table of Contents
Fetching ...

Superposition Yields Robust Neural Scaling

Yizhou Liu, Ziming Liu, Jeff Gore

TL;DR

The paper proposes representation superposition as a central mechanism behind neural scaling laws, using a toy autoencoder model in which weight decay tunes the degree of superposition. It shows two regimes: weak superposition yields power-law loss only under certain data-skew conditions, while strong superposition yields a robust 1/m loss across diverse data distributions, aligning with observations in open-source LLMs and the Chinchilla scaling laws. Empirical results on LLM heads reveal mean-square overlaps scaling as 1/m and losses approximating L = C_m/m^{α_m} + L_no_m with α_m near 1, supporting strong superposition as a driver of width-based scaling. The work offers insights into when neural scaling can improve or break down and suggests architectural and training strategies to leverage superposition, while acknowledging the toy-model limitations and open questions about data structure and safety.

Abstract

The success of today's large language models (LLMs) depends on the observation that larger models perform better. However, the origin of this neural scaling law, that loss decreases as a power law with model size, remains unclear. We propose that representation superposition, meaning that LLMs represent more features than they have dimensions, can be a key contributor to loss and cause neural scaling. Based on Anthropic's toy model, we use weight decay to control the degree of superposition, allowing us to systematically study how loss scales with model size. When superposition is weak, the loss follows a power law only if data feature frequencies are power-law distributed. In contrast, under strong superposition, the loss generically scales inversely with model dimension across a broad class of frequency distributions, due to geometric overlaps between representation vectors. We confirmed that open-sourced LLMs operate in the strong superposition regime and have loss scaling inversely with model dimension, and that the Chinchilla scaling laws are also consistent with this behavior. Our results identify representation superposition as a central driver of neural scaling laws, providing insights into questions like when neural scaling laws can be improved and when they will break down.

Superposition Yields Robust Neural Scaling

TL;DR

The paper proposes representation superposition as a central mechanism behind neural scaling laws, using a toy autoencoder model in which weight decay tunes the degree of superposition. It shows two regimes: weak superposition yields power-law loss only under certain data-skew conditions, while strong superposition yields a robust 1/m loss across diverse data distributions, aligning with observations in open-source LLMs and the Chinchilla scaling laws. Empirical results on LLM heads reveal mean-square overlaps scaling as 1/m and losses approximating L = C_m/m^{α_m} + L_no_m with α_m near 1, supporting strong superposition as a driver of width-based scaling. The work offers insights into when neural scaling can improve or break down and suggests architectural and training strategies to leverage superposition, while acknowledging the toy-model limitations and open questions about data structure and safety.

Abstract

The success of today's large language models (LLMs) depends on the observation that larger models perform better. However, the origin of this neural scaling law, that loss decreases as a power law with model size, remains unclear. We propose that representation superposition, meaning that LLMs represent more features than they have dimensions, can be a key contributor to loss and cause neural scaling. Based on Anthropic's toy model, we use weight decay to control the degree of superposition, allowing us to systematically study how loss scales with model size. When superposition is weak, the loss follows a power law only if data feature frequencies are power-law distributed. In contrast, under strong superposition, the loss generically scales inversely with model dimension across a broad class of frequency distributions, due to geometric overlaps between representation vectors. We confirmed that open-sourced LLMs operate in the strong superposition regime and have loss scaling inversely with model dimension, and that the Chinchilla scaling laws are also consistent with this behavior. Our results identify representation superposition as a central driver of neural scaling laws, providing insights into questions like when neural scaling laws can be improved and when they will break down.
Paper Structure (26 sections, 19 equations, 23 figures)

This paper contains 26 sections, 19 equations, 23 figures.

Figures (23)

  • Figure 1: Superposition leads to robust and fast power-law loss decay with model size. (a) Illustration of no superposition where a three-dimensional space can at most represent three features without any interference (overlap). (b) Toy model results in the regime of weak superposition, where we set data dimension $n = 10240$ (number of features). The toy model will be introduced; more details are in Appendix \ref{['app:fig1']}. (c) Illustration of superposition: there are more features than the dimension of the space. (d) The same toy models in the strong superposition regime show lower losses, which are on power laws with model dimension and have exponents close to 1 (color coding same as panel b). The gray points are from actual LLMs, which have a similar power-law exponent near 1.
  • Figure 2: Toy model of superposition. (a) Architecture and loss of the toy model. (b and c) A row of the matrix $W$, denoted by $W_i$, is the representation of feature $i$. (b) No superposition represented the most frequent features, i.e., the first three ($n=6$ features in $m=3$ dimensional space), without interference. In the frequency-rank plot, height means feature $i$'s frequency $p_i$, and color means the $i$th row vector's norm $\|W_i\|_2$. (c) With superposition, features are all represented, while the representations $W_i$ overlap.
  • Figure 3: Weight decay can tune the degree of superposition. (a) Positive weight decay ($\gamma = 1$ in the figure) has $\|W_i\|_2$ near 0 or 1, with frequent features more likely to be represented (color means $\|W_i\|_2$ in frequency-rank plots). Negative weight decay ($\gamma = -1$) has $\|W_i\|_2$ around $1$. We show results when $\alpha = 1,~m = 100$, yet the claim is generally true. (b) For all models, small weight decays lead to strong superposition, and large weight decays lead to no superposition ($\phi_{1/2}\approx m/n$). More data in Appendix \ref{['app:fig3']}.
  • Figure 4: Loss at weak superposition can be well described by the frequency sum of ignored features. (a) Observation and theory at weak superposition (i.e., Equation (\ref{['eq:unlearn']}) as a function of number of represented features, $\phi_{1/2}n$) agree when weight decay $\gamma$ is positive. (b) For those closest to the ideal no superposition case, we expect $\alpha_m = \alpha -1$, which is close to measured values. Error bars are standard errors. Details in Appendix \ref{['app:fig5']}.
  • Figure 5: Loss scaling at strong superposition is explained via geometry. (a) The row norm distribution is bimodal around $1$. (b) The more frequent the features are, the more likely their norms are greater than $1$. (c) Variance of squared overlaps for features with $\|W_i\|_2>1$ is smaller than that of random unit vectors, i.e., $\frac{2(m-1)}{m^2(m+2)}$. Overlaps are calculated using directions $W_i/\|W_i\|_2$. We show the measured variances ($\gamma = -0.55$) divided by the above theory value for random vectors. (d) The features with $\|W_i\|_2>1$ have $1/m$ mean squared overlaps, where we plotted all the data when $\gamma<0$. (e) At strong superposition ($\gamma<0$), $\alpha_m = 1$ if feature frequencies are flat ($\alpha$ small) due to isotropic vector geometry. But $\alpha_m \approx 2(\alpha -1 )$ if the feature frequencies are skewed ($\alpha$ large). Error bars are standard errors. More details in Appendix \ref{['app:fig6']}.
  • ...and 18 more figures