Superposition Yields Robust Neural Scaling

We thank the reviewer for the nice comments and helpful questions. We will reply to the weaknesses and limitations point by point.

• Weakness 1: We agree, and we may emphasize this weakness more in our discussion of limitations. We tried to use the last paragraph in Section 2 to summarize some key differences between the LLMs and the toy model and justify why some details in LLMs can be ignored, which should also be highlighted or strengthened.

• Weakness 2: This is true. The robustness of $\alpha_m\approx 1$ across $\alpha$ and the transition of $\alpha_m$ are observations, and we try to provide a picture of the underlying mechanisms with more details like vector norms and overlaps observed from the experiments. However, we are limited by the fact that we are not able to analytically solve the directions of the representation vectors. We are actively trying many theoretical analysis methods. One simplified analysis is to assume random directions of the embedding vectors and solve for norms and biases.

• Weakness 3: Thank you for mentioning this. We had a brief discussion on other scaling laws in Appendix B. We are now actively studying the scaling with training steps. The scaling with dataset size should be the same as that with training steps in the online learning setup. The training dynamics in the weak superposition regime have been solved. Yet, the training dynamics in the strong superposition regime is hard to analyze.

• Questions: Given the number of features, $n$, a larger dimension $m$ may lead to less superposition. When we talked about equal performance with fewer parameters, we were imagining a large model with weak superposition and a small model with strong superposition could have similar losses (which is true in the toy models).

• Limitations: Please see reply to Weakness 3.

We thank the reviewer for the helpful questions and suggestions. We will reply to the weaknesses, questions, and limitations point by point as follows.

• Weakness 1 (Structure of the paper): We are sorry for the confusion and inconvenience due to our presentation. The reviewer's suggestion is very helpful, which we will follow. We will also highlight our key definitions, assumptions, and takeaways by the Theorem environment. We attempted to include too many results in the main text. We will emphasize only the most important ones and save space to move Discussion back to the main text.

• Weakness 2 (Weak LLM analysis): Yes, we totally agree with the reviewer, and we should acknowledge this limitation more. Our toy setup is far from actual LLMs, and we used the last paragraph of Section 2 to discuss key differences and why some details in LLMs may not matter. In Discussion (Appendix B), we also discussed other sources of error in LLMs and how to better interpret Figure 8b. Nevertheless, the connection between toy models and LLMs is not rigorous. Regarding the analysis of LLMs, we studied the embedding vectors, loss scaling with width, and token frequency distribution (Figure 22). We tried to analyze properties we have access to and that can be connected to the toy model in some way. The philosophy is that we study all kinds of possible behaviors in the toy model, identify the regime where LLMs are in, and connect LLM behaviors to mechanisms found via the toy model. We can move more details on LLM analysis from the Appendix to the main text to strengthen the last part. We are also willing to conduct mechanistic investigations on LLMs if there is any concrete quantity to look at.

• Nitpick 1: Thank you for this point. We can use "well represented" vs. "poorly represented" or "over represented" vs. "under represented".

• Nitpick 2: Thanks, we will revise.

• Question 1 (Causality?): This is a conceptually important question. Let's think about another example first. Say we can change the temperature to make water ice or liquid. A behavior we found at the lower temperature state is a behavior of the ice. And a behavior at higher temperatures is that of the liquid state. By changing weight decay, we enter the state of "weak superposition" or "strong superposition". We can always say that the scaling law is a property of the strong superposition regime. To be rigorous, "geometry of representations in the strong superposition regime yields a scaling law with width" is the statement, and "superposition causes scaling law" is a simpler version of it.

• Question 2 (Correlation of features): This is a great question. If we consider the correlation between different dimensions of the input vector, the "independent atomic features" should be the eigenvectors of the correlation matrix instead of the dimensions of the input vector. Theoretically, after rotating the weight matrices to another basis, the current analysis still follows. Yet, the non-linearity can introduce extra complexity, which we are not certain about and should study via experiments. We expect the scaling to be robust even if input dimensions are correlated. We should acknowledge this point in our Discussion. And if we consider the eigenvectors of the correlation matrix to be features, the feature frequencies can be different from those of the input dimensions. We measured the co-occurrence matrix of tokens in the Pile, varying context length from 3 to 512. The eigenvalues of the co-occurrence matrix have a scaling exponent $\alpha$ from 1 to 1.2, which does not change much. So, we also do not expect that considering token correlation would change our conclusion.

• Question 3 (Meaning of a sentence): We are sorry for the confusion, "fitting function or manifold" and "learning skills" refer to previous mechanisms proposed to explain the neural scaling laws. We will add citations near these terms. The sentence tries to say that some previous mechanisms conceptually describe what the transformer layers do, while there is another source of error -- representation or embedding -- that needs studying.

• Limitations: We discussed some of our limitations in Appendix B Discussion. We will emphasize what the reviewer pointed out. And we will try to move some results to the Appendix such that we can move the Discussion section back to the main text.

