VB-LoRA: Extreme Parameter Efficient Fine-Tuning with Vector Banks

Dear Reviewer o6ri,

1. Weakness #1 – GPU acceleration and computation overhead

VB-LoRA’s implementation is straightforward, and the proposed factorization approach is also simple to implement in modern deep learning frameworks such as PyTorch, allowing us to fully leverage GPU acceleration. However, the use of subvector decomposition does introduce some computational overhead. Despite this, the reduced number of trainable parameters in VB-LoRA results in lower memory consumption. As shown in Table 8, we compared the training time and memory usage of LoRA and VB-LoRA. VB-LoRA requires approximately 15%-20% more training time than LoRA, while using less memory. It's important to note that this additional overhead is limited to the training phase and does not affect inference, as both LoRA and VB-LoRA merge their parameters back into the original model parameters during this stage.

2. Weakness #2 – comparing with other PEFT methods

Thank you for your suggestion! We will revise the manuscript to include additional baseline comparisons. Below, we compare our method with inserted adapters and BitFit. Our results show that our method outperforms other PEFT methods while using an order of magnitude fewer parameters.

1. Adpt_D: Andreas Rücklé, Gregor Geigle, Max Glockner, Tilman Beck, Jonas Pfeiffer, Nils Reimers, and Iryna Gurevych. AdapterDrop: On the efficiency of adapters in transformers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7930–7946.
2. Adpt_P: Jonas Pfeiffer, Aishwarya Kamath, Andreas Rücklé, Kyunghyun Cho, and Iryna Gurevych. AdapterFusion: Non-destructive task composition for transfer learning. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pp. 487–503.
3. Adpt_H: Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp, 2019.
4. BitFit: Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models, 2022.

3. Question #1 – related work “Deep Fusion”

Thank you for bringing Deep Fusion [1] to our attention! This work provides an efficient method for network training under the network growth mechanism by leveraging pre-trained initializations of smaller networks. It shares a similarity with our approach in that both methods induce structure on the parameters of the original transformer model for improved efficiency. Additionally, initialization can be seen as a way to share computation and reduce costs between training smaller and larger models. We will include a discussion of this work in the related work section of our manuscript.

1. Mazzawi et al., Deep Fusion: Efficient Network Training via Pre-trained Initializations, Forty-first International Conference on Machine Learning, 2024

Dear Reviewer eNJC,

1. Weakness #2 – why MNLI and QQP are discarded?

We adhered to the experimental settings established by our baseline, VeRA, which did not include MNLI and QQP. Despite having access to computational resources, it is a shared resource in university. We chose to focus on Natural Language Generation (NLG) and instruction tuning tasks and evaluation on GPT-2 and Llama-2 models, which we believe is of more interest to the research community.

2. Weakness #2 – why choose CoLA for ablation study?

It is a typical practice to choose smaller sized benchmarks for ablation study, since running small benchmarks across different experimental settings can still be computationally intensive and time-consuming. We chose CoLA for the ablation study due to its relevance and consistency with the prior works, e.g., VeRA. Each experiment on the GLUE benchmark was run five times, and the results are reported as the median of these runs to mitigate variance.

3. Weakness #3 – scalability of VB-LoRA to larger models

Note that there are two distinct differences in the routing module between VB-LoRA and MoE. First, the expert network in VB-LoRA is simply a vector in the vector bank. Second, the gating network does not take any model input; instead, it is freely updated by gradients during fine-tuning. In our experiments (e.g., Figure 3, and Figure 5 in the appendix), we observed that even without noisy logits, our method does not suffer from the load balancing issue seen in MoE.

In terms of scalability, we have conducted experiments with a range of model sizes, from Robota 125M to Llama 13B, and varied the number and length of vector banks from 30 to 2048 vectors and 128 to 1024 dimensions, respectively. Our method demonstrated consistent performance across these variations.

To further demonstrate that our approach is free from load balancing issues that hurt scalability, we present the vector usage in a Gemma-7B model trained on the MetaMathQA dataset, as shown in Figure 2 of the attached PDF file. The vector bank contains 2048 vectors. The distribution of vector usage follows a roughly normal distribution, with most vectors being selected between 40 to 55 times.

4. Weakness #4 – performance of VB-LoRA

VB-LoRA's primary focus is on parameter efficiency rather than predictive performance. Despite the extremely high parameter efficiency, our method still delivers accuracies comparable to or sometimes surpassing baseline models.

