ADAM Optimization with Adaptive Batch Selection

AI summary

Overview

• Research area: Machine learning optimization — specifically stochastic first-order optimizers (Adam) combined with online learning / multi-armed bandit theory.
• Technical level: Advanced. The paper is primarily a theoretical contribution, built around regret analysis, importance sampling, and combinatorial semi-bandits, though the algorithm itself is straightforward to implement.
• Scope: The paper proposes a new Adam variant (AdamCB) that selects mini-batches using a combinatorial bandit, proves a sharper regret bound than prior bandit-based Adam methods, and demonstrates empirical gains.

What This Paper Is About

Standard Adam samples training data uniformly, treating every example as equally useful, which can waste computation on uninformative samples. A prior method, AdamBS, tried to fix this by using a multi-armed bandit to sample "important" examples more often, but its convergence proof contains a critical mathematical error and its algorithm can pick the same sample repeatedly inside one batch. This paper introduces AdamCB, which replaces single-arm bandit sampling with a combinatorial bandit that draws a set of distinct samples per batch, and for the first time provides a valid convergence guarantee for adaptive batch selection in Adam.

Key Contributions

1. A new algorithm, AdamCB, that combines Adam with a combinatorial semi-bandit sampling scheme. It is the first method to correctly adapt bandit-based sampling to the Adam framework, selecting K distinct samples without replacement rather than allowing duplicates.
2. A rigorous regret analysis showing AdamCB achieves a sharper bound than both uniform-sampling Adam and the prior bandit variant AdamBS. The bound explicitly improves as batch size K grows, a property AdamBS could not deliver.
3. A correction of the theoretical errors in AdamBS: the paper identifies that Liu et al. (2020) misapplied Jensen's inequality (the required convexity assumption does not hold), invalidating their claimed regret bound and their claims about batch-size dependence. The authors supply a corrected AdamBS bound for fair comparison.
4. Empirical evaluation across multiple datasets and models showing AdamCB consistently outperforms existing Adam-based methods in convergence rate and final performance, consistent with an independent prior report (Bansal et al., 2022) that AdamBS performs unreliably.

Main Findings

• The prior bandit-based Adam (AdamBS) is theoretically unsound. Its regret guarantee rests on an invalid application of Jensen's inequality, its analysis assumes an impractical doubly heavy-tailed feature distribution, and its sampling design permits the same sample to appear multiple times in a mini-batch — so it cannot benefit from larger batch sizes as claimed.
• AdamCB achieves a sublinear cumulative regret bound, meaning it provably converges to the optimum. The bound is O(d√T + (√d/n^{3/4})·((T/K)·ln(n/K))^{1/4}), where n is dataset size, K is batch size, T is iteration count, and d is parameter dimension.
• Batch size K appears favorably in the bound. Increasing K reduces the regret term, confirming that larger mini-batches genuinely help AdamCB — a direct contrast with AdamBS, where this was claimed but not achieved.
• More data reduces regret. The bound decreases as the number of samples n grows, matching intuition that larger datasets yield faster convergence to optimality.
• Derived under only standard assumptions. AdamCB's guarantee requires only bounded gradients and bounded parameters, unlike AdamBS, which additionally required a restrictive doubly heavy-tailed distribution and offered no analysis for bounded or sub-Gaussian (light-tailed) cases common in practice.
• Gradient magnitude drives sampling. The algorithm assigns higher selection weights to samples whose gradient norms are larger, and lower weights to those with smaller gradient norms, so informative samples are revisited more often.
• Importance sampling preserves unbiasedness. The gradient computed from the biased mini-batch is reweighted by the inverse selection probability, keeping it an unbiased estimate of the full-dataset gradient — essential for the Adam moment updates to remain sound.
• Dependent rounding makes it efficient. Selecting K distinct samples according to per-sample probabilities is done in O(n) time via DepRound, versus the at-least-choose(n,K) cost of a naive approach.

Methodology in Plain English

The researchers reframe training as an adversarial semi-bandit problem in which each data point is an "arm" and pulling a mini-batch means pulling K arms at once. Two design changes distinguish it from prior work:

