Accelerating Block Coordinate Descent for LLM Finetuning via Landscape Expansion

We thank the reviewer for the valuable comments. We now address your concerns below.

Missing comparison with recent memoy-efficient methods. Thank you for the suggestion. We finetune Llama 3.1-8B on MathInstruct dataset under the same setting in Table 2 of the manuscript, using the additional ReLoRA and DoRA. The full benchmark score is shown in the following table:

Notably, both BAdam and BREAD outperform ReLoRA and DoRA. This suggests that BCD-based algorithm might be more suitable for math finetuning task. Due to limited computational resource, we are unable to finish the experiment of LISA during the short rebuttal period. We will provide a more comprehensive comparison with memory-efficient optimization algorithms in the next version of our manuscript.

Some assumptions have not been substantiated to generally hold for LLM finetuning. We acknowledge that the assumptions such as smoothness and deterministic gradient might not hold for general LLM training. However, without assumptions such as smoothness, it is challenging to establish a meaningful convergence result. Theoretical works (e.g., [1, 2]) that analyze Adam's convergence also require certain type of smoothness to derive the descent of the objective for each step. Additionally, while it is possible to remove the deterministic assumption, the analysis will be unnecessarily intricate and does not bring additional insight about the algorithm. The current convergence result based on deterministic gradient can already serve a partial explanatory purpose to show that the algorithm is grounded by classic convergence theory. We remark that it is possible to combine the current analysis with [1, 2] to derive the convergence under stochastic gradient. The important steps would be establishing the approximate descent property (rather descent property) using existing analysis for Adam under stochastic setting, and then providing an overall approximate descent for the BCD loop. Finally, standard telescoping will yield the complexity result. We leave the detailed analysis as future work.

Why choosing LoRA-rank80 and LoRA-rank64? We set LoRA's rank to be 80 and 64 for finetuning 8B and 70B model, respectively. As we explained in Appendix B, we adopt the rank so that the number of trainable parameters of LoRA roughly matches the active block of BAdam/BREAD for a fair comparison.

Suggestions:

1. Statistical evidence for verifying assumptions. Thank you for the suggestion. In the next version, we would like to exam whether the smoothness condition and gradient boundedness condition will be satisfied for a simple and small LLM. Note that verifying smoothness often needs to bound the spectral norm of Hessian, which is computationally prohibitive for large-scale LLMs.

2. Include 2–3 recent memory-efficient methods. Thank you for the suggestion. As we replied previously, we have added comparisons with ReLoRA and DoRA. We will add them in the next version of the manuscript.

We hope that our response is satisfactory and addresses all your concerns. If there are any additional questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications.

References

[1] Zhang et al, "Adam can converge without any modification on update rules", NeurIPS 2022.

[2] Li et al, "Convergence of Adam under relaxed assumptions", NeurIPS 2023.

We thank the reviewer for the valuable comments. We now address your concerns below.

The convergence analysis is based on the deterministic assumption. Since our main focus for this work is to design a practical algorithmic framework (and its practical implementation) rather than establishing new theoretical convergence result, we consider the convergence under the deterministic setting for compact presentation. We remark that it is possible to combine the current analysis with [1, 2] to derive the convergence under stochastic gradient. The important steps would be establishing the approximate descent property (rather descent property) using existing analysis for Adam under stochastic setting, and then providing an overall approximate descent for the BCD loop. Finally, standard telescoping will yield the complexity result. We leave the detailed analysis as future work

BREAD-LoRA requires storing the full model checkpoint. Storing the full parameter checkpoint appears to be inevitable when using full parameter optimization methods such as Adam, SGD, and BCD-based methods. While LoRA has the advantage of storing different types of checkpointed adaptors, full parameter finetuning might achieve higher performance, as we have verified in our experiments.

BREAD-LoRA-SGD might yield better results (W3&Q2). We thank the reviewer for introducing the interesting idea of applying both LoRA and SGD update for inactive blocks. We apply the method in finetuning Llama 3.1-8B under the same setting in Table 1 of the manuscript. The training loss and downstream scores are shown in the following tables:

Table: Training loss.

We observe that BREAD-LoRA-SGD converges evidently faster than both BREAD-LoRA and BREAD-SGD, which supports the hypothesis that applying both memory-efficient optimization algorithms is better than just applying one. For the downstream task, BREAD-LoRA-SGD also slightly outperforms both BREAD-LoRA and BREAD-SGD in terms of average score. Based on these experiments, applying LoRA+ and SGD might yield better result than BREAD-LoRA+. We leave it as an interesting future research direction. In the next version, we will add the suggested variant and BREAD-SGD-LoRA+ (if this LoRA+ version performs well). Thank you again for the nice suggestion.