We appreciate the reviewer’s suggestion to evaluate our method on larger language models, particularly in commonsense and arithmetic reasoning tasks. In response, we have fine-tuned the Mistral-7B and Gemma-7B models on the MetaMathQA dataset and evaluated them on the GSM8K and MATH datasets. We compared our method with LoRA and the concurrent work LoRA-XS [1]. Results show that VB-LoRA outperforms all baselines on the GSM8K dataset, with Mistral-7B utilizing only 0.4% of the parameters compared to LoRA, and Gemma-7B using just 0.3%. Compared to LoRA-XS, our method achieved better results on both evaluation datasets while using only 70% of the parameters with Mistral-7B and 83% with Gemma-7B.

1. Bałazy, Klaudia, et al. "LoRA-XS: Low-Rank Adaptation with Extremely Small Number of Parameters." arXiv preprint arXiv:2405.17604 (2024).

5. Question #1 – memory consumption

We appreciate your insight into the memory consumption concerns. Indeed, the method you mentioned for reducing memory usage was considered and has been implemented already in our submitted supplementary source code, as it integrates seamlessly with our approach. Algorithm 1 serves as a high-level pseudocode to illustrate the formulation of matrices A and B, intentionally omitting the input x for simplicity. Thus, it does not explicitly showcase the memory-saving technique you referenced. We will revise Algorithm 1 as suggested and make it consistent to our implementation.

Dear Reviewer wW2k,

1. Weakness #1 – evaluations on instruction tuning

Thank you for your suggestion. In response, we fine-tuned the Mistral-7B and Gemma-7B models on the MetaMathQA dataset and evaluated them on GSM8K and MATH datasets, and compared them with the suggested work LoRA-XS. Our experimental setup follows the LoRA-XS configuration. The results show that our method outperforms all baselines on GSM8K, with Mistral-7B utilizing only 0.4% of the parameters compared to LoRA, and Gemma-7B using just 0.3%. Compared with LoRA-XS, our method outperforms on both evaluation datasets while using 70% (Mistral-7B) and 83% (Gemma-7B) of parameters. More details can be found in the Global Rebuttal.

2. Weakness #1 – evaluations on NLU and NLG experiments

We chose the GLUE benchmark for evaluation to keep consistency with our baseline methods, LoRA and VeRA, both of which use GLUE in their experiments. Although performance on GLUE is considered saturated, we believe that using this benchmark does not undermine our evaluation, as our primary focus is on parameter efficiency rather than absolute performance metrics. Moreover, we have evaluated our method across a range of model sizes from 125M (Robota-base) to 13B (Llama 2). Overall, these experiments demonstrate the robustness and effectiveness of our method.

3. Weakness #2 – the intuition behind VB-LoRA

We agree that the performance of PEFT methods strongly depends on the task's intrinsic dimensionality [1]. In essence, It's not just the number of training parameters that matters, but also the way they are composed. While both methods use low-rank decomposition, VB-LoRA goes further by breaking layer and module boundaries, employing a less flexible yet expressive reparameterization based on a sparse convex combination of vectors. This approach introduces an inductive bias that can benefit tasks. For more complex tasks like mathematical reasoning, we've included additional evaluations to showcase the effectiveness of our method.

1. Aghajanyan, Armen, Luke Zettlemoyer, and Sonal Gupta. "Intrinsic dimensionality explains the effectiveness of language model fine-tuning." arXiv preprint arXiv:2012.13255 (2020).

4. Weakness #2 – the optimal VB size b and dimensionality of subvectors

Yes, the number of vectors (h) and the length of each vector (b) in the vector bank are task-specific hyperparameters that need to be tuned. Given a fixed budget (h x b), the performance is not highly sensitive to these parameters, as shown in Table 5.

In general, increasing the parameter budget improves performance, as demonstrated in Figure 1. However, the extent of improvement also depends on other factors such as the model architecture and the task difficulty. Performance gains will plateau once these limitations are reached.

5. Weakness #2 – vector selection

In Figure 2 of the attached PDF, we present the vector usage for a Gemma-7B model trained on the MetaMathQA dataset. The distribution of vector usage follows an approximately normal pattern, with most vectors being selected between 40 to 55 times. However, some vectors are chosen more often, with some being selected up to 70 times.

6. Question #1 – adding the multi-index information to the vector selection mechanism can make the TKAM model structure-aware

We apologize for the confusion. Currently, VB-LoRA is not structure-aware, in the sense that the selection of vectors from the vector bank is not informed by factors such as module, layer, and matrix type. However, our framework can be easily extended to become structure-aware. One approach is to make the logits of vector selection conditional on embeddings of the layer, module, and matrix type, which can be implemented through a hypernetwork [1].

1. Mahabadi, Rabeeh Karimi, et al. "Parameter-efficient Multi-task Fine-tuning for Transformers via Shared Hypernetworks." Annual Meeting of the Association for Computational Linguistics, 2021.

7. Question #2 – TKAM and MoE

