CoMERA: Computing- and Memory-Efficient Training via Rank-Adaptive Tensor Optimization

Reminder: discussion window will close in 1 day

Dear Reviwer Rz1r, Thanks a lot for your detailed technical comments and your follow-up discussion.We fully understand that you may super busy with many deadlines at this moment. As the discussion window will close in 1 day, I would highly appreciate it if you confirm your comments have been addressed or not. In our original rebuttal, we have addresed the main weaklness (lack of large-scale experiment) by providing the pre-training result of CodeBERT to show the significant training cost reduction. We also provided details to address your 9 questions (e.g., explanation of various scalarization techniques, difference of our contraction sequence optimization with that in einsum, CuDAGraph optimization, overall training time). In the recent discussion, we have further provided details regarding the comparison with SVDinsTN, result on RTX 4090 (which does not change much from RTX 3090) and code release issue. If our responses have well addressed your comments, we would highly appreciate it if you can acknowledge this. If some of our responses need further clarification, please feel free to let us know! We are staying online to provide potential further clarifications in a timely manner. Thanks again for your detailed technical comments & engagement in the discussion. Warm regards, The authors.

Reminder: discussion window will close in 1 day

Dear Reviewer 6ExQ, Thanks a lot for your constructive comments about our submission CoMERA. We fully understand that you may busy with many deadlines at this moment. As the discussion window will close in 1 day, I would highly appreciate it if you could read our rebuttal. In summary, we have (1) provided a detailed tensor network diagram in the attached PDF file, (2) shown the promising comperssion and 2.5X speedup in pre-training CodeBERT-large, (3）addressed other minior issues such as typos, initial rank setting and accuracy of GaLore and LTE. If our responses have well addressed your comments, we would highly appreciate it if you can acknowledge this. If some of our responses needs further clarification, please feel free to let us know! We are staying online, and we are happy to provide potential further clarification in a timely manner.

follow-up clarification about CodeBERT-large pretraining accuracy and tensor rank/architecture search.

Dear Reviewer 1f6L, Thanks a lot for your discussion and detailed technical discussion. We agree that there are plenty of work about searching tensor ranks and architectures in the machine learning community. Almost all of them are about tensor data compression and post-training model compression. We would like to remark that **searching the architecture and ranks in training is a much more challenging and a rarely investigated task**: the model parameters that we need to compress is not given in advance, and we can not afford a training for each architecture/rank setting due to the huge training cost. This is completely different from tensor data compression and post-training compression, where one can easily evaluate the quality of their rank/architecture setting by doing a cheap forward propagation. This is also why compressed training (even in the simpler low-rank matrix format) is a much more challenging and more important task. Regarding the compression ratio and accuracy in tensor-compressed pre-training, we also would like to remark **three key different features in the LLM domain**: (1) It is natural that compressed models have larger training loss, since its optimization feasible set is smaller than that of an uncompressed training. However, **a slightly larger training loss does not mean worse testing accuracy, because tensorized models have shown smaller generalization gaps (i.e., the difference between training accuracy and testing accuracy) than uncompressed models in many cases**. As an example, let us consider BERT-large pre-training using the same model architecture on the WikiText dataset. Standard pre-training produced a training loss of 1.26, our CoMERA produced a trainining loss of 1.45 (which is worse by 0.19). However, **the compressed model produced by CoMERA outperformed standard BERT-Large in 2 of 3 downstream testing tasks** due to its smaller generalization gaps, as shown below: **Testing accuracy of standard BERT and tensorized BERT on various downstream tasks.** |Models| Accuracy on SST2 |Accuracy on MRPC| Accuracy on SQuAD| |-|-|-|-| |Standard BERT-Large |91.74%| 86.00%|**90.68%**| | Tensorized BERT-Large (ours)|**92.10%**|**86.82%**|88.76%| Since the downstream testing datasets of CodeBERT were not released to the public by its developers, we could not do the same downstream testing at this moment (we are trying to generate a few downstream testing dataset for CodeBERT by ourselves, but it takes time). However, we can make a prediction based on the available training results. Right now, the training loss of CodeBERT using standard pre-training is $0.24$, and training loss produced by our CoMERA is $0.36$. The difference is only **0.12**, which is much better than the loss difference (**0.19**) on the WikiText dataset. Considering the smaller generalization gap of tensorized models, it is reasonable to predict that our compressed CodeBERT-large model will have **better (or at least similar) testing accuracy on downstream code generation tasks**, if those testing datasets are available. (2) Different from regular neural network training which often uses a small dataset and thus the network is often over-parameterized and can have a large compression ratio, LLMs (even small-size LLM like BERT) are pre-trained with a super large dataset and the network is **under-parameterized**. Meanwhile, when those LLMs are released by their developers (often giant tech companies), many tricks have been tried to make sure that these LLMs cannot be dramatically compressed, otherwise customers will quckly swith to much smaller LLMs. As a result, **even a $4.25\times$ overall compression ratio is very large for LLM pre-training**. Note that (i) we only compressed part of the layers (with $\sim 10\times$ compressed ratio on these layers). The embedding tables and output linear layers used a transposted architecture and thus we didn't compress them; (ii) The LLM compression ratio can be further boosted if we combine CoMERA with quantization techniques. (3) **The importance of $2.5\times$ speedup factor and $4.25$ overall compression ratio on CodeBERT-large should not be underestimated**. This means that we can reduce the GPU-hours by $>2.5\times$. **A $2.5\times$ GPU-hour reduction can make a huge impact in the LLM community**, since current LLM pre-training costs millions of US dollars per training run. As a comparison, the recently released and very popular GaLore (oral paper at ICML'2024) has much less memory saving than CoMERA (as shown in our paper) and actually slows down the training. Again, we would like to thank Reviewer 1f6L for his/her technical feedback, which will greatly improve the quality of our work. We hope that the above technical details and potential impact in the LLM community can be considered.