We are grateful for the reviewer's appreciation and helpful questions. We will reply to the Questions and Limitations point by point as follows.

• Question 1 (Phrasing of a sentence): We are sorry that we misused the terminology "representation learning". Our representation-limited error is due to the imperfect token or concept embedding in a low-dimensional embedding space. In language, tokens are atomic features in our mind, and we can conceptually connect them to features in our toy model. And practically, token embeddings are easy to access and analyze. Yes, if the token embedding has some error, the more abstract representations in the middle layers will also have an error. We guess the error of composite features would have the same scaling with width as the atomic features. Our thinking behind the sentence is that there can be two sources of the total error, one is from the imperfect embedding itself (explicit focus of this paper), and the other is from imperfect calculation or transition acting on the embeddings (which is more related to the transformer layers). We propose keeping the phrase "representation" as it could apply to either components of error, although we could also replace it with "embedding" if this would avoid confusion.

• Question 2 (Middle layer superposition): This is a great question! In the AI biology work by Anthropic, they found that in poem completion tasks, the model can plan ahead by representing the next whole sentence at the end of the prompt using the embedding dimension of one token. In the middle layers, the number of concepts to represent may be huge, especially when the context length is large. So, in our mind, for LLMs, superposition may be impossible to avoid. Yet, we do not have more rigorous arguments for this. Token embedding is a subproblem for representation in LLMs and should share the same scaling. There is definitely more to study in representations in the middle layers, and we can add discussions on this topic in the Discussion section. We also imagine that, in terms of new architectures, an adaptive hidden dimension may help performance: if one position needs a lot of previous tokens to understand or is needed by a lot of following tokens for their decoding, it should have a larger embedding dimension.

• Question 3 (Linear representation hypothesis): Thanks for this question. If the linear representation hypothesis (LRH) holds, we expect the parsing loss to have a term that is the number of compositions times the error of one embedding. If LRH does not hold, there are a lot of possible situations. For example, the embedding may be on a manifold. But as long as the manifold is locally Euclidean, the 1/width scaling may still be robust (we are not sure). Another possibility when LRH does not hold is that there are infinitely many independent features, and LLMs need to keep embedding new independent features in the middle layers, which may increase the coefficient but not change the 1/width scaling. We are also sorry for the lack of rigor and preliminary thinking.

• Question 4 (Previous theory): Thanks for mentioning this. Since different theories on neural scaling laws have different setups, we are not very confident in comparing them directly. As far as we can see, previous theories like the phenomenological quanta model and the formal kernel method have the same spirit as our no superposition result. Our finding in the strong superposition regime is therefore complementary to previous theories, suggesting a new mechanism due to the intrinsic model configuration rather than data properties.

• Question 5 (Falsifying our theory): We think a direct experiment is to fix a large depth and only increase the width of LLMs, and study clearly the scaling of loss with width. If the scaling exponent obtained from this experiment is far from 1, then our mechanism cannot explain LLM behaviors. Another approach is about language itself. If one can argue that atomic features of language are not tokens but something much fewer than the number of tokens and embedding dimension, then our theory is false.

• Question 6 (Interpretability and AI safety): Thank you for reminding us about interpretability and safety issues. We agree that encouraging superposition may lead to more powerful but less interpretable models and can have a negative impact on safety. We will add a discussion on these topics.

• Question 7 (Typo): Thank you for the careful reading. We will revise.

• Limitations: Please see reply to Question 6.

We thank the reviewer for the careful reading and professional technical questions. We reply to the questions and concerns point by point below.

• Main weakness 1: We are sorry for the inconvenience due to our presentation. We wanted to include too many points in a limited space (9 pages), which led to dense writing and inadequate details for each point. To solve this problem, we plan to emphasize the core results in the main text by adding more details and moving less important results like Figure 7 to the Appendix. We can also add subsections and use the Theorem environment to clearly list our assumptions and results to improve readability. We are willing to incorporate more suggestions from the reviewers, if any.

