Technical Summary: AdAdaGrad – Adaptive Batch Size Schemes for Adaptive Gradient Methods

1. Problem Statement

In large-scale deep learning, the choice of batch size in minibatch stochastic gradient optimization is a critical hyperparameter that balances computational efficiency against model generalization. While large-batch training is dominant due to hardware advances (GPUs/TPUs), it often leads to a "generalization gap," where models trained with large batches generalize worse than those trained with small batches. This phenomenon is attributed to large-batch training converging to sharp minima in the loss landscape.

Existing strategies to mitigate this gap often rely on heuristic batch-size schedules (e.g., constant sizes or linear warmups) or adaptive learning rate scaling (e.g., AdaScale). However, these approaches frequently lack rigorous theoretical convergence guarantees, particularly for adaptive gradient methods (such as AdaGrad and Adam), which are prevalent in modern deep learning. Furthermore, existing adaptive sampling methods (which dynamically adjust batch sizes based on gradient statistics) have primarily been analyzed for Stochastic Gradient Descent (SGD) under restrictive assumptions (e.g., convexity or fixed step sizes), leaving a gap between theory and practice for adaptive optimizers in nonconvex settings.

2. Methodology

The paper introduces AdAdaGrad and its scalar variant AdAdaGrad-Norm, which integrate adaptive batch size strategies with adaptive gradient updates. The core methodology relies on adaptive sampling methods originally developed for SGD but adapted here for coordinate-wise and scalar adaptive step sizes.

2.1 Adaptive Sampling Tests

The algorithms dynamically increase the batch size $b_k$ during training based on the variance of the stochastic gradients relative to the full gradient. Two primary tests are employed:

1. Norm Test: Checks if the variance of the batch gradient is bounded by a fraction of the squared norm of the full gradient.
   $$\text{Var}_{i \in B_k}(\nabla f(x_k; \xi_i)) \leq \eta^2 \|\nabla F_{B_k}(x_k)\|^2$$
2. Augmented Inner Product Test: A milder alternative that controls the variance of the projection of sample gradients onto the batch gradient, combined with an orthogonality test to ensure the batch gradient is not nearly orthogonal to the full gradient.

2.2 Integration with Adaptive Optimizers

The proposed schemes combine these batch-size rules with:

• AdaGrad-Norm: Updates parameters using a scalar step size $\alpha / \sqrt{v_k}$, where $v_k$ accumulates the squared norms of past gradients.
• AdaGrad: Updates parameters using coordinate-wise step sizes $\alpha / \sqrt{v_k}$, where $v_k$ accumulates squared gradients element-wise.
• Coordinate-wise Adaptation: For AdaGrad, the paper proposes a coordinate-wise exact variance norm test, ensuring the variance condition holds for each parameter dimension individually. This is crucial because AdaGrad's preconditioner is diagonal.

2.3 Implementation Details

The practical implementation uses empirical batch statistics computed from the current minibatch to estimate the variance and determine the next batch size. This contrasts with the theoretical analysis, which assumes an "exact-test" version where the batch size is chosen to satisfy population-level variance conditions. The authors implement these schemes in PyTorch using torch.func and vmap for efficient per-sample gradient computation.

3. Key Contributions

3.1 Theoretical Contributions

• Convergence Guarantees: The paper establishes high-probability convergence rates of $\mathcal{O}(1/K)$ to a first-order stationary point for smooth nonconvex functions within $K$ iterations. This is a significant improvement over the standard $\mathcal{O}(1/\sqrt{K})$ rates often associated with adaptive methods under bounded variance.
• Generalized Smoothness: The convergence analysis for AdAdaGrad-Norm is extended to functions satisfying $(L_0, L_1)$-smoothness (generalized smoothness), a condition more realistic for deep learning architectures (e.g., transformers) than standard $L$-Lipschitz smoothness.
• Expected Strong Growth (E-SG): The authors prove that satisfying the exact adaptive sampling tests implies an Expected Strong Growth condition along the iterates, which is a key driver for the improved convergence rates.
• Novel Analysis for Adaptive Optimizers: The analysis addresses the technical challenge of coupling the stochastic gradient in both the update direction and the adaptive denominator (e.g., $\sqrt{v_k}$), a feature not present in standard SGD analysis.

3.2 Empirical Contributions

• Bridging the Generalization Gap: Experiments on image classification tasks (MNIST, CIFAR-10) demonstrate that the proposed schemes can narrow the generalization gap between small and large batch sizes.
• Efficiency vs. Generalization Trade-off: The methods allow for starting with small batches (for generalization) and progressively increasing to large batches (for efficiency) without manual scheduling.
• Performance:
  • AdAdaGrad (with norm test) consistently outperforms constant-batch baselines and standard adaptive methods in validation accuracy while maintaining high training efficiency.
  • AdAdaGrad-Norm shows strong performance, though in some convex settings, it may underperform compared to SGD with adaptive batch sizes.
  • AdAdam: An extension to Adam (AdAdam) shows similar trends, achieving high accuracy with large effective batch sizes and fewer steps than constant-batch Adam.

4. Results

The paper validates its claims through numerical experiments on:

• Logistic Regression (MNIST): A convex problem where AdAdaGrad achieved the best validation accuracy (90%) and loss, though with a trade-off in computational steps depending on the $\eta$ parameter.
• Three-layer CNN (MNIST & CIFAR-10): On nonconvex objectives, AdAdaGrad significantly outperformed AdaSGD and constant-batch baselines in validation accuracy. For instance, on MNIST, AdAdaGrad with $\eta=0.1$ achieved 96% accuracy in 149 iterations with an average batch size >40,000, whereas constant large batches suffered from poor generalization.
• ResNet-18 (CIFAR-10): The proposed schemes successfully trained larger models. AdAdaGrad with $\eta=0.025$ utilized full batches (50,000) in later stages, achieving high accuracy with only 23 additional steps compared to a constant batch size of 50,000, while avoiding the generalization drop seen in the constant large-batch baseline.

5. Significance and Claims

The paper claims to be the first to provide adaptive batch size schemes based on adaptive sampling methods specifically for adaptive gradient methods that come with:

1. Rigorous convergence guarantees (high-probability $\mathcal{O}(1/K)$ rates) for smooth nonconvex objectives.
2. Strong empirical performance in reducing the generalization gap in large-batch training.
3. Efficient implementations compatible with modern deep learning libraries (PyTorch).

The authors emphasize that their work highlights the potential of adaptive batch size strategies to automate the trade-off between training efficiency and model generalization, particularly in the context of adaptive optimizers which are now standard in large-scale model training (e.g., transformers). They note that while the theoretical results rely on idealized "exact-test" conditions, the empirical results using approximate batch-based tests validate the practical utility of the approach.

Limitations and Future Work: The paper acknowledges that the coordinate-wise tests are computationally expensive for high-dimensional models, leading to the use of standard norm tests in experiments. It also notes that extending these schemes to distributed training paradigms (Data Parallelism, Tensor Parallelism) remains an open area for future research.