ADAM Optimization with Adaptive Batch Selection
Authors: Gyu Yeol Kim
Min-hwan Oh
Venue: ICLR 2025 Poster (2025)
Keywords: ADAM
Combinatorial Bandit
Importance Sampling
Mini-Batch
Optimization
Regret Minimization
Primary area: optimization

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 \textit{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.

Published as a conference paper at ICLR 2025
ADAM OPTIMIZATION WITH ADAPTIVE BATCH SELECTION
Gyu Yeol Kim Scoul National University Seoul, South Korea
Min-hwan Oh Seoul National University Scoul, South Korea minoh@snu.ac ac.
gyuyeolkim@snu ac. kr
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 varying degrees, treating them equally can lead to inefficient convergence. To ad- dress this, a prior work proposed adapting the sampling distribution using bandit framework to select samples adaptively. While promising, the bandit- variant of Adam suffers from limited theoretical guarantees. In this paper, we introduce Adam with Combinatorial Bandit Sampling (AdamCB), which integrates combina- torial bandit techniques into Adam to resolve these issues. AdamCB able to fully utilize feedback from multiple samples at once, enhancing both theoretical guar- antees and practical performance. Our regret analysis shows that AdamCB achieves faster convergence than cdam-based methods including the previous bandit- based variant. Numerical experiments demonstrate that AdamCB consistently outperforms existing methods.
INTRODUCTION
Adam (Kingma & Ba 2015) is one of the most widely used optimizers for training neural networks, primarily due to its ability to adapt learning rates. The standard version of Adam and its numerous variants treat each training sample equally by employing uniform sampling over the dataset. In practice, however, different data samples can influence model updates to varying degrees. As a result, simply performing full dataset sweeps or sampling data with equal weighting may lead to inefficient convergence and, consequently, unnecessary computational overhead when aiming to satisfy given convergence criterion.
To address these challenges, Liu et al. (2020) introduced dynamic approach called AdamBS, which adapts the sampling distribution during training using multi-armed bandit (MAB) framework. In this method, each training sample is treated as an arm in the MAB, allowing more important samples to be selected with higher probability and having a greater influence model updates. This approach was intended to improve both the adaptability and efficiency of the optimization process, presenting promising direction for further advancements.
However, despite its potential benefits, critical issues remain: Specifically, the convergence issue previously identified in the original analysis of Adam (initially reported by Reddi al. 2008 and later resolved by /Zhang et al. (022) also affects its bandit-based variant, AdamBS, as newly discovered in our work. Consequently, the existing theoretical guarantees regarding the efficiency and effectiveness of AdamBS are invalid (see Section based method that can adaptively sample while providing rigorous performance guarantees. This raises a critical question: is il possible to design an algorithm that aduptively adjusts the sampling distribution while ensuring both provable guarantees and practical performance improvements?
As a result, to the best of our knowledge, there is no existing
In this paper, we propose a new optimization method, Adam with Combinatorial Bandit Sampling (AdamCB), which addresses the limitation in both the analysis and implementation of AdamBS by incorporating combinatorial bandit approach into the sample selection process. In this approach, batch selection formulated as a combinatorial action, where multiple arms (samples) are selected simultaneously. This combinatorial bandit framework can take advantage of feedback from multiple samples at once, significantly enhancing the adaptivity of the optimizer. For the first time, we provide provable performance guarantees for adaptive batch selection in base- methods, leading to

---

Published as a conference paper at ICLR 2025
faster convergence and demonstrating both theoretical and practical improvements existing approaches. Our main contributions are summarized as follows:
We propose Adam with Combinatorial Bandit Sampling (AdamCB), a novel optimization algorithm that integrates the Adam method with combinatorial bandit approach for sample selection. To the best our knowledge, AdamCB is not only the first algorithm to successfully combine combinatorial bandit techniques with the Adam framework, but also the first to correctly adapt any bandit techniques to Adam, significantly enhancing its adaptability.
We provide a rigorous regret analysis of the proposed AdamCB algorithm, demonstrating that it achieves sharper regret bound compared to both the original Adam (which uses uniform sampling) and its bandit-based variant, AdamBS (Liu et al. 2020). Additionally, we correct the theoretical errors in the analysis of AdamBS and present revised regret bound (see Table for comparisons).
We perform empirical evaluations across multiple datasets and models, showing that AdamCB consistently outperforms existing Adam based optimization methods in terms of both conver- gence rate and practical performance. Our results establish AdamCB as the first idam-based algorithm to offer both provable convergence guarantees and practical efficiency for bandit- based Adam optimization methods.
PRELIMINARIES
2.1 NOTATIONS
We denote by [n] the set 21, . for positive integer For vector æ E Rd, we denote by ||ا||| the vector's Euclidean norm. For two positive sequences {aa [an)== and {ba 1 On = O(b,) implies that there exists absolute constant C 0 such that on Cb, holds for all 2 1. Similarly, a₂ o(b,,) indicates that lim dn b.. 0.
2.2 EXPECTED RISK AND EMPIRICAL RISK
Expected Risk. In many machine learning problems, the primary goal to develop a model with mean that while models are trained on
robust generalization performance. By generalization, finite sample of data points, we aim for them to perform well on the entire population of data. To achieve this, we focus on minimizing quantity known the expected risk. The expected risk the average loss across the entire population data distribution, reflecting the model's anticipated error if had access to the complete set of possible data samples. Formally, the expected risk defined as:
B(2.)>>P 10(0;2, x v)] == e(0; x(0;x,v/0 w))P(x, y)
(1)
where 0 E Rdd is the model parameter, {(0: x. y) the loss function that measures the error of the model on single data sample (z, v). and P is the true distribution of the data. The gold standard goal is to find the 0 that minimizes the expected risk in Eq. ensuring that the model generalizes well to all data drawn from P.
Empirical Risk. In practice however, the true distribution P is typically unknown. Instead, we only work with finite dataset 'D consisting of samples, which is denoted D = {(T,, y1) ) To approximate the expected risk, we use the empirical distribution P derived from the dataset D. For this empirical distribution P to be a reliable approximation, we assume that the dataset D is representative of the true distribution P. This requires that each sample in the dataset D is equally likely and independently drawn from the true distribution P (i.c., the samples (æ,, 9:) are i..dd according to P). The empirical distribution 1 can be expressed as:
P(x, y; D) = 1 -
8(x 24, W1)
where 8 the Dirac-delta function. With the empirical distribution at hand, the empirical risk the average loss over the given finite dataset D. The empirical risk serves estimate of the expected
2

---

Published as a conference paper at ICLR 2025
risk and is formally defined as:
E ( = ((0) 2. yd y: = - Σ 10 x, V.).
(3)
However, if the dataset is non-unnn non-uniformly distributed, some samples may be over-represented or under- represented, leading to biased estimate of the expected risk. To address this issue, one importance sampling (Katharopoulos & Fleuret 2018), which adjusts the sample weights to ensure the empirical risk remains unbiased estimate of the expected risk.
2.3 OBJECTIVE FUNCTION AND MINI-BATCHES
Objective Function. In the context of optimizing machine learning models, the objective function f(;; D) is often the empirical risk shown (3). Given a dataset D { y:) } the objective function f0; D) is defined as, (0: D)
L(0; z.. Vs). As studied in the relevant literature of Adam optimization (Duchi et al. 2011 Tieleman & Hinton 2012 Kingma & Ba 2015 Dozat 2016 Reddi et al. 2018), we focus on the problem setting where I convex (i.e., l is convex). Then, the goal of the optimization problem is to find parameter 0* E IRd that minimizes the objective function f0; D). This problem is known as empirical risk minimization:
0" E arg min ((0
OERa
The gradient of the objective function / with respect to 0 is denoted by == (0; = E 9i. where g: = (0 Z4s 81) is the gradient of the loss based on the i-th data sample in D.
Mini-Batches. When the full dataset D {xx, y4) ())² very large (i.e., large n), computing the gradient over the entire dataset for each optimization iteration becomes computationally expensive. To address this, mini-batches--mallee subsets of the full dataset-are commonly used reduce computational overhead per iteration. Consider the sequence of mini-batches D1, D....., Dr C D used for training, with corresponding objective functions fo(0) == f(0, D:) for each
(L....,7. Let K be the size the mini miin-bbtchh Dr for all f, then DD == {(x f7 Uu! (Iru.. Uuz).... ..... (o Us )) where J {J]] J?. JKH } S [n]] is the of indices of the samples in the mini-batch D, The objective function f0) for the mini-batch D, is defined as the expected risk over this mini-batch
f((0) = f(:: = e(0) f (ddz,,))
(4)
where P(zx, y,; D:) is the empirical distribution derived from the mini-batch D, The gradient of the objective function fo with respect to is denoted as 9e = Vefit.
Note that the sequence of mini-batches {De) involves choosing mini-batches based on results observed during previous optimization steps, poten- tially adjusting the importance assigned to specific samples. The empirical distribution P(æ: y; Dz) significantly influenced by the method used select the mini-batch D, from the full dataset DD
be selected adaptively. Adaptive selection
2.4 REGRET MINIMIZATION
Cumulative Regret. An online optimization method can be analyzed within the framework of regret minimization. Consider an online optimization algorithm æ that generates a sequence of model parameters 0,,..., .r over T iterations. The performance of can be compared to the optimal parameter 0" E arg minger [0; D), which minimizes the objective function over the full dataset D. The cumulative regret after T iterations is defined as:
RT(T) Σ f(0.; D) T mim f(0: D)
(5)
BERa
where the expectation is taken with respect to any stochasticity in data sampling and parameter estimation. For the optimization algorithm я to converge to optimality, we require the cumulative regret R^ (T) to grow slower than the number of iterations T, specifically R (T) o(T).
3

---

Published as a conference paper at ICLR 2025
Online Regret. An alternative notion of regret is the online regret, defined sequence of mini-batch datasets {D.) or equivalently, over the sequence of functions (fe) Specifically, the online regret of the optimization algorithm π after T iterations is given by:
T Σ BER
Rollin (T)
E
Σ (.(..) min
where the expectation is again taken over any stochasticity in the optimization process. is important to note that the primary focus should not solely be minimizing the online regret. An algorithm might select Dr C ' in a way that allows to perform well on (D+) but it may perform poorly on the full dataset D. Therefore, our ultimate goal remains minimizing the cumulative regret R. (T) Later, in the proof of Theorem | we demonstrate how minimizing the cumulative regret R (2)) in 15
relates to minimizing the online regret R Ronline (T') with respect the sequence (
2.5
RELATED WORK: ADAM AND TECHNICAL ISSUES IN CONVERGENCE GUARANTEES
2.5.1
ADAM OPTIMIZER
Adam (Kingma & Ba 2015) is a widely used first-order gradient-based optimization method that computes adaptive learning rates for each parameter by using both the first and second moment estimates of the gradients.
each iteration t, Adam maintains the accumulated gradients m, <- -1,0m4-1 + (1 $1,1)s and the accumulated squared gradients 24 + f2 + (1 B2)9 where 9i is the gradient at iteration and of represents the element-wise square of gradient 9r. The hyper- parameters 31, 32 E [0,1] 1) control the decay rates of m, and Ut respectively. Since these moment estimates are initially biased towards zero, the estimates corrected as m, + m4/(11111) Bi) and & < 01/(1
mh (3). The Adam algorithm then updates the parameters using 0₂ 0 Or ver+e where e is a small positive constant added to prevent division by zero. The key characteristic of Adam lies in its use of exponential moving average for both the gradient estimates (first-order) and the dlement-wise squares of gradients (second-order). This approach has shown empirical effectiveness for optimizing deep neural networks and has led to many follow-up works, such as Reddi al. (2018), Loshchilov & Hutter 20099) Chen et al. (2020). Alacaoglu et al. (2020), and Chen et al. (2023).
2.5.2 CONVERGENCE OF ADAM
After Adam was first introduced, there was considerable debate regarding its convergence properties. In particular, Reddi al. (2018) provided counterexample demonstrating that Adam might fail to converge under certain conditions (see Section3 of Reddi et al. 2018). In response, numerous variants of adaptive gradient methods have been proposed, such as AMSGrad (Reddi et al. 2018). AdamW (Loshchilov & Hutter 2019). and AdaBelief (Zhuang et al. 2020). to address this issue and ensure convergence.
However, recent studies (Zhang et al. 2022 Défossez et al. 2022) indicate that the standard Adam algorithm itself can achieve convergence with appropriate hyperparameter choices, thereby resolving its earlier theoretical concerns. These recent works provide alternative convergence proofs for the original Adam algorithm without requiring modifications to its update rules, contingent upon hyperparameter conditions being satisfied.
2.5.3 TECHNICAL ISSUES IN ADAM WITH BANDIT SAMPLING
(LIU ET AL.. 2020
The work most closely related to ours is by al. 22220), who proposed AdamBS, extension of the original Adam algorithm and proof framework incorporating bandit sampling approach. However, the initial convergence issue present the original proof of Adam, discussed in Section [.5.2] also affects AdamBS, thus invalidating the convergence guarantee provided by Liu et al. (2020). Moreover, even the alternative convergence proofs of Adam are adapted to the AdamBS framework, several other critical shortcomings persist. We summarize these remaining issues as follows:
AdamBS unfortunately fails to provide guarantees convergence despite its claims, both on the regret bound and on the effectiveness the adaptive sample selection via the bandit approach. Specifically, the claimed regret bound in Theorem I of et (2020) is
4

---

Published as a conference paper at ICLR 2025
incorrect. In particular, Eq.(7) on Page 3 of the supplemental material of Liu et al. (2020) contains an error in the formula expansion This technical error is critical to their claims regarding the convergence rate of AdamBS and its dependence the mini- match size K
Their problem setting also limited and impractical, even if the analysis were corrected. The analysis assumes that feature vectors follow doubly heavy-tailed distribution, a strong and restrictive condition that may not hold in practical scenarios. Importantly, no analysis is provided for bounded or sub- Gaussin (light- (ight-t distributions, which are more commonly encountered in real-world applications.
Despite their claim regarding mini-batch selection of size K, their algorithm design allows the same sample to be selected multiple times within a single mini-batch. This occurs because the bandit algorithm they employ is based single-action selection rather than combinatorial bandit approach. As a result, their method may repeatedly sample the same data points within mini-batch. Moreover, due this limitation, their method fails to achieve performance gains with increasing mini-batch size K, contradicting their claim.
Numerical evaluations (in Section 5) demonstrate poor performance of AdamBS. Our numerical experiments across various models and datasets reveal that AdamBS exhibits poor and inconsistent performance. Additionally, independent evaluation by separate group has also reported inconsistent results for AdamBS (Bansal et al., 2022).
3. PROPOSED ALGORITHM: ADAMCB
3.1 ADAMCB ALGORITHM
Algorithm 1: Adam with Combinatorial Bandit Sampling (AdamCB)
Input: learning rate {at) decay rates ((1,))== 32, batch size K, exploration parameter E [0, 1)
Initialize: model parameters l.. first moment estimate mg +- 0, second moment estimate 0, 6g +- 0, sample weights w1,0 +- 1 for all E [n]] for t= to
Jr, Pt, Smull: + Batch- Batch-Selection( w;-1, K, (Algorithm
Compute unbiased gradient estimate Jr with respect to Jg using Eq.
my + f1,tm411 (1 $1.1).
ve 8201-1 +
B2)99
ô1 + U1 0, max
(1-81.-1) Gc- v2
(1 02
if 22
8.+1 + 0/ a +i+t
wi -
1w Pc, Jr, fosihied Smull, 7) (Algorithm
We present our proposed algorithm, Adam with Combinatorial Bandit Sampling (AdamCB), which described in Algorithm The algorithm begins by initializing the sample weights wo (wi,0, .0, W2.0,-.-, Wn.0 uniformly, assigning an equal weight of to each of n training samples. At each iteration t [T., the current sample weights w (0o0=(wieeem [wi,t-1, W2,1-1
wn,t-1 are used to determine the sample selection probabilities P: := {PI.t. P2.r, Pr.t). where these probabilities are controlled with the exploration parameter 7 (Line 2). subset of samples, denoted by D, S D, is chosen based on these probabilities. The set of indices for samples chosen in the mini-batch DD is denoted by J, : (JH J,....,,..,........... JKK C [n]. Using this mini-batch Dr, an unbiased gradient estimate Ge is computed (Line 3). The algorithm then updates moments estimates ma, Vt, and ê, following the Adam-based based update rules (Lines 4-6). The model parameters O, are subsequently updated based on these moment estimates (Line 7). Finally, the weights - are adjusted to reflect the importance of each sample, improving the batch selection process in future iterations (Line 8).
Liu et al. 2020) apply Jensen's inequality to handle the expectation of the squared norm of the sum of gradient estimates. However, the convexity assumption required for Jensen's 's incquality does not hold in this context, rendering this step their proof invalid.
5

---

Published as a conference paper at ICLR 2025
The following sections describe the detailed process for deriving the sample probabilities pt and selecting the mini-batie batch D, = {(T), (s) from the sample weights w utilizing our proposed combinatorial bandit sampling.
3.2 BATCH SELECTION: COMBINATORIAL BANDIT SAMPLING
In approach, incorporate bandit framework where each sample is treated as arm. Since multiple samples must be selected for a mini-batch, we extend the selection process to handle multiple arms. There two primary methods for sampling multiple arms: with replacement or without replacement. The previous method, AdamBS (Liu et al., 2020), samples multiple arms with replacement. contrast, our proposed method, AdamCB, employs combinatorial bandit algorithm to sample multiple arms without replacement, achieved by Patch-Selecte (Algorithm)
Algorithm 2: Pattch-Selection
Input: Sample weights w1-1, batch size K, exploration parameter 7 E [0, 1)
Set C + (1/K //n)/(1 (1-))
max Wi,t-1 2 C >>= Wi,t1 then
Let / be sorted list of {wi,--1k in descending order
Set S -x=1 wi
for i = to mionon do
Compute 7 + C. 5/(1-i- i
if Wi < then break, else update S +- S Wi,t
Set Small, < {i Wie-- 2 J and Wi,t-1 = T for E Small,
else
10
L Set Smill. < 0
Set Pi, f K ((1-7)
Wi. 1-.111
+ 2) for all
Eis W1s-1
12 Set Jj <- DepRound( K, P... Pn.))) (Algorithm) 6
13 return J,, Pt, Smull.
Weight Adjustment (Lines 2-10). Unlike single-arm selection bandit approach like AdamBS, where Σi= Pt. = 1, because only one sample is selected at time, AdamCB must select K samples simultaneously for a maiih-btcch. Therefore, it is natural to scale the sum of the probabilities to K, reflecting the expected number of samples selected each round Allowing the sum of probabilities to equal K can lead to individual probabilities Pi,t exceeding 1, especially when certain samples are assigned significantly higher weights due to their importance (or gradient magnitude). To ensure valid probabilities and prevent any sample from being overrepresented, AdamCB introduces a threshold т. If a weight w;,--1 exceeds т. the index added to a null set Suullt, effectively removing from active consideration for selection. The probabilities of the remaining samples adjusted to redistribute the excess weight while ensuring the sum of probabilities remains K.
Probability Computation (Line 11). After adjusting the weights, the probabilities for selecting each sample are computed using the adjusted weights W and the exploration parameter y. This computation balances the need to exploit samples with higher weights (more likely to provide useful gradients) and explore other samples. The inclusion of K in the scaling ensures that the sum of probabilities matches the batch size: Σ Pi. = K.
Mini-batch Selection (Line 12). The final selection of K distinct samples for the mini-batch is (2006) and later
performed using DepRound (Algorithm ( originally proposed by Gandhi adapted by [Uchiya et al.) (2010). DepRound efficiently selects K distinct samples from a set of n samples, ensuring that each sample is selected with probability Pi.t- The algorithm has computational complexity of Onx), which is significantly more efficient than naive approach requiring consideration of all possible combinations with complexity of at least (n)
'If the sum of probabilities were constrained to 1, the algorithm would need to perform additional rescaling sampling adjustments. Instead, directly setting Σn Pi,t K aligns the probability distribution with the batch-level selection requirements.

---

Published as a conference paper at ICLR 2025
3.3 COMPUTING UNBIASED GRADIENT ESTIMATES
Given the mini-batch data D, = { Vj Wjlhie. from Algorithm 2 and since is probability over the full dataset D, and D, is sampled according to pt, employ importance sampling technique compute the empirical distribution P for Dr:
P(x, y: Dr) =
1 K
Σ JEJ
8(x
(6)
Tpj.
where 5 is the Dirac-delta function. This formulation ensures that the empirical distribution P for the mini-batch D, closely approximates the original empirical distribution P(x,, y/ D) defined over the full dataset D, as expressed in (q.(). According to the empirical distribution P(xx, y: D,) in (6, the online objective function fe corresponding the mini-batch D, (as defined in Eq. can be computed as
f.(0) = f(O = / ((0:2, v)dP(z y: D.)
{(0;zj:W3]
K
npj
JEJ,
This implies that the gradient gr computed as follows:
Vofe(0) obtained from the mini-batch D. iteration is
Vol(O;Z).15)
yi.t
gx Vofi(0) KAKKK
nPJA
npjj. jEJ.
JEJ.
Here, we denote the gradients for each individual sample in the mini-batch Je is the set of indices for DD, In stochastic optimization methods like SGD and Adam, it is crucial to use an unbiased gradient estimate when updating the moment vectors. We can easily show that gr is an unbiased estimate of the true gradient over the entire dataset by taking the expectation over i.e, Ep, [9c] = g. The unbiased gradient estimate Ir in Eq. 7 is then used to update the first moment estimate my and the second moment estimate in each iteration of the algorithm.
Egg.chjea where
3.4 UPDATE OF SAMPLE W WEIGHTS
The final step in each iteration of Algorithm [ involves updating the sample weights wy- Treating the optimization problem as an adversarial semi-bandit,
partial feedback consists only of the gradients (e.,- The loss li, occurred when the i-th arm is pulled is computed based on the norm of the gradient ||gr,.]|- Specifically, the loss l,2 is always non-negative and inversely related to el This implies that samples with smaller gradient norms are assigned lower weights, while samples with larger gradient norms are more likely to be selected in future iterations.
Algorithm 3: Weight -Update
Input: W;-1,Pr. Jr, {ggatses, Soull,,, ? [0, 1)
for j = to n do
Compute loss Ej.t
Pani 1
(-
|las.ell ? [p,,a)
+ 1.2 plan.
if j E Jr; otherwise bje = -0
j & Smult, then
Wjjt f- Wj,t- exp (KKK41//1)
return We
4 REGRET ANALYSIS
In this section, we present regret analysis for our proposed algorithm, AdamCB. We begin by introducing the standard assumptions commonly used in the analysis of optimization algorithms. Assumption 1 (Bounded gradient). There exists L 0 such that ||9s,ell < L for all [n] and t e [T].
Assumption 2 (Bounded parameter). There exists D 0 such that ||O, -0. D for any s, l (T].

---

Published as a conference paper at ICLR 2025
Discussion of Assumptions. Both Assumptions a 2 the standard assumptions in the relevant literature that studies the regret bounds Adam based optimization (Kingma Ba 2013] Reddi et al. 2018, Luo et al. 2019; Liu et al. 2020 Chen et al. 2020). A closely related work Liu al. 2020 relies on the additional stronger assumption of a doubly heavy-tailed feature distribution. In contrast, the regret bound for AdamCB is derived using only these two standard assumptions.
4.1 REGRET BOUND OF ADAMCB
Theorem (Regret bound of AdamCB). Suppose Assumptions total T iterations with ag = Vt and with $1,c == 3,A1=1, (0, 1). Then, the cumulative regret of AdamCB (Algorithm)
2 hold, and we run AdamCB for
7 with batch size K is upper-bounene hopeer-bounded by
0 TdTT+
Vd
T In
1/4
33/4
K
Discussion of Theorem Theorem establishes that the cumulative regret bound of AdamCB sub-linear in T, i... R. (T) = o(T). Hence, AdamCB is guaranteed to converge to the optimal solution. The first term in the regret bound, dv which is commonly shared by the results in all Adam-based methods (Kingma & Ba 2015 Reddi al., 2018 Liu et al., 2020 The second term, (Vd/n3/4). (//////))((///K batch size K regret. As the number of samples increases, this term decreases, suggesting that having more data generally helps in reducing regret (hence converging faster to optimality). Similarly, increasing the batch size K also decreases this term, reflecting that larger mini-batches can reduce the variance in gradient estimates, thus improving the performance.
((T/K) ) In (n/K))²² illustrates the impact of the number of samples n as well as the
4.2 PROOr SKETCH or THEOREM
In this section, we present the proof sketch of the regret bound Theorem [ The proof start by decomposing the cumulative regret R (T) into three parts: the cumulative online regret R (T) and auxiliary terms (A) and (B), as shown below:
R (T) Rnlinn (T) +
Σ
(J(O,:)) ((D.D.))
+E
min
((, - T min f D) dead
(8)
PERU
(A)
(8)
We now prove the following two key lemmas bound the online regret R R. (T).
Lemma 1. Suppose Assumptions holdd AdamCB (Algorithm 7 with mini-batch of size K, which is formed dynamically by distribution Pt, achieves the following upper-bound the cumulative online regret R Rinsine (T) over iterations.
Mi.z 1
Ro (T) S PIdVT - vdp2
E
+ P3
n²
K
(p,,))²
JEJ,
where PL, P2, and P3 are constants (See Appendix
B.2
Lemma provides an upper bound for the cumulative online regret over iterations. This lemma shows that Pt affects the upper bound of RT Ri (T"). Hence, we wish to choose pt that could lead to minimizing the upper bound. The following key lemma shows that it can be achieved by combinatorial semi-bandit approach, adapted from EXP3 (Auer al. 2002).
Lemma 2. Suppose Assumptions / h we set mimnnn
(n/K)
the batch selection
1)T
(Algorithm 2 and the weight update rule (Algorithm) 3 following AdamCB (Algorithm) 7 implies
1
T Ep.
119,.1 ||
=0
In n
Epp
min t=1
(p
(p), 3
K
1=1
JEJ.
JEJ
8

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of Regret Bounds
Optimizer
Regret Bound
AdamX Tran et al.
2019)
(variant of Adam))
OddTT + v
VT)
a12
AdamBS Liu et al.
2020 (corrected:)
O(dvT + na (Ti)))
AdamCB (Ours)
OldvT w374
va
x In ))
i While the convergence of Adam Kingma & Ba 2015) has been correctly established by Zhang et al. 2022) the results in /Zhang et al. (2022) do not provide regret guarantees. For fair comparisons of regret, we use slight modification of Adam. AdamX (Tran et al. (019), for which correct regret analysis is feasible.
Similarly, for AdamBS(Lin et al. 2020), since the existing proof regret guarantee contains technical errors, we present a corrected prool, resulling in new regret bound (Theorem
Lemma 2 bounds the difference between the expected cumulative loss of the chosen mini-batch
and the optimal mini-batch, showing sub-linear growth in T with dependence the batch size K. Combining Lemma and L we can bound the cumulative online regret R (T'). which also
grows sub-linearly in T. Proofs of Lemma [ and Lemma 2 are in Appendix [B.2 Appendix [B.3]
respectively. The discrepancy terms (A) and (B) in Eq. 8 capture the difference between the full dataset and the mini-batches (Dr) and also bounded sub-linearly in ' (See Lemma
Appendix B B. Since the cumulative regret RD (T) is decomposed into the online regret R
with additional sub-linear terms, we obtain the cumulative regret bound for AdamCB.
4.3 COMPARISONS WITH ADAM AND ADAMBS
Our main goal demonstrate that the convergence rate of AdamCB (Algorithm[] is provably more efficient than those of the existing Adam-based methods including ones that employ uniform sampling and AdamBS (Liu al. 2020) that utilizes combinaonaaa bandit sampling. Since the claimed regret bounds in the original Adam and AdamBS are invalid, we provide new regret bounds for both a variant of Adam, called AdamX (Tran et al. 2019) and AdamBS in Theorems and respectively, which may be of independent interest.
For comparison purposes, we additionally introduce the following assumption:
Assumption 3 (Bounded variance of gradient). There exists O Ouch that Var(Ila.))) S o² for all i E [n]] and E [7]]
Assumption 3 is commonly used in the previous literature (Reddi et al. 2016 Nguyen et al. 2018 Zou et al. 2019: Patel et al., 2022). is important to note that Assumption 3 is not required for the analysis of algorithm in Theorem [ Rather, we employ the assumption to fairly compare with corrected results for the existing Adam-based methods (Tran et al.) 2019 Liu et al. 2020).
Under Assumptions is given by Odd'T + +n
and 3 the regret bound for the Adam variant (AdamX) using uniform sampling 1/2 VdT) (Theorem Appendix C). while the regret bound for corrected version of AdamBS using (mon-combinatorial bandit sampling is OdvT + n -3/1 In n) 1)1/)) I and holld The comparisons of regret bounds are
(Theorem 3 in Appendix D) when Assumptions
outlined in Table
Faster Convergence of AdamCB. The second term in the regret bound of AdamX exhibits a depen- dence on n 1/2, which the rate of regret decrease as the dataset size increases. However, this reduction in regret occurs a slower rate compared to bandit-based sampling methods. Both AdamBS (corrected) and AdamCB achieve an improved n -3/4 dependency, resulting in a faster convergence. When comparing the two bandit-based sampling methods, AdamCB surpasses AdamBS (corrected) in terms of convergence, particularly by the factor of the batch size K. That is, as far regret per- formance concerned, AdamBS does not benefit from multiple samples in batch while our AdamCB enjoys faster convergence. Hence, AdamCB is not only the first algorithm with correct performance guarantees for AdamX with adaptive batch selection, but to our best knowledge, also the method with the fastest convergence guarantees in terms of regret performance.

---

Published as a conference paper at ICLR 2025
Ad
AdamX
Adam8S
AdamCB (ours)
2.25
Tran Lods
5.79
2.1
144
675
Epocts
Epochs
Epocte
2.25
2.0
905t L0G3
177
16
0.3
Epocra
Epochs
Epocts
Figure I: Performances with MLP model on MNIST, Fashion MNIST, and CIFARIO
AdamX
AdamBS
(ours)
(al
3.00
Th3ln Loos
3.7
8:29
0.75
Epochs
Spocts
2.2
Testuate
3.2
1.5
0.2
8.%
Epocho
Epochs
Specte
Figure 2: Performances with CNN model on MNIST, Fashion MNIST, CIFAR1O
NUMERICAL EXPERIMENTS
Experimental Setup. evaluate our proposed algorithm, AdamCB, we conduct experiments using deep neural networks, including multilayer perceptrons (MLP) and convolutional neural networks (CNN), on three benchmark datasets: MNIST, Fashion MNIST, and CIFARIO. Comparisons are made with Adam, AdamX and AdamBS, with all experiments implemented in PyTorch. Performance assessed by plotting training and test losses over epochs, with training loss calculated on the full dataset and test loss calculated the held-out validation data set. Results represent the average of five runs with different random seeds, including standard deviations. All methods use the same hyperparameters:
0.9, B2 0.999, experimental details are provided in Appendix
0.4, K 128, and 0.001. Additional
Results. Figures D and 2 show that AdamCB consistently outperforms Adam, AdamX and AdamBS, demonstrating faster reductions in both training and test losses across all datasets. These results suggest that combinatorial bandit sampling is more effective than uniform sampling for performance optimization. Attempts to replicate the results of AdamBS from|Liu et al. (2020) revealed inconsistent outcomes, with significant fluctuations losses, indicating potential instability and divergence. In contrast, AdamCB exhibits consistent convergence across all datasets. highlighting its superior performance and practical efficiency compared to Adam, AdamX and AdamBS. Additional experimental results in Appendix [F further reinforce the superior performance of AdamCB.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
For each theoretical result, we present the complete set of assumptions in the main paper and the detailed proofs of the main results provided in the appendix, along with experimental details and additional experiments in Appendix Flo to reproduce the main experimental results.
ACKNOW.EDGEMENTS
The authors would like to thank Se Young Chun and Byung Hyun Lee for bringing the previous work (Liu et al. 2020) to attention. This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2022-NR071 1853 and RS RS-2023-00222663) and by AI-Bio Research Grant through Scoul National University.
REFERENCES
Ahmet Alacaoglu, Yura Malitsky, Panayotis Mertikopoulos, and Volkan Cevher. A new regret analysis for adam-type algorithms. In International conference on machine learning, pp. 202- 210. PMLR, 2020.
Peter Auer, Nicolo Cesa-Bianchi, Yoav Freund, and Robert E Schapire. The nonstochastic multiarmed bandit problem. SIAM journal on computing, 32(1):48-1)) -7088877,27 2002.
Aman Bansal, Shubham Anand Jain, and Bharat Khandelwal. Bag of tricks for faster stable image classification. CS231n Course Project Report, 2022. URL httpss//css3In. stanford. edu/
pd
Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu. Closing the generalization gap of adaptive gradient methods in training deep neural networks. In Proceedings of the Twenty-inth International Joint Conference on Artificial Intelligence, 2020.
Yineng Chen, Zuchao Li, Lefei Zhang, Bo Du, and Hai Zhao. Bidirectional looking with a novel double exponential moving average adaptive and non-adaptive momentum optimizers. In International Conference on Machine Learning, pp. 4864-4803 PMLR, 2023.
Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions Machine Learning Research, 2022. ISSN 2835-8856. URL
Timothy Dozat, Incorporating nesterov momentum into adam. International Conference on Learning Representations, 2016.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Rajiv Gandhi, Samir Khuller, Srinivasan Parthasarathy, and Aravind Srinivasan. rependent and its applications to approximation algorithms. Journal of the ACM (JACM), 53(3):324-33 360, 2006.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Angelos Katharopoulos and François Fleuret. Not all samples are created equal: Deep learning with importance sampling. In International conference on machine learning, pp. 2525-2534. PMLR, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimizzaioni 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Rui Liu, Tianyi Wu, and Barzan Mozafari. Adam with bandit sampling for deep learning. Advances in Neural Information Processing Systems, 33:5393-54p0 2020.

---

Published as a conference paper at ICLR 2025
Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976-11986, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Reprexentations, 2019.
Liangchen Luo, Yuanhao Xiong, and Yan Liu. Adaptive gradient methods with dynamic bound of learning rate. In International Conference on Learning Representations, 2019. URL https:
Lam Nguyen, Phuong Ha Nguyen, Marten Dijk, Peter Richtárik, Katya Scheinberg, and Martin Takác. Sgd and hogwild! convergence without the bounded gradients assumption. International Conference Machine Learming, pp. 7750-3758, PMLR, 2018.
Vivak Patel, Shushu Zhang, and Bowen Tian, Global convergence and stability of stochastic gradient descent. Advances in Neural Information Processing Systems, 35:36014 15:36014-36025, 2022.
Sashank Reddi, Ahmed Hefny, Suvrit Sra, Barnabas Poczos, and Alex Smola. Stochastic variance reduction for nonconvex optimization. In International conference on machine learning, pp- 314-323. PMLR, 2016.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
Tijmen Tieleman and Geoffrey Hinton. Rmsprop: Divide the gradient by a running average of its recent magnitude. coursera: Neural networks for machine learning. COURSERA Neural Networks Mach. Learn, 17, 2012.
Phuong Thi Tran et al. On the convergence proof of amsgrad and new version. IEEE Access, 7: 61706-61716, 2019.
Taishi Uchiya, Atsuyoshi Nakamura, and Mineichi Kudo. Algorithms for adversarial bandit problems with multiple plays. In International Conference on Algorithmic Learning Theory, pp. 375-389. Springer, 2010.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi- QuinQuan Luo. Adam can converge without any modification on update rules. Advances neural information processing systems, 35: 19386-28399, 2022.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda. Nicha Dvornek. Xenophon Pa- pademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belicf in observed gradients. Advances in neural information processing systems, 33:18795-18880 2020.
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for conver- gences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 11127-11135, 2019.
12

---

Published as a conference paper at ICLR 2025
Appendix
A AUXILIARY LEMMAS
Definition 1. A functin f: + R is conver for all u, v RR, and all X [0,1].
Xf(u) + (1 A)5(v) 2 s(Au + (1 - ())
Lemma 3. a functinn f IRdd TR is convex, then for all u Rd
[(v) 2 l(u) (v(u)
where (-) denotes the transpose of (-).
Lemma 4 (Cauchy-Schree Schwarz inequality). For all n 2 az b, E R. (1Sisn))
Σ
aab
Σ
Σ by
i=
Lemma (Taylor series). For R. and 0 S a
1
Σ -a'=
and
la
(x)
121
Lemma 6 (Upper bound for the harmonic series). For N E N.
N
N
1
E and
Σ
2V
R=1
Lemma 7. For all E N. and a; b, R such that a, 2 0 and b, 0 all E [n].
Σ Gi
En b
B PROOF FOR ADAMCB REGRET BOUND
In this section, provide proofs of key lemmas, Lemma [ aad Lemma They needed to prove Theorem [ which shows the regret bound for AdamCB. Tn the last of this section, we present the proof for Theorem 1
B.I AUXILIARY LEMMAS FOR LEMMA
We first present auxiliary lemmas and proofs for Lemma Our proofs basically follow arguments in Tran et al. (2019). For the sake of completeness, all lemmas from Tran et al. (2019) restated with our problem setting.
Lemma 8. For all t 2 1, we have
(1 31,5)2 ns for all $
(1 - 31.0)
max
(9)
where 0n AdamCB (Algorithm 7
Proof. Prove by induction on d. Recall that by the update rule er we have éj +- U1. 0z f (1 max 11 0 1, if 2 2. Thus,
(1 31,2)
i = max
01, U2
(1 - 31,1 )²
(1 - }1,2)²
max
U1 U2
(1 - $1,1)²
(1 }1,2)²
max
Vos S 5
)1,
13

---

Published as a conference paper at ICLR 2025
which we proved for the case when t 2 in Eq. Now, assume that
(1 11,0-1)
f = max
(1 - B1,a)² Us for all 1 1
and (q.99 holds for all 1 1 St - 1. By the update rule on i22
(1 - 31,1)²
6 = max
(1 11,c-12 1/-
(1 - 31,²)
(1 (1- 1), U1 for all 1 s1 s 11)
max
(1 11,
max
(1 -,,,²
1/2
(1 )
(1 (1- 11,0-12
(1 12,
max
max
(1 - 01,e) Ug for all 1 ss 1
(1 31,0112 1t
(1 - 11,1-1)2
(1 B1,1²
(1 11,)
max
max
(1 B1,s)²
Us for all S S (1- $1,c-1)² 1/
(1 32,
31 )2 vo for all 1 }
max
which ends the proof.
Lemma 9. For all t 2 1, we have
Vor- S 7(1 31)
where 0e is in AdamCB (Algorithm
Proof. By Lemma 8
(1 - B1,a²²
6 max
U for all $
(1 31,
31,) 1 S 8 S f. such that i = (I- Us. Recall that by the update rule on it.
Therefore, there is some we have V/2 < 12 (1 B2)9 This implies
Uz- (1- - (1-
3 -
k=1
Hence,
(1 )21,)2
62
1 (1 B2,0)² Us
1 - $1.t
V1 - 32
82 gk k=
2 2
1 B1,5
1 - $1.
S SVI 1 32
Σ
82 k
max ISr5s
|1
1 31,a
k=1
Recall the unbiased gradient estimate St in Eq.(7)
1
9j.4
=
K
Πpja
JEJ.
By the triangle inequality property of norms and the fact that Pt, 2 7/n and llg..ll S L for all the unbiased gradient estimate is bounded by T/7y, i.e,
i E [x] and i. [T] from Assumption |
14

---

Published as a conference paper at ICLR 2025
|19/11 L/y. Therefore,
1 $1,
(L/) 32
88
- 31,8
k=1
B1,
((L/)v B2
31,8
vT - B2
(11/) (1- $1,1
$1,0
L
7(1 - 31)
which ends the proof.
Lemma 10. For the parameter settings and conditions assumed in Lemma 7
we have
T
Mi.u I
T T +1
tūt (1 BIIVI 32(1 11...,
Proof. Recall that by the update rule on my Us. we have my + -1 + (1
BLee)e and
Vt B2 + (1 B2)9L. This implies
m E 31.K)
II k=1 Ar=k+1
3
gk: 1s = (1)
3/ gh
1=1
Since for all 2 0e.u Us by Lemma 8 have
mi,u
m
/ //tt,
/BUL,u
Ei - B1,,)) (IT B1.r) gk.ac
Va- B2)t
(1)) Et b $9220
Σt= (1 =1,,)² (II" ((t-x+111x))) B1.)))
(Ek.)) (II- k+ )
/ V(1- B2) Bt)
35-k gk.
El-n 81-1)
EE 81 9h.u
V(1 BB)t E1-
-k=1 B/)tE111
1 oh.w
Ei= - Ik.u
1
< (1 B1)v1 B2 V Ei Bh g2.
where the second inequality is by Lemma 4 the third inequality is from the fact that $1, < and BI. S. B3 for all < k T. and the Tourth inequality is obtained by applying Lemma 5 to
15

---

Published as a conference paper at ICLR 2025
Ei=1 3i Therefore,
m
1
Σ1=
Th
v tōr,u
(1- (1-B:)v1 32v8
Ek= $8 Ph.u
8
81 - Ph
$
(1 Biv 1-
-32v
3/
t Teru
k=1
1
31 k
(1 BIIVI 32vl
1/ k
1
Σ
kk lg
(1-B:)vI 32v
km
where the second inequality is by Lemma and we define
1) = 91 Therefore,
vB2
T
m
n
gk,
(10)
tuta (1- Bi-B
32
It is sufficient to consider
Σ 1 Ek " Firstly, this can be expanded as:
Vi
Σ Σ 1ghuul =1 1911,
1
V2 [n' 19n. + [2
+ V3
I 19 + 17'22 + n 1993.u1]] II]
1 [1TT 10r, + 3 +... ......0 n [or,a]]]]
Changing the role of | as the common factor, we obtain,
T
1 T-1
Ig
1g
17
+
11
+
11
V2"
v3
VT
t=1
1=
+ 1
12110
1
1
"
+...+
VT 11 T-2
v3"
+ 19s.ul
1
V4 n
1
+
VT 1 T-3
V33"
+...
1 0
+ 1gr,u al
/T"
In other words,
T Σ
T
T
k
k-t
Vi
gk
1t
vi"
1=1
k=1
1=1
hut
Moreover, since
Σ
T
T
b-t
Σ
T-t
1
(1)
1
11
n
Vt
Vt
Vt
v8 1-n)
k=t
h=
k=t
k=0
16

---

Published as a conference paper at ICLR 2025
where the last inequality is by Lemma 5
obtain
T
T
19
Mt.
11.u
Vt
1
Furthermore, since
T -
Σ
1 Vt
gttul
1
Σ Li S 5v1++nn2 In
11.T..l
where the first inequality is by Lemma and the last inequality is by Lemma
we obtain
+I+In's In
Ignwul S
|11.T
1 1- 11
k=1
Hence, by Eq.
m
v1 + In
- -1) |191.T,al
(1
12(1--1
which ends the proof.
B.2
PROOF FOR LEMMA !
Lemma 1.
Suppose Assumptions 2 hold. AdamCB (Algorithm D with a mini-batck size K, size
which is formed dynamically by distribution Pt, achieves the following upper-bouse bound for the cumulative online regret Ronline (T') over T iterations,
19 1
Rontime (T) & PIdVT vdp2
Epp
+ P3
n
(Pj.)) 2
JEJ,
where Pr, and are defined as follows:
D
av1+ TnT
11,T22.
P1 = 20y(1 2 P2 (1 (1-31)3
n)
P3 =
20r((1 -1))3(1) 2)2
Note that d is the dimension of parameter space and the inputs of Algorithm follows these conditions: (a) ot = a
Vi (b) 31. 32 E [0, 1), B1, = 3,X for all E (7). X (0, (c)
7 31//BBB and (d) n E 10,
Proof. Recall Lemma
Since fi Rd R is convex, we have, Se(0*) fe(O.) gI 0.). This means that
f((0.) f(())) - g. (0. 0*)
Ic.u gr.u.0r... (0r,u 0%,))
From the parameter update rule presented in Algorithm
o = 0, - a2174// Vo
= 02 - ai m4-1 + (1 32,1)
B1.1
91
v0r
focus on the u-th dimension of the parameter vector 0g E Rdd Substract the scalar 0% and square both sides of the above update rule, have,
(0.41, -0,,)² = (0... 0²,)² 200
Br,t
mi
+ (1 - B1,t)
gt, (0.. - 0%²)) + ag
Mi,u
Vt,
17

---

Published as a conference paper at ICLR 2025
We can rearrange the above equation
VorN
(0,u
(10t,u -0,))² (0c+1,u 07 07,)))
20,(1 =1,1)
+ 2(1 31,e)
a1
m
B1d
m4 1, (8e, fine)
(11)
v U1.u
(1 Bl,a)
Note that,
T
R onnine (T) E
Σ (0))
-inn BER" 1=1
Σ (+)
E
[f((0)) - [s((0))-
f((or))
where E arg minge Σx (0) is defined as the optimal parameter that minimizes the cumu- lative loss over given T" iterations. Hence,
T
T (0,
R onniin (T) = E
(f.(
((00 SE
03)
E
Io (0. Cnn)
(12)
Combining II with Eq
12 T
we obtain
d
Utu
Ro
S
E
B1,e) (10t.u - -0rn²² (0c+1,u - -8())))
20y [
a==1
1=1
+ F
T
04
mi,u
+
d
T
31,t
111-1111
1, (0 Or.u)
2(1 B1,)) 0r
(1 (1,1))
== 1=1
On the other hand, for all 2 2. we have
(0g-1,u)
1/4
Mc- 1,1
m 1.-1.a(0 (0"a -O,,u)=
10% - 01,1)
10--
(81-1,0) 1/4
11-1,u (8*. 0r.a²² )2 + 04-1
m 1,u
201-1
2/ 0--1,M
where the inequality from the fact that pq S p² + q² for any p. q R. Hence.
v (eeu
Rooline (T) s E
((0r, 8,,, - (0 (7,)))
20
S1,1)
d
T
m
0t
+ E
2(1 - B1.1)
Vūeuu
d
T
32,101-1 m²-uu
+ E2
2(1 - B1,5)
Br,t / 11-1, 1,u
+ E
(0 )²
201-1 11 B1,,)
uml 2m
Since 31,1 S (1 < T), we obtain
d T
V
d
T
B1,t s v / 0/-1.u 10", - -2.,)
20r- 1 (1 -1,e) n 02,,)²
<
20,-1 - B1)
==1 1=2
Moreover, we have
11,10-1 m
d T-1
11,441
mi
2(1 - 31,1) / 11-1uu
2(1 B1,441)
==1 1=1
d
T
m
a4
2(1 31,4+1)
4=1 t=1
T
m
2(1 31) 0z,0
w=1 t=1
18

---

Published as a conference paper at ICLR 2025
where the last inequality is from the assumption that 31,t S B1 1(1 ft< T). Therefore,
T
a2
m
31,10-1 mi-lw
Σ
d T
Σ t=1
at 1
2(1
31,1)
V tt.u
2(1 (1,z) V 0
1 31
vt
and we obtain the bound for Ronline (T) as:
V U2N
R. uunline (T) S E
(10r, -0,,²² (0 -
(13)
20r 31,1)
d
mi.u
04
+
E
(14)
31 v Ut.a
d
T
,1, V
11-1,1
+
E
(6% 0e,u)²
(15)
201-1 11 B1)
Now, we start to bound each term:
and 15
Bound for the term
Let us rewrite the term 13 as
d
ft.u
(10e,u - -8,,² (00t1,u - -8(u)2)
2a
B1.e)
11,1
V 10.,0
E
(01,--0 -8,0)
+
(0t,u - 0?,)?
201
(1 B1,1)
20 (1 B1,1)
lu=1 1=
1
Vūr,
E
(0e, - -Prr)²
E
(0 0770)
20
(1,1-1)
200r
-31,r)
Omitting the last term and replacing a: = /2/ vil I1S we obtain
11
E
(10r,u - -0,)² - -00+1,, - -0+t1,u-0001 0%,)))
20
B1,e)
t=1
(01,u -2,)
20x (1
B1,1)
T
Vlogu
V(t 1)82-14
+
(Or, 022)²
2a
((1 - B1,e)
(1 - B1,1-1)
Recall that by the update rule on ét we have Gr + max { (1- Ur - l,us Vt [I-
Therefore,
0z,u (1-8,,)2
(I-Bi.--1)" 1t and hence
(1-31,)2
/ Lur (1)-))e-14
t (1-11, P1,1111 1)2 12-1,
v(t 1)0
31.))
(1 - 31,1-1)
(1 - (1-311)
(1 31,-1)
v tt
/ - -1)0t-
B1,c-1)
91,--1)
0
19

---

Published as a conference paper at ICLR 2025
Now by the positivity of the essential formula VtR... V(1-1)0 1,u we obtain (1-Bi,t-1)
(1-B1,e)
d
T
E
((0r,u
00,)²
(0c+1,u 11:+1,u-00 -0,)0)
2cve
(1
B1,e)
Lu=1
t=1
D
01,0
D²
T
Vtt- - 1)0c-111
+
E
20
B1)
2a
(1,)
(1 -1,0-1)
Lu=l t=2
D
d
T
$
(ary(1-50)) VT
2a
BITT)
where the last inequality is by Lemma 9
Bound for the term
14
T
2
d
T
m Vtte,a
a2
7.a
E
E
B1
Ut
$1
α
E
d
B1) - B2(1 17) 11
v +1
B1
aV In T +
3 /I
B2(1 7)
u=1
where the last inequality is by Lemma 10
Bound for the term
15 By Assumption 2 that |10m.
0all
D for any m, n E [T), a4 a/vi, Vt.
and B1,t= = 31
I 31 S 1, we obtain
BI 1111,1 (8", 02,0)2
D22
d
T
E
S
E
Σ
B (t
1)82-1,
2a 31)
2a(1 - 31)
Therefore, from Lemma 9 we obtain
31,1 1-- (0'a - -,,,)²
dD2L
E
5
E
Σ
BL (t (t-1)
201-1 (1 31)
20 3112
2=2
1==2
Note that
T
I Σ 3 8 1
(t (=2
3
V (t-
1)
<
Σ Bit
<
81 (1- A)²
1=2
1=2
1=2
where the first inequality is from the fact that B1 1, and the last inequality is from Lemma Thus, the bound for the term 15 is
31,1 11-1,u
d3, d3
E
(6 0c.u)²
5
201-1 (1 -31)
20
(1) X]2
We bounded for terms
14 and 15
dD2L
av In +1
d Σ
R. (7) <
VT
E [11gIT,1.1]
20 - -1132
(1- 31
I 35(1 17) m=1
d3
+ 20ry(1 =31 2)2
Hence,
R. S S1dpp + P2
Σ
[lI +
(16)
20

---

Published as a conference paper at ICLR 2025
where P1, P2. and Ps defined as the following:
D P1 20 - 81)2 (1 -1- 3.) 32(1 - -))
+
d31 D2L
20y(1
B1) 2
Now, we consider Σ
E
[1lg.T,,ull]
which in the right-hand side of Eq. (6).
T
T
Σ Ellg
E
Σ r=1
S d
Σ mim
vd
Σ
E [1laa
where the first inequality is due to the concavity of square root. Recall that the unbiased gradient estimate is 9t Ejed, 9id, Hence, npsit"
Realine
+ P2 Σ Ep. IlIgu:T.all + Ps
S pidvT +pvd
Σ
py
|llee1
+ Ps
t=1
S pPdd +
Σ
Eyg
K JEJ,
np,,t
Uj.t
P3
The last inequality uses Jensen's inequality to the convex function
||
112. Therefore,
Runtine (T) SidvT P2Vd
Ep.
Σ oEJ
yis
+ P3
n²
K2
Ps.a
t=1
T
1 2 (ps..
S SpidVT + pavd
Ep
Σ
+ P3
n K
2
t=
JEJ
where the last inequality is by Lemma
This completes the proof of Lemmal
B.3 PROoF LEMMA
Lemma 2. Suppose Assumptions hold. we set min 1,
In (n//)
the batch selection
(e- 1)TK
(Algorithm 2 and the weight update rule (Algorithm) 3) following AdamCB (Algorithm) ! implies
T
|1aH. 11
(V
n Kni In
Ep.
min
Ep..
(p
(p),a)² 0z
K
t=1 JEJ.
JEJ
Proof. We set Ej Payga (-
12
in Algorithm Since ||9,, S L and Pi,e Pmánn
(p
for all E [n] and E [I]] by Assumption I
we have lee E [0.I]
Let W, Ei= Wi. Then, for any (TT.
We
Σ
Wi.t
Wa.t W
W
W.1
Sma.
Ky KT
Σ
Wi.1-1
Whi-1 Wi-1
exp
W
n
21

---

Published as a conference paper at ICLR 2025
The last equality is by the weight update rule in Algorithm Algorithm 2 we have
3
From the probability computation in
Pi,e = K (1-
Wi,t
+
2
2
Ky
Ei==1 ==1 w
11
n
Thus, we obtain the following bound,
Ky lis
Kylse S lipe S 1
0 5
npirr
By the fact that e I 1- - x + (e 2)x² for all x E [0, 1]. and considering Ky ls,e as x, we have
W₁
wi,t-1
K7 bit + (e (e-22 (2)
Ky 2 (1,t)
W1,1-1 1
W
Wi-1
W. 1
sE(m]
Sma.
iESun
Wi,t-1
K7 Ri,t + (e-2)( K7 P3.t)
= 1 +
W
vE(n Veln
PaL
Ky 1 + (c (2)( Ky li,1)
=1+
Σ
K
7
n
K(e - 2)72 21)2 n²(1 -))
< 1
7
Σ
Puclet +
Σ piolls,e)2
n21 7)
iE(n](SSur
K(e - K(2-2)7 Σ lee
< 1 7 n01 VEJiYSun.
la,t
n²(1 -?) iE(m]
The last inequality uses the fact that Pi, = l, < for i E Jj and pr. 0 for i & J.. Taking logarithms and using the fact that In(1 + +2)= S x for all x > -1 gives
We In W
7 n(1 ((1-))
latt + K(e 2)n²
Σ les iE(m
S
n² - 7)
By summing over f, we obtain
WT I In W1 n(1
T
K(e - 2) T
Σ le+
Σ Σ lus
n² (1- [=1 [=1ie(m ie(m)
[=1 FlliJe)ae S.mmi
On the other hand, for the sequence {JA] of batches with the optimal Σt= EJEs, li among all subsets Jj containing K elements,
WT In W1
In Eje; WST
Ejer; In
K
2
+ In
W1
K
K
Σ
b In
n
JEJT tijg
The first line above uses the fact that
Σ
Wjr 2 W1
lEd
and the second line uses Wj.r exp (-(K7/n) Eezgsman 2zt). From combining results,
ΣΣ
n
K
T
T lett K)
Σ lae
In
<
7
7)
n11 (1)
JEJ; E.f&S
1=1 JeJe
1=1 (E(m))
22

---

Published as a conference paper at ICLR 2025
Since Eye;; Eugesmart S4 S 1- Et= Eiesun la trivially holds, we have
Σ jEJ; t:j&Sun,
t
Σ E
ln
K
1
lse + (e k) lep
2
7)
n(1 7) ic(n)
jEJz t.jESuil
tm iEJ;
Let LMIN-K(T) i= ET= Ejer; la and LEXP3xX of both sides and using the properties of fie,, we obtain,
= = Eje,, Cae. Taking the expectation
K K (T)
(e (e-2)K7 n(1 Am iE(n]
7 In n (1 7)
([LExpAK(T)]+
Σ
Expt
This is because the expectation of EGe is Le2 from the fact that DepRound selects i-th sample with probability pi,t- Since Σ Ei=1 lift pT) have the following statement,
K
LMIN-K(T)) - E[(LEx3-x S (e 1)
K
Using the fact that (T) TK and choosing the input parameter as Y
nln (n/K) min 1, (e- 1)TK Jo we obtain the following,
LMIN-K (T) - B[LEEPAA((T) S 2VE KnT In n < 2.63 K
7 KnT' In
K
Therefore, considering the scaling factor, we have:
1gg
11g,ee1?
L
min
(LMIN-K(T) - EE| LEXP3X33 (T)))
Pr
(p).e)²
Pt
(py.c)2
Prii
JEJ.
jEJ.
2.63L2
S
KnTIn
Pim
(O(K |KITin n
o
K
This completes the proof of Lemma
23

---

Published as a conference paper at ICLR 2025
B.4 PROOF FOR THEOREM (REGRET BOUND ADAMCB)
In this section, we present the full proof of Theorem [ Recall that the online regret only focuses on the minimization over the sequence of mini-batch datasets {Di) Thus, the online regret of the algorithm at the end of T iterations is defined as
Σ BER
Romline (T) = Σ ((0D))) min
However, ultimate goal is find the optimal selection of the parameter under the full dataset. Consider an online optimization algorithm π that computes the sequence of model parameters 01...., 0r. Then, we can compare the performance of with the optimal selection of the parameter mingeR f(;; D) under the full dataset. The cumulative regret after 7 iterations is
Rs (T) = E
f(0;; T min f(0: D)
dER"
where the expectation is taken with respect to any stochasticity data sampling and parameter estimation. Before we prove Theorem [ we first prove the following lemma.
Lemma 11. The cumulative regret R² (T) can be decomposed into sub- which includes the cumulative online regret ontine (T) and additional terms that are sub-linear in T:
R (T) - (T) O(VT)
Proof. First, rewrite (T) by expanding the terms inside the expectations. We add and subtract the sum Et= f(0;; Dj) inside the expectation:
T
Rd (T) = E
Σ [(0;;)) - -T min f(0; D)
DER4
T ((0;)) Σ f(e,: D.) + Σ D.) T min f(0; D)
T
T
E
Σ
BE
We also add and subtract the term mingeRd the following.
Σt f(0; D,) inside the expectation. Then, we have
T
T
T
R" (T)=
((0rD)
Σ
f(0r;D.)+ +
Σ f(0,; T min f(0; D)
Oel
1=1
T
T
= E
F(0::D) Σ ((0::Di)
+E
f(0;:
> 5(0:D.) 0
min
1=1
E
min BERd
f(0; D:) - T min f(0 D)
Since the second term of the right-hand side in above equation is equal the online cumulative regret R (T), we can rewrite R# (T) as:
RRd (T) = Rs Ru (T))
T Σ F(O::D))
+ E
f(0::)
(17)
1=1
E
min Σ
f(0:Di) -T
min f(0: BERd
(18)
BERd
1=1
Now, let consider each term in detail.
24

---

Published as a conference paper at ICLR 2025
Bound for the term 17 Recall the expression of f(0; D) and fu; = f(0; Di):
s(0: D) 1 Σ ((0:.. f(0:
1
Σ
f(0;T,,
K
nPj.
jEJ.
where Jr the set of indices in the subset dataset (mini-batch) at iteration t, De S D. For any 0 E IRd, we have
((0:I,.
E(0; Mj)
((f(0: =E
E
K
nPjr
K
TiPj.
jEJ,
JEJ,
8(; Ti, V:)
Pi,t
E(0; v:) f(;
K
nPat
n
jEJ,
Note that, by linearity of expectation, we can interchange the expectation and the summation, Since [0; D,)] = f(0; D), we have for the term as:
f(0;;)) D) Σ f(e:
[f(O;:D) f(00: D.Jl
B[f(0;: f.).: =0
Bound for the term 18 Let 0* be the parameter that minimizes the cumulative loss over the full dataset D, i.e, E arg mingeRs (0; D). Since is optimal for the full dataset, we have:
min [0; D) = ;
BER4
Similarly, denote the optimal parameter for the cumulative regret for mini- mini-batch datasets by 07 arg mingeR Σt= f(0; Dz). Given these notations, we can write the term (88 as:
118
E min Σ [(0,D.) -T min [(0 E Σ [(0; T ((00)) BER4
BER
We can add and subtract the term ET= f(0"; D.)
inside the expectation.
T
E
D,) - T ;()":D)
E
f(0;; D,) - 5(0:
1=1
E
Σ ;(0::))) T 5(0;
Note that [[ = (() holds as we have shown when bounding the term
17 By the
linearity of expectation, have
T
T (() = T(0;)) D)
E
(()):
t=1
Since E IT= f(0 D.) T -((0; D)
0 holds, the term (8) reduces to
T
18
= E
D.) f(0::
Σ (f(())) - 1((077)
E
25

---

Published as a conference paper at ICLR 2025
By the convexity of ft, have:
f((0)) fe(0)) S g. (0: - 0*)
Therefore,
T
T
f((0)))
S
Σ oe (0 - 0*)
Using bounded gradients assumption (Assumption Cauchy-Schwarz inequality (Lemma we have
Lee, lIgell I L/y (Proof in Lemma and
2'
T
18
E
Σ
ge O0 - 0*)
Elllelllo; - -00 S (L/)) Σ
Billo 0-
Recall the parameter update rule, dst1 < 0g a.m+/ + c). Then
10 - 0- S 110 0r + an / h/me/(/// + +
(19)
Now, we claim that [|m,]] is bounded. The update rule for the first moment estimate:
my + -1 (1 $1.e)9t
Then, the expression for my is:
m E B1,k)
II Birr 9k
k=1
r=#++
where BI.; = B1At-1 with B1 < and < Note that || is bounded by L/> for all This
implies that:
ml 5
31,1
II 31.
|g
1=1
1==++1
2 II
(L/))
1 B1Ab=11
3,
(1hh) Σ
81 1
MC6 13 Acck 13
((1/) Σ
81
k=1
L
7(1 B1)
The last inequality is due to Lemma
Therefore, the step size in Eq. 19 aL B1) Vley(1 - 31)
bounded by:
ael|my |I
ey(1 a,L
V0g +
We use the fact that ar = a./v.. By summing over T iterations, we obtain
T
aL
2 67(1 - 31)
Σ
Blllo; 0*
e7(1 31)
The last inequality is by Lemma Finally, get
TTL2VT
18
(L/r)
Σ E|10 0-1]]
21 B1) = -
In summary, the cumulative regret Rd RT"T") is decomposed by the following:
R. (I') R (T) + 17 + 18
where
17 0 and 18
O(VT). Thus, this completes the proof of Lemma saying Rd (T) = Romin (T) + O(VT)
26

---

Published as a conference paper at ICLR 2025
Now, prove the main Theorem
Proof. From Lemma 11 we have shown that the cumulative regret R (T) can be decomposed into the online regret R R (I') with the additional sub-linear terms. Hence, we are left to bound the cumulative online regret R oaline (I"). Recall the first key lemma (Lemmal 1
19 12 (P5,,) )
R. oaline (T) S PIdvT vdpa
Ppp
P3
K
jEJ,
Recall also the second key lemma (Lemma
2
Σ
T
Ep
Σ
1 (pj.))2
min Σ
T
Ep
19 11 (P),a) )
-
KnT In n
K
JEJ.
JEJ
Let we denote M minp,
Ei=1 Ey,
|1 Eser
Then by Lemma 2 we have
(ps.e) 1
119 (P,,)2
P1
M + C KnTin K
JEJ.
where C >> 0 is a constant. By plugging above equation to Lemmal D we obtain
Vd
oaline (T) S pidV + P2
M + C KnTTn n
K P3
nvK
Vd
Vd
SpddT + nvK K
VM + pr
CV
CVKNT
KnTTn
In
n
+ Ps
nvK K
K
pidVT P2vd
P4Vd
nT
1/4
M
In K
+ P3
RvK
We use the fact that V S Va+ vō in the second inequality and we define pa == P2VC.
Now, should consider M. Using the tower property, can express M as,
M min
Ep
(2),)²
)2
Pe
jEJ
|| 11?
min
Bps
Pi
(pj.e) )2
| Pt.
JEJ.
T
11ah.el1
min
E
(pr.e)² Pi,e
==1
JEJ
T
1gr.ell2
min
Epp.
ps
Pi.
1=1
jEJ.
T
l1gr.
11?
K min
Pi.
For this minimization problem, it can be shown that for every iteration é, the optimal distribution PF is proportional to the gradient norm of individual example. Formally speaking. for any t, the -ll...ll for |lo..e1" is (P5..)" >m |ls.all
optimal solution P the problem arg min,, ΣT. Ep. all j E [n]. By plugging this solution,
Σ
Pi.e
T
M = K
Σ
Iga.ell
t=1
27

---

Published as a conference paper at ICLR 2025
By plugging M to the online regret bound expression,
vd
nf
1/4
Reinine (T) S pidvT + P2
M
P4
In
+ P3
nv
K
K
Vd
T
2
Vd
nT
1/4
= + nvK
1
+ P1
In
+ P3
K
K
p + Vdp2 VVp2
1
T
2
+ P1
Vd
nT
In
n
1/4
+ Ps.
n
K
K
By Assumption 1 |Iahell I for i E [n) and t E [T]. Then, the second term in the right-hand side of above inequality bounded by Lp2vdT, which diminishes by the first term that have order of (dv/)) Hence, the online regret R Rutime (T) after T iterations is,
vd nf n
1/4
Rolline (T) O(dVT) o
K
Finally, by Lemma 11 can bound the cumulative regret using the bound of the online regret
R7 (T) R R (T)
+
S O(dVT)+++
Vd
nT
ln
n
1/1
+ O(VT)
K
K
Vd
T
#
o d dVT+ +
In
3//
K
K
This completes the proof of Theorem
28

---

Published as a conference paper at ICLR 2025
C PROOF FOR CONVERGENCE RATE WHEN USING UNIFORM SAMPLING
To compare the convergence rate between using uniform sampling and bandit sampling, we will now prove the following Theorem 2 It is important note that Theorem 2 includes an additional condition- AAstiionnnssumttoonnooooooo which was not present in Theorem
This assumption plays key role 1
in distinguishing the results between these two theorems.
Theorem 2. Suppose Assumptions 72 and 3 hold.
The convergence rate for AdamX (variant of
Adam) using uniform sampling given by:
o
Vd ((aVT=
VT) n2/2
Proof. We start the proof from the first key lemma (Lemma
Suppose Assumptions Z hold. AdamCB (Algorithm 7 with mini-batch of size K, which is formed dynamically by distribution Pt. achieves the following upper-bound the cumulative online regret Ronnime (T) over T iterations,
Lemma 1.
19 1 (P,,)) )
R (T) S PıdvT Vdpz
n²K
pp
+ Ps
(20)
JEJ,
where PL p2, and ps are defined as follows:
D
aV av1+tn? IT
D2L
P1 20 -1)2
P2 (1 3. -)
P3 200 31) X)2
Note that d is the dimension of parameter space and the inputs of Algorithm ffolows these a
conditions: (a) ai
Vp (b) Sr. 32 E [0, 1). BI, == 3, for all E [T]. E (0,1), (c)
7 31/v and (d) n 10.
Consider the second term in the right-hant side of (Q.
1
T
|| 1 (Pj.e)2
Exp
E
Epp
Pe
n n2K K
(pi,))
n K
ps
JEJ.
jEJ.
T
119s,ell
Epp
(pi.e) )22 Pt.e
2 K
JEJ,
Ep
n
K
Pi.t
jEJ,
T
11?
n
Pia
The tower property is used in the first equality. Since mini-batch size K is multiplied in the last equality. Therefore, express the cumulative online regret R (T) as:
Σ
los is independent to j Jr, the
Pit
1
T
11gt, 1|2
Relline (T) DidVT Vdpa
n2
pt
Pi.e
+ P3
In the case when we select samples uniformly, we can set the probability distribution pt to satisfy Pa, = 1/n for all E [T] and E [n]. By plugging it, we obtain
R
- pidvT= Vapp
Σ n
Σ
l1 12
+ Ps
29

---

Published as a conference paper at ICLR 2025
Now, recall Assumption 3
Assumption 3. There exists o such that Var(Ilgi, (l) S o² for all E [n] and E [T]
|sscell) + o²
I 11
n
Therefore, the online regret bound R Ratine
(T) for uniform sampling is,
2
0²
R online (T) = = +
Vd.
T
n
Applying the fact that va+b Va vb, we obtain,
T Vā
R.T Roaline OT)=
Vd
119
+ O
n
By Assumption |ga.ell < L for [n] and t E [T]. Then, the second term in the right-hand side of above inequality bounded by O(vdT), which diminishes by the first term that have order of (dvTT). Hence, the online regret R (T) after T iterations is given by
Vd VT
R unline = +0
n 1/2
we bound the cumulative regret using the online regret, which completes the regret analysis for uniform sampling.
Finally, by Lemma 11
Vd
Rs (T) Roaline (T) =)++) +0
VT + +O(VT)
n/2
dVT + Vd
VT
n/2
30

---

Published as a conference paper at ICLR 2025
D CORRECTION OF ADAMBS (Liu ET AL.
2020)
This section introduces the corrected analysis for AdamBS Liu et al. 2020). We use Algorithm and Algorithm for modified AdamBS.
Algorithm 4: (Corrected) Adam with Bandit Sampling (AdamBS)
Input: learning rate {ar)?= decay rates (,1,,)/m 82, batch size K, exploration parameter ? E [0, 1)
Initialize: model parameters Oo: first moment estimate mg 0: second moment estimate 00 - 0, in f- 0: sample weights w/ < for all E [x]
for t to T do Compute sample distribution Pr for all E [n]]
2im1 Wirr + 2
Pj. (1 - 7)
Σ
Select a mini-batch D. := {zz, Uy DeJ. by sampling with replacement from Pt Compute unbiased gradient estimate with respect to the mini-batch Dr using Eq.
my + m-1 1 + 01-111111 -B11)
vt - 32 + (1 B2)9
i1 - U1 0g - maxx o1.2
(1-31--1) 1 v1 ift 2 2
s.t1 t 0: - a,my/ / (V +
Wr Weight-Update(waa Pcs Jr, iogneiseee 7) (Algorithm
Algorithm 5: (Corrected) Weight- for AdamBS
Input: W4-1, Pt, Jr, {ggatiEJ. and E [0, j = to n do
Compute loss fj.r
P
(- ||ss.e11 + I.² [Ps.a)²
if Jr, otherwise, Cj
Pinn
Compute unbiased gradient estimate E,, 83.0 E 1()=())
Kpja
Update sample weights Whe + Wj,- 1ep-1expppppp exp(->ll.e/n)
return Wi
At iteration t [T]. AdamBS chooses mini-batie D. (iey.:))sse. of size K according to probability distribution with replacement. We denote Jr as the set of indices for the mini-batch D,. Then, the algorithm receives the loss, regarding losses from all chosen samples in the mini-bater D loss, & Ejen, lg denote as li [0, 1]. The unbiased estimate of the loss lz is,
E,4 L/ Σk=1 I(j )(()
K
KP3
We have a following key lemma concerning the rate of convergence AdamBS. Lemma 12 (Corrected version of Lemma 1 Ln al. (2020). Suppose Assumptions 7 hold. If set 7 min
1,
nlnn (e-1)T
the weight update rule (Algorithm following AdamBS (Algorithm 49
implies
1 (p
19 (p
Epg.
min
O(KVnTInn)
jeJ.
Algorithm for all E [ [n] by Assumption have la E [0, 1).
Proof. We set b,,t = -liyaa
! (P
Since, 1|gr,ll2 L and Pi,t 2 Pinin
Pinn
31

---

Published as a conference paper at ICLR 2025
We use the following simple facts, which are immediately derived from the definitions,
Prilii Σ f/+==
== E.¹
(21)
K JE4
( D k== 1(i 1(i= Jti (()
" Eh. 1=1
2K=1 1(i 1t)
Σ pr.((er.)²= Σ pil
les S Σ
la
(22)
Kpt.t
K
1m1
Let W₁ == Ei= Wr. Then, for any l E [T].
W
n
Wiet
W
We-1
Wi,t-1 1
(-2
exp
Wi-1
n
The last equality is by the weight update rule in Algorithm Algorithm 4 we have
5
From the probability computation in
W1
2 > 2
Pi. (1 - 7)
+
==1 Wjt-1
n
n
Thus, we obtain the following bound,
41 Σ k= I(i= JoC) J()=
0 -l. i,t
2'
< Sli $
n
Kpi,t
By the fact that e x S 1 - x + (e 2)x²²
for all x E [0, 1]. and considering n 2 as x, we have
W
n
W
-- + (e (e-2) (2)))) n1
W
W
12
1=1
n
Pi 7/n
7 Ml (e-2) 2) (a,))
1
1 1-7 7
n
1=1
1- 7 Σ
7/n
pirlee +
(e (e-2)(7//m3² 1 1-7
Σ
prolls.)2
$ 1- el + (e - (2)2(///)3
7/n
n
1-7
The last inequality uses Eq. for all > - gives
21 and Eq. 22
Taking logarithms and using the fact that (1 + x)
S x
W.
7/n p/r + (e (e-2) (2))////))
S
4
W
1 -
1 1-7
By summing over I, we obtain
WT W
7/12
(e - 2)
de +
On the other hand, for any action j.
In
W W
In
J.
7 T
8 Inn
W
t=1
From combining results,
T
T
T
Σ
c
2 (1 7) Σ
n Inn
(e
2)
7
l
t=1
32

---

Published as a conference paper at ICLR 2025
We next take the expectation of both sides with respect to probability distribution pt and since Ep, [3,1] Est have
2 7) 53 n nn
T
E ps
(e 2)
Ei
1=1
Since j Je were chosen arbitrarily, can choose the best J for every iteration Let (T) Eiz EjEr; E/e and (T) = Ei= Eje,, lyt. Summing over j E Jf, and using the fact that Eí Lisl lepe mlvn((T) we have the following statement,
K
nK Inn
LEXx (T) 2 7) [TIN(T)
(e (MI(N
Then, get the following.
LMIN (T) EXxx (T) (e I)IINI((T) nk Iun
?
Using the fact that LMIN (T) < TK and choosing the input parameter as 7 = min 1, v
In (e-
obtain the following.
(T) EXLExss (T)) < - & 2.63K VTT
Therefore, considering the scaling factor, we have:
19,4 12
19 12
L
Epp
min P4
Ep
Σ
(LMIN(T) - EE Lexes (T)))
(p),) 2
Pt
(Pj,.)
Pinn
JEJ
jEJ
2.6312
K KVnT'nn
Punn OKK VnTlnn)
Theorem (Corrected version of Theorem 4 in Liu et al. (2020)). Suppose Assumptions 72 hold. The convergence rate for (corrected) AdamBS using bandit sampling is given by:
o ( d Vd 33/4 (Tlm))1/4
Proof. From Lemma 11 we have shown that the cumulative regret R (T" can be decomposed into the online regret R Ro (I') with the additional sub-linear terms. Hence, are left to bound the cumulative online regret Roaline (T). Recall the first key lemma (Lemma
T
9 2
R (T) S PIdVT - vdpz
Bpp
+ P3
n²K
(p..))
jEJ
We can apply Lemma to AdamBS as AdamCB, since both AdamBS and AdamCB follow the same model parameter update rule. However, we use the corrected lemma (Lemma (2) for AdamBS, rather than applying the key lemma (Lemma used for AdamCB. Recall Lemma 12
T
1gg (2
T
19 (p
E
min
E
=O(KV (nTlnn)
pt
JEJ.
1=1
JeJ.
33

---

Published as a conference paper at ICLR 2025
Let we denote M == minp,
Ei=1 Ep.
EjEre
| (pg,a)²
Then by Lemma 12
we have
T
19 1 (p),,)² )2
Σ
M + C'KVTTnnn
pt
1=1
JEJ
where C² > O is a constant. By plugging above equation to Lemma] 1 we obtain
Vd C+CKKVA C'K M'KyTThnn Inn P3
R.
oaline
(T) S pidVT + P2
M
ny K
Vd
Vd
ppddT + nvK
VM
+
P2
nvK C'KvnTIn +
p2vd
PsVd
ppddT
VM
(nTInn) + +B3
RVK K
We use the fact that +b+bb < va + vb in the second inequality and we define Po == P2vC" Now, we should consider M. Using the tower property and applying the optimal solution for pt at each iteration, we can express M as,
T
n
M K
This follows the same argument as in the proof of Theorem (see Appendix B.4). M to the online regret bound expression,
Then, by plugging
Vd nv K
Vd M
RsT online Spidi
(nT (nT (n
02
PS
n
vd
T
+ Pss vd
-iddT + P2
K
Σ gi.eell
(nT In + P3
nvK K
1=1
Ps Vd
= PIdVT vdpa
1ss.ell
(nT Inn) 1/4 + Ps
n
By Assumption I |gg.ell for [n] and [T]. Then, the second term in the right-hand side of above inequality bounded by Lp2v dT, which diminishes by the first term that have order of (dvT). Hence, the online regret R* R (T) after T iterations is,
Roniine = Odvv) +o (nTmmn)//)
(Y
Vi
Finally, by Lemma I1 we can bound the cumulative regret using the bound of the online regret as
Vd (nT (nT)m))/)) + O(VT))
(T) R. valine
(T) +0VT) O(dVT)
Vd
o dvT
3/4 (Thm))1/4
This completes the proof of Theorem
3
34

---

Published as a conference paper at ICLR 2025
E ADDITIONAL ALGORITHM
E.1 DEPROUND ALGORITHM
Algorithm 6: DepRound
Input: Natural number K(< n),, sample distribution = pp'.p" .p..... ......,pp)) with Ei=1 K
Output: Subset [n] with distinct K elements
while there is i with 0 < p² < do
Choose distinct i,j with 0 < p < 1 and 0 < p <1
Set a min - p'.p' } and [ = min(p', (o', p²}
Update y and p³ as:
(p' + ta,p' -a)
with probability o+3
(,'
(p' 3.p' + +5) with probability a+3
return (i:p' 1,1,
The DepRound (Gandhi et al. 2006) (Dependent Rounding) algorithm is used to select a subset of elements from a while maintaining certain probabilistic properties. ensures that the sum of probabilities is preserved and elements are chosen with the correct marginal probabilities.
F
MORE ON NUMERICAL. EXPERIMENTS
F.1 DETAILS EXPERIMENTAL SETUP
We compared our method, AdamCB, with Adam, AdamX, and corrected AdamBS. The experiments measured training loss and test loss, averaged over five runs with different random seeds, and included 1-sigma error bars for reliability. Throughout the entire experiments, identical hyper-parameters are used with any tuning as shown in Table
Table 2: Hyper-parameter used for experiments
parraee
Value
Learning rate a; Exponential decay rates for momentum 31.1, ,2
0.001
0.9, 0.999
Decay rate for 31,1 for convergence guarantee À for non-zero division
1-le-8
1e-8
Loss Function Batch Size K exploration parameter y
Cross-Entropy
128
0.4
Number of epochs
10
We trained MIP models on the MNIST, Fashion MNIST, and CIFAR-10 datasets. The detailed architectures of the ML. models for each dataset are provided Table
3
Table 3: MLP Architecture for MNIST/Fashion MNIST (left) and CIFARIO (right)
Layer Type
Input
Output
Layer Type
Input
Output
Flatten
(N, 28281)
(N, 28281)
Flatten
(N, 32323)
(N. 32323)
Dense ReLU
(N, 28281)
(N, 512)
Dense ReLU
(N, 32323)
(N. 512)
Dense ReLU
(N, 512)
(N, 256)
Dense ReL.U
(N, 512)
(N, 256)
Dense
(N, 256)
(N, 10)
Dense
256)
(N, 10)
We also trained CNN models on the same datasets. The detailed architectures of the CNN models for each dataset presented in Table
35

---

Published as a conference paper at ICLR 2025
Table 4: CNN Architecture for MNIST/Fashion MNIST (left) and CIFARIO (right)
Layer Type
Input
Output
Layer Type
Input
Output
Conv ReLU
(,, 32, 32)
(N, 64, 32, 32)
Conv + Rel.U
(N, I, 28, 28)
(N, 32, 28, 28)
MaxPool
(N, 64, 32)
(N, 64, 16, 16)
MaxPool
(N, 32, 28, 28)
(N, 32, 14, 14)
Conv + ReLU
(N, 64, 16, 16)
(N, 128, 16, 16)
Conv + ReLU
(N, (N, 14, 14)
(N, 64. 14, 14)
MaxPool
(N, 128, 16. 16)
(N. 128, 8,
MaxPool
(N, 64, 14, 14)
(N, 7.
Conv ReLU
(N, 128, 8,
(N, 256, 8,
Flatten
(N, 64, 7)
(N, 3136)
MaxPool
(N, 256,
(N, 256, 4,
Dense
(N, 3136)
(N, 128)
Flatten
(N, 256, 4,
(N, 25644)
Dense + Softmax
(N, 128)
(N, 10)
Dense
(N, 25644) (12)
(N, 512)
Dense Softmax
(N, 10)
Table 5: VGG Architecture for MNIST/Fashion MNIST (left) and CIFARIO (right)
Layer Type
Input
Output
Layer Type Input
Output
Conv ReLU
(N, I, 28, 28)
(N, 64, 28, 28)
Conv ReLU
(N, 3, 32, 32)
(N, 64, 32, 32)
Conv + ReLU
64, 28.
64, 28, 28)
Conv + ReLU
(N, 32, 32)
64, 32, 32)
MaxPool
(N, 64, 28, 28)
(N, 64, 14, 14)
MaxPool
(N, 64, 32, 32)
(N, 64, 16, 16)
Conv + ReLU
(N, 64, 14, 14)
(N, 128, 14, 14)
Conv + ReLU
(N, 64, 16, 16)
(N. 128, 16, 16)
Conv + ReLU
(N, 128, 14, 14)
(N, 128, 14. 14)
Conv ReLU
(N, 128. 16, 16)
(N. 128, 16, 16)
MaxPool
(N, 128, 14)
(N, 128, 7,
MaxPool
(N, 128, 16, 16)
128, 8)
Conv + ReLU
(N, 128, 7)
(N, 256, 7,
Conv + ReLU
(N, 128, 8, 8)
(N, 256, 8, 8)
Conv + ReLU
(N, 256, 7)
(N, 256. 7.
Conv + ReLU
(N, 256, 8, 8)
(N, 256, 8, 8)
Conv + ReLU
(N, 256, 7)
(N, 256, 7,
Conv + ReL.U
(N, 256, 8, 8)
(N, 256, 8, 8)
MaxPool
(N, (N, 256, 7)
3,3)
MaxPool
(N, 256,
256, 4,
Flatten
(N, 256, 3)
(N, 2304)
Flatten
(N, 256. 4)
(N. 4096)
Dense
(N, 2304)
(N, 512)
Dense
(N, 4096)
(N. 512)
Dense
(N,
(N, 512)
Dense
(N, 512)
(N, 512)
Dense
(N, 512)
(N, 10)
Dense
512)
10)
We also evaluated the AMSGrad optimizer and the corrected AdamBS algorithm (Algorithm 4 on the CIFAR-10 dataset using both MLP and CNN models. The results are presented in Figures and 4 From these plots, it is evident that our AdamCB algorithm outperforms the other Adam-based algorithms. To further assess performance, we conducted experiments using the VGG model, which is a larger architecture compared to the MLP and CNN models. The detailed structure of the VGG architecture provided in Table and the results are shown in Figure
Adam
AdamBS
AMSGrad
AdamX
AdamBS (corrected)
AdamCB (ours)
MNISTT
[b] Rashion MNIST
CBAR-10
2.
Tran Lo06
20-
0.3
1.4
Spochs
22-
2.3
Epocis
Figure 3: Comparison Adam-based optimizations on MLP model
36

---

Published as a conference paper at ICLR 2025
Adam AdamX
AdamBS
AMSGrad
AdamBS (corrected)
AdamCB (ours)
MNIST
Rashion MNIST
CIFAN-10
223
2.-
Tran L:55
175
3.00
032
0:2
Eposta
spocts
spachs
2.0
232
Tock L203
150
125
822
Epocts
Epects
Epachs
Figure 4: Comparison of Adam-based optimizations on CNN model
Adam AdamX
AdamBS
AMSGrad
AdamBS (corrected)
AdamCB (ours)
(a) MNIITT
[b] Fashion MNITT
(c) CIPAR-10
3.2
Tran La05
10-
Epocts
Epects
Epochs
225
20-
272
2.-
250
00
Epoche
Epocts
Spochs
Figure 5: Comparison of Adam-based optimization on VGG model
F22 ADDITIONAL EXPERIMENTS
To further evaluate the effectiveness of our proposed method, we conducted additional experiments
using logistic regression, R8 (He et al.] 2016), ConvNeXt-Bare (Liu et al. 2022). and
Adam
AdamBS
AdamCB (ours)
26
Loss
Test Loos
Epachs
Epochs
Figure 6: Comparison of Adam-based optimizations on the logistic regression model (MNIST)
37

---

Published as a conference paper at ICLR 2025
AdamBS
AdamCB (ours)
MINIST
Reshion MINST
CIFAR-10
20-
1.2
Epochs
3.2
50-
0.3
4.9
33-
Eoocto
Epects
Epochs
Figure Comparison of Adam-based optimizations on ResNet-18 model
Adam
AdamBS
AdamCB (eurs)
(a) MNIST
[b] Rashion MNIST
CIFAR-10
20-
Tatking L250
7.0
1.5
1.3
03-
12-
1.5
0.5
80
Epocis
Epects
Spochs
Figure 8: Comparison of adammbbased optimizations ConvNext-base model
Adam
AdamBS
AdamCB (ours)
(a) MINIST
(b) Fashion MNST
(c) CIFA--00
3.5
Epochs
Spochs
Epochs
K.15
4015
6.5
Epachs
Spocha
Figure 9: Comparison Adam-based optimizations on ConvNext-large model
38

---

Published as a conference paper at ICLR 2025
LonvNeXt-Large (Liu et al. 2022 networks. The archecture of The logistic regression model was employed to assess the performance of our algorithm in convex optimization settings.
For general non-convex optimization, we tested our method the ResNet-18, ConvNeXt-Base, and ConvNeXXt-Laree models. Notably, ResNet-18 (11.4 million parameters), ConvNeXt-Base (89 million parameters), and ConvNeXt-Large (198 million parameters) architectures compared to the simple MLP and CNN models evaluated in the previous section. These experiments demonstrate the scalability and efficiency of our algorithm on larger, more complex models.
substantially larger
In all experiments, proposed algorithm, AdamCB, consistently outperformed existing methods, reaffirming its effectiveness across both convex and non-convex optimization tasks and on models of varying complexity.
39

---

Published as a conference paper at ICLR 2025
G WHEN L IS NOT KNOWN
Algorithm 7: Wei th Weihhtuaaaaee (with unknown L)
Input: Wt. J,, (gisteede S, 7 [0, 1), -1-1 Set Lt + max Lt- 1-11 maxjEJ, (Ig)eel)
for j = to do
Pinga (- Compute loss f,, L.
11 [Psa)=
12 Piin
ifj Jji otherwise E/e
j & Smil, then
Wjja t Wja-1 exp (- -Klia//))
return Wi, Lt
Lemma 13. (Lemma when L is unknown) For all 2 1, we have
L
Vo S 7(1 B1)
where 0n is in AdamCB (Algorithm
Proof. The argument follows the same reasoning as presented in Lemma 9 with the modification that L is replaced by L,, reflecting the condition that I L, for all i [ at any t.
Lemma 14.
(Lemma 7 when I, is unknown) Suppose Assumptions 7 hold. AdamCB (Algorithm with mini-batch of size K, which is formed dynamically by distribution pt. achieves the following upper-bound for the cumulative online regret R ontine (T)) over T iterations,
T
Mitt 12 (Pj,))
Roniime (T) S P/dVT Vada
+ P/3
n2k
Pd
JEJ,
where phy Pby, and & are defined as follows:
Dº (a) 0s V. (b) Br. B2 [0, 1), 31, == 31 -or- all (T), X E (0, 1), (c) n = 31//B2 S 1,
0V Tn
d3, D2 LT
A 20 -))
P/2 -= - (1-3,)3V1) -1
AS = 20 31) X)2
Note that d is the dimension of parameter space and the inputs of Algorithh follows these conditions:
and (d) 1).
Proof. The proof is the same as Lemma until bounding the terms
13
14 and
15
Bound for the term 13). Following the same reasoning Lemma have
Dt,a BI.)) ((0r.
D2
TŪT Bi.T)
d.D2
0220² - - 0,033)
-
5 Σ
S 200 - 3,,
VT
20
2x
a
where the last inequality is by Lemma 13
Bound for the term
Nothing changes here.
Bound for the term
15 Following the same reasoning as Lemma
we obtain
d T
BI, 0t 1.
D2
d T
31) (0", 0e.))
S 2a(1 - 31)
E
Σ
SI,
(t
1100-11
2a
1=2
Therefore, from Lemma 13 we obtain
T
31, 0z-1, (8t, 0l,u)²
dD>
T Σ
E
5
E
(23)
201-1(1 31)
20 B1)²
1=2
40

---

Published as a conference paper at ICLR 2025
Since Le is running max. {Lo)== is non-decreasing sequence, i.e., L1 $ $ the inequality (23
Lr. Thus,
becomes
B1,e 11-1,u (0 0r,u)²
dD
E
E
Σ 1=
B (t- (t-1)
BI.
20h-1
31)
207 31)²
1=2
Note that
31 (1-A) A)2
[ (t (t-
3
Σ 31
Σ
B1t1
1=2
where the first inequality is from the fact that the bound for the term 15 is
1, and the last inequality is from Lemma
Thus,
31,1 01-1,
d3,
E
31) (0", 0c,aa²
S
201-1
207(1 - 31) 2)2
We now bounded three terms:
14) and 15p. Hence,
dD²
ay In +1
d Σ
R. (T) S 20 - B1)2
VT
E [111.T
(1 (1- 31) 32(1 17) a=1
+ 200 1 - 31) X)2
Thus, express R
online (T) as
E (1l..r + Ro (T) S p/dvV 1/ Σ
where p1, p2z and Pl are defined as the following:
D²Lr
av I
d33 L2 P's 200 - 3:) (00(1-11)200
p 2a -,)22
P5 - (1-3))V1 32 (01--)) n)
2)2
The subsequent proof process same as Lemma !
Note that, by Assumption 1 Lr is always less than or equal to the theoretical upper bound of the maximum gradient norm across all iterations (T.). Hence, we have / S pl. P/ -22, and pla = Pa- This implies that Lemma I holds even when L is not known.
Lemma 15. (Lemma when L unknown) Suppose Assumptions
2 hold. If we set y
min-
1,
nln (n/K) [e 1y-K
the batch selection (Algorithm 2 and the weight update rule (Algorithm 7
following AdamCB (Algorithml 7 implies
|| (pi.e)²
T Bpp
|1gh 2| (Pj,.)² 2
co(VV K
E
min
JEJ.
JEJ.
Proof. The proof is the same as Lemma]
2 However, at the last part, where we scale,
T
Ep JEJs
1
min Pt
T
Epg
Σ
L 1 P
LMINo E(LEExPx((T) (24) (P5,)²
(P,,)2
yEJ.
Since L, is running max, (L.TT is non-decreusing sequence, ie., L1 s L2 < < Lr-
Hence, Fq. 24 becomes
2.63LT Pmin
LMIN- (T) E[LExP1 (TT)
KnT In K
By Assumption 1 Lr is always less than or equal to L. which implies Lr = O(1. This completes the proof of Lemma 15
41

---

Published as a conference paper at ICLR 2025
Lemma 13 implies that Lemma hollds even when L is not known.
Theorem 4. (Regret bound of AdamCB (Theorem 7 when L is unknown) Suppose Assumptions
1
hold, and we run AdamCB for a total 'I' iterations with Ox = - and with B1,, := X X (0, 1). Then, the cumulative regret of AdamCB (Algorithm)
Vt D with batch size is upper-bounded by
o dvT +
Vd
T K
1/4
In
(25)
3/4
K
Proof. The overall proof is similar to the proof of Theorem (when L is known) detailed in Appendix B.4 The part that is different is when bounding the term Eq. 18) in Lemma []
20 27 31)
(18) S (LT/T)
Σ t=1
Bli0
0-1ll
By Assumption D LT is always less than or equal to the upper bound of the maximum gradient norm across all iterations (L), which implies LT O(1). Therefore, we have 18 (VVT). This implies that Lemma [1] still holds. Since both Lemma and Lemma [ hold even when L is not known according to Lemma 14 and Lemma 15 we complete the proof of Theorem 4 by following same proof process as Theorem
42


Reviews:
[{"review_id": "hboindyR5A", "reviewer": "Reviewer_797D", "summary": "This paper introduces the AdamCB algorithm, aiming to address inefficiencies in the Adam optimizer by improving data sampling. AdamCB integrates combinatorial bandit techniques, enabling adaptive batch selection to focus on informative samples. The authors claim enhanced theoretical guarantees, providing a rigorously analyzed regret bound that surpasses those of standard Adam and its bandit-based variant, AdamBS, which are identified here as having flawed guarantees. Experimental results across diverse datasets demonstrate the practical benefits of AdamCB, indicating faster convergence and greater efficiency.", "questions": "1. In the weight adjustment section, it is unclear why AdamCB requires the sum of probabilities to equal $K$ instead of 1. This choice appears to necessitate an additional operation in the sampling strategy, specifically the introduction of a threshold $\\tau$. Clarification on the rationale behind this requirement would be beneficial.\n2. The paper could further explain why increasing the sample size $n$ leads to faster convergence in AdamCB. A theoretical or intuitive justification for this relationship would strengthen the understanding of the algorithm’s convergence behavior.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper rigorously addresses and corrects the theoretical flaws in convergence guarantees for both Adam and AdamBS, presenting refined proofs that offer independent value to the community.\n2. The proposed AdamCB algorithm is shown to be both theoretically robust and empirically effective, with rigorous analysis and extensive experimental validation.\n3. The paper presents the method and supporting claims clearly, facilitating reader comprehension and enhancing accessibility of the technical content.\n4. Through a fair and thorough comparison, this paper evaluates AdamCB alongside Adam and AdamBS, incorporating corrected theoretical guarantees to provide a balanced assessment.", "weaknesses": "1. The paper's motivation could benefit from further clarification and depth. In the abstract and introduction, the authors state that uniform sampling leads to inefficiencies in Adam, but they should specify the type of inefficiency (e.g., memory, computational, time, or convergence efficiency; presumably the latter). Additionally, the authors should provide evidence or discussion showing that alternative sampling methods indeed improve Adam's efficiency, strengthening the case for the proposed approach.\n2. Algorithm 3 requires prior knowledge of $L$ for the weight update rule, which may be limiting in practical applications. It would be valuable to discuss potential ways to relax this requirement or clarify how the authors manage this constraint in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R2N0AST88b", "reviewer": "Reviewer_7DWV", "summary": "This paper proposes a variant of the Adam optimizer with a combinatorial bandit approach (AdamCB) for adaptively selecting batches to train on.  Using a combinatorial bandit to select a batch of examples addresses the limitation of a prior approach called Adam with Bandit Sampling (AdamBS) which failed to improve performance with larger batch size due to myopic approach of selecting a single sample at a time with replacement.  Another core contribution of the paper is new theoretical analysis of Adam, AdamBS, and the proposed AdamCB that relaxes prior assumptions as well as fixes errors in the prior proofs.  These convergence guarantees show AdamCB to have faster convergence than Adam and AdamBS and faster convergence with increasing batch size.  Experimental studies show AdamCB to outperform Adam and AdamBS on small scale MLP and CNN tasks.", "questions": "The convergence rate of Adam, Adam with Bandit Selection, and the proposed Adam with Combinatorial Bandits are all dominated by the $d\\sqrt{T}$ term.  In practice, what is the expected speedup of using combinatorial bandit in training settings with >100 million parameters?\n\nIn the LLM setting, not only are models over 100 billion parameters in some cases, we also rarely loop through the full dataset if at all during pretraining.  What benefits if any do you expect adaptive batch selection to provide in this setting?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The primary strengths of this paper are as follows:\n- The writing is clear and easy to understand.  \n- The assumptions required for the convergence analysis is more general than previous work and the theory also covers Adam and AdamBS.  \n- The paper identifies and addresses incorrect assumptions made in the analysis of AdamBS.", "weaknesses": "- The experiments are fairly limited in scale and not very reflective of practical settings that we are in these days with large models and large datasets.\n- The benefits of adaptive selection will likely be limited in settings with large model dimensionality since then the $d\\sqrt{T}$ term will dominate the $\\sqrt{d}/n^{3/4}T^{3/4}$ term controlled by adaptive selection.  The potentially marginal improvement of adaptive selection is not discussed as far as I can tell in the main paper (it is discussed in the AdamBS paper).\n- It is unclear how useful bandit selection will be in settings where we see the entire dataset just a few times as with LLM pretraining.\n- It seems like the theoretical analysis has a lot of overlap with Tran et al. 2019 but this is mainly mentioned in the Appendix and not stated in the main text.  It also detracts from the novelty of the theoretical analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lRUouTkIG9", "reviewer": "Reviewer_fG67", "summary": "This paper proposes an extension of the Adam optimizer by integrating adaptive batch selection. It also identifies a flaw in the proof presented in previous papers. Based on extensive theoretical analysis and some experiments, the proposed method demonstrates better convergence and improved performance.", "questions": "1. **Semi-supervised learning**: Semi-supervised learning is a popular machine learning task where the testing dataset is given without labels. In this context, I wonder if the adaptive sampling method can still assign weights to both training and unlabeled testing samples.\n\n2. **Extra regularizations**: Many regularization techniques are sample-agnostic, such as weight decay and dropout. Sometimes, regularization is only related to the output instead of the input, for example, using a total-variation loss to encourage smoothness in image generation. How does the adaptive sampling method work in these cases?\n\n3. **Adversarial samples/noises**: The paper mentions that \"samples with a low gradient norm are assigned a low weight, whereas samples with larger gradient norms are more likely to be chosen in future iterations.\" What if there is noise in the dataset, which is common in real-world datasets? Is it beneficial for the model to learn more from noisy samples that are difficult to fit?\n\n4. **Data augmentation**. Does this method compatible with data augmentation? The augmented data will share strong similarity among samples. What should be considered when using data augmentation? Will the training be biased in an unexpected manner?\n\n5. **Comparison with other methods for data sampling**: Given the popularity of large language models (LLMs), it is important to fine-tune/pretrain LLMs with a wise combination of data. Is there any potential to demonstrate this method in such a setting? How does this method compare with currently adopted reinforcement learning-based methods for sampling, such as \"DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining, NeurIPS 2023\"?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper clearly points out the flaw in the proof of previous papers.\n2. This paper provides a new convergence rate of Adam/AdamBS in the new perspective.\n3. The results on some simple datasets are better than baselines.", "weaknesses": "I appreciate the theoretical contribution, and my major concern is the applicability of this method in real-world applications. Since the optimizer is one of the most fundamental components of the entire machine learning process, I wonder if this method can be directly integrated into existing practices. Additionally, I question whether the performance remains robust given the extra effort required to assign sampling bias towards certain samples during batch construction. Please see the question part for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)