Are the gradient of shallow layers computed? As shown in equation (5), there are two backward modes. The full backward requires computing the gradient of all the layers, while the efficient backward only computes the gradient up to the active layer.

Suggestion on the choice of BREAD variants. BREAD-LoRA variant offers faster convergence and better downstream performance over BREAD-SGD. The observation is consistent with the fact that LoRA can converge faster than SGD for finetuning task, especially when SGD is performed in pure bfloat16 precision. We thereby suggest to use BREAD-LoRA or BREAD-LoRA+ for finetuning tasks.

Suggestions:

• Acknowledge LOMO. We will acknowledge and mention LOMO when introducing BREAD-SGD in the next version.
• Notation definition. Thank you for the suggestion. The $s_J^t$ refers to the Adam optimizer states of expansion matrices. We will add a definition in the manuscript.
• Discussion on Addax. We thank the reviewer for referring the related work. Addax applies both first-order and zeroth-order update with different batch sizes, where all the layers are equally treated. In contrast, BREAD distinguishes between active and inactive blocks; the active block drives convergence, while the inactive block facilitates landscape expansion. We will discuss the relationship between BREAD and Addax in the next version of our manuscript.

We hope that our response is satisfactory and addresses all your concerns. If there are any additional questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications.

We thank the reviewer for the valuable comments. We now address your concerns below.

Limited novelty: Using LoRA for inactive block is relatively straightforward. We would respectfully argue that BREAD is a well-motivated algorithmic framework, rather than a trivial combination of BCD and LoRA. The framework naturally incorporates BCD with any light weight optimizers (not just LoRA technique) for alleviating the sub-optimality issue and accelerates the convergence of BCD. Specifically, the effectiveness of BREAD framework has been justified through theory, efficient design, and strong empirical results under tight memory constraints. We list the details below.

Theoretical insight. In Proposition 3.1, we identified the suboptimal landscape issue of BCD, where optimizing a single block may lead to suboptimal solution. Proposition 4.2 further reveals that low-rank landscape expansion can mitigate this suboptimality and improve the convergence of BCD. To our knowledge, this phenomenon has not been carefully examined in prior work.

Computational-efficient algorithmic design. By leveraging the intermediate derivatives computed during the backpropagation of BCD, BREAD calculates the gradient of the expansion matrices without requiring additional backpropagation steps.

Competitive performance under limited memory budget. Similar to LoRA and BAdam, BREAD is able to finetune Llama 3.1-70B model with only A100-80GB GPUs. Compared with memory-efficient baselines, the model finetuned by BREAD attains the highest MT-bench score and average math benchmark score in instruction tuning tasks. Notably, BREAD outperforms Adam in supervised finetuning on Alpaca-GPT4, despite its 80% less memory consumption.

In summary, we believe that BREAD is a well-motivated algorithmic framework with theoretical insights, computational-efficient design, and competitive empirical performance, rather than a trivial combination of BCD and LoRA.

No comparison with memory-efficient optimizers such as Galore, DoRA. We have actually provided the comparison with Galore in terms of memory cost and downstream performance; see our Table 1-3 in the manuscript. To include more methods, we finetune Llama 3.1-8B using DoRA and ReLoRA under the same setting in Table 2 of the manuscript. We present the results in the following two tables, where we also list other methods for completeness.

Notably, both BAdam and BREAD outperform ReLoRA and DoRA. This suggests that BCD-based algorithm might be more suitable for math finetuning task. We will provide a more comprehensive comparison with memory-efficient optimization algorithms in the next version of our manuscript.

Proposition 3.1 only applies to specific conditions and Theorem 4.2 is for deterministic case. Let us mention that the primary focus of this work is to design a practical algorithmic framework rather than focusing on theoretical analysis. The derived propositions are mainly intended to motivate the algorithm framework design, and the convergence result serves a partial explanatory purpose to show that the algorithm is partly grounded by classic convergence theory. We consider the convergence under the deterministic setting for compact presentation. We remark that it is possible to combine the current analysis with those of [1, 2] to derive the convergence under stochastic gradient, which we leave as a future work.

Generalization to different model architectures. Thank you for the suggestion. We finetune the Qwen 2.5-7B model using the Alpaca-GPT4 dataset for 3 epochs. The MT-bench scores (with GPT-4 as evaluator) are displayed in the following table:

The BREAD-LoRA method clearly outperforms baseline methods, which demonstrates the effectiveness of BREAD in finetuning models other than Llama series. We will add these results to include the Qwen model in the next version of our manuscript.

