Provable Scaling Laws of Feature Emergence from Learning Dynamics of Grokking

Introduction: What is Grokking? An Intuitive Analogy

Before diving into the technical intricacies, let's establish a common understanding. Imagine you're preparing for a particularly challenging exam (say, the CFA Level III Exam).

1. Initial Phase (Memorization): You start by tackling practice questions without fully grasping the underlying principles. Your only recourse is to rote memorize various question types and their solutions. While you achieve 100% accuracy on questions you've already encountered, you find yourself utterly stumped by a new one. This scenario perfectly illustrates memorization: performing flawlessly on the "training set" but failing miserably on the "test set."
2. Epiphany (Generalization): Undeterred, you persist in solving problems and synthesizing concepts. At a certain juncture, everything suddenly clicks; you achieve a holistic understanding of the economic and financial principles underpinning all the knowledge points. From then on, regardless of how the questions are framed, you can apply your understanding to solve them with ease. This "epiphany"—the transition from rote memorization to genuine comprehension—is precisely what we call Grokking.

The core objective of this paper is to leverage rigorous mathematical language to unravel what truly happens inside neural networks, leading to this fascinating moment of "grokking."

The Core Framework: Li2's Three-Stage Learning Journey

The Li2 framework, proposed by the authors, dissects the entire learning process into three sequential stages. We can visualize this process with the following diagram:

• Three Learning Stages of the Li2 Framework

These three stages are:

1. Stage I: Lazy Learning (Memorization): The model takes shortcuts, opting for the simplest way to memorize training answers.
2. Stage II: Independent Feature Learning: The model is compelled to "think," with each neuron independently learning truly valuable patterns (features) within the data.
3. Stage III: Interactive Feature Learning: Neurons begin to "collaborate as a team," communicating and complementing each other to form a more comprehensive and efficient understanding of the problem.

Now, let's delve into what transpires during each stage, step by step.

Stage I: Lazy Learning - The "Lazy" Memorizer

This marks the initial phase of training. Picture a simple two-layer neural network as a two-tier reporting team: junior staff (hidden layer W) process raw data, then report their findings to the manager (output layer V), who makes the final decision.

• What Happens?
  • Initially, the junior staff (W) are novices; their data processing methods are random and disorganized.
  • The manager (V) is exceptionally "smart" yet "lazy." Instead of training the staff, they opt to directly memorize the correspondence between each staff member's random reports and the correct answers. In finance, this is akin to performing ridge regression, where one finds coefficients for an existing set of (even random) explanatory variables to minimize the error between predicted and actual values.
  • Result: The model performs progressively better on the training data (training loss decreases, accuracy increases), but this is a false sense of prosperity. It has merely memorized the answers, without acquiring any generalizable knowledge.
• Key Player: Weight Decay
  • This constitutes the first crucial insight of this paper. You can conceptualize weight decay as an "anti-cheating mechanism" or a "cost constraint." It penalizes overly complex, large-magnitude weights.
  • Without weight decay: The manager (V) can effortlessly and perfectly memorize all training answers. Consequently, they would perceive, "I am already perfect; there's nothing left to improve," leading to a zero "feedback signal" (i.e., gradient ) transmitted to the junior staff (W). Without effective feedback, the staff would never grow, and the model would perpetually remain in the memorization phase.
  • With weight decay: The manager (V) cannot perfectly cheat; they will always make a slight error. This very "imperfection" compels them to transmit a meaningful, non-zero feedback signal  to the junior staff, instructing them: "Among your random reports, some elements correlate with the correct answers; you must extract this useful information!"
  • The following figures vividly illustrate this point:
    • With Weight Decay (Left Three Figures)
    • Without Weight Decay (Right Three Figures)

The output layer V quickly learns and perfectly fits, but the feedback signal  rapidly decays to zero, and the hidden layer W remains static. The model consistently stays in the memorization phase, unable to generalize.

In the early training phase, the output layer V learns first (orange line moves first), followed by the hidden layer W (blue line moves later). The critical feedback signal  (gray dashed line) first rises then falls, signifying the transmission of effective information. Ultimately, the model successfully groks (the red line for test accuracy eventually surpasses).

Key Takeaways from Stage I: The first step of Grokking necessitates an initial memorization phase. This stage, seemingly useless overfitting, plays the crucial role of creating a meaningful "learning signal" for the subsequent phase of genuine learning, with weight decay acting as the spark that ignites this signal.

Stage II: Independent Feature Learning - The Onset of "Epiphany"

Now, the junior staff (W) have finally received an effective feedback signal  from the manager (V). The true curtain of learning rises.

• Core Concept: Energy Function 
  • This represents the second, and most crucial, insight of this paper. While this concept might seem abstract, let's use another analogy.
  • Imagine a topographic map of undulating mountains; this is the "energy landscape." Each neuron (junior staff) is like a climber, whose goal is to find and ascend the highest possible peak.
  • What Do the Peaks Represent? They represent meaningful, generalizable patterns (features) within the data. For instance, in the modular addition task (e.g., 7 + 8 mod 15 = 0) studied in the paper, a "peak" might correspond to an underlying mathematical rule such as "identifying parity" or "identifying whether it's a multiple of 3."
  • What is the Climbing Process? It is the neuron's learning process. It will climb along the steepest ascent (gradient ascent), ultimately settling at the summit of a particular peak.
  • Why is it called "independent" learning? Because in this stage, each climber (neuron) focuses solely on their own path, independently seeking their own peak, without interaction among themselves yet.
