🤖 machine learning

Scaling the Memory of Balanced Adam

This paper proposes that the momentum parameter $\beta$ in balanced Adam should be treated as a memory-scale variable rather than a fixed constant, introducing a refresh rule ( $R_\beta \approx 1000$ ) that significantly improves training robustness across 11 vision and language experiments by aligning the optimizer's statistical memory horizon with the effective learning horizon.

Original authors: Alberto Fernández-Hernández, Cristian Pérez-Corral, Jose I. Mestre, Manuel F. Dolz, Enrique S. Quintana-Ortí

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alberto Fernández-Hernández, Cristian Pérez-Corral, Jose I. Mestre, Manuel F. Dolz, Enrique S. Quintana-Ortí

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

The Big Picture: Tuning the "Memory" of a Learning Robot

Imagine you are teaching a robot to recognize cats, write poems, or solve math problems. To do this, the robot uses a tool called Adam (an "optimizer"). Think of Adam as the robot's internal coach. Every time the robot makes a mistake, the coach looks at the history of past mistakes to decide how to adjust the robot's brain for the next try.

For a long time, this coach had two different "memory settings" (called $\beta_1$ and $\beta_2$ ). One setting remembered the direction of the mistakes, and the other remembered the size of the mistakes.

The Discovery:
Recent research (including this paper) found that you don't actually need two different settings. If you set them to be exactly the same ( $\beta_1 = \beta_2$ ), the robot learns just as well, but the system becomes much simpler. Now, the coach only has one memory knob to turn.

The Problem:
The big question is: What number should we turn that knob to?
Most people just pick a standard number (like 0.9) and stick with it, no matter what. The authors of this paper argue that this is like using the same pair of shoes for a 5-mile walk and a 100-mile marathon. It doesn't make sense because the "distance" of the training changes.

The Core Idea: The "Refresh Count"

The authors propose a new way to think about that memory knob. Instead of thinking of it as a fixed number, they say it should be thought of as a memory horizon.

The Analogy: Imagine the robot is reading a book to learn a language.
- If the robot has a short memory, it only remembers the last few sentences.
- If the robot has a long memory, it remembers the whole chapter.

The paper introduces a concept called the Refresh Count ( $R_\beta$ ). This answers the question: "How many times does the robot completely refresh its memory of the past during the entire training session?"

The authors found that the robot learns best when this "refresh count" stays roughly the same, regardless of how long the training session is.

Short Training Session: The robot needs a short memory (a lower number) so it can refresh its memory about 1,000 times.
Long Training Session: The robot needs a long memory (a higher number) so it still refreshes its memory about 1,000 times.

The Solution: A Simple Rule

The paper tested this idea on 11 different tasks, ranging from teaching a robot to recognize images (like cats and cars) to teaching it to write text (like LLMs).

They discovered a simple "Golden Rule":
Set the memory knob so that the robot refreshes its memory exactly 1,000 times during the useful part of training.

If the training is short (e.g., 6,000 steps), the rule picks a lower number (like 0.82).
If the training is long (e.g., 40,000 steps), the rule picks a higher number (like 0.97).

Why Does This Matter? (The Results)

The authors compared their "Refresh Rule" against the standard "Fixed Rule" (where everyone just uses 0.944).

Average Performance: Both methods were very close in terms of average success. The standard fixed number is actually quite good.
The "Safety Net" (Robustness): This is where the new rule shines.
- Imagine you are driving a car. The "Fixed Rule" is like driving at a speed that works well on most roads, but sometimes you might hit a pothole and crash.
- The "Refresh Rule" is like a smart suspension system. It adjusts to the road length.
- The Result: The new rule reduced the "worst-case" failures by 33%. In other words, it made the training much more reliable. It ensured that every single experiment (all 11 of them) performed almost perfectly, whereas the old method had a few experiments that struggled significantly.

The Takeaway

The paper argues that we shouldn't treat the memory setting in AI training as a static, unchangeable constant. Instead, we should treat it as a scale.

Just as you wouldn't use the same map scale for a city and for a whole country, you shouldn't use the same memory setting for a short training run and a long one. By adjusting the memory based on how long the training lasts (keeping the "refresh count" at 1,000), we make the AI training process more robust and less likely to fail in unexpected ways.

In short: Don't just pick a number and hope for the best. Adjust the robot's memory based on how long the journey is, and you'll get a much smoother ride.

Technical Summary: Scaling the Memory of Balanced Adam

Problem Statement