Statistical significance of the result. Due to the high computational cost of training LLM, it is almost infeasible to perform multiple runs within the limited rebuttal period. We remark that in Figure 1, BREAD consistently and evidently converges faster than BAdam across different choices of rank and block switch frequency. The consistent faster convergence partially justifies BREAD's performance over BAdam in LLM's finetuning.

Sensitivity of rank, block switch frequency, and learning rate ratio. First, he ablation study in Figure 1(c) has shown that BREAD with even rank-1 expansion matrix achieves significantly faster convergence than BAdam. As the rank increases, the convergence speed increases as well. Second, the result in Figure 1(b) shows that BREAD consistently achieves faster convergence as the block switch frequency $K$ increases (under a reasonable number, say 512). This may attribute to that when using larger $K$, the Adam optimizer will aggregate more historical information, leading to a better search direction. Third, we observe that using the same learning rate for the active module and expansion matrices consistently lead to stable convergence for all BREAD variants, and achieve acceleration over BAdam. However, when applying a 10$\times$ larger learning rate in SGD's update of BREAD-SGD, we observe that training will diverge. We conclude that BREAD will safely converge by applying a universal learning rate, and suggest not to assign highly imbalanced learning rate for active and inactive blocks.

We hope that our response is satisfactory and addresses all your concerns. If there are any additional questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications.

References

[1] Zhang et al, "Adam can converge without any modification on update rules", NeurIPS 2022.

[2] Li et al, "Convergence of Adam under relaxed assumptions", NeurIPS 2023.

We thank the reviewer for the response. In the next version of our manuscript, we would emphasize more the novelty of our algorithmic framework, provide comparisons with other memory-efficient methods, include experiments on alternative architectures, discuss the assumptions used in our theoretical results, etc. Once again, we appreciate your valuable comments and your active participation in the review of our manuscript.

We thank the reviewer for the valuable comments. Overall, let us mention that our primary focus lies in proposing the BREAD algorithmic framework (and its practical implementation) that accelerates BCD for LLM finetuning, rather than focusing on theoretical analysis. The derived propositions are mainly intended to motivate the algorithm design, and the convergence result serves a partial explanatory purpose to show that the algorithm is grounded by classic convergence theory. Below, we address your concerns regarding our theoretical results.

Confusion in Proposition 3.1. The objective of Proposition 3.1 is to show that fixing $W_{\text{out}}$ may potentially exclude the optima, which can happen in the BCD update scheme. As stated in Proposition 3.1, we require $W_{\text{out}}$ in the left hand side to be not full column rank for the strict inequality to hold. Namely, if one fixes $W_{\text{out}}$ at a degenerate point, optimizing other variables only (such as the scheme of BCD) will attain a strictly higher function value compared to the optimal value (i.e., the left hand side). Then, based on this result we conjecture that this issue may potentially slow down the convergence of BCD's subproblem, as it may not include the optima at all. This observation motivates the design of the BREAD algorithmic framework.

Since the feature learning depends on the model parameterization, characterizing the sub-optimality of BCD's subproblem for other layers (other than the last layer) for an arbitrary model is almost impossible. Instead, let us use the following much simplified regression model to demonstrate the effect of fixing one of the other layers. Since the sub-optimality holds for such a simple model, we would then extrapolate that the same issue will likely remain for the training of the much more complicated language models.

Consider the following regression problem $\min_{W} ||H(W, X) - Y||\_2^2$ where the input is $X \in \mathbb{R}^{B\times m}$ and the ground truth label $Y=PX \in \mathbb{R}^{B\times n}$, with rank($P$)=$r$. Suppose H(W, X) is a multi-layer model: $H(W, X):= f_{W_L} \circ \cdots f_{W_2} \circ f_{W_1}(X)$ $f_{W_\ell}(Z):=\sigma(W_\ell Z)$, where $\sigma(\cdot)$ is the ReLU activation function. Suppose one fix any $W_\ell$ with rank($W_\ell$)<r, then by the rank inequality there always exists $(X, Y$) such that $\min_{\{W \backslash W_\ell\}} ||H(W, X) - Y||\_2^2 > \min_{W} ||H(W, X) - Y||\_2^2$

In addition, we will change all the $W$ in the paper to boldface to keep consistency.

Proposition 4.1’s validity. When L=0, the problem essentially becomes finding $Cx=e_y$. One rank-1 solution is let $C_{:, i}=0 \ \forall i \neq y$, $C_{:, y}=\frac{x}{||x||_2^2}$. Therefore, Proposition 4.1 correctly holds under the scenario mentioned in the review.

