LoRA: Low-Rank Adaptation of Large Language Models

Helllo, lets now look into the famous LoRA, which is now one of the goto PEFT adaptation methods.

Fine tuning means training a pre-trained network on new data to improve its performance on a specific task. The problem of this full-finetuning task is:

1. It’s expensive to both train and store the weights.
2. Doing full parametric update on the new data which the distributions different from the one which is trained on causes catastrophic forgetting and loss of neuro plasticity.

LoRA is one of the parameter efficient adaptation methods.

The main idea of LoRA is, not all parameters in the original weight is relevant for the task and instead of doing adaptation on the full matrix, lets create other matrices which when they’re multiplied give us the full matrix. This will yield us a smaller matrix as B and A. This will significantly reduce our token space. Rank is the number of linearly independent vectors, which can’t be combined to produce another one.

For example: Let 𝑑 = 1000, 𝑘 = 5000, (𝑑 × 𝑘) = 5,000,000; using 𝑟 = 5, we get (𝑑 × 𝑟) + (𝑟 × 𝑘) = 5,000 + 25,000 = 30,000. Less than 1% of the original.

This also allow us We can easily switch between two different fine-tuned model (one for SQL generation and one for Javascript code generation) just by changing the parameters of the A and B matrices instead of reloading the W matrix again.

we scale the decomposed matrix using alpha / r(hyper parameters), where ‘r’ tell us the amount of compression to be done and ‘alpha’ is a scaling factor which tell us the amount of parameter change added to the original balancing between the knowledge of the original pre trained model and the new added knowledge. We can increase the value of ratio can be increased if we want to increase the weight of the fine tuning and reduce it when we want to decrease the weight.

When it comes to the two matrix, A which removes the linearly dependent weights and B gives us back the dimension, we initialize A from a gaussian distribution and we initialize B as 0 and we do back propagation to figure out the right values for them. This means, instead of training the full weights W we adapt the low ranked A and B.

The great thing about LoRA besides having less parameters, training time, RAM usage and weight merging is we can switch between different LoRA weights to switch tasks.

We can implement LoRA from scratch like this

1
class LoRA(nn.Module):
2
    def __init__(self, input_dim, output_dim, rank=1, alpha=1, device=device):
3
        super().__init__()
4

5
        self.A = nn.Parameter(
6
            torch.empty(rank, input_dim, device=device)
7
        )
8
        self.B = nn.Parameter(
9
            torch.zeros(output_dim, rank, device=device)
10
        )
11

12
        nn.init.normal_(self.A, mean=0, std=1)
13

14
        self.scale = alpha / rank
15
        self.enabled = True
16

17
    def forward(self, original_weight):
18
        if self.enabled:
19
            # h = w0 + BA *(alpha/rank)
20
            return original_weight + (self.B @ self.A) * self.scale
21
        return original_weight

References:

Original Paper: Hu, Edward J., et al. “Lora: Low-rank adaptation of large language models.” arXiv preprint arXiv.09685 (2021).

Code implementation: https://github.com/yeabwang/LoRA