TLDR: New research reveals that the implicit bias of the Adam optimizer fundamentally changes when moving from full-batch to mini-batch training. Unlike full-batch Adam, which favors ℓ∞-max-margin solutions, mini-batch Adam’s bias becomes data-dependent and can shift towards â„“2-max-margin solutions. The study introduces a proxy algorithm and a fixed-point formulation to characterize this behavior. In contrast, the Signum optimizer is shown to retain its ℓ∞-max-margin bias across various batch sizes, highlighting a key difference between adaptive optimizers.
The Adam optimizer, a cornerstone in the world of deep learning, has long been celebrated for its efficiency and widespread applicability. Despite its practical success, a complete theoretical understanding of its behavior, particularly concerning its ‘implicit bias,’ has remained a complex challenge. Implicit bias refers to the tendency of optimization algorithms to favor solutions with specific structural properties, even without explicit regularization. This phenomenon is crucial for explaining why over-parameterized models generalize well.
Previous research into Adam’s implicit bias primarily focused on the ‘full-batch’ regime, where the entire dataset is used to compute gradients at each step. These studies suggested that full-batch Adam tends to converge towards solutions aligned with an ℓ∞-geometry, often referred to as the ℓ∞-max-margin classifier. This behavior is similar to Sign Gradient Descent (SignGD), which is also known for its ℓ∞-max-margin direction.
A New Look at Mini-Batch Adam
However, modern deep learning predominantly relies on ‘mini-batch’ training, where only a small subset of data samples is used for each gradient update. This paper, titled “Implicit Bias of Per-sample Adam on Separable Data: Departure from the Full-batch Regime”, delves into this crucial, yet under-explored, mini-batch setting. The authors, Beomhan Baek, Minhak Song, and Chulhee Yun, pose a fundamental question: Does Adam’s characteristic ℓ∞-bias persist under mini-batch training?
The surprising answer, as revealed by their research, is no. The implicit bias of Adam can significantly deviate from its full-batch behavior when trained with mini-batches. Through experiments, the researchers observed that while full-batch Adam converges to the ℓ∞-max-margin direction on Gaussian data, mini-batch Adam (specifically, variants with a batch size of one, termed ‘incremental Adam’) converges closer to the â„“2-max-margin direction. This marks the first theoretical evidence demonstrating that Adam’s implicit bias is fundamentally altered in the mini-batch context.
Understanding the Shift in Bias
To illustrate this departure, the researchers constructed a specific class of structured datasets called ‘Generalized Rademacher (GR) data.’ On these datasets, incremental Adam was provably shown to converge to the â„“2-max-margin classifier. This stands in stark contrast to the ℓ∞-max-margin bias observed with full-batch Adam. The GR data effectively eliminates the coordinate-wise adaptivity of Adam’s preconditioner, simplifying the dynamics and highlighting the core difference.
For more general datasets, the team developed a ‘uniform-averaging proxy’ algorithm, dubbed ‘AdamProxy,’ which captures the limiting behavior of incremental Adam as a specific momentum parameter (β2) approaches one. They characterized its convergence direction using a novel data-dependent dual fixed-point formulation. This framework allows for predicting the algorithm’s preferred solution, which, depending on the dataset, can be different from both the conventional â„“2- and ℓ∞-max-margin solutions. Numerical experiments confirmed that this fixed-point solution accurately predicts the limiting directions for both AdamProxy and incremental Adam on various datasets, including Gaussian and shifted-diagonal data.
Signum: An Invariant Alternative
In light of Adam’s changing bias, the researchers also investigated other optimization algorithms. They found that Signum (a variant of SignSGD with momentum) behaves differently. Unlike Adam, Signum consistently maintains its bias toward the ℓ∞-max-margin solution, regardless of the batch size, provided its momentum parameter is sufficiently close to one. This invariance makes Signum a compelling alternative for scenarios where a consistent ℓ∞-geometry preference is desired across different batching schemes.
Also Read:
- Optimizing Federated Learning in Wireless Networks: A New Approach to Bias and Variance
- Understanding the Parameter Cost of Robustness in Neural Networks
Implications and Future Directions
The findings underscore that the implicit bias of Adam is critically dependent on both the batching scheme and the specific characteristics of the dataset. This research offers a potential explanation for the observed ‘Adam–SGD gap,’ where Adam’s performance advantage over SGD is often more pronounced in large-batch training and diminishes with smaller batches. The ℓ∞-adaptivity, often cited as a source of Adam’s strength, may diminish in the mini-batch regime, leading to this gap.
While this work provides significant theoretical insights, the analysis for general datasets currently relies on the asymptotic regime where the β2 momentum parameter approaches one and a batch size of one. Future research aims to extend this framework to finite β2 values, larger batch sizes, and other common sampling schemes like random reshuffling, to provide a more comprehensive understanding of Adam’s intricate behavior in real-world deep learning applications.