We chose TKAM for its simplicity and strong performance. In our ablation study (Section 4.4), TKAM outperforms other baselines like Gumbel-Softmax and Straight-Through Gumbel-Softmax. While alternatives like DSelect-k [1] and Expert Choice routing [2] are worth exploring, we focused on introducing the general idea in this paper.

1. Hazimeh, Hussein, et al. "Dselect-k: Differentiable selection in the mixture of experts with applications to multi-task learning." Advances in Neural Information Processing Systems, 2021.
2. Zhou, Yanqi, et al. "Mixture-of-experts with expert choice routing." Advances in Neural Information Processing Systems, 2022.

8. Question #3 – VB-LoRA and Compacter

Compacter and our work are conceptually similar, as both are PEFT methods that operate in an efficiently reparameterized space and reduce redundancies by globally sharing information. However, Compacter designates shared and adapter-specific parameters, while our approach introduces a vector bank with a learned sharing mechanism. This mechanism offers more flexibility by allowing dynamic, context-dependent utilization of shared parameters instead of static designation. We will include this discussion of the relationship between VB-LoRA and Compacter in the related work.

9. Limitations

We agree. We will include these points in the limitations section.

Dear Reviewer e7Nd,

1. Weakness #1, Question #2 – the parameter efficiency when rank is high

First, it’s important to note that rank may not be directly comparable across different methods. In many approaches, rank dictates the number of independent trainable parameters. However, in our method, the majority of the trainable parameters reside in the vector bank, allowing us to use a smaller rank compared to other methods. For instance, when training LLaMA2-13B for instruction tuning, VB-LoRA performs well with r=6, whereas LoRA requires r=64 and VeRA requires r=1024. New experiments on Mistral-7B and Gemma-7B also show that the rank in our method (rank=4) can be significantly lower than baseline methods (r=64), while achieving better performance. More details can be found in the Global Rebuttal.

If high rank is needed, VB-LoRA can reduce the number of parameters more significantly compared to other PEFT methods due to the global sharing mechanism. In this case, the gap in parameter efficiency between LoRA and Full FT becomes smaller, whereas VB-LoRA can sustain a high rank while maintaining a compact vector bank.

1. Bałazy, Klaudia, et al. "LoRA-XS: Low-Rank Adaptation with Extremely Small Number of Parameters." arXiv preprint arXiv:2405.17604 (2024).

2. Weakness #2, Question #1 – decomposed cumulative gradient parameters are more sensitive than the original model parameters during the training process

The reason for this increased sensitivity is the decomposition process. Similar to hypernetwork, the parameters are generated by another parameterized model, instead of directly being updated through gradients. In our case, the cumulative gradient updates ΔW are decomposed into matrices A and B, and further into sub-vectors in VB-LoRA. The vectors in the vector bank and logits are the parameters directly updated through gradients. This can be related to the training instability issues observed in hypernetworks [1], where heuristics such as gradient clipping are typically used to manage these instabilities [2]. In TKAM, we found that adding noise to the logits [3] can lead to difficulties in training stability.

1. Ortiz, Jose Javier Gonzalez, John Guttag, and Adrian V. Dalca. "Magnitude Invariant Parametrizations Improve Hypernetwork Learning." The Twelfth International Conference on Learning Representations. 2024.
2. Ha, David, Andrew Dai, and Quoc V. Le. "Hypernetworks." arXiv preprint arXiv:1609.09106 (2016).
3. Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).

3. Weakness #3 – vector specialization

Figure 3 shows the footprint of the entire training process. It is important to note that active exploration predominantly occurs in the early stages of training. As training progresses, each sub-vector starts to focus more on much fewer (specialized) vectors within the vector bank. To highlight this pattern, we have plotted the footprint at different training periods in Figure 1 in the attached pdf file. This visualization demonstrates that updates become progressively sparser in the later stages of training.

4. Weakness #4 – virtual dimension b

The number of vectors (h) and the length of each vector (also known as the virtual dimension b) in the vector bank are task-specific hyperparameters that need to be tuned. One constraint on b is that it needs to be a common factor of all hidden dimensions to ensure compatibility across the entire model, as the hidden dimensions of FFN layers may differ from those of attention layers.

Table 5 shows that given a fixed budget (h×b), the performance is not highly sensitive to the exact values of these parameters. However, selecting a moderate b may help achieve higher performance. Given the budget of the vector bank, a larger b reduces the number of vectors in the vector bank, potentially making each vector less specialized. On the other hand, a smaller b increases the number of trainable parameters (as discussed in Sec. 3.4) and complicates the vector selection process.

We will include this discussion in the paper.

5. Limitation #2 – evaluations with GPT-4

GPT-4 is only used for evaluating instruction tuning for Llama2. All other experiments do not rely on GPT-4, thus remaining consistent over time.