Thanks for your follow-up, and further discussion about pre-training performance.

Dear Reviewer 1f6L, Thanks for your quick follow-up and insightful discussion. Your comment on Galore is insightful. Indeed the comparision should not be just focused on compression ratio and GPU-hours. We completely agree that Galore has better convergence gaurantee. Specifically, it will keep the same accuracy with standard training, under the assumption that the SVD compression in gradient compression does not lose much gradient information. Meanwhile, we also would like to gentlely point out that Galore will not gaurantee convergence and may face divergence issue when it has the same high compression ratio as CoMERA, since the gradient information loss is huge. It is also worth noting that practical LLM pre-training iterations are normally **bounded by the given limited budget (e.g., dollar amounts or GPU-hours)**. As a result, many large-size LLMs are pre-trained for only 1 or 2 epochs in practice. **With the same limited budget of GPU-hours, CoMERA can train $2.5\times$ more iterations than standard pre-training and GaLore, achieving better acuracy**. Regarding the testing accuracy of CoMERA: the testing accuracy of LLMs are evaluated on **multiple downstream tasks**, which are quite different from the pre-training dataset. As a result, a model (even if it has better training loss) often behave very differently on various downstream tasks: it may has better testing accuracy on some tasks, but meanwhile has worse testing accuracy on other tasks. This is very normal, and is exactly the case in our tensorized BERT-large model: our model outperformed the baseline BERT-large on two downstream tasks (SST2 and MRPC), and under-performed basline BERT-large on one downstream task SQuAD. Currently our produced model beats standard BERT-large on **two out of three** downstream tasks (**at the cost of $2.5\times$ less GPU-hours** in pre-training), and we look forward to its future performance as our pre-training methods keep evolving. Thanks agian for your timely discussion. Your participation in the discussion is highly appreciated: it will enable a fair and thorough evaluation of our work. Please feel free to let us know if you have further questions. Best regards, The authors.

Thanks, and further information about pre-training (including LlaMA models)

Dear Reviewer 1f6L, Thanks a lot for your further follow-up discussion: we highly appreciate it, especially considering that yourself may also have many deadlines at this critical moment. Your conern is very valid. Meanwhile, we would like to share more details and our own experience about the evaluation and pre-training of LLMs. Regarding the SQuAD: you are VERY right, this is a more challenging downstream task. But the reason of performance difference across various downstream tasks is not just caused by the diffculty of a specific task like SQuAD. For instance, if we want to achieve a different compression ratio, we find that the accuracy on SQuAD can be higher than the basline, but the acuracy on the easier task MRPC may drop a little bit. As a result, we think that the main reasons are as follows: --(1) **We tested the pre-trained model on downstream tasks without any fine-tuning**. The performance on various downstream tasks definitely can be further improved if fine-tuning is performed using each downstream dataset. --(2) Since we did not conduct task-specific fine-tuning, we are atually evaluating a pre-trained model directly on multiple very differerent downstream tasks. Mathematically, this is equivallent to solving an optimization problem with objective function $f_0(x)$ (in pre-training), then evaluating its solution quality on three different objective functions $f_1(x), f_2(x)$ and $f_3(x)$ (on downstream tasks). Without fine-tuning (i.e., optimizing $f_1(x), f_2(x)$ and $f_3(x)$ **individually**), it is very normal that some tasks show superior performacne and some tasks show slightly degraded performance. Regarding pre-training **larger LLMs** like LLaMa-1B (probably you mean LlaMA-7B?). It is indeed our goal to test our method on large models like LlaMA. However, as an academic group, we are constrainted by the computing budget. Here is one budget data sample: pre-training GPT-2 (with 762 Million parameters) would need 100,000 A100 GPU-hours (equivalent to 250,000 USD based on market cloud computing price). Based on the scaling law released by Google Deepmind (Hoffmann, et al., “Training compute-optimal large language models,”arXiv:2203.15556, 2022), we estimate that the pre-training budget of LlaMa-1B would be around **430,000 USD**, which is far beyond the financial capability of most (if not all) academic groups. Pre-training LlaMA-7B will cost **millions of USD**. We are trying very hard to get more computing resources to validate our ideas on LlaMA-scale models, which is super challenging as an academic group. To be honest, starting up a company may be a better option to achieve this goal (that's why we said that we had some IP concerns when replying to a reviewer's question about a full code release). We will be very happy to share our result with the community in the future, if we can get the testing results on LlaMA-scale models. Thanks again for your insightful discussions. Your active paticipation in the disucssion is highly appreciated.

Thank you for your reply

Thank you so much for your rebuttal and address my questions. This paper formulates an optimization framework to search for ranks inside the tensor training format. I still think this question/problem setup is well exposed in efficient ML works (like search pruning ratio, etc.). Although the optimization description is huge, it is not new to me. Moreover, the claimed benefits is not validated on a large model, and I suppose on the large model (especially pre-training), we need to keep a tight compression ratio to have a comparable accuracy (I saw unmatched accuracy on CodeBert) However, I am very happy to see this work deliver a real training speedup via contraction path order optimization and some CUDA implementation optimization. It would be great to see codes open-sourced. I would like to keep my rating to 6. Thank you so much.