Adapprox: Memory Efficient Optimization via Adaptive Randomized Low-Rank Approximation

Response to Question

Thank you for your question! Our method leverages low-rank matrix approximation to compress the second moment, reducing memory usage associated with storing the second moment matrix. This compression not only benefits single-node training but also has notable implications in multi-node training:

1. In multi-node settings, where memory constraints are often exacerbated due to the large model sizes and batch splits, our approach retains its memory-saving advantage. By storing only low-rank approximations of second moments, we reduce the memory footprint across all nodes, enabling more efficient utilization of distributed resources.
2. Multi-node training typically involves communication of optimizer states (e.g., gradients, moments) between nodes. With our low-rank approximation, the second moment matrix representation is inherently compact, leading to reduced communication overhead compared to full-matrix methods. This could accelerate the synchronization process, especially in bandwidth-limited environments.
3. In particular, for tensor parallelism (e.g., splitting a large weight matrix across multiple nodes and thus the corresponding optimizer states), our "low-rank approximation" becomes even more accessible and applicable. Since the rank of a submatrix cannot exceed the rank of the original matrix (see detailed explanation below), applying low-rank approximation via Adapprox at each node for the corresponding submatrix has the potential to achieve higher accuracy.

However, there are potential limitations to consider. Adapprox requires additional matrix operations during the S-RSI process. While these operations are efficient and scalable on modern hardware, in a distributed system, their execution time may slightly impact synchronization compared to simpler operations. Furthermore, if the synchronization interval is short (e.g., after every training step or a small number of steps), even minor latencies can accumulate over time, potentially increasing the overall training time.

That said, there are viable mitigation strategies. For instance, reducing the frequency of low-rank matrix synchronization can help minimize latency. Additionally, performing low-rank projections in parallel with other computations (e.g., forward or backward passes) can effectively hide the associated latency.

Overall, we believe our optimizer offers substantial benefits in multi-node training scenarios, particularly in terms of memory and communication efficiency. We will explore this aspect further in future work and include more detailed analysis in subsequent studies.

(Explanation for "the rank of a submatrix cannot exceed the rank of the original matrix")

Without loss of generality, consider a large weight matrix $W$ used in a model, partitioned column-wise for tensor parallelism:

$$W = [W_1 W_2 \cdots W_p].$$

For each submatrix $W_i$, there is

$$\mathrm{rank}(W_i)\leq \mathrm{rank}(W).$$

This property follows directly from the definition of matrix rank: any set of linearly independent columns in $W_i$ must also be linearly independent in $W$. Therefore, the rank of $W_i$, being a subset of the columns of $W$, cannot exceed $\mathrm{rank}(W)$.

Response to Weaknesses

Thank you for your thoughtful suggestion. We appreciate your feedback and have carefully taken it into account.

To address your concern, we conducted a series of new comparison experiments, including the baselines SGD and SGD with momentum (SGDM), as well as similar methods not previously included, such as GaLoreAdamW [1]. The training process was conducted from scratch on the GPT-2 125M model using the OpenWebText dataset. All experiments were performed on a single NVIDIA RTX 4090 GPU. The parameter settings for the compared methods are summarized in the table below.

(Those parameters not mentioned are set as default.)

We present the training loss recorded at every 1,000 steps along with the overall training time for each method in the table below:

(In the initial submission of this rebuttal, GaLoreAdamW's rank decomposition mechanism was not enabled, and the results corresponded to AdamW. We have now corrected this oversight and included the accurate data for rank = 128.)

The results above align with the conclusions presented in our manuscript. While SGD and SGDM require the least memory and computational cost, they converge significantly slower. This reinforces our belief that incorporating the second moment is essential for effectively training LLMs. CAME may reduce the training loss quickly in the initial stages but slows down as training progresses. Adafactor demonstrates convergence behavior similar to that of AdamW. As GaLoreAdamW compresses the gradient, which is crucial for guiding the optimization direction, its performance drops a lot and becomes unstable. Our Adapprox method demonstrates stable performance and achieves the lowest final training loss, even when the rank for approximation is aggressively fixed at a constant value of 1. This also indicates that, like Adafactor and CAME, we compress the second moment matrix using two vectors, effectively eliminating nearly all memory usage associated with the second moment.

We will include these additional experimental results as figures and tables in our revised version to enhance the clarity and comprehensiveness of our experiments. Thank you again for your suggestion!

[1] Zhao et al., 2024. GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection.

Response to Weakness 2

Thank you for your insightful comment. We agree that techniques orthogonal to low-rank matrix factorization, such as CAME or quantization, have significant potential to complement our proposed method. While our primary focus was to develop and analyze the core low-rank approximation approach, we recognize the value of integrating these techniques to further enhance efficiency and scalability.

In particular, CAME inspire us to develop additional mechanisms to address approximation errors in the low-rank approximation of the second moment matrix. Combining our method with quantization could reduce memory usage and computational overhead further. Exploring such synergies is a promising direction for future work, and we will consider this in subsequent research. We appreciate your suggestion and will include a discussion of this point in the manuscript to highlight these opportunities.

Response to Weaknesses 1

Thank you for your thoughtful comment. The parameters $k_0$, $k_{min}$, $k_{max}$, $l$, and $p$ are indeed selected based on the trade-offs between approximation accuracy, memory usage, and computation efficiency.

In terms of approximation accuracy, Equation (10) in our manuscript illustrates how these parameters influence the low-rank approximation:

$$\mathbb{E} \| A - QU^\top\| \leq \left[ \left( 1 + \sqrt{\frac{k}{p - 1}} \right)^{2l+1} \sigma_{k+1}^{2l+1} \right. + \frac{e\sqrt{k + p}}{p} \left. \sqrt{\sum_{j > k} \sigma_j^{2(2l+1)}} \right]^{1/(2l+1)}.$$

As shown, larger values of $k, l$ and $p$ generally result in smaller approximation errors. However, increasing these parameters also results in higher computational cost or memory consumption.

First, $l$ represents the number of power iterations in our S-RSI method. While increasing $l$ improves approximation accuracy, it also inherently increases computation cost, as power iterations are sequential. Based on our experiments, $l=5$ offers a good balance between accuracy and efficiency. Below are the results of an ablation study on $l$, where $A$ is a $m \times n$ matrix sampled from a normal distribution:

As shown, $l=5$ achieves sufficient accuracy, with further increases yielding only marginal improvements. Therefore, we recommended $l=5$ in our manuscript.

Similarly, we recommended $p=5$ in our manuscript:

Besides, $k_0$, $k_{max}$, and $k_{min}$ are critical parameters that govern both approximation accuracy and memory usage. We suggest setting $k_{min} = 1$ to ensure that our method could fully exploit the low-rank structure of the second moment matrix. Meanwhile, we recommend setting $k_0 = k_{max}$ to effectively control the upper bound of memory consumption. The specific values for $k_0$ and $k_{max}$ depend on the available memory resources in practice.

Regarding the need to validate our design choices on larger-scale cases. To address this, we conducted additional experiments on a 1.1B parameter GPT-2 model using 4 NVIDIA RTX 4090 GPUs. The model configuration is as follows:

• n_positions = 1024
• n_ctx = 1024
• n_embd = 1408
• n_layer = 36
• n_head = 22

The parameter setting is as follows:

(Those parameters not mentioned are set as default.)

The results are as follows:

As shown, our parameter choices remain effective for Adapprox even on larger models. We will include these additional results in the revised manuscript to address the concern and further demonstrate the robustness of our approach.

In practice, optimal hyperparameter settings may be different on the specific use case. As indicated in Equation (10), increasing $p$, $l$, and $k$ reduces approximation error but also will increase memory and computational cost. Therefore, we recommend a trade-off based on the available hardwares and the requirements of the training tasks.

Response to Question 3

Thank you for your suggestion. Table 3 presents the zero-shot results of pre-trained GPT-2 models using various optimizers. The pre-training times for these optimizers are shown in Table 4, where Adapprox incurs a latency increase of approximately 6% compared to Adam on GPT-2 117M. Memory usage for Adapprox is detailed in Table 2 and depends on the selected rank $k$. For Adafactor and CAME, as they utilize a fixed rank-1 ($k=1$) approximation, their memory usage is nearly half that of Adam. We will clarify this information further in the revised manuscript.

In summary, Adapprox typically requires more time than AdamW but uses less memory. While it requires more memory than Adafactor and CAME, it achieves better training performance. Overall, our method provides a more balanced trade-off between memory utilization and training efficiency.

Response to Question 4

Thank you for your suggestion. Below, we provide both a theoretical analysis and experimental results regarding memory utilization. We will incorporate this information into the revised manuscript accordingly.

Theory Analysis of Memory Usage

In low-rank approximation for the second moment matrix, memory utilization is dominated by the storage of the decomposed matrices $Q$ and $U$. For a second-moment matrix $V$ of size $m\times n$, the memory usage under low-rank approximation is:

$$\text{Memory Usage} = k \cdot (m+n),$$

where $k$ is the rank of the approximation. In contrast, storing the full matrix $V$ requires $m\times n$ elements. The ratio of memory saved can be expressed as:

$$\text{Memory Savings} = 1 - \frac{k \cdot (m+n)}{m\cdot n}.$$

For larger models (with higher $m$ and $n$), the relative savings increase as long as $k$ remains small compared to $m$ and $n$. However, increasing $k$ leads to higher memory usage, and for sufficiently large $k$, the savings diminish. For example, consider a second moment matrix $V$ of $1024\times 1024 (m=n=1024)$:

• with $k=16$, the memory savings are $1-16\times(1024+1024)/1024^2 = 96.88\%$;
• with $k=256$, the memory savings are $1-256\times(1024+1024)/1024^2 = 50.00\%$.

In addition, to achieve memory savings in this example, $k$ must be less than $512$.

Experimental Results

To validate the efficiency, we measured the memory utilization of optimizer states for various ranks $k$ across GPT-2 models of different sizes during active training. The memory usage includes both the low-rank approximation and other optimizer states (e.g., first moment). Below are the results (in MB):

As shown in the results above, when $k$ is small (e.g., 1 or 8), memory utilization for optimizer states is reduced by nearly half compared to AdamW, demonstrating significant savings due to compressed second moment. This reduction is particularly beneficial for larger models, such as GPT-2 1.1B.

As $k$ increases, memory usage grows, approaching that of AdamW at higher ranks. For example, at $k=256$, memory utilization for GPT-2 1.1B is 38% lower than AdamW. This underscores the importance of selecting an appropriate $k$ to balance memory savings and approximation accuracy.

These results highlight the efficiency of low-rank approximations in reducing memory utilization for optimizer states, particularly when $k$ is small. We will include these findings in the revised manuscript to further illustrate the memory-saving potential of our method.

Response to Question 1 and Weakness 2

Thank you for your observation regarding the validation curves in Figure 4. The similarity in the validation curves arises because all experiments were conducted under identical training conditions, including model initialization, data ordering, and random seeds. This ensures a controlled comparison and eliminates variability that could obscure the differences introduced by the optimizers. The jumps or dips in the curves reflect specific points in the training process where the model interacts with particularly challenging or easy batches, which remain consistent across all runs due to the shared conditions.

Despite this, we also suspect that the optimization landscape is highly non-convex, characterized by large flat regions interspersed with dense clusters of local minima. This structure could contribute to the consistent behavior of the validation curves across different optimizers, as they may follow similar trajectories through the landscape under identical initialization and training conditions. We are interested in further exploring this suspicion in future work.

Response to Question 2 and Weakness 1

Thank you for your suggestion to examine the impact of rank on decomposition time during training. To address this, we conducted additional experiments measuring the training speed across different ranks and model sizes during active training on 4 NVIDIA RTX 4090 GPUs. For this study, we selected GPT-2 models with 125M, 345M, and 1.1B parameters. We fixed the rank at 8, 16, 32, 64, 128 and 256, respectively. We recorded the training speed (iteration per sections) for each rank and record the total training time over 1,000 steps. The parameters for S-RSI are $p=5$ and $l=5$.

Preliminary findings suggest that while higher ranks lead to increased computation time, the scaling remains sub-linear due to efficient parallelization and the overlap of computations with other training tasks. Naturally, the exact statistics may vary depending on the specific hardware used. For example, the reduction in training speed when increasing the rank from $k=8$ to $k=64$ is mostly within 10%. Besides, in this setting, the training speed only shows a significant reduction when $k=256$. When using the same rank $k$ across models of different sizes,the number of layers must be considered, since layers are processed sequentially. For example, GPT-2 125M has 12 layers, while GPT-2 345M has 24 layers.

This observation underscores the importance of our adaptive rank mechanism, which is designed to minimize computational time and memory usage by dynamically adjusting the rank as needed. As mentioned in Section 4.1, with the adaptive rank mechanism, Adapprox maintains a higher rank only during the initial training stages and transitions to a relatively low rank for the remainder of training. For instance, as shown in Table 2, the mean rank during training is approximately 1.4 for GPT-2 117M and 7.7 for BERT 345M. This will accelerate the overall training process.

In addition, for the primary purpose of saving memory, the rank $k$ is typically not set too high. Otherwise, there may be little to no memory savings, defeating the goal of memory efficiency.

Response to Weakness 2 and Question 2

Thanks for your question. We define the approximation error as $||A-QU^\top||_F/||A||_F$, which measures the normalized error ratio. Since this metric captures the relative error, it is minimally influenced by the size or magnitude of $A$, making it adaptable to models of varying sizes.

Regarding the specific value of the threshold $\xi$ for $||A-QU^\top||_F/||A||_F$, which controls the adaptive rank direction (whether to increase or decrease the rank), we currently set it empirically to $1\times 10^{-3}$. This value was selected after testing several options within the range $[1 \times 10^{-5}, 1 \times 10^{-2}]$, as it provides a good balance between accuracy and rank adjustment. A threshold that is too small results in slower rank reduction and increased computational cost, while a threshold that is too large reduces the rank too aggressively, potentially impacting accuracy.

For rank modifications, we were inspired by the gradient descent framework and designed a rank adjustment mechanism in the form $k \pm \Delta k$. This approach allows $k$ to adapt incrementally within a local neighborhood, enabling the algorithm to maintain approximation performance comparable to the previous step while effectively exploring and adjusting the rank. Currently, we empiriclly set $\Delta k$ as $\lceil 0.02k_{max} \rceil$.

Response to Weakness 3 and Question 3

Thank you for your question.

First, the ratio of additional latency introduced by Adapprox (due to S-RSI or AS-RSI) over the overall latency during training depends on factors such as the model size, batch size (which affects forward pass time), and the hardware used. Below, we provide detailed record for GPT-2 117M trained on 8 NVIDIA Tesla V100 GPUs with a global batch size of 128 and a sequence length of 1024. The timings were recorded using TensorBoard’s PyTorch Profiler, with the start of the step set to 0 seconds:

As results shown above, the backward pass constitutes the majority of the training time. From this perspective, although Adapprox introduces additional computations, the relative increase in time remains acceptable.

We also conducted experiments with different model sizes and ranks to investigate how latency varies. Due to limitations in our current hardware and the constrained rebuttal period, we used 4 NVIDIA RTX 4090 GPUs with a batch size of 8. For this study, we selected GPT-2 models with 125M, 345M, and 1.1B parameters. The rank was fixed at 8, 16, 32, 64, 128, and 256, respectively. We recorded the training speed (iterations per second) for each rank and the total training time over 1,000 steps. The parameters for S-RSI were set to $p=5$ and $l=5$.

Preliminary findings suggest that while higher ranks lead to increased computation time, the scaling remains sub-linear due to efficient parallelization and the overlap of computations with other training tasks. When using the same rank $k$ across models of different sizes,the number of layers must be considered, since layers are processed sequentially. For example, GPT-2 125M has 12 layers, while GPT-2 345M has 24 layers.

Note that the results above represent a relative comparison within our specific setting. Furthermore, as the batch size and number of devices increase, the proportion of additional time introduced by Adapprox's S-RSI or AS-RSI relative to the total training time will be further reduced.

Response to Weakness 4 and Question 4

Thank you for your question.

To elaborate on the influence of rank adaptation on convergence, it is essential to revisit the intuition and motivation behind our approach. Our goal is to provide a more flexible method for approximating the second moment during training, balancing the trade-off between training performance, memory usage, and computational cost.

The underlying intuition is that higher accuracy in the approximation (achieved by using a higher rank $k$) results in behavior more closely aligned with the counterpart optimizer, AdamW. Consequently, a higher rank $k$ typically leads to more stable convergence during training.

It is important to emphasize the term "stable" here, as optimizing large language models involves navigating a highly non-convex landscape. In such scenarios, occasional noise or fluctuations (resulting from less precise approximations) can sometimes help escape local minima and achieve better training loss, highlighting the complex interplay between accuracy and optimization dynamics.

As for our rank adaption mechanism, the error threshold $\xi$ in AS-RSI plays a key role in controlling rank adjustments. A smaller $\xi$ tends to maintain higher ranks for longer, ensuring greater accuracy but at the expense of higher computational overhead. Conversely, a larger $\xi$ leads to more aggressive rank reductions, potentially compromising convergence stability.

To demonstrate the influence of rank adaption mechanism and the choice of $\xi$ on convergence performance, we conducted the following ablation studies on GPT-2 345M using 4 Nvidia RTX 4090 GPUs and with a batch size of 8 over 1,000 steps. We tested three scenarios:

• (1) Fixing $k=128$ to represent a case where $\xi$ is too small;
• (2) Fixing $k=1$ to represent a case where $\xi$ is too large;
• (3) Using $\xi=1\times 10^{-3}, k_0=k_{max}=128, k_{min}=1$, which represents the adaptive rank mechanism.

The training loss results are as follows:

The results above align with our analysis. The adaptive rank mechanism achieves performance between the two extremes, demonstrating its ability to balance accuracy, memory efficiency, and computational cost. The mechanism could effectively adjust the rank to meet the demands of each training stage, avoiding unnecessary computational overhead while maintaining stable convergence.

Response to Weakness 5

Thanks for your comment. Our current focus on NLP downstream tasks aligns with prior work in the field and reflects the natural applicability of Adapprox to transformer-based architectures. We agree that extending Adapprox to other domains, such as vision tasks or alternative model architectures, could further strengthen its generalizability. The adaptive rank mechanism and memory-efficient optimization strategies used in Adapprox are not inherently tied to NLP models and could be adapted for other domains. We will explore these broader applications in future work.

Response to Question 5

Thanks for your question. We will release it publicly upon acceptance of the manuscript. The detailed experimental settings are provided in Section 4 of the manuscript. Specifically, our choice for rank setting is $k_0=k_{max}=256$. The rank adjustment is $\Delta k = \lceil 0.02 k_{max} \rceil$. The error threshold is $\xi = 1\times 10^{-3}$.

For specific use cases, we recommend adjusting these parameters based on the available memory and training requirements. Selecting an appropriate rank is especially crucial to achieve a good trade-off between memory savings and computational efficiency while maintaining training performance.

Thank you for your thoughtful and detailed feedback, as well as your positive evaluation of our work. Below, we address your insightful suggestions and provide further clarification on the points you raised.

Response to Weakness 1 and Question 1

Thank you for your comment.

Comparison with Adaptive Rank Techniques

The compared methods, Adafactor [1] and CAME [2], rely on fixed rank-1 low-rank approximation techniques and do not include an adaptive rank mechanism. To the best of our knowledge, our work is the first to explicitly introduce an adaptive rank mechanism in memory-efficient optimization algorithms via low-rank matrix approximation.

The motivation behind this innovation is that our S-RSI enables flexible rank approximation, allowing the algorithm to dynamically select the rank for approximating the second-moment matrix as needed. This adaptability ensures a balance between approximation accuracy, computational cost, and memory utilization.

Comparison with GaLore

GaLore [3] emphasizes memory efficiency through low-rank gradient projections, which target gradient compression (and consequently both the first and second moments), rather than focusing solely on the second moment as Adapprox does. Considering its mechanism, GaLore offers both potential benefits and limitations compared to our methods:

• Benefits: GaLore theoretically enables higher memory savings. While Adapprox (along with Adafactor and CAME) focuses solely on compressing the second moment, which provides a memory saving upper bound equivalent to the storage of the second moment, GaLore projects gradients into a compact space. This allows for memory savings with respect to both the first and second moments, potentially achieving greater overall memory efficiency. In addition, by projecting gradients into a compact space, GaLore also achieves greater computational efficiency in the corresponding operations.
• Limitations: While GaLore compresses gradients to achieve memory efficiency, gradients are critical as they provide the direction for updates during optimization. Any inaccuracies in the gradient direction could lead to instability in the optimization trajectory, potentially affecting both convergence speed and overall performance. Therefore, this makes it more challenging to achieve a suitable trade-off between memory utilization and training performance.

To compare GaLore with our method in depth, we conducted a series of experiments over GPT-2 125M on NVIDIA RTX 4090 GPU. The configurations for the experiments are as follows:

The recorded training losses at every 1,000 steps are presented below:

The memory utilization are presented below:

The results above align well with our prior analysis. GaLoreAdamW demonstrates greater computational efficiency and could achieve higher memory savings for the same rank setting. However, Adapprox achieves better and more stable training performance, showcasing its robustness in optimization tasks. (Notably, when memory utilization is comparable, such as $k=1$ for Adapprox and $k=256$ for GaLoreAdamW, our Adapprox outperforms GaLoreAdamW.)

[1] Shazeer, Noam, and Mitchell Stern. "Adafactor: Adaptive learning rates with sublinear memory cost." International Conference on Machine Learning. PMLR, 2018.

[2] Luo, Yang, et al. "CAME: Confidence-guided Adaptive Memory Efficient Optimization." Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2023.

[3] Zhao, Jiawei, et al. "GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection." Forty-first International Conference on Machine Learning.

Response to Major Question 1

Thank you for your detailed observations and valuable feedback. We acknowledge that our presentation may have caused some confusion. Below, we provide clarifications and will ensure these points are addressed more clearly in the revised manuscript:

• The term "Adafactor approximation" refers specifically to the low-rank matrix approximation method proposed and employed within the Adafactor algorithm (this method is also used in CAME). We compare this approximation with our S-RSI, focusing on approximation accuracy rather than the broader optimization process. Specifically, for a matrix $A \in R_{\geq 0}^{m \times n}$, Adafactor uses two vectors $R \in R_{\geq 0}^{m \times 1}$ and $S \in \mathbb{R}_{\geq 0}^{1 \times n}$, where $R = A1_n$, and $S = 1_m^\top A/1_m^\top A1_n$. Here, $R$ represents the row sum of $A$, and $S$ represents the normalized column sum of $A$. In Figure 2(a), we present the mean approximation error, calculated as $1/s\sum_i^s||A_i - R_i S_i||_F$, where $s$ denotes the number of samples.
• The "Adafactor time" reflects the time required to compute the low-rank approximation within Adafactor's implementation. Specifically, as mentioned in the first point, this includes the time taken to compute $R$ and $S$. To ensure a fair comparison, we implemented this single module independently for evaluation.
• To address this, we will enhance clarity by adding a table in our revised version to summarize the numerical results, making the comparisons easier to interpret. We acknowledge that S-RSI is slower than Adafactor's low-rank approximation method, which involves only summation and element-wise division. However, Figure 2(b) also highlights that S-RSI’s runtime on GPUs scales sub-linearly with rank $k$, contrary to theoretical expectations (as outlined in Table 1 of our manuscript).

Response to Major Question 2

Thank you for your insightful feedback regarding Figure 4. We appreciate your thorough observations and provide the following clarifications:

• In our experiments, the sequence length was set to 1024 tokens, with a batch size of 128. Therefore, the total batch size (measured in tokens) is 131,072 tokens per iteration. We will clarify this in the revised manuscript.
• There are three points:
  • The term "uniform training parameters" refers to the use of identical hyperparameters (e.g., batch size, number of training iterations, number of warmup iterations, peak learning rate, minimum learning rate, learning rate scheduler, weight decay, etc.) across all methods to ensure a fair comparison.
  • Regarding the "empirical testing" methodology, we determined hyperparameters (primarily the peak learning rate and minimum learning rate) through grid search, prioritizing stability and competitive performance for the baseline methods. The hyperparameter search space for learning rates was $[10^{-5}, 10^{-2}]$.
  • For CAME, we also carefully considered its behavior when selecting the learning rate. As you pointed out, its convergence pattern often show rapid reduction in training loss during the initial stages but slower progress in later stages. We attribute this to CAME's inherent design. Specifically, in Algorithm 2 of CAME (lines highlighted in blue), a "confidence score" mechanism is employed to adjust the step size of updates, where the scale factor is $1/\sqrt{S_t}$ in their notation. In the experiments, we observed that $1/\sqrt{S_t}$ is almost always greater than 1, effectively enlarging the learning rate. This explains why CAME converges more quickly in the early stages but slows down in later stages. Moreover, this mechanism can sometimes cause instability in CAME's convergence. Our experimental results reflect these characteristics, which we believe highlight an inherent feature of CAME's design.
• In the revised manuscript, we will simplify the presentation by reporting only perplexity, which is more interpretable for practitioners.
• Thanks for your suggestion. To address this, we will include a supplementary table or a detailed illustration summarizing the final performance metrics within the context.

Response to Minor Question 3

Thank you for your detailed observations. We address your points as follows:

• The correct model is BERT 345M, and we will revise Table 5 to ensure consistency with the rest of the manuscript.
• Thank you for pointing out the inconsistency in the dimensions of matrix $U$. We intended $U$ to be $n \times k$, with $U^\top$ as $k \times n$. We will review and correct this discrepancy throughout the manuscript and explicitly include the dimensions of the matrices in Algorithm 1 to improve clarity.
• The units are seconds, and we will update the figure caption to reflect this explicitly.
• Thank you for this insightful suggestion. We will revise the inputs list of Algorithm 3.

We greatly appreciate your detailed and valuable feedback and will incorporate these revisions to improve the clarity, accuracy, and overall readability of the manuscript.

Response to Weakness 1

Thank you for your detailed and valuable feedback. We have carefully addressed your concerns and suggestions in the responses above.

Response to Weakness 2

Thank you for highlighting this important point. We plan to release the code upon acceptance of the manuscript. In the meantime, the pseudocode provided in the manuscript can be used to reproduce our results.

Response to Weakness 3

Thank you for your suggestion to expand the comparison baselines. We have conducted additional experiments to include comparison results with GaLore [1]. Regarding Adam-mini, as it is also a submission to ICLR 2025, we plan to conduct a comprehensive comparison with it in future work. Below, we provide details of the additional results on GPT-2 125M.

(Those parameters not mentioned are set as default. All experiments were performed on a single NVIDIA RTX 4090 GPU.)

The training loss recorded at every 1,000 steps and the overall training time for each method are presented in the table below:

(In the initial submission of this rebuttal, GaLoreAdamW's rank decomposition mechanism was not enabled, and the results corresponded to AdamW. We have now corrected this oversight and included the accurate data for rank = 128.)

The results above align with the conclusions presented in our manuscript. While SGD and SGDM require the least memory and computational cost, they converge significantly slower. This reinforces our belief that incorporating the second moment is essential for effectively training LLMs. CAME may reduce the training loss quickly in the initial stages but slows down as training progresses. and Adafactor demonstrates convergence behavior similar to that of AdamW. As GaLoreAdamW compresses the gradient, which is crucial for guiding the optimization direction, its performance drops a lot and becomes unstable. Our Adapprox method demonstrates stable performance and achieves the lowest final training loss, even when the rank for approximation is aggressively fixed at a constant value of 1. This also indicates that, like Adafactor and CAME, we compress the second moment matrix using two vectors, effectively eliminating nearly all memory usage associated with the second moment. However, we acknowledge that Adapprox introduces additional time overhead, a common trade-off for memory-efficient optimizers that involve extra computation.

[1] Zhao et al., 2024. GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection.

Response to Major Question 3

• We chose zero-shot testing for pre-trained models to ensure a consistent comparison across models without relying on task-specific fine-tuning or prompt engineering. This evaluation metric is inspired by a series of post-training compression studies, such as SparseGPT [1] and Wanda [2], which also employ zero-shot accuracy alongside perplexity to measure the performance of pre-trained models and their compressed variants.
• We conducted zero-shot evaluation using the LM Harness library [3], which facilitates standardized testing of language models. The saved pre-trained models were loaded into this library, and the zero-shot evaluation mode was activated to generate the corresponding results.

[1] Frantar, Elias, and Dan Alistarh. "Sparsegpt: Massive language models can be accurately pruned in one-shot." International Conference on Machine Learning. PMLR, 2023.

[2] Sun, Mingjie, et al. "A Simple and Effective Pruning Approach for Large Language Models." The Twelfth International Conference on Learning Representations.

[3] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, et al. A framework for few-shot language model evaluation. Version v0. 0.1. Sept, pp. 8, 2021.

Response to Major Question 4

• The objective of this experiment was to evaluate the optimizer's performance in a complete end-to-end training pipeline, encompassing both the pre-training and fine-tuning stages. To align with this goal, we performed fine-tuning on pre-trained models optimized with the respective optimizers.
• We agree that PEFT methods are relevant baselines for a study focused on memory efficiency in fine-tuning, and we acknowledge the valuable contributions of methods like LoRA [1]. However, our primary focus was on comparing full-model optimization approaches leveraging low-rank matrix approximation. Therefore, we selected baselines such as Adafactor and CAME for this study.
• The primary hyperparameter in this stage is learning rate, which can vary across optimizers and tasks. To ensure a fair comparison, we conducted all experiments on BERT-345M using learning rates from the set [1e-5, 5e-5,1e-4, 2e-4] respectively and retained the best results for each optimizer. Similarly, for GPT-2 117M, we tested learning rates from the set [2e-5, 4e-5, 5e-5], again keeping the best results for each optimizer. For batch size, we set 128. Other optimizer specific hyperparameters are set the same as pre-training stage as we mentioned in Section 4. The batch size was fixed at 128 for all experiments, and other optimizer-specific hyperparameters were set the same as in the pre-training stage, as described in Section 4. We will add these details to the revised manuscript to ensure clarity.

[1] Hu, Edward J., et al. "LoRA: Low-Rank Adaptation of Large Language Models." International Conference on Learning Representations.

Response to Minor Question 1

Thank you for your insightful comment. We perform a projection step at this stage, which was inadvertently omitted from our pseudocode. This step clamps negative elements of $QU^\top$ to zero. This projection is necessary not only to ensure non-negativity but also to potentially reduce the approximation error as $V_t \geq 0$. Mathematically, there is:

$$||V_t - [Q_tU_t^\top]_+||_F \leq ||V_t - Q_tU_t^\top||_F$$

where $[x]_+$ denotes the operation of clamping negative values to zero.

We will clarify this procedural detail and its justification in the revised version of our manuscript.

Response to Minor Question 2

Thank you for your comment. Dynamic rank adaptation was introduced to balance memory efficiency and approximation accuracy. While memory for $k_0=k_{max}$ is allocated upfront, maintaining a fixed rank throughout training could lead to unnecessary computational overhead and memory usage, particularly in scenarios where a smaller rank suffices (e.g., during the later stages of training or for less complex structures). Dynamic adaptation allows the rank to adjust based on the training requirements, optimizing both resource utilization and performance. In contrast, a fixed-rank approach may allocate excessive resources uniformly, which might not translate into proportional gains in accuracy.

Regarding Response to Weakness 3

I have several concerns about these tables:

• As previously mentioned, different methods require different optimal learning rates. This concern is particularly relevant for SGD and GaLore, where the original paper [4] used a learning rate of 1e-2 (100 times larger) for a model of comparable size and similar number of steps.

• Generally, the learning rate selection appears questionable. For models of this scale, 1e-4 is notably too small (see, e.g. [6], Table 2.1).

• Other hyperparameter choices also raise concerns. The batch size, in particular, is extremely small (see [6], Table 2.1). It's also known that an inadequately small batch size can severely impact transformer performance (see, e.g. [7] Figure 5, 6).

• With batch size = 8, sequence length 1024, and 10k steps, the total number of tokens processed is approximately 82M, which is substantially below even the Chinchilla-optimal token count (which should be ~20x the parameter count) [8]. This means the authors are only comparing performance during the very early stages of model training.

• The scheduler type and number of warmup steps used in the experiment are not specified. Proper warmup is known to be crucial for training transformers (see, e.g., [9] Table 1 and [7] Figure 8).

To summarize, I believe this experimental setup does not provide a sufficient basis for comparing optimizer performance. For a fair comparison of optimization algorithms, I would recommend that the authors conduct experiments with GaLore using the same setup as for other baselines in original paper.

[1] Chen, et al. "Symbolic discovery of optimization algorithms.", NeurIPS, 2023.

[2] Pagliardini et al. "The ademamix optimizer: Better, faster, older.", arXiv:2409.03137, 2024.

[3] Zhang et al. "Adam-mini: Use fewer learning rates to gain more.", arXiv:2406.16793, 2024.

[4] Zhao et al. "GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection.", ICML, 2024.

[5] Hu et al. "LoRA: Low-Rank Adaptation of Large Language Models.", ICLR, 2022.

[6] Brown et al. "Language models are few-shot learners." arXiv:2005.14165, 2020.

[7] Popel et al. "Training tips for the transformer model." arXiv:1804.00247, 2018.

[8] Hoffmann et al. "Training compute-optimal large language models." arXiv:2203.15556, 2022.

[9] Shazeer & Stern. "Adafactor: Adaptive learning rates with sublinear memory cost." ICML, 2018.