Clarification on Theorem 4.2. The sample complexity is established by telescoping the inequality over block epoch $t$ rather than the block subproblem step $K$. We now sketch the derivation of the sample complexity based on the inequality of Theorem 4.2. Let us recall the following definitions:

• $W^t$: Weight at the start of block epoch $t$.
• $W_i^t$: Weight at the start of block epoch $t$, block sub-problem $i$.
• $W_{i}^{t, k}$: Weight at block epoch $t$, block sub-problem $i$, update step $k$.

The Theorem 4.2 gives $H(W_{i}^{t, k+1})-H(W_{i}^{t, k}) \leq \mathcal{O}(\alpha) ||\nabla_i H(W_{i}^{t,k})||\_2^2 - \mathcal{O}(\beta) \sum_j ||\nabla_{C_j} H(W_i^{t,k})||\_2^2$ Sum over $k=1, \cdots, K$ steps and apply the assumption of smoothness and bounded gradient gives $H(W_{i+1}^t) - H(W_{i}^{t}) \leq -K \mathcal{O}(\alpha) ||\nabla_i H(W_i^t)||\_2^2 - K \mathcal{O}(\beta) \sum_j ||\nabla_{C_j} H(W_i^t)||\_2^2 + \mathcal{O}(K^2(\alpha^3+\alpha^2\beta))$ Sum over blocks $i=1, \cdots, D$ gives $H(W^{t+1}) - H(W^t)$ $\leq-K\mathcal{O}(\alpha) \sum_{i=1}^{D} ||\nabla_i H(W_i^t)||\_2^2 - DK\mathcal{O}(\beta) \sum_j ||\nabla_{C_j} H(W_i^t)||\_2^2 + \mathcal{O}(DK^2(\alpha^3+\alpha^2\beta))$ $\leq -K\mathcal{O}(\alpha)||\nabla_W H(W^t)||\_2^2 - DK\mathcal{O}(\beta) ||\nabla_C H(W^t)||\_2^2 + \mathcal{O}(D^2K^2(\alpha^3+\alpha^2\beta)),$ where the last inequality follows from the assumption of smoothness and bounded gradient.

Let $\alpha=\mathcal{O}(1/\sqrt{T}), \beta=\mathcal{O}(1/\sqrt{T})$ and telescope the inequality over $t=1, \cdots, T$, we have $\frac{1}{T} \sum_{t=1}^{T}||\nabla_W H(W^t)||_2^2 = \mathcal{O}(\frac{1}{\sqrt{T}K} + \frac{K^2}{T^{1.5}}).$

We will derive the detailed convergence rate in the next version of our manuscript.

Why classical Adam's convergence result cannot be applied, since $t$ is not changing. We clarify that $t$ is the block epoch and changes over training. While one can follow the Adam's analysis to analyze the descent of one block subproblem, the gradient norm of the converged point will contain constant error of $\mathcal{O}(1/\sqrt{K})$, as $K$ is a prefixed constant value. To establish the convergence of the algorithm, one needs to jointly analyze the descent of a full block epoch rather than a single block sub-problem.

Why standard coordinate descent result cannot be used? The update of the active block uses Adam update rule rather than gradient descent update. Therefore, the bias between Adam update and the true block gradient has to be bounded using a different analysis; see equation 12-14 in Appendix D for the error analysis of Adam update.

The convergence result is for deterministic case. Since our main focus for this work is to design a practical algorithmic framework rather than establishing new theoretical convergence result, we consider the convergence under the deterministic setting for compact presentation. We remark that it is possible to combine the current analysis with [1, 2] to derive the convergence under stochastic gradient, which we leave as a future work.

Typos. We thank the reviewer for the careful reading of our manuscript, and will revise all the typos in the next version of our manuscript.

Choice of K beyond 512. Our preliminary experiments show that once $K$ reaches above 512, the convergence acceleration becomes marginal. Hence, we only present the convergence result up to $K=512$. This may due to that the algorithm will overly optimize one block before moving to the next block and slow down the convergence.

We hope that our response is satisfactory and addresses all your concerns. If there are any additional questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications.

References

[1] Zhang et al, "Adam can converge without any modification on update rules", NeurIPS 2022.

[2] Li et al, "Convergence of Adam under relaxed assumptions", NeurIPS 2023.

We thank the reviewer for acknowledging our response and raising the score. In the next version of our manuscript, we would include more detailed explanation on Proposition 3.1, and discuss the scenarios where the strict sub-optimality holds when intermediate layers are frozen. We would also rewrite the convergence statement, provide detailed derivation for the sample complexity, and discuss the challenge of our analysis compared to previous works. Additionally, we will correct all the typos you noted. We appreciate the reviewer's valuable comment and active participation in the review of our manuscript.