SwitchLoRA: Switched Low-Rank Adaptation Can Learn Full-Rank Information

## Abstract

Abstract In the training of large language models, parameter-efficient techniques such as LoRA optimize memory usage and reduce communication overhead during the fine-tuning phase. However, applying such techniques directly during the pre-training phase results in poor performance, primarily because the premature implementation of low-rank training significantly reduces model accuracy.
Existing methods like ReLoRA and GaLore have attempted to address this challenge by updating the low-rank subspace. However, they still fall short of achieving the accuracy of full-rank training because they must limit the update frequency to maintain optimizer state consistency, hindering their ability to closely approximate full-rank training behavior.
In this paper, we introduce SwitchLoRA, a parameter-efficient training technique that frequently and smoothly replaces the trainable parameters of LoRA adapters with alternative parameters. SwitchLoRA updates the low-rank subspace incrementally, targeting only a few dimensions at a time to minimize the impact on optimizer states. This allows a higher update frequency, thereby enhancing accuracy by enabling the updated parameters to more closely mimic full-rank behavior during the pre-training phase.
Our results demonstrate that SwitchLoRA actually surpasses full-rank training, reducing perplexity from 15.23 to 15.01 on the LLaMA 1.3B model while reducing communication overhead by 54% on the LLaMA 1.3B model.
Furthermore, after full fine-tuning the SwitchLoRA pre-trained model and the full-rank pre-trained model on the GLUE benchmark, the SwitchLoRA pre-trained model showed an average accuracy gain of about 1% over the full-rank pre-trained model. This demonstrates enhanced generalization and reasoning capabilities of SwitchLoRA.

## 1 Introduction

The size of large language models (LLMs) has increased rapidly due to the advent of the transformer architecture []. To support the training of large models, distributed training techniques such as data parallelism [, ], tensor parallelism [], pipeline parallelism [, ] and the Zero Redundancy Optimizer [] have been employed.
However, distributed training of trillion-scale models incurs significant inter-node communication overhead from synchronizing extensive parameter gradients across multiple nodes.

To address these challenges, various parameter-efficient strategies have been proposed. Techniques such as model sparsification [, ] and progressive model pruning during training [] have shown promise. Additionally, methods leveraging Singular Value Decomposition (SVD) to approximate full-rank matrices in low-rank spaces have been explored [, , ].
Beyond the entire training process, several techniques improve adaptability and efficiency during the fine-tuning phase. For example, methods such as the Adapter [, ] and Prefix-tuning [] introduce additional trainable layers while freezing the remaining parameters.

Another noteworthy fine-tuning strategy is Low-Rank Adaptation (LoRA) [], which introduces no computational overhead during inference while maintaining training accuracy.
However, previous studies [, , ] have observed that parameter-efficient methods such as LoRA perform less efficiently during the pre-training phase because the premature use of low-rank training leads to a considerable loss in model accuracy.
To increase the rank of updated parameters and benefit from low-rank training, ReLoRA [] applies the structure of LoRA and periodically resets LoRA adapters. Similarly, GaLore [] projects gradients onto a subspace, updating this subspace periodically.
These approaches update the descent direction of trainable parameters to mimic the behavior of full-rank training, thereby overcoming the limitations observed in existing implementations of low-rank adaptation.
However, we find that the intervals between resetting/updating steps in ReLoRA and GaLore are set to relatively large values because too frequent changes in the updating direction can cause inconsistency in optimizer states, which may not sufficiently approximate the behavior of full-rank training, resulting in a loss of accuracy.

To address this challenge, as illustrated in Figure 1, we introduce SwitchLoRA, which enables smooth and frequent adjustments to the trainable parameters of the LoRA matrices while introducing negligible additional computational overhead.
SwitchLoRA maintains a set of candidate vectors for each matrix within the LoRA adapters. At each training step, it replaces portions of the column or row vectors with these candidate vectors, subsequently training the LoRA adapters.
This process minimizes the impact on optimizer states, thus allowing for a higher update frequency compared to ReLoRA and GaLore. By more closely approximating full-rank parameter updating behaviors during the pre-training phase, this approach enhances overall accuracy.

Figure: Figure 1: SwitchLoRA: An enhanced LoRA with dynamic vector switching for pre-training. In traditional LoRA, an adapter $\mathbf{B}\mathbf{A}$ is added to the matrix $\mathbf{W}$ of linear layers. $\mathbf{B}$ and $\mathbf{A}$ are trained while $\mathbf{W}$ is kept frozen (as depicted in the left part of the figure). SwitchLoRA enhances this by dynamically switching vectors within $\mathbf{B}$ and $\mathbf{A}$. The figure illustrates an example of this process: when the third column(labeled as black \raisebox{-0.9pt}{3}⃝) of $\mathbf{B}$ is switched, the corresponding third row(labeled as white \raisebox{-0.9pt}{3}⃝) of $\mathbf{A}$ is temporarily frozen. Similarly, when the second row(labeled as black \raisebox{-0.9pt}{2}⃝) of $\mathbf{A}$ is switched, the corresponding second column(labeled as white \raisebox{-0.9pt}{2}⃝) of $\mathbf{B}$ is also temporarily frozen.
Refer to caption: x1.png

#### Our contribution:

- •
We propose SwitchLoRA to facilitate smooth and frequent adjustments to the trainable parameters of the LoRA matrices through low-rank adaptation, maintaining the accuracy of full-rank training while reducing communication overhead.
- •
To mitigate inconsistencies in optimizer states when parameters are switched, SwitchLoRA resets the corresponding optimizer states and temporarily freezes the affected parameters. Additionally, SwitchLoRA employs a different initialization rule for LoRA adapter parameters and their associated candidate vectors, thereby improving the overall efficiency of the training process.
- •
We experimentally validate SwitchLoRA on various sizes of the LLaMA model. SwitchLoRA shows significant perplexity improvements when compared to ReLoRA [] and GaLore []. For the 1.3B model, SwitchLoRA achieves a perplexity of 15.01, surpassing the 15.23 perplexity obtained with full-rank training.
Furthermore, by performing full fine-tuning on the resulting 1.3B model using the GLUE [] tasks to validate the reasoning capabilities, we demonstrate that SwitchLoRA enhances model accuracy by approximately 1% on average, compared to the full-rank training method.

## 2 Methodology

A substantial body of research, such as various pruning methods [, ], has demonstrated that neural networks tend to exhibit low-rank characteristics after certain stages of training. Techniques for parameter-efficient fine-tuning, such as LoRA, capitalize on this observation.
Concurrently, studies like [, ] have revealed that overparameterization in neural networks can lead to implicit regularization, thereby enhancing generalization. These findings underscore the importance of training with full parameters during the initial phase. Further empirical evidence supporting this phenomenon is provided in works like [, , , ].
Based on these insights, this section proposes a method designed to train a substantial number of parameters while selectively updating only a portion of the parameters at any one time to reduce communication overhead.

### 2.1 Low-Rank Adaptation (LoRA)

Introduced in [], LoRA is designed specifically for the fine-tuning stage of model training.

Consider a pre-trained model with a weight matrix $\mathbf{W}\in\mathbb{R}^{m\times n}$ from a specific linear layer. LoRA proposes an innovative modification: transforming $\mathbf{W}$ into $\mathbf{W}+\frac{\alpha}{r}\mathbf{B}\mathbf{A}$. Here, $\mathbf{B}\in\mathbb{R}^{m\times r}$ and $\mathbf{A}\in\mathbb{R}^{r\times n}$ are newly introduced matrices, where $r$ is a positive integer significantly smaller than both $m$ and $n$. And $\alpha$ is a constant hyperparameter, set to $r$ in the following description to clarify the algorithm’s mechanics. Then during fine-tuning, $\mathbf{W}$ is kept frozen while matrices $\mathbf{B}$ and $\mathbf{A}$ are trained. At the inference stage, $\mathbf{B}\mathbf{A}$ is added to $\mathbf{W}$ which preserves the model’s original structure. The matrix $\mathbf{A}$ is initialized using Kaiming initialization [], while $\mathbf{B}$ is initially set to a zero matrix to ensure consistency.

### 2.2 SwitchLoRA

Below, we detail our proposed SwitchLoRA algorithm, the steps of which are outlined in Algorithm 1 and Algorithm 2.

#### Switching process

Now, let us delve deeper into the linear system $(\mathbf{W}+\mathbf{B}\mathbf{A})\mathbf{x}=\mathbf{y}$. As illustrated in Figure 1, we decompose the matrix $\mathbf{B}$ into its column vectors $\mathbf{b}_{k}\in\mathbb{R}^{m\times 1}$ for $k=1,\ldots,r$, represented as $\mathbf{B}=[\mathbf{b}_{1},\ldots,\mathbf{b}_{r}]$.
Similarly, we decompose matrix $\mathbf{A}$ into its row vectors $\mathbf{a}_{k}^{T}\in\mathbb{R}^{1\times n}$ for $k=1,\ldots,r$, leading to $\mathbf{A}^{T}=[\mathbf{a}_{1}^{T},\ldots,\mathbf{a}_{r}^{T}]$. Hereafter, we call these vectors $\mathbf{b}_{k}$ and $\mathbf{a}_{k}$ as LoRA vectors.

