🔬 condensed matter

Tunneling the Loss Landscape: Bypassing Memorization with Monte Carlo Parameter Swapping

This paper interprets the "grokking" phenomenon in neural networks as a glassy relaxation process characterized by kinetic arrest and low parameter mobility, and proposes a State-Aware Monte Carlo Parameter Swapping (SAM-Swap) optimization method that accelerates generalization by introducing random exploration to bypass memorization states.

Original authors: Lai Shun Chan, Xiaotian Zhang, Yue Shang, Ge Zhang, Entao Yang

Published 2026-08-04

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Lai Shun Chan, Xiaotian Zhang, Yue Shang, Ge Zhang, Entao Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are watching a student cram for a math test. At first, they are just memorizing the answers to every single practice problem. They get a perfect score on the practice sheet, but if you ask them a slightly different question, they have no idea what to do. This is "memorization." Then, suddenly, after a long period of staring blankly at the same problems, something clicks. The student stops just reciting answers and starts understanding the logic behind them. They can now solve brand-new problems they've never seen before. This sudden shift from rote memorization to true understanding is a mystery that has puzzled scientists studying artificial intelligence.

In the world of AI, this phenomenon is called "grokking." It happens when a computer model trains on a dataset, gets perfect at memorizing the data, but fails to generalize to new data for a very long time. Then, abruptly, it figures out the pattern and becomes brilliant. Scientists have wondered: Why does the model get stuck in this "memorization mode" for so long? Is it just slow, or is it trapped? To answer this, researchers are borrowing ideas from physics, specifically the study of "glass." You know how glass is a liquid that has cooled down so fast it gets stuck in a rigid, disordered state? It's not quite a solid, but it can't flow either. Scientists think AI models might get stuck in a similar "glassy" state during training, where they are frozen in place, unable to find the better solution even though it's right there.

This new paper dives deep into that idea. The researchers, led by Lai Shun Chan and colleagues, wanted to prove that grokking is indeed a form of "glassy" behavior and, more importantly, to find a way to break the model out of that frozen state. They didn't just watch the model; they built a special toolkit to measure exactly how the model's internal "brain" (its parameters) was moving. They discovered that during the long memorization phase, the model's movements become incredibly slow, repetitive, and stuck in a narrow path, much like a person trying to walk through a crowded room who is forced to shuffle in a straight line without ever turning.

To fix this, the team invented a clever trick called "SAM-Swap." Inspired by how physicists help glassy materials relax by swapping the positions of particles, they created a method that randomly swaps the values of the model's internal numbers. It sounds like it would mess everything up, but instead, it acts like a gentle shake that wakes the model up. This simple swap allows the model to escape its frozen state and find the "generalization" solution much faster. The paper suggests that by understanding the geometry of how these models move, we can stop them from getting stuck and help them learn faster, turning a long, frustrating wait into a quick "aha!" moment.

The Story of the Frozen Brain

The Mystery of the "Grokking" Pause
Imagine you are training a robot to solve a puzzle. You show it thousands of examples. At first, the robot is a superstar at memorizing the exact answers. It gets 100% on your practice test. But when you give it a new puzzle it hasn't seen, it fails miserably. It keeps failing for a long, long time. Then, suddenly, after thousands of more tries, it figures out the rule and gets 100% on the new puzzles too. This is "grokking."

For a while, scientists thought this was just a quirk of how the robot learned. But a new theory suggests the robot gets "frozen." Think of it like a river that freezes over in winter. The water (the learning process) stops flowing. The robot is stuck in a "glassy" state. It's not that the solution doesn't exist; it's that the robot can't move its internal parts enough to find it.

The Three Tools to Measure the Freeze
To prove this, the authors built a three-part toolkit to watch the robot's brain in action. They wanted to know: Is the robot moving? Is it stuck in a loop? And is it exploring new paths?

