Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, Zhi-Qin John Xu
Abstract
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 is the standard deviation of parameters of the layer with input neurons.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 4e3a78ab-c0f0-4a6b-99f8-0990deec939cCited by top-tier papers4
- Adaptive Preconditioners Trigger Loss Spikes in AdamZhiwei Bai, Zhangchen Zhou, Jiajie Zhao, Xiaolong Li et al.ICML 2026 · 9 citations
- Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data?Yutong Yin, Zhaoran WangICLR 2025
- Swing-by Dynamics in Concept Learning and Compositional GeneralizationYongyi Yang, Core Francisco Park, Ekdeep Singh Lubana, Maya Okawa et al.ICLR 2025
- An Analysis for Reasoning Bias of Language Models with Small InitializationJunjie Yao, Zhongwang Zhang, Zhi-Qin John XuICML 2025
Builds on15
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- Faith and Fate: Limits of Transformers on CompositionalityNouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li et al.NeurIPS 2023 · 728 citations
- Improving Transformer Optimization Through Better InitializationXiao Shi Huang, Felipe Pérez, Jimmy Ba, Maksims VolkovsICML 2020 · 181 citations
- Understanding the Difficulty of Training TransformersLiyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen et al.EMNLP 2020 · 158 citations
- Compositional Abilities Emerge Multiplicatively: Exploring Diffusion Models on a Synthetic TaskMaya Okawa, Ekdeep Singh Lubana, Robert P. Dick, Hidenori TanakaNeurIPS 2023 · 113 citations
Related papers
- Attention as a HypernetworkSimon Schug, Seijin Kobayashi, Yassir Akram, João Sacramento et al.ICLR 2025
- In-Context AlgebraEric Todd, Jannik Brinkmann, Rohit Gandikota, David BauICLR 2026 · 3 citations
- Reasoning with Latent Thoughts: On the Power of Looped TransformersNikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar et al.ICLR 2025
- Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable TasksRahul Ramesh, Ekdeep Singh Lubana, Mikail Khona, Robert P. Dick et al.ICML 2024 · 17 citations
- Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of GeneralizationBoshi Wang, Xiang Yue, Yu Su, Huan SunNeurIPS 2024 · 48 citations