The product $\mathbf{B}\mathbf{A}$ can be expressed using these LoRA vectors as follows:

$$ $\displaystyle\mathbf{B}\mathbf{A}=\sum_{k=1}^{r}\mathbf{b}_{k}\mathbf{a}_{k}^{ T}.$ (1) $$

Let $\mathcal{C}(\mathbf{B})$ denote an ordered set containing $\min(m,n)$ vectors, each having the same dimensions as $\mathbf{b}_{k}$. Furthermore, ensure that $\{\mathbf{b}_{1},\ldots,\mathbf{b}_{r}\}\subset\mathcal{C}(\mathbf{B})$. Similarly, define $\mathcal{C}(\mathbf{A}^{T})$ as an ordered set containing $\min(m,n)$ vectors, each having the same dimensions as $\mathbf{a}_{i}$. Also, let $\{\mathbf{a}_{1}^{T},\ldots,\mathbf{a}_{r}^{T}\}\subset\mathcal{C}(\mathbf{A}^
{T})$.
Moving forward, we will refer to $\mathcal{C}(\mathbf{B})$ and $\mathcal{C}(\mathbf{A}^{T})$ as the candidate vectors for $\mathbf{B}$ and $\mathbf{A}$, respectively.

It is known for $k$ matrices $\mathbf{W}_{1},\mathbf{W}_{2},\ldots,\mathbf{W}_{k}$, the following inequality holds:

$$ $\displaystyle rank(\sum_{i=1}^{k}\mathbf{W}_{i})\leq\sum_{i=1}^{k}rank(\mathbf {W}_{i}).$ (2) $$

If we adopt the strategy in LoRA to add $\mathbf{B}\mathbf{A}$ to $\mathbf{W}$ and only update $\mathbf{B}$ and $\mathbf{A}$ from the pre-training stage, according to (2), the rank of updated parameters of the local linear system through the entire training process will be limited to $2r$.
This limitation can potentially impede the training efficacy.
To mitigate this issue, we alter the values of $\mathbf{b}_{k}$ and $\mathbf{a}_{k}$ to $\mathbf{b}^{\prime}_{k}\in\mathcal{C}(\mathbf{B})$ and $\mathbf{a}^{\prime}_{k}\in\mathcal{C}(\mathbf{A}^{T})$ at appropriate frequencies, respectively, with these new values randomly selected from predefined candidate vectors list $\mathcal{C}(\mathbf{B})$ and $\mathcal{C}(\mathbf{A}^{T})$(one of $\mathbf{b}_{k}^{\prime}$ or $\mathbf{a}_{k}^{\prime}$ can be the same as $\mathbf{b}_{k}$ or $\mathbf{a}_{k}$).
To maintain the consistency of the model’s output, we adjust $\mathbf{W}$ by adding the difference between the old and new LoRA components.
To be more precise, when $\mathbf{b}_{k}$ and $\mathbf{a}_{k}$ are updated to $\mathbf{b}^{\prime}_{k}$ and $\mathbf{a}^{\prime}_{k}$, we accordingly adjust $\mathbf{W}$ with the equation $\mathbf{W}\leftarrow\mathbf{W}+\mathbf{b}_{k}\mathbf{a}_{k}^{T}-\mathbf{b}_{k}
^{\prime}\mathbf{a}_{k}^{\prime T}$.

When implementing these updates, the updated parameters of both $\mathbf{B}$ and $\mathbf{A}$ are derived from $\min(m,n)$ distinct candidate vectors, which ensures updated parameters are full-rank.
Readers can refer to [, , ] for more details.

When selecting candidate vectors, we have the option to choose randomly from $C(\mathbf{B})$ or $C(\mathbf{A}^{T})$. Alternatively, we can select candidate vectors sequentially from $C(\mathbf{B})$ or $C(\mathbf{A}^{T})$, restarting from the beginning once the end of the set is reached. We find that varying the matching orders of vectors $\mathbf{b}_{k}$ and $\mathbf{a}_{k}$ yields only minor differences in outcomes. A theoretical explanation for this phenomenon is provided in Appendix A. Additionally, to conserve GPU memory, spare candidate vectors can be offloaded to the CPU.

#### Switching frequency

As mentioned in [, , ], the model initially exhibits full internal rank during pre-training, and the internal rank of each layer decreases progressively over time. Consequently, we have adopted an exponential decay function for the switching frequency, namely $frequency=Ce^{-\theta step}$, where the coefficients are determined empirically.
Besides, the selection of LoRA rank $r$ for $\mathbf{B}\mathbf{A}$ is influenced by the final internal rank of the layers, which has been extensively explored in [, , ].

#### Optimizer states resetting

Currently Large Language Models (LLMs) predominantly utilize Adam [] and AdamW [] optimizers over SGD, which rely on optimizer states. It is crucial to note that after switching LoRA vectors, the gradients associated with these parameters are also changed, which prevents the reuse of optimizer states. To address this issue, when $\mathbf{a}_{k}$ is switched, we reset the optimizer states of $\mathbf{b}_{k}$. And conversely, when $\mathbf{b}_{k}$ is switched, we reset optimizer states of $\mathbf{a}_{k}$. Note that we reset optimizer states of counterpart pair rather than optimizer states of the switched parameters itself. This approach will be further explained in Appendix A.
Additionally, when the optimizer states are reset to zero, we freeze corresponding parameters for $N$ steps to maintain the robustness of the training. In this study, $N$ is set to $5$.

Figure: Algorithm 1 Switch algorithm: $\mathbf{W},\mathbf{P},\mathbf{Q}=\text{switch}(\mathbf{W},\mathbf{P},\mathbf{Q },i,j)$. $\mathcal{C}(\mathbf{P})[i]$ is $i$-th predefined candidate vectors for $\mathbf{P}$.

Figure: Algorithm 2 SwitchLoRA training process. switch_num$(step,r,interval_{0},\theta)$ is an integer generator function which yields $\lfloor s\rfloor+X$ numbers sampled from $1$ to $r$ where $s=r/(interval_{0}e^{\theta step})$ and random variable $X\sim\text{Bernoulli}(s-\lfloor s\rfloor)$, i.e. $P(X=1)=1-P(X=0)=s-\lfloor s\rfloor$.

#### Initialization of SwitchLoRA

Results in [, ] have demonstrated the importance of initialization of LoRA matrices $\mathbf{B}$ and $\mathbf{A}$ to the training effects.
Unlike these works, which are applied only during the fine-tuning stage, our method is utilized throughout the entire training process.
To achieve appropriate initialization for matrices $\mathbf{B}$ and $\mathbf{A}$ along with their candidate vectors, we follow the idea of Xavier initialization [] and Kaiming initialization []. The derived initialization values are presented as follows:

$$ $\displaystyle std[\mathbf{B}]=std[\mathbf{b}]=(\frac{r}{\sqrt{mn}})^{\frac{1}{ 4}}gain^{\frac{1}{2}}\quad\forall\mathbf{b}\in\mathcal{C}(\mathbf{B}),$ $\displaystyle std[\mathbf{A}]=std[\mathbf{a}]=(\frac{\sqrt{m}r}{\sqrt{n}n})^{ \frac{1}{4}}gain^{\frac{1}{2}}\quad\forall\mathbf{a}\in\mathcal{C}(\mathbf{A}^ {T}),$ (3) $$

where $gain$ is a constant dependent on the type of activation function used.

A detailed analysis of the above results can be found in Appendix A.

## 3 Related work

#### Direct low-rank factorization method

Numerous studies [, , , ] have demonstrated the effectiveness of using low-rank factorization to approximate the weights of linear layers in deep neural networks. They employ methods such as SVD to achieve a factorization $\mathbf{U}\mathbf{V}$ that minimizes $\|\mathbf{W}-\mathbf{U}\mathbf{V}\|$.
Later on, Pufferfish [] and subsequent work in Cuttlefish [] employ full-rank training prior to low-rank training to enhance efficiency. Additionally, they introduce adaptive strategies to determine the necessary duration of full-rank training and to select the appropriate rank for each linear layer for SVD.
Further developments in this field include InRank [], which proposes a low-rank training approach based on greedy low-rank learning []. Additional research such as [] integrates orthogonality into the low-rank models to enhance training accuracy, while [] introduces low-rank ordered decomposition, a generalization of SVD aimed at improving low-rank training efficiency.
These innovations mainly focus on convolutional neural networks (CNNs) and smaller-scale language models.

#### LoRA variants

