Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing
Transformers have shown impressive capabilities across various tasks, but their performance on compositional problems remains a topic of debate. In this work, we investigate the mechanisms of how transformers behave on unseen compositional tasks. We discover that the parameter initialization scale plays a critical role in determining whether the model learns inferential (reasoning-based) solutions, which capture the underlying compositional primitives, or symmetric (memory-based) solutions, which simply memorize mappings without understanding the compositional structure. By analyzing the information flow and vector representations within the model, we reveal the distinct mechanisms underlying these solution types. We further find that inferential (reasoning-based) solutions exhibit low complexity bias, which we hypothesize is a key factor enabling them to learn individual mappings for single anchors. We validate our conclusions on various real-world datasets. Our findings provide valuable insights into the role of initialization scale in tuning the reasoning and memorizing ability and we propose the initialization rate $γ$ to be a convenient tunable hyper-parameter in common deep learning frameworks, where $1/d_{\mathrm{in}}^γ$ is the standard deviation of parameters of the layer with $d_{\mathrm{in}}$ input neurons.
Paper
Similar papers
Peer review
Summary
The work mainly investigates whether transformers indeed learn primitive functions or simply learn surficial correlations (i.e., "memorizing") when trained on a mannually-designed toy compositional task, which contains a set of composite numerical functions (e.g., $g_1(x)=x+5, g_2(x)=x+1,...$). Authors find that the learned mechanism has a phase transition (which correlates with the parameter initialization scale) : smaller initialization scale leads to learning the composite solution ("inferential solution"), where models learn the exact meanings of primitive functions and compose them to solve the task; while larger initialization scale leads to learning solutions based on the surficial correlation in the training data ("symmetric solution"). Centering on this finding, the paper also conducts empirical analysis from the perspectives of information transmission, the similarity of vector representations learned by transformers and model complexity to offer interpretation. Besides, the paper also finds that there is a positive correlation between the the complexity of the training data (i.e., the number of the data that do not meet the composition rule.) and the difficulty of fitting the training data.
Strengths
1. The topics of this paper, whether the solutions learned by transformer-based models are inferential or just memorizing, are very important and of interest to the NeurIPS community. Besides, the most part of the paper has clear logic structure, thus being easy-to-follow. 2. (Based on the setting discussed in the paper) Authors conduct analysis experiments from multiple aspects, including tracing the information flow within the trained transformers, inspecting the similarity of learned representations and the model complexity, all of which lead to a converging interpretation to the initial observation: smaller initialization scale leads to inferential solutions where larger initialization scale leads to memorizing surficial patterns.
Weaknesses
My main concerns are two-fold: 1. The experiment setting of this work might be over-simplified (as the authors discussed in the limitation part): the synthetic data with little variance (four simple primitive functions, sixteen combinations in total). Though I believe the conclusions in the paper are correct for this setting, I question the generality of the conclusions and the analysis results in other (more complex) settings (especially in this paper, there is only empricial evidence but no explicitly theoretical guarantee for the conclusion). If we train transformers on other classic compositional tasks (e.g.,: SCAN [1] or COGS [2], both of which are also synthetic but with more natural language variance, thus being ideal testbeds), will we have similar findings? Besides, the further discussion and validation (at least some initial experiments) on how to use the conclusion to guide the training (fine-tuning) of real pre-trained language models are also lacking. All these factors weaken the contribution of this work. 2. The research question discussed in the paper (i.e., inference v.s. memorization) is highly related to another line of works ("grokking") [3, 4, 5, 6] (i.e., generalization v.s. memorization). The mainly conlusion in the paper (initialization scale correlates to the learned pattern, and even the weight decay part in the Appendix E) is also similar with the conclusion in this "grokking" work [5]. I find the discussion of the relationship (maybe the difference) between this work and the previous "grokking" works are missed in the paper, which weaken the novelty of the work. [1]: Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks, https://arxiv.org/abs/1711.00350 . [2]: COGS: A Compositional Generalization Challenge Based on Semantic Interpretation, https://arxiv.org/abs/2010.05465 . [3]: Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, https://arxiv.org/abs/2201.02177 . [4]: Towards Understanding Grokking: An Effective Theory of Representation Learning, https://arxiv.org/abs/2205.10343 . [5]: Omnigrok: Grokking Beyond Algorithmic Data, https://arxiv.org/abs/2205.10343 . [6]: Explaining grokking through circuit efficiency, https://arxiv.org/abs/2309.02390.
Questions
Q1. Regarding the first point in the Weaknesses part: Q1.1. Can the main conclusion and the analysis results introduced in the paper generalize to other well-studied compositional tasks (e.g., SCAN [1] or COGS [2]) ? Q1.2. From my point of view, a direct usage for the conclusion introduced in this paper might be initializing the LoRA [3] parameters (or Prompt-Tuning [4] learnable parameters) when fine-tuning pre-trained language models (PLMs). I am wondering if we can observe similar phenomenon when we use PLMs + LoRA/Prompt-Tuning to learn solutions for the task used in this paper or other tasks (e.g., SCAN or COGS)? Q2. Regarding the second point in the Weaknesses part: Q2.1. How to understand the difference and the correlation between the conclusions of this work and [5] ? Q2.2. For larger initialization scale, if we train the transformers far beyond over-fitting, will the learned patterns change (from memorization to inferential solutions)? just as the "grokking" paper described. Q2.3. What will the learned solution be? if we use larger initialization scale and larger weight decay? Q2.4. I suggest the authors discuss "grokking" works in the Related Work section. Q3. An additional question: In Section 6, it is claimed that "Based on our analysis of the mechanisms underlying different solutions, we can predict that models with different initialization scales will exhibit varying degrees of difficulty in learning data with different complexities.". I do not see clear relationship between the initialization scale and data complexity discussed in the previous sections. Could you please explain more about this relationship (between Section 6 and previous sections)? Besides, I think instead of using "epochs for 100% training acc", it is more meaningful to use the "epochs for 100% testing acc" (e.g., using inferential solution as the ground-truth). Q4. A minor suggestion: the abbreviations are supposed to be introduced when they are used in the paper first time. (e.g., Eq(6), FNN, by the way, I think it should be FFN (Feed-Forward Network) here.) [1]: Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks, https://arxiv.org/abs/1711.00350 . [2]: COGS: A Compositional Generalization Challenge Based on Semantic Interpretation, https://arxiv.org/abs/2010.05465 . [3]: LoRA: Low-Rank Adaptation of Large Language Models, https://arxiv.org/abs/2106.09685 [4]: The Power of Scale for Parameter-Efficient Prompt Tuning, https://arxiv.org/abs/2104.08691 [5]: Omnigrok: Grokking Beyond Algorithmic Data, https://arxiv.org/abs/2205.10343 .
Rating
6
Confidence
4
Soundness
2
Presentation
2
Contribution
2
Limitations
The authors do discuss the limitations of this work in Section 7.
Summary
This paper studies the role of transformer initialization scale on its ability to compositionally generalize. Using a simple arithmetic task, the authors find that when initialization scale is low, the model learns inferential (or compositional) solutions, whereas when the initialization scale is high, models tend to learn memorized solutions, leading to poorer generalization performance. These core findings are supplemented with additional analysis of their model, e.g., through mapping the information flow of each task via analysis of the attention maps. They find that small scale initializations learn token embeddings that are interpretable; that is, the token embeddings generally learn the correct ordinal structure. Finally, they show an analysis that demonstrates that the learning time scales with the inferential complexity of the dataset/task.
Strengths
The study pairs an interesting set of subfields in deep learning: compositionality with neural tangent kernel learning theory. The results are unambiguous and demonstrate strong effects. I have little doubt that these effects are replicable, and that insights from this line of work can be transferred to other domains.
Weaknesses
The primary weaknesses: 1) Other studies demonstrate that this type of arithmetic generalization can be achieved without specific scale-dependent initializations on the arguably more difficult task of length generalization by manipulating positional encodings (Zhou et al., 2023; Zhou et al., 2024; Ruoss et al., 2024; McLeish et al., 2024). In some ways, this paper addresses a simpler form of compositional generalization – systematic compositional generalization – which focuses on generalization on novel combinations of previously seen elements (rather than novel structures of previously seen elements (productive generalization)) (Hupkes et al., 2020). 2) Clarity: I found the paper difficult to understand in some parts. This was partially due to writing, but also partially due to lack of labeling (e.g., figures). For example, the paper often uses the token ID of the anchor pairs (e.g., 3 and 4), rather than the actual referent values of those anchors (-2 and -8), which makes it difficult to follow exactly what figures are showing. Unfortunately, I did not understand the motivation, results, and conclusions of the analyses presented in Section 5.2 and Figure 4. It could be that it was already redundant with results from Section 5.1 (which I did find interesting and easy to follow).
Questions
1) Is there an error in Fig. 1a? It appears the anchor for ‘4’ says ‘8’, when I think it should read ‘-8’. 2) It would be helpful to attach some of the references included in related work to the Introduction section, where they reference the NTK work as guiding the rationale and hypotheses. 3) Is a non-inferenetial composite anchor considered to not follow the traditional rules of arithmetic? That is, [3|4] -> -6, even though $-2 - 8 \neq -6$? The terminology of non-inferential composite anchors was a bit confusing to me. Relatedly, in the dataset, how many non-inferential composite anchors are given? Just 1? 4) Are key items sampled from 20-99 uniformly? Is there a motivation for why this specific set of keys? What are the noise items? The examples provided didn’t seem to provide any examples with noise. 5) In the transformer model initialization, are *all* parameters initialized from the same scale distribution, including: embedding matrix, PE vector, Q/K/V matrices, and MLPs? Or is it a subset of these parameters? 6) I would be curious to see the learning trajectories for Fig. 2c to get a sense of convergence rate across initializations. 7) In general, it would be helpful if anchor pairs were labeled (perhaps with parentheses) of the actual values the correspond to. For example, it is difficult to verify that Fig. 3a/c are doing correctly without having first memorized what the values the anchor pairs refer to. 8) I had trouble understand the motivation and corresponding results for section 5.2, and what the general conclusion/takeaway was. I found the Figure (fig 4) confusing to understand. This is likely related to the above question 7, where it is hard to keep track what the anchor values refer to without having to refer back to Fig 1. 9. Please label the rows and columns of Fig. 5a. 10. Condensation (line 262) is never defined. 11. Regarding Fig 6. There appears to be an effect that adding layers in the small initialization regime improves the sample efficiency of the model. Is this effect true? 12. Lines 132, 133: From the definitions it wasn’t immediately clear what generalization on data meant. Does this mean different key values were used on generalization? Whereas task generalization is specific to anchor pairs?
Rating
6
Confidence
4
Soundness
3
Presentation
2
Contribution
3
Limitations
The dataset/task is limited and focuses only on a specific type of ‘systematic compositional generalization’ (i.e., novel combinations of anchor pairs), and does not fully address other forms of compositional generalization, such as length generalization. Also, the clarity of the paper is a bit difficult to follow, which could potentially limit its impact.
Summary
Transformers show remarkable capabilities but struggle with compositional tasks, raising questions about their true understanding versus input-output mapping. The authors investigates transformers' generalization to unseen compositional tasks using anchor functions, revealing that initialization scale impacts solution types. With smaller initialization scales, models tend to infer solutions, while larger scales lead to symmetric solutions. Analysis shows symmetric solutions combine anchor information without compositional understanding, whereas inferential solutions capture underlying primitives. These insights suggest model complexity influences the ability to learn compositional primitives, with initialization scale being crucial.
Strengths
- The investigation into the effect of initialization on compositional abilities of transformers is well-motivated and practically useful. - The discovery of phase transitions from symmetric memorization to generalization based on initialization scale is compelling.
Weaknesses
- While the problem and findings are new, the approach lacks novelty. - The focus on symmetric solutions is novel, but the generalizability to real data and a broader range of models remains uncertain.
Questions
- Do you have a solid definition for "Phase Transition"? If so, how did you define it? - Do you have any additional results with other data/models (no need to run new experiments for the rebuttal)?
Rating
6
Confidence
4
Soundness
3
Presentation
4
Contribution
3
Limitations
Experiments were conducted primarily on a simple synthetic setup. Future work should verify the extent to which these findings can be generalized and explore the implications of these findings on real-world applications of transformers.
Thanks for the authors' response
Dear Authors, Thanks a lot for the detailed response. As for the new results on SCAN and COGS, I appreciate the results which demonstrate that initialization scales to some extent affect the generalization performance on unseen (out-of-distribution) compositions. I still have questions regarding this point: Do the model's o.o.d. performance exhibit the phase change or gradual increase with the changing of initialization scale? Moreover, I value the mechanistic insights of two phases discussed in the paper a lot. Could you please validate these findings (i.e., Figure 3 and Figure 4) on SCAN (or COGS)? I will keep thinking about the rest parts in the rebuttal. I write this comment to give you feedback ASAP. Reviewer MVAX.
Thanks for the reviewers' response
Dear Reviewer MVAX, Thank you for your appreciation of our results on the SCAN and COGS datasets and for your insightful questions. We immediately conducted further experiments to address your concerns. Here are our responses: $\textbf{Point 1}$ Do the model's o.o.d. performance exhibit a phase change or gradual increase with the changing of initialization scale? $\textbf{Reply}$ Due to time constraints, we re-ran experiments on SCAN and COGS with $\gamma=0.3$ and $\gamma=0.7$, using $\gamma=0.5$ as a reference point. Preliminary results indicate that $\gamma=0.3$ does not lead to a complete absence of OOD generalization (the model still shows some performance on OOD data), which might initially suggest the absence of a phase change. However, detailed analysis revealed that $\gamma=0.3$ models tend to generalize to only a few specific types of OOD data while failing to generalize to others. Here is a detailed explanation: In simple terms, OOD data can be subdivided into various types, each with nearly identical generalization difficulty within its type but different difficulties across types. For example, in the COGS dataset, OOD data is divided into 21 categories, each constructed differently and with varying difficulties. If we closely observe the generalization accuracy of each category concerning the initialization scale, we can see a clear phase change. For the SCAN dataset, with a training set of 5,000 data points, the attached PDF shows that the $\gamma=0.7$ model generalizes well on OOD data, achieving an accuracy of 23% on OOD data for the $\gamma=0.3$ model. Specifically, if the OOD data consists solely of commands like "I_TURN_LEFT" (e.g., "turn left twice"), the model can achieve nearly 100% accuracy. However, if the output commands include adjacent "I_TURN_LEFT" and "I_TURN_RIGHT," the model almost always fails. Thus, considering OOD data with adjacent "I_TURN_LEFT" and "I_TURN_RIGHT" commands, a phase change appears. For the COGS dataset, trained with 20,000 data points, the attached PDF shows that the $\gamma=0.7$ model achieves 65% accuracy on OOD data, while the $\gamma=0.3$ model achieves only 6%, primarily driven by the "pp_dative_to_do_dative" OOD data type. Excluding this specific type, a clear phase change can be observed. In summary, for OOD data with uniform difficulty, the phase change in generalization accuracy concerning the initialization scale is evident. However, due to varying difficulties among different OOD data types, the phase change position may slightly vary. This variance can blur the phase change when averaged over multiple data types, creating an appearance of a gradual increase. We also recommend examining Figure 2 in the ONE-PAGE PDF. Here, the test data difficulty is uniform (with inference chains of length 3), making the phase change more apparent. As expected, at a data size of 10,000, the large initialization accuracy is only 50% (random guess), while the small initialization accuracy is nearly 100%. $\textbf{Point 2}$ Could you please validate these findings (i.e., Figure 3 and Figure 4) on SCAN (or COGS)? $\textbf{Reply}$ Thank you for appreciating our mechanistic insights. We studied the attention matrices in SCAN. For the $\gamma=0.7$ model, we found that later machine instructions tend to have higher attention weights to earlier corresponding human instructions. Different attention heads in the same layer often focus on different instructions within the same sentence. We observed that shallow attention heads usually focus on a single instruction, while deeper heads associate multiple instructions. Some heads only attend to tokens a few positions ahead. Combining the information of the two types of heads, the machine instruction part can intuitively predict the next instruction. We believe that transforming human language into machine instructions may require the help of MLP layers. For the $\gamma=0.3$ model, the attention matrices do not show clear structures, suggesting that the high regularity of attention matrices is crucial for the model's generalization in compositional tasks. Additionally, we examined the generalization mechanisms on PrOntoQA. As the layers deepen, the model's attention matrices transmit information about the next object along the inference chain to the last position in the sequence (the output True or False position). This step-by-step reasoning enables the model to judge the correctness of a Question and output True or False. Similarly, the large initialization model's attention matrices lack such clear structures, resulting in random guessing. We again thank the reviewer for the prompt and detailed feedback and hope our rebuttal meets your satisfaction. Sincerely, The Authors
Thanks for the response.
Dear Authors, Thanks for the additional results and analysis. My concerns around the (Weaknesses@1 and Q@1) have now been cleared up. I think it is good to discuss the new experiment settings and results in the revised version of the paper. Regarding Q3, I would like to clarify my question: since you want to connect the proposed metric "inferential complexity" with the degrees of the difficulty of learning data (I think here we both focus on the solution that models learned, memorizing or inferential). training acc can not unveil which type of solution models learned (both memorizing and inferential can lead to high acc). Hence using testing acc (or a validation acc with o.o.d. data) might be more reasonable. If the inferential complexity positively correlates with "epochs for achieve a high validation/testing acc"? Reviewer MVAX.
Dear Reviewer MVAX, Thank you for your recognition of our new experimental work, and we are pleased to have addressed your concerns. Below, I will explain the rationale behind our design of Figure 6 (which relates to your Q3) and hope to discuss the presentation of this experiment with you. $\textbf{Point 1}$ training acc can not unveil which type of solution models learned (both memorizing and inferential can lead to high acc). Hence using testing acc (or a validation acc with o.o.d. data) might be more reasonable. $\textbf{Reply}$ The primary intent behind Figure 6 was to illustrate the difficulty of fitting data with varying inferential complexity using models with different initializations, specifically in terms of the time required for training or generalization on i.d. data. In practice, when there is sufficient data, training accuracy often can reflect test accuracy, and the time required for a model to fit the data becomes an important metric for researchers. We aimed to demonstrate that the same data, when fitted by models with different initializations, can exhibit significantly different fitting speeds, which naturally affects the model's performance under a fixed number of training steps. In essence, the core message we intended to convey with fitting speed aligns with the display method you suggested. When a model is initialized with small values, it tends to learn patterns. Therefore, for data with low inferential complexity, the model achieves a faster fitting speed and learns the inferential solution. Conversely, for data with high inferential complexity, the model's fitting speed decreases, and it is more challenging to learn the inferential solution. On the other hand, for models with large initializations, because the data of varying inferential complexities all disrupt the symmetry in the data, the learning speed across different complexities remains the same, resulting in the model learning a symmetrical solution. Why didn't we initially use the method you suggested (i.e., displaying inferential accuracy) to present our results? There are two main reasons: First, for data with very high inferential complexity, we altered multiple composite mappings in the training data to non-inferential mappings. This made it difficult to define a universally accepted form of inferential mapping under these settings. In an extreme scenario, if all mappings were randomly disrupted into non-inferential mappings, it would be meaningless to compare unseen mappings with the original inferential mappings. The second reason is that we believe the earlier sections of the paper already provide readers with the conclusion that small initialization tends to learn inferential solutions for low-complexity data, while large initialization learns symmetrical solutions for higher-complexity data. Thus, further characterization of data with higher inferential complexity might not add significant value. We acknowledge that your suggestion may address concerns shared by other readers. One possible revision could be to move all the content of Section 6 and the above discussion to the appendix to avoid potential misunderstandings from readers. We would then include the newly added experiments in the main text for a thorough discussion, highlighting the practical significance of small initialization and similar setups in real-world problems. The results of the acc on the inference or symmetry solution are as follows: For inferential complexities of 0 and 1, models with small initialization can learn the inferential solution (inferential complexity 1 is similar to the experimental setup in the main text). For inferential complexity 5, models with small initialization fail to learn either the inferential or symmetrical solution because the form of the inferential solution on unseen anchors becomes unclear based on seen anchors. For models with large initialization, they can learn the symmetrical solution across all inferential complexities, as the symmetry in the data was not disrupted. $\textbf{Point 2}$ If the inferential complexity positively correlates with ... $\textbf{Reply}$ If we focus on i.d. generalization, the answer is yes for small initializations. In fact, the test accuracy on i.d. data follows a similar trend to the training accuracy, which is the expected outcome we mentioned in our first rebuttal. However, for o.o.d. data, as mentioned above, there may be cases where neither the inferential nor the symmetrical solution is learned. In these cases, the situation does not improve with additional training steps, making the definition of "epochs to achieve a high validation/testing acc" challenging. If we consider "epochs to achieve a high validation/testing acc" as inf if the model fails to achieve perfect generalization on o.o.d. data, then this conclusion would also hold true. We hope this clarifies our approach and look forward to your feedback. Sincerely, The Authors
I update my score.
Dear Authors, As most of my concerns are addressed in the discussion phase, I would like to update my score to this paper (from 4 to 6). Please update the new experiment results (provided in the rebuttal phase) to the revised version of the paper. Reviewer MVAX.
Dear Reviewer MVAX, Thank you very much for your feedback and for taking the time to reconsider your score based on our responses during the discussion phase. We appreciate your constructive comments and are pleased that our additional experiment results have addressed your concerns. We will update the revised version of the paper with the new experiment results as requested. Authors.
Reviewer response to author rebuttal
I thank the authors for their response to my questions. I think their additional experiments are a good addition to the study, and demonstrate robustness of their results for compositional generalization. A quick clarifying question on Fig 1 of the PDF: How is dataset size varied in the SCAN dataset? I was under the impression the dataset size is fixed (see the SCAN github).
Thanks for the reviewer's response
Dear Reviewer LUj2, Thank you for your appreciation of our additional experimental work, and we're glad to have addressed your concerns. Below, I will explain how we constructed our dataset: $\textbf{Point 1}$ How is dataset size varied in the SCAN dataset? I was under the impression the dataset size is fixed (see the SCAN github). $\textbf{Reply}$ The official dataset size provided by the authors is indeed fixed. For our SCAN dataset construction (as well as for other datasets like COGS), we based our work on the dataset provided by the original authors. To better illustrate the significant differences in performance between models with different initializations on varying dataset sizes, we conducted experiments using different dataset sizes. The smaller datasets were all derived by subsetting larger ones. For example, the 1w-sized dataset was derived from the full dataset available on the original authors’ GitHub, and the 0.5w-sized dataset was further subset from the 1w-sized dataset, and so on. Additionally, the datasets used by different initialization models with the same dataset size are identical. This approach allows us to observe the relationship between model performance and dataset size while maintaining the official data distribution. We hope our rebuttal satisfies your concerns. Please feel free to reach out with any further questions, and we look forward to your feedback. Sincerely, The Authors
Thanks
I thank the author for their clarification -- the approach makes sense. I keep my score at a 6, due to existing concerns regarding the paper's clarity and presentation. Nevertheless, I think the results are sound and interesting.
Dear Reviewer LUj2, Thank you for your thoughtful response and for recognizing the soundness and interest of our work. We appreciate your feedback and understand the concerns regarding the paper's clarity and presentation. As mentioned in our rebuttal, we will be making further revisions to the manuscript based on your suggestions and those of the other reviewers. The Authors.
Thank you for the detailed response and the additional experiment. I don't have any further questions, and I'll maintain my score toward acceptance.
Dear Reviewer 6WHp, Thank you for your detailed feedback and for taking the time to review our work. We greatly appreciate your recognition of our efforts, and we are pleased that our additional experiments have addressed your concerns. The Authors.
Decision
Accept (poster)