• How Does Data Volume Affect "Grokking"?
  • The paper indicates that the morphology of this "energy landscape" is not fixed; it is shaped by the training data.
  • When Data is Abundant: Generalizable "regularity peaks" will stand tall, while "memorization peaks" (corresponding to rote learning shortcuts) will be relatively low. Climbers will naturally choose to ascend the higher peaks representing true underlying rules.
  • When Data is Scarce: The landscape becomes rugged, and "memorization peaks" might appear higher and easier to climb than "regularity peaks." Climbers could then be misled into learning incorrect, non-generalizable patterns.
  • The Critical Point: This explains why Grokking often occurs at a "critical data volume." Only when the data volume surpasses a threshold, sufficient to shape a "favorable" energy landscape (where regularity peaks are sufficiently prominent), will "grokking" occur.
  • The phase transition diagram below clearly illustrates this: The x-axis is problem size M, and the y-axis is the required training data ratio. The color represents test accuracy (red = generalization, purple = memorization). A clear boundary line (black dashed line) exists, with the generalization region above the line and the memorization region below.
    • Generalization/Memorization Phase Transition for Different Tasks
• How Does the Learning Rate Affect "Grokking"?
  • The learning rate is analogous to the climber's "step size."
  • Small Learning Rate: With small steps, the climber can steadily ascend nearby peaks, and even if this peak isn't the highest, it might be a good "regularity peak."
  • Large Learning Rate: With large steps, the climber might leap over several small hills, directly jumping to the globally highest peak. However, if data is insufficient, this highest peak might ironically be a "memorization peak."
  • This accounts for another empirical observation: when the data volume is in the critical zone, a smaller learning rate is more conducive to Grokking, as it helps the model stabilize within the "basin of attraction" of the generalization solution.
    • Generalization/Non-Generalization Phase Transition at Different Learning Rates

Key Takeaways from Stage II: Genuine learning occurs in Stage II. Its essence lies in neurons performing "hill-climbing" (optimization) on the "energy landscape" shaped by data. Whether "grokking" (generalization) occurs depends on the landscape's morphology, which in turn is determined by data volume. The learning rate dictates the "hill-climbing" strategy, influencing which peak the model ultimately settles on.

Stage III: Interactive Feature Learning - Efficient "Team Collaboration"

In Stage II, each staff member independently acquired certain skills. However, an efficient team requires collaboration.

• What Happens?
  1. Avoiding Redundancy (Repulsion of Similar Features): If two staff members (neurons) coincidentally acquire identical skills (e.g., both learn to identify parity), the system generates a "repulsive force," prompting one of them to learn something else equally useful (e.g., identifying multiples of 5). This resembles team division of labor, preventing wasted human resources and enhancing overall "feature diversity."
  2. Filling Gaps (Top-Down Modulation): After synthesizing the preliminary learning outcomes of all staff, the manager (V) might realize, "Our team now knows how to identify parity and multiples of 5, but no one yet knows how to identify multiples of 3." Consequently, they would issue new "feedback signals." These signals would dynamically alter the "energy landscape," marking the path to the "identifying multiples of 3" peak more clearly, encouraging the remaining staff to conquer this stronghold.
• Key Player: Muon Optimizer
  • Standard optimizers like Adam might encounter difficulties in Stage III, as they treat all staff uniformly, potentially leading to homogenized skill acquisition among team members.
  • The Muon optimizer acts like a smarter project manager. It actively identifies which skills have already been mastered, then suppresses the motivation of other staff to learn these existing skills, while simultaneously enhancing their drive to explore entirely new ones.
  • Result: Teams utilizing Muon can acquire the entire skill set required to solve a problem more rapidly and comprehensively, thereby achieving superior generalization performance.
    • Comparison of Test Loss for Adam and Muon Optimizers in Modular Addition Task

(In this figure, the orange Muon optimizer generally shows lower test loss than the blue Adam optimizer across different settings, indicating its stronger generalization capability.)

Key Takeaways from Stage III: The final stage of learning involves the integration and refinement of features. The model enhances feature diversity and completeness through internal interaction, ultimately forming an efficient, comprehensive solution. The advantages of advanced optimizers like Muon are particularly evident in this stage.

Conclusion and Insights

Having journeyed through the core interpretation of this paper, let's now summarize:

1. Grokking is not magic, but an ordered three-step process: From "lazy memorization" to "independent epiphany" and finally to "collaborative mastery."
2. Memorization is a prerequisite for generalization: The seemingly useless Stage I (overfitting) is a necessary step; it creates the essential "driving signal" for genuine learning.
3. Weight Decay is the igniter: Without it, the learning engine cannot start.
4. The essence of learning is "hill-climbing": Neurons search for "peaks" representing universal principles on the "energy landscape" shaped by data.
5. Data volume and learning rate are key regulating valves: They jointly determine whether the model ultimately learns "rules" or "shortcuts."
6. Team collaboration improves efficiency: The model ensures that the learned knowledge is comprehensive and non-redundant through internal interaction and advanced optimizers.

The value of this paper lies in its pioneering effort to provide a computable, verifiable mathematical framework for the enigmatic phenomenon of Grokking, derived from the first principles of gradient dynamics. It unifies the roles of hyperparameters (such as weight decay, learning rate, sample size), the impact of data structure, and the behavior of optimizers under a coherent theoretical umbrella, profoundly deepening our understanding of the deep learning "black box."