After the introduction of LoRA in [], which facilitated fine-tuning with very few trainable parameters, numerous works are proposed to improve the performance of LoRA.
Improvements include better initialization strategies for LoRA matrices as demonstrated in [, , ]. Additionally, [, ] have adjusted learning rates for $\mathbf{B}$ and $\mathbf{A}$ to optimize training outcomes.
Other research efforts, such as those in [, ], have modified the training process of LoRA.
Moreover, some studies, such as [, ], focus on training models from scratch within a sparse model structure.

Similar to our approach, various LoRA variants employ strategies to increase the rank of updated parameters by merging parameters of adapters into $\mathbf{W}$. For instance, Chain of LoRA [] and ReLoRA [] merge $\mathbf{B}\mathbf{A}$ into $\mathbf{W}$ and restart training at regular intervals. ReLoRA enables low-rank training during the early phases, yet it still requires 33% of the steps to be full-rank training.
Delta-LoRA [], another variant, targets the fine-tuning phase by updating the matrix $\mathbf{W}$ using the gradients from the LoRA matrices $\mathbf{B}$ and $\mathbf{A}$ as they are updated, enhancing accuracy for fine-tuning.

#### Other compression methods

In addition to previously discussed techniques, there are many other methods to compress models during training.
For instance, several studies have introduced quantization to LoRA [, , ], effectively reducing memory overhead during fine-tuning.
Other research employs iterative pruning and growth techniques during training [, , , ].
Additionally, some works focus on compressing gradients through quantization [, ] or gradient projection [].
Notably, [] presents a recent method for training from scratch that utilizes SVD to project gradients into a periodically updated subspace.
This approach also enables the addition of quantization, offering enhanced memory efficiency compared to LoRA.

## 4 Experiments

### 4.1 Experimental setup

Our studies are carried out on the LLaMA model [], with model sizes reduced to 130M, 250M, and 350M.
We designed our experiments based on the settings described in [] to benefit from established hyperparameter configurations.
The specific hyperparameters for these models are detailed in Table 1. We use Adam optimizer to train the model with $\beta_{1}=0.9,\beta_{2}=0.999$.
We use a cosine learning rate schedule with 100 warm-up steps and a total of 40,000 training steps.

The pre-training experiments utilize the C4 dataset [], with the first 46M samples of the training dataset serving as our training data, and samples from the entire validation dataset used for testing. The evaluation of validation loss is performed on 10M tokens for all our experiments, with evaluations conducted every 1,000 steps.
Additionally, we utilize some of tasks from the GLUE benchmark [] to assess the reasoning capabilities of the models.
All experiments are conducted using 8xNVIDIA Tesla A100 80GB PCIe GPUs. Gradient accumulation is applied when GPU memory reaches its limit.

We have conducted ablation studies to assess the impact of various configurations, detailed in Appendix B.

**Table 1: Model sizes and architectures used in our experiments**
| Params | Hidden | Heads | Layers | Batch size | Batch size per GPU | Seq. len. |
| --- | --- | --- | --- | --- | --- | --- |
| 130M | 768 | 12 | 12 | 600 | 150 | 256 |
| 250M | 768 | 16 | 24 | 1152 | 72 | 512 |
| 350M | 1024 | 16 | 24 | 1152 | 72 | 512 |
| 1.3B | 2048 | 32 | 24 | 1536 | 16 | 512 |

To ensure fairness across all experiments, the initialization method described in Section 2 is applied to both LoRA and SwitchLoRA experiments. We deploy LoRA adapters across all attention layers and fully connected layers in these experiments.

For the hyperparameters in Algorithm 2, we initiate with $interval_{0}=40$ and set $N=5$. The parameter $\theta$ is adjusted to ensure that the switching frequency is one-third of its initial frequency at the $1/10$ of total steps.

All experiments were repeated multiple times to select the best results. The learning rates for pre-training experiments were selected from a predefined set $\cup_{n=2,3,4}${1e-n 2e-n, 5e-n}.
We have determined that the optimal learning rate remains consistent across different model sizes for all methods. Specifically, the learning rate for full-rank training is set at 0.001, while the learning rate for the LoRA method is 0.01. For SwitchLoRA, the learning rate is slightly higher at 0.02.

### 4.2 Basic experiments

Figure: Figure 2: Loss results for 130M, 250M, and 350M models with a LoRA rank of $128$.
Refer to caption: x2.png

Figures 2 displays the experimental results for the 130M, 250M, and 350M models, respectively, with the LoRA rank set to $128$. The data reveal that while LoRA alone does not yield satisfactory training results, SwitchLoRA approaches the performance of full-rank training. The performance gap continues to grow as model size increases. This suggests that the low-rank training approach, such as LoRA, might cause models to become trapped in local minima, while SwitchLoRA mitigates this issue by dynamically changing trainable parameters.

**Table 2: Perplexity results at step 38,000 for 130M, 250M and 350M.**
|  | 130M | 250M | 350M |
| --- | --- | --- | --- |
| Full-rank | 27.71 | 20.19 | 18.72 |
| LoRA(rank$=128$) | 34.74 | 29.56 | 31.87 |
| SwitchLoRA(rank$=128$) | 30.26 | 20.97 | 19.96 |
| SwitchLoRA(rank$=256$) | $\backslash$ | 19.82 | 18.70 |

**Table 3: Perplexity results at step 38,000 for 1.3B models.**
|  | 1.3B |
| --- | --- |
| Full-rank | 15.23 |
| SwitchLoRA(rank$=256$) | 15.89 |
| SwitchLoRA(rank$=512$) | 15.01 |

Figure: Figure 3: Loss results for 250M, 350M and 1.3B models using higher LoRA ranks.
Refer to caption: x5.png

As shown in Figure 3, additional experiments conducted on the 250M, 350M and 1.3B models using higher LoRA ranks demonstrates improved performance compared to those with the rank set at $128$, achieving outcomes close to those of full-rank training.
Although utilizing a higher rank yields better outcomes, it may not be more economical to increase the LoRA rank instead of increasing the model size for larger models for several reasons. First, the method still has potential for further refinement. Second, a lower LoRA rank enables training on devices with limited memory capacities.
Furthermore, in the context of 3D parallelism, inter-node communication is predominantly influenced by data parallelism, where communication overhead is proportional to trainable parameters. The trainable parameters for each model are detailed in Table 4. For further discussions on potential ways to enhance the SwitchLoRA strategy, refer to Section 5. Additionally, the impact of distributed training is detailed in Appendix F.

**Table 4: Comparison of trainable parameters: full-rank models vs. LoRA and SwitchLoRA.**
| Full-rank | 247.5M | 247.5M | 368.2M | 368.2M | 1339.5M | 1339.5M |
| --- | --- | --- | --- | --- | --- | --- |
| (Switch)LoRA | $r=128$ | $r=256$ | $r=128$ | $r=256$ | $r=256$ | $r=512$ |
| 98.9M | 148.4M | 125.6M | 185.4M | 370.7M | 609.7M |  |

### 4.3 Comparison with other methods

Among all related methods, the works which are most close to ours are ReLoRA [] and GaLore []. We do comparison experiments on these two methods to further validate the effectiveness of our algorithm.
The learning rates for all methods are tuned in $\cup_{n=2,3,4}${1e-n 2e-n, 5e-n}.

#### Comparison with ReLoRA

Figure: Figure 4: Comparison between ReLoRA and SwitchLoRA. In the figure, red circles denotes the steps at which the parameters of the LoRA adapter are reset.
Refer to caption: x8.png

Since ReLoRA requires full-rank pre-training as warm-up, we do full-rank pre-training on SwitchLoRA too to do a fair comparison. We train 250M LLaMA model specified in Section 4, with detailed settings available in Appendix C. In the Figure 4, we compare ReLoRA and SwitchLoRA with different full-rank pre-training steps. It shows that our method can still perform better when ReLoRA uses 5,000 steps full-rank pre-training and SwitchLoRA uses 200 steps full-rank pre-training. Furthermore, when both algorithms are subjected to the same 1,000 steps of full-rank pre-training, SwitchLoRA shows significant improvements on ReLoRA.

The frequency for resetting the LoRA adapters in ReLoRA is set to 1/5,000, significantly lower than the initial switching frequency of 1/40 in SwitchLoRA experiments. As illustrated in Figure 4, we observe a rapid decrease in loss at each resetting step in the ReLoRA experiments. In contrast, the loss reduction in SwitchLoRA experiments is steady and more rapid.

#### Comparison with GaLore

In the comparison experiments with GaLore, we strictly follow the setup in Galore []. Detailed setup can be found in Appendix C.
For the 350M LLaMA model, GaLore achieves a perplexity of 20.29, whereas SwitchLoRA performs slightly better, with a perplexity of 19.58.
In addition, we conducted additional experiments on the 350M model, changing only one hyperparameter to assess its impact. The perplexity results are shown in Table 5.

When further reducing the rank, as shown in Table 5, our method performs significantly better.
This improvement may be because GaLore’s use of SVD focuses on the most significant directions, whereas SwitchLoRA covers all update directions, including less important ones that still require training.

