ADAM Optimization with Adaptive Batch Selection

Thank you for taking the time to review our paper and for your thoughtful and valuable feedback. We appreciate your positive recognition of our work and the constructive comments you have provided. Below, we address each of your comments and questions in detail:

---

[W] Applicability of the Method in Real-World Applications

We would like to clarify that our paper primarily focuses on optimization with adaptive batch selection using a bandit method, and our contributions are centered on this specific optimization framework. While the reviewer raises important questions regarding the type of task, model dependence (e.g., regularization techniques like weight decay and dropout), and data-related concerns (e.g., noise and augmentation), we emphasize that the proposed method is task-agnostic and model-independent, as long as the loss function can be computed.

Below, we specifically address each aspect while reiterating the broader applicability of our method:

1. Type of Task: The method is not restricted to any particular task or domain. It dynamically adapts to the task at hand by leveraging the computed loss and gradient norms, making it applicable to supervised learning, semi-supervised learning, and even self-supervised learning setups.
2. Model Dependence (Regularization Techniques): Our approach does not interfere with standard or task-specific regularization techniques such as weight decay, dropout, or total variation loss. These regularizations function independently of how batches are constructed. Our method focuses solely on selecting the most informative samples for efficient optimization, complementing these regularization strategies.
3. Data Samples (Noise and Augmentation): Regarding data augmentation, the approach is fully compatible and can incorporate augmented samples effectively without biasing the training process, provided the augmentation strategy is balanced.

In summary, while the reviewer raises interesting questions, the scope of our paper is the general optimization method rather than the intricacies of specific tasks, models, or data manipulations. We hope this clarification resolves any questions and underscores the broad applicability of our combinatorial bandit-based (adaptive batch selection) optimization method.

---

Answers to Questions

[Q1] Semi-Supervised learning

Once a specific loss function is defined, the feedback is derived from the gradient of that loss function. Therefore, even in semi-supervised learning settings, as long as a loss can be computed, our proposed method can be applied. For unlabeled data, pseudo-labeling techniques could be used to assign temporary labels, allowing the computation of a loss. This enables our adaptive sampling approach to function effectively, regardless of whether the setting is fully supervised or semi-supervised.

---

[Q2] Extra regularizations

Our method does not conflict with sample-agnostic regularization techniques such as weight decay or dropout, as these operate independently of batch construction. For regularizations related to the output, such as total variation loss, the adaptive sampling method would indirectly align by prioritizing samples that lead to higher gradient norms, which often correlate with areas requiring regularization (e.g., regions with high variance). The flexibility of our approach ensures compatibility with a wide range of regularization techniques.

---

[Q3] Adversarial Samples/Noises

Our method can be uOur method does not inherently prioritize noisy samples unless they consistently produce high gradient norms over multiple iterations. In practice, gradient norms for noisy samples often decrease once the model learns to ignore the noise, reducing their selection probability. However, for datasets with substantial noise, adding noise-robust loss functions or noise-detection mechanisms can complement our method to prevent overfitting to such samples.

---

[Q4] Data augmentation

Data augmentation can indeed affect the adaptive sampling process, as augmented samples often exhibit strong similarity to original samples. Our method could be adapted to treat augmented samples as lower-priority if they exhibit lower gradient relevance, thereby ensuring that the model focuses on diverse samples. Alternatively, adaptive sampling could prioritize augmented samples when they introduce new informative patterns.

---

[Q5] Comparison with Other Sampling Methods (e.g., DoReMi)

Our method leverages a combinatorial bandit framework focused on adaptive batch selection, which is theoretically grounded in gradient-based relevance. In contrast, reinforcement learning-based methods like DoReMi optimize data mixtures with LLMs as the specific use case. There is a clear difference between the two methods. We are confident that its ability to prioritize informative samples enhances efficiency in many large-scale applications.

Answers to Questions

---

[Q1] Clarification on Setting $\sum_{i=1}^{n} p_{i,t}=K$.

