📊 statistics

The Effect of Mini-Batch Noise on the Implicit Bias of Adam

This paper introduces a theoretical framework demonstrating that mini-batch noise fundamentally alters the implicit bias of Adam's momentum hyperparameters, revealing that while the standard $(\beta_1, 0.999)$ configuration is optimal for small batches, larger batches require adjusting $\beta_1$ closer to $\beta_2$ to avoid anti-regularization and improve generalization.

Original authors: Matias D. Cattaneo, Boris Shigida

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Matias D. Cattaneo, Boris Shigida

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 trying to teach a student (an AI model) to solve a complex puzzle. You have two main tools to help them learn:

The "Memory" Tool (Momentum): This is like a student who remembers their past mistakes and smooths out their learning path. If they stumbled yesterday, they adjust their steps today to avoid stumbling again. In the paper, this is controlled by a setting called $\beta_1$ .
The "Noise" Tool (Mini-batch Size): This is like how many practice problems the student sees at once.
- Small Batch (High Noise): They see only a few problems at a time. The feedback is "noisy" or jumpy because it's based on a tiny, unrepresentative sample.
- Large Batch (Low Noise): They see thousands of problems at once. The feedback is smooth, clear, and very accurate.

The paper investigates a third tool, $\beta_2$ , which controls how much the student trusts their recent history of mistakes versus their distant history.

The Big Discovery: The "Goldilocks" Switch

For years, the standard advice for training AI has been to set the "Memory" ( $\beta_1$ ) to 0.9 and the "Recent History" ( $\beta_2$ ) to 0.999. This means: "Trust your recent history way more than your older history."

The authors of this paper discovered that this standard advice is only half-right. It depends entirely on how much "noise" (how small your batch size) you are using.

Here is the twist they found:

1. The "Small Batch" Scenario (High Noise)

The Situation: You are giving the student very few practice problems at a time. The feedback is jumpy and chaotic.
The Problem: The student's "Memory" tool (the default setting) is actually making things worse. It's holding onto old, noisy information that confuses the student, making them get stuck in "sharp" corners of the solution space (places where the answer is fragile and changes easily).
The Fix: In this noisy environment, you need to increase $\beta_2$ (trust the recent history even more).
The Analogy: Imagine walking through a foggy forest with a shaky map. If you rely too much on your memory of where you were 10 minutes ago (which might have been a wrong turn), you'll get lost. You need to rely heavily on what you see right in front of your feet right now. The paper shows that in high-noise settings, the standard setting ( $\beta_1 \ll \beta_2$ ) is actually the correct one because the "noise" helps cancel out the bad effects of the memory.

2. The "Large Batch" Scenario (Low Noise)

The Situation: You are giving the student thousands of practice problems at once. The feedback is crystal clear and smooth.
The Problem: Now, the "Memory" tool is acting like a stubborn habit. Because the feedback is so clear, the student's memory of past steps starts to push them toward "sharp" corners again, but this time, increasing $\beta_2$ makes it worse.
The Fix: In this clear environment, you should make $\beta_1$ and $\beta_2$ equal (e.g., both 0.9).
The Analogy: Now imagine walking through a sunny, open field with a perfect map. If you rely too much on your recent steps (high $\beta_2$ ) compared to your older steps, you start zig-zagging unnecessarily. Instead, you want a balanced memory where your recent steps and your older steps work together in harmony. The paper predicts that when the data is large and clean, setting $\beta_1 \approx \beta_2$ leads to a "flatter," more stable solution that generalizes better.

The "Switching Point"

The paper calculates a specific "tipping point" based on the size of your data batch.

Below the tipping point (Small Batches): The "Noisy" rules apply. Keep $\beta_1$ low and $\beta_2$ high (the default).
Above the tipping point (Large Batches): The "Clear" rules apply. Bring $\beta_1$ and $\beta_2$ closer together.

Why Does This Matter? "Sharpness" vs. "Flatness"

The authors use a metaphor of terrain:

Sharp Minima: Imagine a needle standing on its tip. It's a solution that works perfectly for the specific practice problems you saw, but if you change the problem slightly (like a new test question), the needle falls over. This is bad for generalization.
Flat Minima: Imagine a wide, flat valley. You can walk around a bit, and you stay in the valley. This solution is robust and works well even when the problem changes slightly.

The paper argues that the interaction between your "Memory" settings and your "Batch Size" secretly pushes the student toward either the needle (sharp) or the valley (flat).

In small batches, the noise naturally pushes you toward the valley, but only if you use the standard settings.
In large batches, the lack of noise means the memory settings take over. If you don't adjust them (by making $\beta_1$ and $\beta_2$ similar), the memory pushes you back toward the needle.

Summary of the Paper's Claims

The Default isn't Universal: The famous setting of $\beta_1=0.9$ and $\beta_2=0.999$ is great for small batches but suboptimal for very large batches.
The Reversal: As you increase the batch size, the "best" relationship between $\beta_1$ $β_{1}$ and $\beta_2$ $β_{2}$ flips.
- Small Batch: Keep $\beta_1$ much smaller than $\beta_2$ .
- Large Batch: Make $\beta_1$ and $\beta_2$ roughly equal.
The Evidence: They proved this mathematically by analyzing how the "noise" in the data interacts with the "memory" of the optimizer. They also tested this on language models (like Llama 3) and found that the validation performance (how well the model does on new data) followed their predictions: the "best" settings changed as the batch size grew.

