Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing

2024-05-08

AI Summary

This study investigates why Transformer models adopt either reasoning or memorization strategies in compositional generalization tasks. Method: We identify parameter initialization scale as the decisive factor: smaller initialization standard deviations induce low-complexity biases, encouraging models to learn compositional primitives rather than memorizing input-output mappings. We propose a tunable hyperparameter—the “initialization rate” γ—to uniformly characterize layer-wise initialization scaling, and integrate information-flow analysis, representation visualization, complexity-bias quantification, and evaluation across multiple real-world datasets. Contribution/Results: We demonstrate that precise control of initialization scale systematically shifts model behavior from memorization toward structured reasoning, substantially improving out-of-distribution generalization on unseen compositional problems. Our findings provide a principled, initialization-based mechanism for enhancing mathematical reasoning and compositional generalization in large language models.

Technology Category

• Knowledge Representation and Reasoning: Computational Complexity of Reasoning
• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Large Vision Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience
• Web Mining and Content Analysis: Large pretrained models with web data

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 $gamma$ to be a convenient tunable hyper-parameter in common deep learning frameworks, where $1/d_{mathrm{in}}^gamma$ is the standard deviation of parameters of the layer with $d_{mathrm{in}}$ input neurons.

Problem

• Transformer Models
• Unseen Combinations
• Mathematical Functions

Innovation

• Transformer Models
• Initialization Parameters
• Strategy Control