**Table 5: Perplexity comparison for GaLore and SwitchLoRA with different experimental setup.**
|  | Standard | Model size=130M | Rank=128 | Rank=32 | Seq. len. = 512 |
| --- | --- | --- | --- | --- | --- |
| GaLore | 20.29 | 26.17 | 22.52 | 34.09 | 19.03 |
| SwitchLoRA | 19.58 | 25.93 | 20.93 | 25.26 | 18.19 |

The gradient projection subspace update frequency in GaLore is set at 1/200, while the initial switching frequency for SwitchLoRA is 1/40. Additionally, since updates in GaLore are performed via SVD, the subspace changes are less frequent compared to approaches that randomly select a new subspace. Consequently, the subspace changes in GaLore are, in fact, less efficient.

### 4.4 Reasoning ability comparison

Current works on low-rank training for LLMs, such as [, ], primarily evaluate models based on perplexity and lack validation of reasoning abilities.
To validate the reasoning abilities, we also conducted full fine-tuning using the resulting checkpoints from the aforementioned experiments. We fine-tuned the models on GLUE tasks []. For the checkpoints trained using SwitchLoRA, all LoRA adapters are merged into the original weights such that $\mathbf{W}\leftarrow\mathbf{W}+\mathbf{B}\mathbf{A}$ before the fine-tuning process. Detailed experiment settings are provided in Appendix C.

**Table 6: GLUE benchmark of the full-rank, SwitchLoRA and GaLore pre-trained 350M models. The metric for STS-B is the Pearson correlation, while Matthew’s correlation coefficient is used for CoLA. Accuracy is reported for the other tasks.**
|  | CoLA | STS-B | MRPC | RTE | SST2 |
| --- | --- | --- | --- | --- | --- |
| Full-rank pre-trained | 42.95$\pm$5 | 87.26$\pm$0.2 | 79.16$\pm$1 | 59.86$\pm$1 | 90.88$\pm$0.5 |
| SwitchLoRA pre-trained | 23.13$\pm$15 | 87.71$\pm$0.5 | 76.86$\pm$2 | 56.24$\pm$5 | 90.83$\pm$0.3 |
| GaLore pre-trained | 40.23$\pm$2 | 86.14$\pm$0.5 | 72.70$\pm$4 | 54.66$\pm$4 | 89.35$\pm$0.5 |

We first perform full fine-tuning on the pre-trained 350M models. The full-rank pre-trained model is from Section 4.2. Similarly, the SwitchLoRA pre-trained model, with a LoRA rank of 256, is also from Section 4.2. The GaLore pre-trained model originates from newly conducted experiments, where the batch size, sequence length, and rank are the same as in the SwitchLoRA experiment. This GaLore pre-training experiment resulted in a perplexity of 21.61.

The full fine-tuning results for these three models are shown in Table 6. From the results, we observe that for the 350M models, except for the CoLA task, SwitchLoRA outperforms GaLore by [insert performance difference], and outperforms the full-rank model by [insert performance difference].

We also conduct fine-tuning experiments on the 1.3B LLaMA models, one pre-trained using full-rank and the other pre-trained using SwitchLoRA with a LoRA rank of 512.
Both pre-trained models are from Section 4.2.
As shown in Table 7, SwitchLoRA performs slightly worse in some tasks and better in others compared to the full-rank results. Overall, the average score of SwitchLoRA exceeds the full-rank results by approximately 1%.

**Table 7: GLUE benchmark of the Full-rank pre-trained and SwitchLoRA pre-trained models. The metrics for all tasks except STS-B are accuracy. The metric for STS-B is Pearson correlation coefficient.**
|  | CoLA | STS-B | MRPC | RTE | SST2 |
| --- | --- | --- | --- | --- | --- |
| Full-rank pre-trained | 48.60$\pm$2 | 87.64$\pm$0.1 | 78.43$\pm$1 | 58.05$\pm$3 | 91.93$\pm$1 |
| SwitchLoRA pre-trained | 47.43$\pm$3 | 88.49$\pm$0.3 | 80.15$\pm$2 | 61.37$\pm$3 | 92.39$\pm$0.5 |

## 5 Limitations and future work

While our results are promising, there are several areas for future exploration.
In our experiments, we have demonstrated that selecting a larger LoRA rank is necessary to achieve accuracy comparable to full-rank training. Additionally, finely tuning the switching frequency of the LoRA vectors presents significant challenges.
To address these limitations, we propose the following directions for future work, as illustrated in Figure 5.

- •
In our experiments, we simply used exponentially decreasing switching frequencies, which may not be the optimal approach. Guidelines should be developed to help set appropriate switching frequencies throughout the training process.
- •
Going further, a more detailed idea is to examine each layer of the model to adjust the switching frequencies.
For instance, LoRA-drop [] evaluates whether the rank is sufficient using a norm of $\Delta\mathbf{W}\mathbf{x}$. This is rational because different types of layers, such as the $Q,K,V$ matrices in transformer layers, exhibit significantly varied behaviors.
- •
In our work, we simply chose candidate vectors at random or sequentially. However, during training, all candidates are updated separately, leading to significant differences among them. The selection of these candidates may improve the training outcomes.

Figure: Figure 5: Future work roadmap.
Refer to caption: x10.png

## 6 Conclusions

In this work, we introduce SwitchLoRA, a novel training strategy designed for parameter-efficient pre-training. Our approach achieves comparable accuracy to full-rank training while reducing the trainable parameters to approximately 50% to 60% of those in traditional full-rank training. Moreover, the computational overhead and memory usage are nearly identical to those of LoRA when using the same number of trainable parameters.
We further validate the reasoning abilities of models trained with SwitchLoRA using the GLUE benchmark. The results from the 1.3B model indicate that SwitchLoRA not only matches but also slightly outperforms full-rank training by about 1% in accuracy.

## Appendix A Theoretical analysis

In this section, we conduct a thorough discussion of our algorithm and address the following key aspects:

- 1.
Demonstrating that the order of LoRA vectors does not impact performance;
- 2.
The effectiveness of our algorithm;
- 3.
Discussion on resetting optimizer states;
- 4.
Detailed process to deduce the values for initialization.

First, we take a closer look at the properties of the local linear system. Assume that the loss function of the model is denoted by $\mathcal{L}$. Our discussion focuses on the scenario where the input $\mathbf{x}$ and output $\mathbf{y}$ are vectors, satisfying the equation:

$$ $\displaystyle\mathbf{y}=(\mathbf{W}+\frac{1}{r}\mathbf{B}\mathbf{A})\mathbf{x},$ (4) $$

where the bias term is omitted for simplicity.

Next, we calculate the gradients of the column vectors of $\mathbf{B}$.
For a function $f(\mathbf{x})$, we denote $\nabla_{\mathbf{x}}f$ as the partial derivative of $f$ with respect to $\mathbf{x}$.
Recall the decomposition of $\mathbf{B}\mathbf{A}$ as defined in the previous equations. For $k=1,\ldots,r$, the gradient of $\mathbf{b}_{k}$ with respect to the loss function $\mathcal{L}$ is given by:

$$ $\displaystyle\nabla_{\mathbf{b}_{k}}\mathcal{L}=(\mathbf{a}_{k}^{T}\mathbf{x}) \nabla_{\mathbf{y}}\mathcal{L}.$ (5) $$

Note that when the input $\mathbf{x}$ is a vector, $\mathbf{a}_{k}^{T}\mathbf{x}$ becomes a scalar. Consequently, the gradients of $\mathbf{b}_{k}$ are proportional to the gradients of $\mathbf{y}$.

We can also derive the gradients of the row vectors of $\mathbf{A}$ as follows:

$$ $\displaystyle\nabla_{\mathbf{a}_{k}}{\mathcal{L}}=((\nabla_{\mathbf{y}} \mathcal{L})^{T}\mathbf{b}_{k})\mathbf{x}.$ (6) $$

In this expression, $(\nabla_{\mathbf{y}}\mathcal{L})^{T}\mathbf{b}_{k}$ is a scalar, indicating that the gradients of $\mathbf{a}_{k}$ are aligned in the direction of the input activations.

In fact, the gradients expressed in (5) and (6) and be derived as follows:

Consider the expression for $\mathbf{y}_{i}$ given by $\mathbf{y}_{i}=\sum_{j,k}\mathbf{B}_{ij}\mathbf{A}_{jk}\mathbf{x}_{k}+\sum_{j}
\mathbf{W}_{ij}\mathbf{x}_{j}$ for $i=1,\ldots,m$. The partial derivative of the loss function $\mathcal{L}$ with respect to $\mathbf{B}_{ij}$ is computed as

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial\mathbf{B}_{ij}}=\sum_{k}\frac {\partial\mathcal{L}}{\partial\mathbf{y}_{k}}\frac{\partial\mathbf{y}_{k}}{ \partial\mathbf{B}_{ij}}=\frac{\partial\mathcal{L}}{\partial\mathbf{y}_{i}} \frac{\partial\mathbf{y}_{i}}{\partial\mathbf{B}_{ij}}=\frac{\partial\mathcal{ L}}{\partial\mathbf{y}_{i}}\sum_{k}\mathbf{A}_{jk}\mathbf{x}_{k},$ (7) $$