• Main weakness 2:

  • Yes, features being represented are those having non-zero representation vectors. Features being well-represented means they can be correctly recovered or predicted. We can highlight this by using a "Definition" environment.
  • Based on the definition of "represented or not", the threshold should be 0. In practice, we need to use a (small) non-zero threshold. We chose 0.5 since it is equally far away from 1 and 0, potentially minimizing misclassification. But this choice is very conservative since 0.5 is not small. When $\gamma \geq 0$, the norm distribution concentrates near 0 or 1 (Figure 11), and choosing a threshold from 0.05 to 0.95 does not change the outcome. When $\gamma < 0$, all norms are around 1 but may have large variance (Figure 12). In this case, choosing 0.5 as a threshold will give a conservative lower bound for the number of features being represented. Our goal is to establish a relation between $\gamma$ and superposition. If a lower bound is close to 1, it suffices to say that $\gamma<0$ can robustly generate strong superposition. We should add a remark on this issue. We can also provide a scan or sensitive study of the choice of thresholds in the Appendix.
  • The number of ETF-like vectors has a bound $\sim m^2/2$. When $m^2/2<n$, the less important features cannot have small overlaps and have to decrease their norms to have less interference. The two distinct modes are being in ETF-like configuration (with mutually small overlap) or not. This particular behavior is irrelevant to LLMs since $m^2/2>n$ is always true. Please see Figure 19, LLM embedding vector norms is not bimodal. Then, LLM loss should be dominated by the strongly represented features (ETF-like), more likely leading to $\alpha_m = 1$, which seems to be true. We scan all kinds of behaviors of toy models, then identify which regime LLMs are in, and use the mechanism found from the toy model in that regime to explain the LLM behaviors.

• Main weakness 3: We are sorry for our oversight. The assumption is that all biases are zero. In our experiments, since both $m$ and $n$ are large, the expected value of $x_i~(i>m)$, proportional to the probability of $p_i$, is extremely small, approximately satisfying the assumption.

• Main weakness 4: "ETF-like" has two main points: (i) Overlaps distributed closely around $\kappa$ in Eq. (5); (ii) The number of such vectors has an upper bound. We verified point (i) by Figure 6d (i.e., the mean of overlaps agrees with $\kappa$) and Figure 15 (i.e., the variance is small at least for small $\alpha$). And yes, weakly represented features behave differently, as they have much larger mean overlap (compare Figure 16 and Figure 6d). We verified point (ii) in Figure 6c, where the bound is not exactly $m^2/2$ but around. This discrepancy is part of why we refer to the vectors as being "ETF-like" (we do not claim they form an ETF). What we can rigorously say is that when $\alpha = 0$, the optimal configuration is an ETF. However, when $\alpha$ increases, one would expect important features to occupy larger angular space and break the evenness. Yet, we still observe that the vectors are similar to ETF for a wide range of $\alpha$ and yield a robust $\alpha_m \approx 1$. In our Discussion, we admit our main weakness is that we cannot analytically solve the toy model. What we did was to connect empirical findings and simple theoretical examples to generate a zeroth-order understanding.

• Main weakness 5: This problem is also due to the fact that we cannot solve the toy model now. Specifically, we found that solving the directions of the embedding vectors is unexpectedly hard. What we are actively trying is to assume random directions and solve vector norms and biases. We are curious how many phenomena can be explained by such an ansartz. What we presented in the paragraph starting the line 217 is a hypothesized explanation for the transition observed. Some results in Figure 16 are confusing to us, too. The proposed picture seems to agree more with observations in the small $\alpha$ regime, but not the regime after the transition. Our understanding is that when $\alpha$ is large, the overlap variance is also getting large (e.g., Figure 15). So, the decreasing mean overlap over weakly represented features does not mean they are not squeezed. The weakly represented ones may be squeezed into different positions in the space (and have lower mean but higher variance) such that the most important features can have lower overlap with each other. As long as each weakly represented feature contributes to constant loss (e.g., has one large overlap), our proposed explanation may be valid. Again, the study of direction configurations is hard in high-dimensional space for both theory and experiment. We should highlight that our explanation for the transition is one hypothesis. We believe that the empirical results are correct and worth sharing, and many theoretical studies can be done later by us or others with various tools.

• Additional concerns 1: We defined "the absence of superposition" in line 59. We will rewrite to make it clear what "weak superposition" is.

• Additional concerns 2: We will use "input sample" or "input".

• Additional concerns 3: Thank you, we will explain more in the caption.

• Additional concerns 4: Only weight decay. Since we used AdamW, weight decay is a separate line of code.

• Additional concerns 5: Thanks for the suggestion! This would help a lot.

• Additional concerns 6: $m = 10, 15, 39, 63, 100$. They are roughly uniform in the log scale. Will make this clear.

• Additional concerns 7: Thanks! We will revise.

• Additional concerns 8: You are right, we will revise.

• Additional concerns 9: The second line is copied from Figure 5b. Therefore, we can summarize all behaviors of $\alpha_m$ in different regimes on one plot. We tested 10 $\gamma$ evenly from $-1$ to $1$. The points with high weight decay have $\alpha_m \ll 1$. And the reason is provided around line 166: if weight decay is too large, the number of learned features will be fewer than $m$, which leads to slow decay of loss and is not ideal.

• Additional concerns 10: For models with sizes smaller than 3 billion, the embedding and LM head matrices are usually tied. Larger models usually do not tie the matrices. Analyzing either matrix leads to the same conclusion. However, analyzing the embedding matrix will lead us to a lengthy discussion on how the representation error would propagate through transformer layers, which may distract from the focus of this paper. We can provide the analysis result of the embedding matrices in the Appendix.