While recent evidence suggests that the Adam optimizer performs robustly when its momentum parameters are tied ( $\beta_1 = \beta_2 = \beta$ ), reducing it to a single-parameter optimizer, the optimal value for this remaining parameter remains poorly understood. Treating $\beta$ as a fixed, dimensionless constant (e.g., the standard 0.9 or 0.99) is insufficient because the same value represents vastly different dynamical regimes depending on the total duration of training. A value of $\beta$ that is appropriate for a short training run may be ineffective for a long one. The core problem addressed is how to choose the single memory scale of balanced Adam in a way that is aware of the training scale.

Methodology

The authors propose reframing $\beta$ not as a static constant, but as a variable controlling a "statistical memory horizon."

Definitions:
- Memory Horizon ( $H_\beta$ ): Defined as $H_\beta = (1-\beta)^{-1}$ , representing the approximate number of recent optimization steps that dominate the optimizer's internal estimates.
- Effective Learning Horizon ( $T_{ES}$ ): An estimate of the useful duration of training, derived from early-stopping times on validation trajectories.
- Refresh Count ( $R_\beta$ ): A dimensionless quantity defined as $R_\beta = T_{ES} / H_\beta = (1-\beta)T_{ES}$ . This measures how many times Adam "refreshes" its internal statistics during the useful phase of training.
The Refresh Rule:
The central hypothesis is that balanced Adam behaves robustly when the refresh count is kept approximately constant, regardless of the training scale. The proposed rule is:
$R_\beta \approx R_0 \implies \beta = 1 - \frac{R_0}{T_{ES}}$
where $R_0$ is a target constant.
Experimental Setup:
- Scope: 11 experiments spanning vision (ResNet50, ViT-B/16, EfficientNet-B0, Swin-T) and language (Llama60M, NanoGPT, T5-small) models trained on diverse datasets (e.g., ImageNet, CIFAR-100, C4, BookCorpus).
- Protocol: All runs enforce the diagonal constraint $\beta_1 = \beta_2 = \beta$ . A grid of 13 beta values (ranging from 0 to 0.999) was swept.
- Calibration: 8 experiments served as a development set to calibrate $R_0$ , while 3 held-out experiments were used for validation.
- Metric: Performance was measured by the relative validation loss gap compared to the "oracle" $\beta^*$ (the best $\beta$ found in the grid for that specific experiment). Metrics included Mean Gap, Maximum Gap, and Conditional Value at Risk (CVaR) over the worst 25% of cases.

Key Results

The study found that selecting $\beta$ to maintain a constant refresh count significantly improves robustness compared to fixed baselines.

Optimal Target: The development set calibration identified $R_0 = 1000$ as the target that minimizes the maximum gap. This rule selects different $\beta$ values depending on $T_{ES}$ (ranging from 0.822 to 0.968 across the experiments).
Robustness vs. Average Performance:
- The fixed baseline $\beta = 0.944$ (the strongest fixed choice in the grid) achieved the best mean global gap (0.290%).
- However, the refresh rule ( $R_0 = 1000$ ) significantly improved worst-case robustness. It reduced the global maximum validation gap by 33.4% (from 1.328% to 0.885%) and lowered the global CVaR by 32.7% (from 1.059% to 0.713%).
- Crucially, the refresh rule brought all 11 experiments within 1% of their respective validation oracles, whereas the fixed baseline only achieved this for 9 out of 11.
Stability: The method proved robust to moderate errors in estimating $T_{ES}$ . Even with multiplicative noise up to $\sigma=0.20$ in the horizon estimate, the refresh rule maintained a lower global CVaR than the fixed baseline.

Significance and Claims

The paper argues that the remaining hyperparameter of balanced Adam is better understood as a memory-scale variable rather than a fixed constant.

Scale-Aware Perspective: The work provides a "budget-aware perspective on optimizer scaling." It posits that $\beta$ should be viewed as a control for the optimizer's statistical memory clock, which must be matched to the effective learning horizon.
Dynamics over Defaults: The results suggest that treating Adam's momentum as part of the learning dynamics (specifically, the number of statistical refreshes) is superior to treating it as a static default.
Practical Implication: While a fixed $\beta$ (like 0.944) is a strong practical default, the refresh rule offers a principled mechanism to adjust $\beta$ as the training horizon changes, ensuring the optimizer's internal memory is neither too short (causing instability) nor too long (causing slow adaptation) relative to the training duration.

The authors acknowledge limitations, noting that $T_{ES}$ must be estimated (making the rule scale-calibrated rather than zero-shot) and that the specific value $R_0=1000$ is an empirical finding for the tested suite, not necessarily a universal constant. Future work is suggested to disentangle architecture-specific effects and study the interaction between the memory clock and learning rate schedules.