where we use the fact that $\frac{\partial\mathbf{y}_{k}}{\partial\mathbf{B}_{ij}}=0$ when $k\neq i$. This derivation confirms (5). Similarly, the derivative with respect to $\mathbf{A}_{jk}$ is

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial\mathbf{A}_{jk}}=\sum_{i}\frac {\partial\mathcal{L}}{\partial\mathbf{y}_{i}}\frac{\partial\mathbf{y}_{i}}{ \partial\mathbf{A}_{jk}}=\sum_{i}\frac{\partial\mathcal{L}}{\partial\mathbf{y} _{i}}\mathbf{B}_{ij}\mathbf{x}_{k}.$ $$

This calculation leads to (6).

#### Independence of vectors updating

In our algorithm, candidate vectors are either randomly selected or chosen sequentially to replace vectors in $\mathbf{A}$ and $\mathbf{B}$, which alters the matching pairs of $\mathbf{b}_{k}$ and $\mathbf{a}_{k}$. A natural question arises: Does the matching order of these vector pairs influence the training effects?

In the following discussion, we will use the notation $\tilde{\mathbf{v}}$ to denote trainable parameters that are initialized with the value of $\mathbf{v}$.

For the sake of clarity, we focus on one linear layer without a bias term for our discussion. We denote $\mathcal{L}(\tilde{\mathbf{W}}\mathbf{x})$ as the loss when the weight matrix of the linear layer under study is $\mathbf{W}$, with the vector $\mathbf{x}$ as input activations. This formulation intentionally omits contributions from other layers and the bias term, as they are beyond the scope of our subsequent analysis.

To integrate the LoRA matrices while preserving the initial loss value, we reformulate $\mathcal{L}(\tilde{\mathbf{W}}\mathbf{x})$ as $\mathcal{L}((\mathbf{W}-\sum_{k}{\mathbf{b}}_{k}{\mathbf{a}}_{k}^{T}+\sum_{k}
\tilde{\mathbf{b}}_{k}\tilde{\mathbf{a}}_{k}^{T})\mathbf{x})$.
Further, we simplify this expression to $\mathcal{L}(\mathbf{a}_{1},\ldots,\mathbf{a}_{r};\mathbf{b}_{1},\ldots,\mathbf
{b}_{k};\mathbf{x})$. A simple observation is

$$ $\displaystyle\mathcal{L}(\mathbf{a}_{1},\ldots,\mathbf{a}_{r};\mathbf{b}_{1}, \ldots,\mathbf{b}_{k};\mathbf{x})=\mathcal{L}(\mathbf{0},\ldots,\mathbf{0}; \mathbf{0},\ldots,\mathbf{0};\mathbf{x}).$ (8) $$

Recall that the gradient $\nabla_{\mathbf{b}_{k}}\mathcal{L}=(\mathbf{a}_{k}^{T}\mathbf{x})\nabla_{
\mathbf{y}}\mathcal{L}$.
We derive the following expression:

$$ $\displaystyle\Delta\mathbf{b}_{k}\mathbf{a}_{k}^{T}=(c(\mathbf{a}_{k}^{T} \mathbf{x})\nabla_{\mathbf{y}}\mathcal{L}+\text{opt_state}(\mathbf{b}_{k})) \mathbf{a}_{k}^{T},$ (9) $$

where $c$ is a negative value from optimizer and opt_state$(\mathbf{b}_{k})$ is optimizer state of $\mathbf{b}_{k}$, determined by the value of $(\mathbf{a}_{k}^{T}\mathbf{x})\nabla_{\mathbf{y}}\mathcal{L}$ of previous steps.
Moreover, the value of $\nabla_{\mathbf{y}}\mathcal{L}$ will remain unchanged, as indicated by (8).
Consequently, the component $\Delta\mathbf{b}_{k}\mathbf{a}_{k}^{T}$ is influenced solely by $\mathbf{a}_{k}$ and not by other LoRA vectors.
Similarly, the value of $\mathbf{b}_{k}\Delta\mathbf{a}_{k}^{T}$ is influenced only by $\mathbf{b}_{k}$ when switching $\mathbf{a}_{k}$.
Note that the updated weight can be expressed as

$$ $\displaystyle(\mathbf{b}_{k}+\Delta\mathbf{b}_{k})(\mathbf{a}_{k}^{T}+\Delta \mathbf{a}_{k}^{T})-\mathbf{b}_{k}\mathbf{a}_{k}^{T}=\Delta\mathbf{b}_{k} \mathbf{a}_{k}^{T}+\mathbf{b}_{k}\Delta\mathbf{a}_{k}^{T}+\Delta\mathbf{b}_{k} \Delta\mathbf{a}_{k}^{T},$ (10) $$

where $\Delta\mathbf{b}_{k}\Delta\mathbf{a}_{k}^{T}$ represents a minor term that can generally be disregarded. Hence, the updates derived by $\mathbf{b}_{k}$ and $\mathbf{a}_{k}$ are nearly independent.

From this discussion, we can conclude that the order of vectors $\mathbf{a}_{k}$ and $\mathbf{b}_{k}$ does not influence the parameter updates in the current step. For instance, for $1\leq i,j\leq r$, back propagation of $\mathcal{L}(\mathbf{a}_{1},\ldots,\mathbf{a}_{j},\ldots,\mathbf{a}_{i},\ldots,
\mathbf{a}_{r};\mathbf{b}_{1},\ldots,\mathbf{b}_{k};\mathbf{x})$ and $\mathcal{L}(\mathbf{a}_{1},\ldots,\mathbf{a}_{i},\ldots,\mathbf{a}_{j},\ldots,
\mathbf{a}_{r};\mathbf{b}_{1},\ldots,\mathbf{b}_{k};\mathbf{x})$ yield almost the same parameters updating to the weight matrix of the linear layer.

#### Effectiveness of SwitchLoRA

Consider the following modification to the original model. For the weight matrix $\mathbf{W}\in\mathbb{R}^{m\times n}$ of a specific linear layer in the model, replace $\mathbf{W}$ with the product of matrices $\mathbf{B}^{0}\mathbf{A}^{0}$, where $\mathbf{B}^{0}\in\mathbb{R}^{m\times\min(m,n)}$ and $\mathbf{A}^{0}\in\mathbb{R}^{\min(m,n)\times n}$. This modification results in a full-rank weight matrix $\mathbf{B}^{0}\mathbf{A}^{0}$ and introduces more parameters than the original model. Consequently, it is anticipated to achieve results that are at least as good as those of the original model when the full parameters of this modified model are trained.

We now compare the modified model with another model that implements the SwitchLoRA strategy. Define $\mathbf{B}^{0}_{:,i}=\mathcal{C}(\mathbf{B})[i]$ and $\mathbf{A}^{0}_{i,:}=\mathcal{C}(\mathbf{A}^{T})[i]^{T}$ for $i=1,\ldots,\min(m,n)$. It becomes apparent that the two models are quite the same except that the model applying SwitchLoRA strategy updates only subsets of parameters incrementally.

In optimization, it is well-established that for problems with separable objective functions, the parameters of each separable group can be optimized independently. Although the loss function of the SwitchLoRA model is not separable, the preceding discussion has demonstrated the independence between the LoRA vectors. Consequently, we can infer that the inseparable components of the loss function concerning parameters within the same linear layer are modest. Therefore, this suggests that training subsets of parameters incrementally, as in the SwitchLoRA model, is likely more effective than other methods, such as the layer-wise training approach [, ].

#### Reset of optimizer states

Let us discuss whether it is reasonable to zero out the optimizer states of LoRA vectors and temporarily freezing them when switching their counterpart LoRA vectors.

Consider a scenario where $\mathbf{b}_{k}$ is switched while $\mathbf{a}_{k}$ is not.
Note that, according to (8), the forward propagation remains unaffected after the switching occurs. During the initial step after switching $\mathbf{b}_{k}$, with $\mathbf{a}_{k}$ being frozen, the only term contributing to the weight matrix update is $\Delta\mathbf{b}_{k}\mathbf{a}_{k}^{T}$ according to (10).
We previously established that this term, $\Delta\mathbf{b}_{k}\mathbf{a}_{k}^{T}$ in (9), is not influenced by other LoRA vectors apart from $\mathbf{a}_{k}$. Consequently, changes made to $\mathbf{b}_{k}$ or any other recently switched LoRA vectors do not impact the accuracy of the optimizer states for $\mathbf{b}_{k}$. This substantiates the rationale behind resetting the optimizer states.