Parameter Mobility (PM): The "Step Size" Meter.
Imagine the robot is walking. In the beginning, it takes big, confident steps as it learns. But once it hits the memorization phase, the authors found that the robot's steps become tiny, almost invisible. It's like the robot is standing on one foot, barely twitching. The paper shows that the distance the robot's internal numbers move drops by huge amounts (from about 1 to 0.0001). It's physically stuck.
Replica Correlation (RC): The "Twin Test."
This is the most fun part. The researchers created "twins" of the robot. They started two identical robots from the exact same spot and gave them tiny, random nudges. If the robots were free to explore, they would quickly go in different directions. But during the memorization phase, the twins stayed glued together. Even with nudges, they followed the exact same path. This "history dependence" means the robot is so stuck in its current state that it can't even imagine a different way to solve the problem. It's like two people trying to walk through a narrow tunnel; no matter how they try to wiggle, they are forced to stay in the same line.
Fractal Dimension (FD): The "Wiggle" Detector.
This measures the shape of the robot's path. If the robot is exploring, its path is messy, winding, and full of turns (high dimension). But during the grokking pause, the path becomes a straight, boring line. The authors found the "Fractal Dimension" dropped to almost 1. This means the robot is moving in a straight, "channel-like" tunnel, unable to turn left or right. It's trapped in a narrow corridor.

The "Glass" Theory Confirmed
The paper suggests that this combination—tiny steps, twins that can't separate, and a straight-line path—is the signature of a "glassy" system. The robot isn't just slow; it's kinetically arrested. It's trapped in a state where it can't escape, even though the "generalization" solution (the exit) is right there. The authors argue that the training process cools the system down too fast, locking it into this frozen state before it has a chance to relax into the better solution.

The Magic "Swap" Trick
So, how do you unfreeze a glassy robot? In physics, scientists use a method called "Swap Monte Carlo" to help glass relax. They imagine swapping the positions of particles to help them find a better arrangement. The authors applied this to the AI model.

They created a new tool called SAM-Swap. Here's how it works:

They watch the robot's path.
When they see the path becoming too straight (the Fractal Dimension drops below 1.1), they know the robot is stuck.
They then randomly swap the values of the robot's internal numbers within the same layer.

You might think, "Wait, if you swap the numbers, won't you destroy what the robot learned?" Surprisingly, no. The authors found that this swap doesn't break the model; it acts like a gentle shake. It breaks the "channel" the robot is stuck in and lets it explore new paths.

The Results: A Faster "Aha!" Moment
The results were dramatic. In the standard training (using a method called AdamW), the robot took about 3,000 epochs (training cycles) to finally generalize. But with the SAM-Swap trick, the robot generalized in just 650 epochs. That's a massive speedup!

The paper also compared this to other methods, like adding random noise (shaking the robot with static electricity). While adding noise helped a bit, the structured "swap" was more effective. The key finding is that to escape the glassy state, the robot needs to do something that looks like "diffusion"—moving randomly and exploring, rather than just shuffling in a straight line.

What This Means
This paper doesn't just say "grokking is weird." It gives us a way to measure why it happens and a tool to fix it. It suggests that the timing of when a model learns isn't just about the math of the problem; it's about the journey the model takes. If the journey gets too straight and narrow, the model gets stuck. By introducing a little bit of organized chaos (the swap), we can help the model break free and learn faster.

The authors are careful to note that this was tested on a specific math task (modular arithmetic) and a specific type of model. They suggest that while this works well here, we need to see if it works for bigger, more complex models like those used for language or vision. But the core idea—that we can "tunnel" through the loss landscape by understanding the geometry of movement—is a powerful new way to think about how AI learns.

In short, the paper tells us that sometimes, to learn faster, you don't need to push harder. You just need to shake things up a little bit and let the model take a different path.

Technical Summary: Tunneling the Loss Landscape with Monte Carlo Parameter Swapping

Problem Statement

The paper addresses the phenomenon of grokking in neural network training, where models achieve perfect training accuracy (memorization) but fail to generalize to test data for a prolonged period before abruptly generalizing. While previous theories have framed this through weight norms, circuit competitions, or representation alignment, a recent statistical physics perspective suggests grokking is a form of computational glass relaxation. This theory posits that training involves a "fast cooling" into a glassy, memorized state followed by a "slow relaxation" toward a generalizing equilibrium.

However, this glass analogy has remained largely theoretical and macroscopic. There is a lack of direct empirical validation regarding the training dynamics that characterize this transition. Specifically, it is unclear how the optimization trajectory behaves during the memorization phase and whether interventions inspired by glass physics (such as particle swapping) can accelerate the transition to generalization.

Methodology

1. Experimental Setup

The authors replicate the setup from previous glass theory work (Zhang et al., 2025), training a 1-layer transformer on the arithmetic task $(x^2 + y) \mod 67$ . Models are trained for 10,000 epochs using AdamW. To isolate dynamical effects, the authors use full-batch training without dropout, controlling noise levels explicitly.

2. Three-Component Framework for Dynamics