In short: If you are training with small chunks of data, stick to the defaults. If you are training with massive chunks of data, you should tweak your settings to balance your memory more evenly.

Problem Statement
Adam and AdamW are standard optimizers in deep learning, yet the selection of their momentum hyperparameters, $\beta_1$ and $\beta_2$ , often lacks a principled basis beyond convention (typically $\beta_1=0.9, \beta_2=0.999$ ). While recent large-scale training practices sometimes lower $\beta_2$ (e.g., to 0.95) to improve stability, the interaction between these hyperparameters, batch size ( $b$ ), and the resulting implicit bias toward loss landscape sharpness remains poorly understood. This is critical because loss sharpness serves as a mechanistic proxy for generalization in multi-epoch training and influences post-training performance (e.g., quantization degradation and catastrophic forgetting) even in single-epoch pretraining. The central question is whether the preferred relationship between $\beta_1$ , $\beta_2$ , and batch size changes across different noise regimes.

Methodology
The authors develop a theoretical framework to analyze the implicit bias of Adam by isolating the effects of mini-batch noise and memory. The approach proceeds in three steps:

Memory Removal: Utilizing a technique from Cattaneo and Shigida [10], the authors convert the memory-dependent Adam iteration into a memoryless iteration with an explicit correction term. This allows the dynamics to be approximated as a function of the current state only, with a global error bound of $O(\eta^2)$ .
Noise Expansion and Averaging: They apply this framework to mini-batch Adam, expanding the correction term up to second-order monomials in noise derivatives. By averaging over all permutations of samples (without-replacement sampling), they derive an explicit correction term that separates full-batch effects from mini-batch noise effects.
Sharpness Proxy Interpretation: The dominant terms of the resulting correction are interpreted through a sharpness proxy. Specifically, the analysis focuses on how the coefficients of these terms depend on $\beta_1$ , $\beta_2$ , and the "simple noise scale" $B_{simple}$ (related to the critical batch size). The authors analyze the monotonicity of the sharpness-bias coefficient $C_{total}$ as a function of the betas under two regimes: fixed $\beta_1$ with varying $\beta_2$ , and fixed $\beta_2$ with varying $\beta_1$ .

Key Contributions and Results
The paper provides the first theoretical explanation for a batch-size-dependent reversal in the preferred choice of Adam's betas, driven by an implicit sharpness-bias mechanism:

Reversal of Monotonicity: The relationship between $\beta_2$ $β_{2}$ and sharpness penalization is not monotonic across all batch sizes.
- High-Noise, Small-Batch Regime ( $b \ll 2B_{simple}$ ): Mini-batch noise reverses the monotonicity. In this regime, a larger $\beta_2$ compensates for the anti-regularizing memory term. Consequently, the default ordering $\beta_1 \ll \beta_2$ (e.g., $0.9$ vs $0.999$) is theoretically well-motivated to improve generalization and reduce model sensitivity.
- Low-Noise, Large-Batch Regime ( $b \gg 2B_{simple}$ ): Increasing $\beta_2$ strengthens the anti-regularization induced by memory. Here, choosing $\beta_1$ close to $\beta_2$ (e.g., $\beta_1 = \beta_2 = 0.9$ ) is predicted to be superior for generalization.
Symmetric Effect for $\beta_1$ : A similar reversal occurs when sweeping $\beta_1$ with $\beta_2$ fixed at 0.999. For large batches, larger $\beta_1$ is preferred; for small batches, smaller $\beta_1$ is preferred to maintain the $\beta_1 \ll \beta_2$ structure.
Transition Scale: The transition between these regimes occurs at a batch size scale controlled by $2B_{simple}$ , which is closely related to the critical batch size.
Empirical Validation: Experiments on a small Transformer overfitting on WikiText-2 and a 280M-parameter Llama3 pretraining run on DCLM support these qualitative predictions. The results show that in small-batch settings, increasing $\beta_2$ reduces validation loss and sharpness, whereas in large-batch settings, this trend reverses, with lower $\beta_2$ (or $\beta_1 \approx \beta_2$ ) yielding better performance.

Significance
The paper claims to provide the first theoretical argument based on generalization that explains why the optimal relationship between Adam's momentum parameters depends on the batch size. It challenges the notion that default hyperparameters are universally optimal, suggesting instead that the "sweet spot" for $\beta_1$ and $\beta_2$ is determined by the interplay between memory-induced anti-regularization and mini-batch noise. The findings offer a rule of thumb for practitioners: use the default $\beta_1 \ll \beta_2$ for small batches and consider $\beta_1 \approx \beta_2$ for very large batches to optimize for generalization and model sensitivity. The theory aligns with substantial empirical work but provides a mechanistic derivation rooted in implicit bias and loss landscape sharpness.