If we choose not to freeze $\mathbf{a}_{k}$, we derive the following from a similar equation to (9):

$$ $\displaystyle\mathbf{b}_{k}\Delta\mathbf{a}_{k}^{T}=c((\nabla_{\mathbf{y}} \mathcal{L})^{T}\mathbf{b}_{k})\mathbf{x}+\mathbf{b}_{k}\text{opt_state}( \mathbf{a}_{k}).$ (11) $$

This formula demonstrates that without resetting $\mathbf{a}_{k}$, the update direction would be completely incorrect.

The reasoning for switching $\mathbf{a}_{k}$ and its implications can be deduced in a similar manner.

#### Derivation of parameters initialization

The initial values of $\mathbf{B}$ and $\mathbf{A}$ were specified in Section 2. In this section, we present the derivation process.

The main idea of [] and [] is to maintain a balance in the variance of the activation and gradients across layers during forward and backward propagation. In this study, we focus on balancing the variance of activations. Furthermore, we aim to ensure the updated parameters derived from $\mathbf{B}$ are of the same amount as those derived from $\mathbf{A}$:

$$ $\displaystyle\Delta\mathbf{B}\mathbf{A}\sim\mathbf{B}\Delta\mathbf{A}.$ (12) $$

Consider two matrices, $\mathbf{W}_{1}$ and $\mathbf{W}_{2}$, both characterized by zero mean and uniform distribution. The standard deviation (std) of the elements of their product is given by:

$$ $\displaystyle std[\mathbf{W}_{1}\mathbf{W}_{2}]=\sqrt{k}std[\mathbf{W}_{1}]std [\mathbf{W}_{2}],$ (13) $$

where $k$ represents the output dimension of the matrix $\mathbf{W}_{1}$.
To ensure the stability of forward propagation, it is crucial that the output of each layer maintains a standard deviation of $1$. However, when the matrix $\mathbf{W_{2}}$ represents activation values, its standard deviation, denoted as $\text{std}[\mathbf{W}_{2}]=\text{gain}$, differs from $1$ due to the influence of the activation function. For ReLU activations, $gain=\sqrt{2}$.
Following this principle, we derive:

$$ $\displaystyle std[\frac{1}{r}\mathbf{B}\mathbf{A}\mathbf{x}]=\frac{\sqrt{r}}{r }std[\mathbf{B}]std[\mathbf{A}]{\sqrt{n}}=gain.$ (14) $$

The standard deviation of the gradients for LoRA vectors is given by:

$$ $\displaystyle std[\nabla_{\mathbf{b}_{k}}\mathcal{L}]=\sqrt{n}std[\mathbf{a}_{ k}]std[\mathbf{x}]std[\nabla_{\mathbf{y}}\mathcal{L}],$ $\displaystyle std[\nabla_{\mathbf{a}_{k}}{\mathcal{L}}]=\sqrt{m}std[\mathbf{b} _{k}]std[\mathbf{x}]std[\nabla_{\mathbf{y}}\mathcal{L}].$ (15) $$

Assuming the updated parameters are solely influenced by the gradients of the current step, to obtain (12), the following condition must be met:

$$ $\displaystyle std[\nabla_{\mathbf{B}}\mathcal{L}\mathbf{A}]=std[\mathbf{B} \nabla_{\mathbf{A}}\mathcal{L}].$ (16) $$

From this, we derive:

$$ $\displaystyle std[\nabla_{\mathbf{B}}\mathcal{L}\mathbf{A}]=\sqrt{r}std[\nabla _{\mathbf{b}_{k}}\mathcal{L}]std[\mathbf{A}],$ $\displaystyle std[\mathbf{B}\nabla_{\mathbf{A}}]=\sqrt{r}std[\mathbf{B}]std[ \nabla_{\mathbf{a}_{k}}\mathcal{L}].$ (17) $$

By combining (14)-(A), we achieve the following standard deviations:

$$ $\displaystyle std(\mathbf{A})=({\frac{\sqrt{m}r}{n\sqrt{n}}})^{\frac{1}{4}} gain^{\frac{1}{2}},\quad std(\mathbf{B})=(\frac{r}{\sqrt{mn}})^{\frac{1}{4}} gain^{\frac{1}{2}}.$ (18) $$

## Appendix B Ablation study

In this section, we mainly use the 130M model with a LoRA rank of $128$ and a batch size of 128. For hyperparameters not explicitly mentioned, we follow the configurations detailed in Section 4.

In Figure 6, we did two experiments. In the first experiment, we evaluate the model’s performance with varying descent rates for frequencies while maintaining a constant initial switching interval of $40$. In the second experiment, we maintain a consistent descent rate for frequencies as detailed in Section 4, but we vary the initial switching interval across different experiments. It is evident from our results that both hyperparameters significantly impact training accuracy.

Figure: Figure 6: Loss comparison for the 130m model with different $interval_{0}$ and $ratio$, where the parameter $ratio$ determines the point at which the switching frequency is reduced to one-third of its initial value, occurring at the step $total_step\times ratio$.
Refer to caption: x11.png

In Figure 7, we conducted a series of experiments with various frequency settings. The results indicate that the choice of frequency settings plays a crucial role in the model’s effectiveness. Specifically, we find that setting both the initial frequency values and the descent rates to moderate levels is essential for achieving optimal performance. Extremely high or low frequency settings tend to degrade the model’s performance, indicating a sensitive balance that must be maintained.

Figure: Figure 7: Perplexity comparison for the 130m model with different switching frequencies. Each point in the figure has a triple label $(interval_{0},{ratio},{perplexity})$, with its size corresponding to the perplexity value. The parameter $ratio$ determines the point at which the switching frequency is reduced to one-third of its initial value, occurring at the step $total_step\times ratio$.
Refer to caption: x13.png

In Figure 8, we conduct experiments to investigate the impact of the number of frozen steps $N$. The results indicate that the choice of $N$ influences the loss outcomes. This phenomenon can be explained as follows: when $N$ is excessively large, the training parameters may become biased towards different subsets of the data. Conversely, if $N$ is too small, at the moment the freezing is canceled, the gradients will have a larger contribution to the parameter updates due to the nature of momentum-based optimizers. This leads to potentially abrupt changes in model behavior. However, selecting an optimal value for $N$ is relatively straightforward, as this value is robust across different model since it simply determines how many steps are needed to warm up switched LoRA vectors. Therefore, this hyperparameter does not require frequent adjustments across various experiments.

Figure: Figure 8: Comparison of loss for the 130m model at different values of $N$.
Refer to caption: x14.png

In Figure 9, we present the results from a focused comparative study where we evaluated our initialization strategy against the traditional LoRA initialization method through two distinct experiments. The results indicate that our initialization method outperform traditional approach for initialization. Notably, the loss curve for LoRA initialization reveals a slower decrease in initial loss compared to that of SwitchLoRA initialization. This phenomenon in LoRA initialization can be attributed to the slow warm-up of matrix $\mathbf{A}$ and its associated candidate vectors due to (6). In contrast, our method modifies the initialization values to allow for more rapid adjustments, enabling the model to adapt more effectively to the training data.

Figure: Figure 9: Loss comparison for the 130m model between traditional and our enhanced initialization methods.
Refer to caption: x15.png

## Appendix C Experimental setting details

### C.1 Experimental settings of ReLoRA

We adhere strictly to the setup described in ReLoRA [] for our comparative experiments with ReLoRA. Specifically, the warm-up steps for the scheduler are set to 1,000. The learning rates are as follows: 5e-4 for full-rank pretraining, 1e-3 for ReLoRA, and 1e-2 for SwitchLoRA. The total batch size is established at 20,000. All other settings remain consistent with our previous experiments as detailed in Section 4.

### C.2 Experimental settings of GaLore

Continuing in the same vein, we also strictly adhere to the setup outlined in GaLore [] for our comparison experiments with GaLore. Specifically, we set the warm-up steps for the scheduler at 6,000. The total batch size is adjusted to 60,000. The learning rate is standardized at 1e-2 for all GaLore experiments, while for SwitchLoRA, it is set at 2e-2. All other experimental settings remain consistent with those detailed in our previous experiments, as described in Section 4.

### C.3 Experimental settings of fine-tuning

The hyperparameters for fine-tuning used in the experiments described in Section 4.4 are presented in Table 8 and Table 9.

**Table 8: Hyperparameters for different GLUE tasks for the 350M models.**
|  | CoLA | STS-B | MRPC | RTE | SST2 |
| --- | --- | --- | --- | --- | --- |
| $lr$ (Full-rank) | 8e-6 | 1e-5 | 1e-5 | 8e-6 | 3e-6 |
| $lr$ (SwitchLoRA) | 3e-5 | 5e-5 | 5e-5 | 5e-5 | 1e-5 |
| $lr$ (GaLore) | 2e-6 | 5e-6 | 5e-6 | 3e-6 | 8e-7 |
| Batch size |  |  | 16 |  |  |
| Epochs |  |  | 30 |  |  |
| Sequence length |  |  | 512 |  |  |