• Why $\sum_{i=1}^{n} p_{i,t}=K$ instead of $\sum_{i=1}^{n} p_{i,t}=1$?

We thank the reviewer for providing us with the opportunity to clarify this point. AdamCB uses a combinatorial bandit framework to sample multiple arms (samples) without replacement in each mini-batch. Unlike single-arm selection bandit algorithms like AdamBS, where$\sum_{i=1}^{n}p_{i,t}=1$ because only one arm is selected at a time, AdamCB must select $K$ simultaneously for a mini-batch. Therefore, it is natural to scale the sum of the probabilities to $K$, reflecting the expected number of samples selected in each round.

If the sum of probabilities were constrained to 1, the algorithm would need to perform additional rescaling or sampling adjustments to ensure $K$ samples are drawn, which would unnecessarily complicate the sampling process. Instead, directly setting $\sum_{i=1}^{n}p_{i,t}=K$ aligns the probability distribution with the batch-level selection requirements. By setting $\sum_{i=1}^{n} p_{i,t}=K$, AdamCB simplifies the sampling process and ensures compatibility with mini-batch training.

• Rationale for the Threshold $\tau$

Allowing the sum of probabilities to equal $K$ can lead to individual probabilities $p_{i,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 $\tau$. If a sample's probability $p_{i,t}$ exceeds $\tau$:

1. Its index is added to a null set $S_{null,t}$, effectively removing it from active consideration for selection.
2. The probabilities of the remaining samples are adjusted to redistribute the excess weight while ensuring the sum of probabilities remains $K$.

This adjustment ensures that no single sample dominates the mini-batch while maintaining the proportional relationship between the weights $w_{i,t}$ and the probabilities $p_{i,t}$. We again thank the reviewer for the chance to clarify this key aspect.

---

[Q2] Relationship Between Sample Size $n$ and Convergence in AdamCB

We thank the reviewer for their thoughtful feedback, which provides us with an opportunity to enhance the understanding of our algorithm’s convergence behavior by further explaining the relationship between the sample size $n$ and the faster convergence of AdamCB.

For theory, the second term of the regret bound of AdamCB (Theorem 1) is given as, $(\sqrt{d} / n^{3/4}) \left( (T/K) \ln{(n/K)}\right)^{1/4}$. From this term, it is evident that the regret decreases as $n$ increases. This implies that larger $n$ leads to smaller regret, improving the algorithm’s convergence rate.

For intuitions, increasing the sample size $n$ expands the pool of data samples from which mini-batches are drawn. This broader pool allows AdamCB (as well as other algorithms albeit slower convergence rate) to access a wider range of informative samples during each update step. As a result, the selected mini-batches are more representative of the overall data distribution, reducing variance in the gradient estimates and leading to more accurate updates. This relationship underscores the advantage of AdamCB in leveraging large datasets to achieve faster convergence, making it particularly effective in modern machine learning tasks with abundant data.

Thank you for taking the time to review our paper and for your thoughtful and valuable feedback. We deeply appreciate your recognition of our work and the constructive comments you have provided. Below, we address each of your comments and questions in detail:

---

[W1] Clarification Regarding Inefficiency

We thank the reviewer for the opportunity to clarify this point. The inefficiencies of Adam (with uniform sampling) can be categorized into two main issues:

1. Convergence Inefficiency: As explained in Section 2.4.2, the analysis of Adam reveals a technical error in the original framework, preventing it from providing convergence guarantees. This issue implies that Adam can potentially diverge under certain conditions.
2. Algorithmic Limitation with Uniform Sampling: Adam's uniform sampling nature limits its ability to fully leverage the feedback provided by multiple samples in each mini-batch. This constraint leads to slower convergence (even with corrections) compared to our proposed algorithm, AdamCB, which utilizes a combinatorial bandit sampling mechanism to address these inefficiencies (see Table 1).

Our proposed AdamCB algorithm overcomes these challenges by dynamically adapting the sampling distribution to prioritize informative samples, resulting in faster convergence. This improvement is supported by rigorous theoretical guarantees, as outlined in Theorem 1. The provable efficiency of our method is verified through comparisons with existing methods, as shown in Table 1.

AdamCB achieves provable regret bounds and demonstrates superior convergence performance compared to both the original Adam and its bandit-based variant, AdamBS. Furthermore, our experiments highlight that AdamCB consistently outperforms its counterparts across various models and datasets. We will refine the paper to ensure these clarifications are more explicit and accessible. Thank you for the opportunity to highlight these points.

---

[W2] Knowledge of $L$ is NOT needed.

We sincerely thank the reviewer for the comment regarding the knowledge of $L$ in Algorithm 3. We are happy to clarify that the presentation of the algorithm requiring knowledge of $L$ is actually without loss of generality, meaning that prior knowledge of $L$ is not necessary.

In fact, we can relax this requirement and adopt a dynamic approach where the upper bound $L$ is replaced with a running maximum based on the gradient norms observed during training. Specifically:

• Dynamic update of $L$: In practice, we can replace $L$ with $L_{t}=\max_{t' \leq t} \max_{i \in [n]} \|g_{i,t'}\|$, where $L_{t}$ maintain the running maximum gradient norm observed during training up to iteration $t$. This ensures $L_{t}$ is non-decreasing and provides a sufficient upper bound for the gradient norms.
• Implementation in Algorithm 3: The weight update rule is modified to incorporate $L_t$, which is updated at every iteration: $L_t \leftarrow \max(L_{t-1}, \max_{i \in [n]} \| g_{i,t} \|).$This ensures that $L_t$ remains a valid upper bound for all gradients observed during training.
• Impact on Theoretical Guarantees: The use of $L_t$ still preserves the theoretical analysis of the algorithm, as the bounded gradient assumption is satisfied by construction. The convergence guarantees and performance remain unaffected.

This modification enhances the practicality of our method while retaining the rigor of our theoretical results. We sincerely thank the reviewer again for their valuable feedback and for providing us the opportunity to improve the practicality of our approach.

[W2] In high dimension

We appreciate the reviewer’s comment and are happy to clarify any potential misunderstandings. The convergence guarantee shown in our paper is a regret comparison between the true optimal solution and the algorithm's performance, where the first term, $\mathcal{O}(d \sqrt{T})$, represents the leading term. But, that is about the difference between the performance of an optimization algorithm and the optimality.

However, if the reviewer’s point pertains to the "marginal improvement of adaptive selection" (as it appears from the comment), this requires examining the difference between adaptive and non-adaptive (uniform) sampling strategies.

As illustrated in Table 1, AdamCB achieves a convergence rate of $\mathcal{O}\left(d\sqrt{T} + \frac{\sqrt{d}}{n^{3/4}}\left(\frac{T}{K}\ln{\frac{n}{K}}\right)^{1/4}\right)$, while corrected Adam (with uniform sampling) achieves a convergence rate of $\mathcal{O}\left(d\sqrt{T} + \frac{\sqrt{d}}{n^{1/2}}\sqrt{T}\right)$. As shown in the regret analysis, the constant factor in the leading term $\mathcal{O}(d \sqrt{T})$ is the same for both AdamCB and corrected Adam. Therefore, the regret gap arises from the difference in the second terms: $\mathcal{O}\left(\frac{\sqrt{d}}{n^{1/2}}\sqrt{T}\right) - \mathcal{O}\left(\frac{\sqrt{d}}{n^{3/4}}\left(\frac{T}{K}\ln{\frac{n}{K}}\right)^{1/4}\right)$. This gap becomes larger as $d$ or $T$ increases, highlighting the greater comparative improvement provided by AdamCB in high-dimensional settings. This behavior is both theoretically demonstrated and empirically validated in our experiments with neural network models, as presented in the experiments of Appendix G.

We sincerely thank the reviewer for raising this question, as it provides an opportunity to clarify this point in detail. However, we respectfully disagree with the notion (if the implication was intended) that there is a diminishing marginal improvement of adaptive selection in high dimensions. The theoretical results in our paper clearly demonstrate that adaptive selection retains its advantages in both low- and high-dimensional settings. Additionally, we have reviewed the discussion in the AdamBS paper mentioned by the reviewer. We believe that the discussion there does not necessarily align with the theoretical results in our work as well as the comparison mentioned above, and as shown in our analysis, the results in AdamBS themselves are invalid. We would be more than happy to include this discussion in the revised manuscript to further highlight our contributions.

In conclusion, adaptive selection consistently provides a distinct advantage in both low- and high-dimensional settings, as demonstrated by our theoretical analysis and empirical results.

---

[W3, Q2] LLM pretraining

We appreciate the reviewer’s point regarding the applicability of our method to settings such as LLM pretraining, where the dataset may only be swept a few times. However, our method remains valuable in such scenarios. Each gradient update during LLM pretraining is computationally expensive, and any unnecessary gradient computation becomes significantly more wasteful compared to smaller models. Even in cases where the dataset is seen only a limited number of times, the adaptive batch sampling introduced by our method can still provide meaningful efficiency improvements by prioritizing the most informative gradients. Exploring the full potential of this adaptive strategy in LLM pretraining contexts is indeed an intriguing direction for future work. That said, we respectfully believe this aspect should not be considered a weakness of our work as a general optimization method.

---

[W4] On Tran et al. 2019

We would like to clarify that while some techniques from Tran et al. 2019 are utilized to address the technical error in the original Adam analysis, we have explicitly acknowledged and properly credited Tran et al. 2019 in our work.

It is important to emphasize that the primary goal of our research is not to fix errors in the analysis of Adam (or AMSGrad). Rather, the main theoretical contribution of our work lies in designing and proving the provable efficiency of adaptive batch sampling (leveraging a combinatorial bandit approach) for Adam-based optimization. This contribution is independent of Tran et al. 2019.

Our theoretical improvement in convergence efficiency, achieved through the rigorous integration of a combinatorial bandit framework into Adam optimization, is a core novelty of our work. This innovation is not present in Tran et al. 2019 and represents a significant advancement in the field.

Therefore, the novelty of our work is not detracted in any sense, as it addresses a distinct problem and provides contributions that go beyond the scope of Tran et al. 2019. We are happy to clarify this.

Thank you for taking the time to review our paper and for your thoughtful and valuable feedback. We appreciate your positive recognition of our work and the constructive comments you have provided. Below, we address each of your comments and questions in detail:

---

[W1] Experiments and practical relevance

We would like to emphasize that the primary focus of our work is to propose a provably efficient and practical optimization algorithm that addresses longstanding inefficiencies—and even incorrectness—in Adam-based methods. AdamCB is a general optimization algorithm with convergence guarantees, making it applicable not only to large models but also to a wide range of models and optimization tasks. We sincerely hope that this fundamental focus is appropriately considered, as our algorithm is not solely designed for large models such as LLMs.

That said, in response to the reviewer's feedback, we have performed and included additional experiments with larger-scale models such as ResNet-18 (11.4 million parameters), ConvNeXt-Base (89 million parameters), and ConvNeXt-Large (198 million parameters), in addition to the previously included experiments with MLPs, CNNs, and the VGG network (Figure 5). These new results are presented in the supplementary material (in Appendix G) of the updated manuscript, with some snapshots shown in the tables below (all results are test errors on CIFAR-10). These experiments clearly demonstrate that AdamCB consistently outperforms baselines such as Adam and AdamBS, even for models with larger parameter counts and higher complexity.

ResNet-18 (11.4 million parameters)

ConvNext-base (89 million parameters)

ConvNext-large (198 million parameters)

The inclusion of new experiments on these larger models complements the MLP, CNN, VGG network results already presented in the main paper, bridging the gap between simpler architectures and larger models. This expanded evaluation further substantiates the effectiveness and scalability of AdamCB across a diverse range of architectures. For more details, please refer to Appendix G in the updated manuscript.

With these new results, alongside the findings already presented both in theory and experiments, we strongly believe that the efficacy of our method is well supported. Furthermore, our core contribution—the provable efficiency of AdamCB—is further strengthened by this comprehensive evaluation. We are also more than willing to conduct additional evaluations, time permitting, to further strengthen our results. We sincerely and respectfully request the reviewer to recognize the potential impact of our work in light of our main contributions and the additional evidence provided.

We thank the reviewer for their continued discussion and for assigning a positive score.

However, there appears to be a misunderstanding in the reviewer's comment regarding the performance gap between AdamCB and Adam, as our key results indicate the opposite of what the reviewer's comment suggests. We sincerely hope the reviewer recognizes our genuine intention to ensure that the evaluations are based on what our results truly represent. To this end, we are happy to provide further clarification on this point.

As explained in our previous response (see the section "In high dimension"), the performance gap between AdamCB and Adam does NOT diminish as $d$ increases. On the contrary, the performance gap increases with $d$ with $O(\sqrt{d})$ -- the gap is given by:

$$\mathcal{O}\left(\frac{\sqrt{d}}{n^{1/2}}\sqrt{T}\right) - \mathcal{O}\left(\frac{\sqrt{d}}{n^{3/4}}\left(\frac{T}{K}\ln{\frac{n}{K}}\right)^{1/4}\right).$$

This reflects the fact that AdamCB retains its comparative advantage even in high-dimensional settings (or for any $d$, both in low- and high-dimension). We had clearly stated this above and the reviewer may refer to our response for details (e.g., the distinction between regret compared to the optimality and comparative advantage of AdamCB over Adam).

To further illustrate, consider the example you provided. The relative difference in regret between Adam and AdamCB (the larger this value, the greater the advantage of AdamCB) is shown below:

• 89 million parameters, training examples $n = 50,000$, iterations $T = 1000$, batch size $K = 128$

$$C \cdot \left( \sqrt{\frac{89,000,000}{50,000}}\sqrt{1000} - \frac{\sqrt{89,000,000}}{50,000^{3/4}}\left(\frac{1000}{128} \ln{\frac{50,000}{128}}\right)^{1/4} \right) \approx \mathbf{1328} \cdot C$$

• 89 million parameters, training examples $n = 50,000$, iterations $T = 10,000$, batch size $K = 128$

$$C \cdot \left( \sqrt{\frac{89,000,000}{50,000}}\sqrt{10,000} - \frac{\sqrt{89,000,000}}{50,000^{3/4}}\left(\frac{10,000}{128} \ln{\frac{50,000}{128}}\right)^{1/4} \right) \approx \mathbf{4208} \cdot C$$

• 198 million parameters, training examples $n = 50,000$, iterations $T = 10,000$, batch size $K = 128$

$$C \cdot \left( \sqrt{\frac{198,000,000}{50,000}}\sqrt{10,000} - \frac{\sqrt{198,000,000}}{50,000^{3/4}}\left(\frac{10,000}{128} \ln{\frac{50,000}{128}}\right)^{1/4} \right) \approx \mathbf{6276} \cdot C$$

where $C > 1$ is some constant. (By the way, this gap is in cumulative sense as the regret was defined in cumulation.) These calculations demonstrate that AdamCB maintains a significant comparative advantage over Adam in all cases.

Furthermore, this advantage is observed not only in theory (as shown in the regret bound) but also in experiments. Please see Figure 8 and Figure 9 in Appendix G for experimental results using ConvNext models, where AdamCB consistently maintains its comparative advantage over Adam.

We sincerely hope this clears up any misunderstanding, and we would be happy to provide further clarifications if needed. In light of this clarification (and given our earnest effort to address your feedback comprehensively), we kindly and respectfully ask the reviewer to reconsider the rating.