1. Sampling without replacement. Instead of drawing K samples independently (which can repeat a sample), they compute a selection probability for every sample, then use dependent rounding (DepRound) to pick exactly K distinct indices while honoring those probabilities. A threshold τ caps any sample's probability at the batch-size-derived limit, with the excess redistributed; samples whose weight exceeds τ are temporarily placed in a "null set" so they don't distort the distribution. An exploration parameter γ mixes the weighted distribution with a uniform one to avoid over-committing to early-appearing "good" samples.

2. Reweighted gradient estimates. Because the batch is drawn non-uniformly, each sampled gradient is divided by its selection probability, yielding an unbiased gradient for the Adam moment updates.

After each step, sample weights are updated multiplicatively based on the observed gradient norms — an EXP3-style update — so informative samples gain weight over time and are likelier to be chosen later. The exploration rate γ is then set to a specific formula involving n, K, and T to balance the regret from exploration against the benefit of exploitation. The proof decomposes cumulative regret into an online-regret term plus two discrepancy terms between the full dataset and the mini-batches, bounds each separately, and combines them.

Why This Matters

This paper closes a gap that has quietly persisted in the adaptive-optimization literature: many methods claim to speed up training by sampling "important" data, but lack valid convergence guarantees. By identifying and correcting a specific error in an influential prior paper and by supplying the first sound guarantee for bandit-based batch selection in Adam, this work puts data-adaptive sampling on firm theoretical footing. It also shows the theory is not merely cosmetic — the batch-size dependence in the bound corresponds to realizable performance gains.

Real-world applications:

• Large-scale model training / fine-tuning, where each gradient step is expensive and prioritizing high-information examples reduces wall-clock time to a target loss.
• Data-efficient learning regimes (medical imaging, scientific simulation, rare-event data) where only limited labeled samples exist and controlling which are used per step matters.
• Distributed and streaming pipelines, where batches must be composed from heterogeneous shards and per-sample utility varies widely.
• Cost-sensitive cloud training, where reducing the number of iterations needed to hit a convergence threshold translates directly into hardware cost savings.

Industry relevance: Optimizer behavior sits at the core of every major ML framework. A drop-in modification to batch selection that is provably convergent and empirically faster is directly relevant to teams running training jobs at scale, and the algorithm's O(n) selection cost keeps the overhead negligible relative to a forward/backward pass.

Future Directions

• Extending the analysis to non-convex objectives. The current regret bound assumes a convex loss, which does not hold for deep networks; a non-convex convergence analysis would broaden applicability.
• Adaptive or dynamic batch size. The theory shows larger K helps, but not how to choose K on the fly; an adaptive schedule for K could be a natural follow-up.
• Generalizing beyond Adam. The combinatorial bandit sampling layer could plausibly be attached to SGD with momentum, AdamW, or other adaptive optimizers.
• Empirical robustness at scale. Systematic comparison across very large models, long training horizons, and distribution-shifted or noisy data would test whether the theoretical advantage persists where bounded-gradient assumptions strain.
• Interaction with privacy and fairness constraints. Non-uniform sampling changes which examples are seen, raising questions about differential privacy guarantees and representation of minority subgroups.

Target Audience

Optimization and learning-theory researchers working on stochastic gradient methods, bandit algorithms, and importance sampling will find the regret analysis and the correction of AdamBS most valuable. Practitioners building or maintaining deep-learning optimizer libraries and large-scale training pipelines will benefit from the concrete, implementable algorithm and its batch-size efficiency. Graduate students in machine learning with some background in online learning or convex optimization will find it a useful case study in turning a heuristic sampling idea into a provably convergent method.

Authors’ abstract

Adam is a widely used optimizer in neural network training due to its adaptive learning rate. However, because different data samples influence model updates to varying degrees, treating them equally can lead to inefficient convergence. To address this, a prior work proposed adapting the sampling distribution using a bandit framework to select samples adaptively. While promising, the bandit-based variant of Adam suffers from limited theoretical guarantees. In this paper, we introduce Adam with Combinatorial Bandit Sampling (AdamCB), which integrates combinatorial bandit techniques into Adam to resolve these issues. AdamCB is able to fully utilize feedback from multiple samples at once, enhancing both theoretical guarantees and practical performance. Our regret analysis shows that AdamCB achieves faster convergence than Adam-based methods including the previous bandit-based variant. Numerical experiments demonstrate that AdamCB consistently outperforms existing methods.