**Table 9: Hyperparameters for different GLUE tasks for the 1.3B models.**
|  | CoLA | STS-B | MRPC | RTE | SST2 |
| --- | --- | --- | --- | --- | --- |
| $lr$ (Full-rank) | 8e-6 | 1e-5 | 2e-5 | 1e-5 | 5e-6 |
| $lr$ (SwitchLoRA) | 1e-5 | 1e-5 | 1e-5 | 2e-6 | 5e-6 |
| Batch size |  |  | 16 |  |  |
| Epochs |  |  | 30 |  |  |
| Sequence length |  |  | 512 |  |  |

## Appendix D Implementation of LoRA vector switching

We discuss the code implementation of SwitchLoRA, focusing on its efficiency and memory consumption.

#### Implementation Adjustments in Optimizer

The primary distinction in the implementation of SwitchLoRA from conventional approaches lies in its handling of gradients and optimizer states at the granularity of row or column vectors within matrix parameters. Consider the scenario when using the AdamW optimizer: typically, each trainable parameter group in AdamW is associated with a “step” state which is implemented as a float scalar value in the code. To facilitate the resetting of specific rows or columns in matrices, we modify the type of "step" in the optimizer to a 32-bit float matrix with the same shape as the corresponding parameters. In fact, this modification does incur some extra memory overhead. An alternative approach would be to implement “step” as a row vector for $\mathbf{A}$ and a column vector for $\mathbf{B}$. However, this would require more complex code management, and thus, we have not adopted this strategy in our implementation.
With the capability to manipulate optimizer states and gradients at the level of rows and columns, we can now execute operations such as resetting optimizer states and freezing specific rows or columns of parameter matrices.

#### Implementation of the Switching Process

We can either randomly select or sequentially select candidate vectors. However, fragmented operations on a GPU can’t fully utilize its capabilities. Since several candidate vectors are switched at each step, this will impact training efficiency. As an example, during the initial phase of SwitchLoRA training for the 1.3B LLaMA model with a LoRA rank of 512, approximately $\frac{512}{40}\approx 13$ candidate vectors are switched for each LoRA matrix at every step.

By organizing a list of candidate vectors into a matrix and selecting vectors sequentially, we can perform operations on multiple vectors simultaneously. For example, consider a scenario where we need to set the values of candidate vectors $\mathcal{C}(\mathbf{B})$ at indices 4, 5, 6 to the values of $\mathbf{B}$ at indices 7, 8, 9, respectively. Let $\mathcal{C}^{\mathbf{B}}$ be a matrix defined as $\mathcal{C}^{\mathbf{B}}\in\mathbb{R}^{m\times\min(m,n)}$, where each column $\mathcal{C}^{\mathbf{B}}{:,i}=\mathcal{C}(\mathbf{B})[i]$ for $i=1,\ldots,\min(m,n)$. We can then directly assign $\mathcal{C}^{\mathbf{B}}_{:,4:7}=\mathbf{B}_{:,7:10}$. This arrangement enables us to consolidate operations on multiple contiguous indices into a single operation, enhancing efficiency.
Consequently, we employ a sequential selection approach and apply this technique. By implementing this approach, the switching process now occupies only about 1/40 of the training time during the initial training phase.

#### Memory offloading for candidate vectors

The use of candidate vectors leads to additional GPU memory usage. This memory overhead can be reduced by offloading it to the CPU. The offloading process can be decoupled from other training processes. By utilizing non-blocking CPU offloading, we can handle both offloading and other training processes in parallel, which can be readily implemented using frameworks like PyTorch.

The amount of parameters offloaded at each step is approximately $switch_freq\times lora_rank/hidden_dim\times total_param$. For the 1.3B LLaMA model, using 16-bit precision for model parameters, this translates to:
$1/40\times 512/2048\times 1.3e9\times 2$ bytes $\approx 16.25MB$.

## Appendix E Distribution of Singular values

Figure: Figure 10: Rank distribution of LoRA on different types of linear layers.
Refer to caption: x16.png

Figure: Figure 11: Rank distribution of full-rank training and SwitchLoRA on different types of linear layers.
Refer to caption: x19.png

Given that the rank distribution significantly influences the training efficacy of models [, ], we conducted experiments to examine the rank distribution of SwitchLoRA.
As outlined in Section 4, experiments were conducted on the 350M model to analyze the rank distribution of linear layers after 40,000 training steps.
Figure 10 demonstrates that the singular values of weight matrices converge within a limited range when trained with LoRA, indicating dominance of LoRA adapters in the linear layers.
This dominance is expected, as the singular value distribution of weight matrices during the pre-training phase exhibits a form of illness, due to updates being limited to the low-rank adapter $\mathbf{B}\mathbf{A}$.
In contrast, as illustrated in Figure 11, the rank distribution of SwitchLoRA closely approximates that of full-rank training, suggesting a more robust and more effective adaptation process.

## Appendix F Impact on distributed training

As demonstrated in [], for a transformer model with $n$ layers and a hidden dimension of $h$, the memory required for model parameters scales proportionally with $nh^{2}$. Assuming these parameters are stored in $fp16/bf16$ format occupying $\Psi$ parameters, the memory footprint for optimizer states would be approximately $12\Psi$ bytes when using the Adam optimizer as stated in [].
Additionally, when the batch size is $b$ and the sequence length is $s$, the memory consumption for activations scales with $bshn$. To manage memory demands for large models, gradient accumulation can be utilized to adjust the batch size per GPU to $1$. Moreover, activation checkpointing can be implemented to reduce memory consumption, though it comes with a trade-off: a 33% increase in computational overhead.

In this work, we primarily focus on the memory consumption associated with optimizer states, which constitutes a significant portion of the overall memory usage for models with tens of billions of parameters. Assuming that full-rank training requires $knh^{2}$ bytes of memory, where $k$ is a constant. Our algorithm, as well as LoRA, reduces memory usage from $knh^{2}$ to $2knhr$, with $r$ representing the LoRA rank.

In addition to memory usage, parameter-efficient training also reduces communication overhead. When implementing 3D parallelism to train large language models, tensor parallelism is typically limited within a single machine due to its substantial communication demands. Pipeline parallelism introduces some idle “bubble” time, which cannot be eliminated even with fast communication. And its communication overhead remains relatively low. The main part of inter-node communication stems from data parallelism, where the same amount of gradients as parameters is communicated at every training step. Consequently, having fewer trainable parameters can significantly decrease communication overhead. Moreover, reduced memory consumption allows a larger portion of the model to reside on a single GPU, potentially decreasing the degree of pipeline parallelism needed and consequently reducing the associated “bubble” time.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - lines 470-472: there are missing numbers at "[insert performance difference]". 
- line 193-198: The authors implemented an exponential decay function for switching frequency during training, defined as frequency = Ce^(-θ * step), with coefficients determined empirically.
  - This reliance on empirical tuning limits the method’s generalizability across various models and datasets, as optimal values may vary depending on specific training scenarios.
  - The approach assumes a progressive decrease in each layer’s internal rank during training, yet this behavior may not be consistent across all models or tasks.
  - Although the exponential decay function is inspired by observed trends, the paper does not provide a theoretical framework to justify this specific form of decay.
  - A comparative analysis is lacking in the current approach to support the superiority of this decay function.
  - The predetermined exponential decay schedule does not account for the dynamic nature of training, potentially reducing its effectiveness in varied scenarios.

Furthermore, the authors provided some critics on previous works but the reviewer has a different perspective. the switchLoRA brings about specific scenarios without generalization. while, based on previous works, the reviewer can argue the following process to enhance:
  - The low-rank adaptation matrices are initialized by performing SVD on the pretrained weights. This method selects only the top singular vectors, retaining task-relevant information while significantly reducing the number of trainable parameters.
  - The low-rank matrices derived from the pretrained weights are kept frozen during training. Only a small, trainable matrix, positioned between these frozen matrices, is updated in fine-tuning. This approach reduces computational and memory overhead, as adaptation occurs through a single small matrix rather than full-rank updates.
  - In contrast to methods where parameter count scales with model dimensions, this approach keeps a constant trainable parameter count by using the small matrix with fixed dimensions. This design is highly efficient for large-scale models, where maintaining a low parameter count and memory efficiency is essential.