To characterize the training dynamics, the authors introduce a framework combining three metrics derived from glassy dynamics:

Parameter Mobility (PM): Measures the net Euclidean displacement of weights over a training window. It quantifies the magnitude of stepwise parameter motion.
Replica Correlation (RC): Measures the similarity between independently evolving training trajectories (replicas) branched from the same initial state at a waiting time $t_w$ . High RC indicates that perturbations do not lead to divergent paths, suggesting strong history dependence and confinement.
Fractal Dimension (FD): Characterizes the geometry of the training trajectory in weight space. It relates the accumulated path length to the net displacement.
- $D_f \approx 1$ : Indicates straight, ballistic, or channel-like motion (low-dimensional).
- $D_f > 1$ : Indicates tortuous, diffusive, or indirect motion.

3. Intervention: State-Aware Monte Carlo Parameter Swapping (SAM-Swap)

Inspired by the swap Monte Carlo algorithm used to accelerate equilibration in glass-forming systems, the authors propose SAM-Swap.

Mechanism: When the Fractal Dimension ( $D_f$ ) drops below a threshold (indicating entry into a constrained, low-dimensional regime), the algorithm randomly swaps parameter values within the same layer.
Constraints: To prevent catastrophic failure, swaps are rejected if the new training loss exceeds a cap (set to the initial training loss).
Rationale: This structured perturbation aims to disrupt the "kinetic arrest" of the memorized state without destroying learned features, effectively "tunneling" through the loss landscape.

Key Results

1. Characterization of Grokking as Kinetic Arrest

The authors' framework reveals that standard optimization (AdamW) exhibits clear signatures of glass dynamics during the grokking phase:

Collapsed Mobility: Once memorization is reached, PM drops by orders of magnitude ( $\sim 1$ to $\sim 10^{-4}$ ), indicating the network becomes nearly immobile despite poor test performance.
Two-Step Relaxation & Aging: RC exhibits a two-step decay. During memorization, RC plateaus, indicating trajectories are confined to a restricted region. Furthermore, RC shows aging behavior: replicas branched later (larger $t_w$ ) remain highly correlated for longer periods, demonstrating that the system becomes increasingly resistant to perturbations the longer it stays in the memorized state.
Channel-Like Geometry: FD collapses to $\approx 1$ after memorization, suggesting the optimization trajectory becomes effectively one-dimensional and "channel-like," preventing exploration of alternative solutions.

2. Effectiveness of SAM-Swap

The introduction of SAM-Swap significantly accelerates generalization:

Reduced Grokking Time: Models using SAM-Swap generalize at approximately $t \approx 650$ epochs, compared to $t \approx 3000$ epochs for the baseline AdamW model.
Restoration of Dynamics: SAM-Swap restores PM, reduces RC (breaking the history dependence), and increases FD (making trajectories more tortuous and exploratory).
Comparison with Noise: The authors compare SAM-Swap against large-amplitude additive Gaussian noise. Both interventions accelerate generalization, suggesting that random exploration in the parameter space (diffusion) is the key mechanism for escaping the kinetic arrest, rather than the specific structure of the perturbation.

Significance and Claims

The paper claims to provide the first direct dynamical diagnosis of grokking, moving beyond macroscopic loss analysis to characterize the microscopic behavior of the optimization trajectory.

Validation of Glass Analogy: The results provide quantitative evidence that grokking corresponds to a glass-like kinetic arrest in parameter space, characterized by collapsed mobility, strong history dependence (aging), and geometric confinement (channel-like paths).
Intervention via Disruption: The work demonstrates that disrupting constrained trajectories via structured perturbations (SAM-Swap) can bypass the memorization phase. This supports the view that generalization delay is not solely due to the scarcity of generalizing solutions, but rather the inaccessibility of these solutions due to dynamical constraints.
Unified Framework: The three-component framework (PM, RC, FD) offers a practical tool to distinguish between general slow motion and true kinetic arrest. It clarifies that low mobility alone is insufficient to explain grokking; the combination of low mobility, high history dependence (RC), and low-dimensional geometry (FD) is required to identify the arrested state.

Limitations

The authors acknowledge that their Fractal Dimension estimation relies on a fixed temporal window (20 epochs), which may obscure multiscale dynamics or transients outside this window. They note that the current formulation is a local diagnostic rather than a full multiscale characterization. Additionally, the empirical validation is currently limited to modular arithmetic tasks, and the generalizability to larger vision or language models remains an open question.