soundness: 3
strengths: - In empirical tests, SwitchLoRA demonstrates performance improvements, achieving lower perplexity than full-rank training, especially on the LLaMA 1.3B model.
- Despite frequent updates, SwitchLoRA keeps computational and memory overhead low by using pre-trained candidate vectors.
- When fine-tuned on GLUE tasks, SwitchLoRA shows a slight improvement in accuracy over full-rank models, indicating enhanced generalization.
summary: - the authors proposed a method that leverages frequent parameter updates in LoRA (Low-Rank Adaptation) matrices during pre-training. 
- LoRA techniques in related works shown in the paper optimize memory and communication overhead during fine-tuning but underperform during pre-training due to low-rank constraints. 
- Existing methods (ReLoRA and GaLore) address this by periodically resetting low-rank subspaces, but their large intervals result in accuracy loss according to the perspectives of the authors. 
- In constrast, the method in this paper (SwitchLoRA) allows smooth, frequent parameter updates by switching LoRA vectors without significantly impacting the model’s optimizer states. 
- Key features of the idea:
  - SwitchLoRA frequently replaces portions of column and row vectors in LoRA matrices with pre-defined candidate vectors, allowing it to approximate full-rank training behavior more closely.
  - For each matrix in the LoRA adapter, a set of candidate vectors is maintained. The system can switch vectors dynamically, keeping model output consistent and effectively increasing the adaptability of low-rank spaces.
  - The frequency of switching is controlled by an exponential decay function, allowing the model to dynamically adjust update rates.
  - To manage the switch in parameters while maintaining optimization stability, SwitchLoRA resets certain optimizer states, allowing the model to stabilize quickly.
  - SwitchLoRA employs a specific initialization scheme for candidate vectors, enhancing stability and effectiveness during training.
weaknesses: - Dynamic parameter adjustment impedes scalability for very large models or environments with limited resources due to additional overhead and computational costs of scaling factors.
- Broader applicability is limited since the paper primarily evaluates SwitchLoRA within language tasks, leaving its performance and adaptability in other domains.
- SwitchLoRA assumes that task-appropriate configurations can be achieved simply by adjusting scaling factors on existing model parameters. While effective within the tested scenarios, this approach may not generalize across models with different architectures or to tasks where such adjustments cannot capture necessary model changes.
- The SwitchLoRA paper provides limited comparison to other parameter-efficient fine-tuning techniques, such as adaptive pruning or selective parameter updates.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - The main claim of the paper must be backed with experiments, as noted in the Weaknesses section.
- The authors don’t compare validation loss curves with GaLore, which, along with point 2 in Weaknesses, casts doubt on their claims of outperforming GaLore. Including these validation loss curves would better substantiate their claims.
- Can the optimizer states be updated layer-wise, instead of updating the entire model? This could lead to further memory saving.
- The authors can include Tables/Figures comparing the memory usage of SwitchLora compared to other methods,

Minor
- Table 5 needs clarity. Are the metrics in the first two columns for a rank of 256? The configurations linked to each value are unclear and difficult to follow.
soundness: 2
strengths: - The overall switching methodology - selecting candidate vectors to reset the optimizer states - is novel and enables the use of high switching frequencies. 
- The evaluations and experiments against LoRA and full-rank training are extensive and clearly show the benefits of using SwitchLoRA against them.
- The proposed method maintains performance against full-rank training while reducing the number of trainable parameters to 50-60% to full-rank training, with minimal communication overhead.
summary: SwitchLoRA addresses limitations in low-rank adaptation methods like ReLoRA and GaLore, which restrict update frequency to maintain optimizer state consistency, thus limiting their approximation of full-rank training. SwitchLoRA frequently and smoothly alternates LoRA adapter parameters, updating only a few dimensions at a time to reduce the impact on optimizer states. This approach allows for higher update frequency, achieving accuracy improvements by closely approximating full-rank behavior. The authors validate SwitchLoRA on various LLaMA model sizes, comparing against full-rank training, ReLoRA, and GaLore. They further perform full fine-tuning of the pre-trained model on GLUE to the model’s validate reasoning abilities.
weaknesses: - The paper claims that high intervals between reset/update steps in ReLoRA and GaLore are needed to avoid inconsistency in optimizer states, which otherwise wouldn't approximate full-rank training well. SwitchLoRA, on the other hand, uses a default highest switching frequency of 40, which then decays exponentially. GaLore reports that this frequency is close to optimal and does not cause issues for them. The core motivation presented in the paper is GaLore's inability to handle high switching frequencies, yet these frequencies are never actually tested in this paper, contradicting the stated motivation - no scenario is shown where SwitchLoRA uses such high switching frequencies.
- The source of the claimed improvements in the paper is unclear. The authors attribute it to resetting optimizer states and temporarily freezing parameters, but there’s no evidence supporting this. It could instead be due to (1) the exponential decay switching rule or (2) using a random subspace instead of SVD-based updates, which would substantially reduce the novelty of the work. Experiments are needed to confirm these claims.
- In Table 5, it appears the first two columns use a rank of 256. The GaLore paper reports perplexity values of 18.95 and 25.36 for these configurations, while this paper lists them as 19.58 and 25.93, despite claiming to use identical settings. This discrepancy is unclear, and if the original GaLore values are correct, GaLore outperforms SwitchLoRA. (The remaining configurations presented in this table were not covered in the GaLore paper.)

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Qu 1: Did you experiment with different frequencies of resetting in ReLoRA? (and for GaLore?)

Qu 2: Do you, or can you, provide a direct comparison of training times between the different approaches?

Qu 3: Could you say more about how the approach would scale in terms of memory/compute requirements for much larger models?
soundness: 3
strengths: The technique appears to offer significant gains over previous approaches. It achieves similar levels of accuracy to full-rank training with only 50-60% of the trainable parameters. The idea seems to intutively make sense.
summary: The paper explores the use of LoRA adapters during pre-training. The authors detail a new technique that is able to often replace vectors of the trainable parameters of LoRA adapters during training. This smooth and frequent adjustment of the trainable parameters provides a better approximation to full-rank training. Comparisons are made to full-rank training, ReLoRA and GaLore. 

When an existing vector in the LoRA adapter is replaced, W is adjusted by adding the difference between the old and new LoRA components. Vectors are chosen from a predefined set of options. 

The switching frequency using this technique is must higher than for the other approaches compared. A high-frequency of switching is employed early in training and this is reduced over time (exponentially decreasing in frequency).
weaknesses: It is currently a little unclear to me if the approach would scale or not to larger models. Could you detail the memory and compute implications of training larger models in more detail please. Can you extropolate from your current experiments to give us more confidence of the scalability of the approach?

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: Please refer to the Weakness.
soundness: 1
strengths: This work identifies the limited restarting frequency in ReLoRA, and proposes SwitchLoRA to address this problem by replacing the LoRA adapters frequently in the form of vectors of trainable parameters.
summary: This paper proposes SwitchLoRA, a parameter-efficient training method that periodically switches the vectors in low-rank terms with a set of vectors of trainable parameters. SwitchLoRA aims to achieve full fine-tuning model performance with much less memory footprint.
weaknesses: - This submission is wordy and incomplete. Here are some examples:
    - In line 472, there are two placeholders for experiment results, "`[insert performance difference]`". Please fill in all placeholders.
    - The oversimple flowchart in Fig. 5 does not add information beyond the Section 5 of future work. 
    - Inconsistent spelling like "pre-train" vs "pretrain". Please ensure consistency in terminology and spelling.
- Problematic experiment design.  
    - Section 4.2, in line 342, it is unclear to me how to use LoRA and SwitchLoRA to pretrain a language model from scratch. What is the value of the frozen weights, i.e., $W$ in Fig. 1(b)? Please consider providing a clear explanation of how LoRA and SwitchLoRA are initialized and used for pretraining from scratch.
   - Section 4.4 fine-tunes the resultant checkpoint from Section 4.2 for each baseline and SwitchLoRA, which means the baselines and SwitchLoRA have different start points when the fine-tuning starts. In other words, this is not a controlled experiment. 
- Questionable results. 
    - The standard deviation annoted in Tab 6 and 7 are questionable. For GLUE, such large std values like $23.13\pm 15$ and $72.70\pm 4$ in Tab 6 and $47.43\pm 3$ in Tab 7 cannot be seen in other works. Please verify the correctness of these values and explain any potential sources of such high variability.
- Insufficient experiments
    - In line 465, Table 6 presents the fine-tuning experiments on half of the subsets of the GLUE benchmark without MNLI, QNLI, QQP, and STSB. In most literature they are included. Please consider expanding the GLUE benchmark experiments to include the missing subsets.
    - Lack of large models. This work claims that SwitchLoRA improves the training efficiency of LLMs, but only small models (up to 1.3B) are included. Conduct experiments on larger models (e.g., models with tens of billions of parameters) to better support the claims about improving LLM training efficiency.

Meta-review: In this paper, the authors claimed that a higher update frequency for parameters of low-rank approaches can approximate full-rank behavior and thus achieve better performance for LLM training. Therefore, they proposed  SwitchLoRA, which frequently and smoothly alternates LoRA adapter parameters. 

The reviewers raised Major concerns, which were not addressed well in the authors' rebuttal: 1, the claims made in the paper that can gain performance improvement for LLM training are not well supported by explanations or experiments; 2, the experimental setup and results are not convincing; 3, the presentation needs to be significantly improved.

Therefore, based on its current shape, this paper is not ready for publication in ICLR.

Decision: Reject