Fine-tuning neural networks with LoRA (Low-Rank Adaptation)

Low-rank adaptation (LoRA) is a popular method of fine-tuning neural networks that trains a relatively small number of parameters while achieving comparable performance to full fine-tuning of all parameters (Hu et al., 2021). Instead of fine-tuning the original weights from the base model, an additional set of weights (that are much lesser in number) are trained instead. This reduces computational cost, allowing for cheaper and faster fine-tuning.

Figure 1: Illustration of low-rank adaptation (LoRA) during training (left) and inference (right).


LoRA works by using two low-rank matrices, the down-projection matrix \( \mathbf{A} \in \mathbb{R}^{d \times r} \) and the up-projection matrix \( \mathbf{B} \in \mathbb{R}^{r \times d} \), to represent a original weight matrix from the base model \( \mathbf{W} \in \mathbb{R}^{d \times d} \) (Figure 1). \(\mathbf{W}\) represents the connections between two layers in a neural network . In LoRA, \(2dr\) parameters are fine-tuned, as opposed to \(d^2\) parameters in full fine-tuning. By using \(r << d\), the number of parameters involved in fine-tuning is drastically reduced.

During training (i.e. fine-tuning), on top of multiplying the input vector \(\mathbf{x}\) with \(\mathbf{W}\) , \(\mathbf{x}\) is multiplied with \(\mathbf{A}\) and \(\mathbf{B}\) before summing the resulting vectors to obtain the output vector, \( \mathbf{h} = \mathbf{Wx} + \mathbf{BAx} \). Gradients are computed and applied only to \(\mathbf{A}\) and \(\mathbf{B}\), with \(\mathbf{W}\) remaining frozen during fine-tuning. During inference, there is no additional latency as the number of parameters is kept the same by merging the low-rank matrices with the original weight matrix, \( \mathbf{W'} = \mathbf{W} + \mathbf{BA} \), and performing the forward pass with the merged matrix instead, \( \mathbf{h} = \mathbf{W'} \mathbf{x} \).

LoRA's approach of fine-tuning a separate set of weights yields other benefits as well. Firstly, we can switch dynamically between multiple LoRA models by performing the forward pass through all matrices (\( \mathbf{h} = \mathbf{Wx} + \mathbf{BAx} \)) and not merging the original and LoRA weight matrices. This allows for the replacement of only LoRA weight matrices when switching between models and avoids expensive I/O operations involving large matrices. Secondly, as LoRA works orthogonally from other other parameter-efficient fine-tuning (PEFT) techniques, it is able to be used together with them.

References

[YouTube] Edward Hu - What is Low-Rank Adaptation (LoRA)


Comments

Post a Comment

Popular posts from this blog

Training an LLM (In a Nutshell!)

A large language model (LLM) learns how to reply to a conversation by learning how to predict the next token (akin to a word), given the preceding conversation. This is framed as a multi-class classification problem, where each token represents a different class. The LLM outputs the likelihood of each token in the vocabulary (i.e. set of all possible tokens) of being the next token (which represent the probability parameters of a categorical distribution). Prediction of the next token occurs by sampling from the categorical distribution of all possible tokens. After a token is predicted, it is used along with tokens from the preceding conversation to predict the next token. Training an LLM typically consists of three main stages (different LLMs may have different training schemes): Unsupervised pre-training Supervised fine-tuning Reinforcement learning Stages 2 and 3 are commonly termed collectively as the fine-tuning stage. Unsupervised pre-training In the unsupervised pre-tra...
Read more

Multi-head Latent Attention (In A Nutshell!)

Image
In this post, I will dive into the Multi-head Latent Attention (MLA) mechanism, one of the innovations presented by the DeepSeek Team! This post assumes prior knowledge of the attention mechanism and the key-value cache. For a quick refresher on these topics, refer to my previous post on self-attention! One of the main problems with multi-head self-attention is the memory cost associated with the size of the key-value cache. MLA reduces the size of the key-value cache and speeds up LLM inference. The core idea is to cache latent embeddings that are  shared across all heads (and for both keys and values)  instead of different key and value embeddings for each head like in multi-head self-attention (Figure 1). The latent embeddings are multiplied with different key and value up-projection matrices for each head to produce different key and value embeddings unique to each head. Having unique key and value embeddings for each head maintains the expressivity of the attention mech...
Read more

Self-Attention and the Key-Value Cache (In A Nutshell!)

Image
The Transformer architecture underpins most modern large language models (LLMs). In the seminal paper  "Attention is All You Need" , Vaswani et al. propose a Transformer architecture that relies solely on the multi-head self-attention mechanism to learn global dependencies (i.e. relationships) between words in a sentence. In this post, I will first explain the multi-head self-attention mechanism that is used in LLMs such as the original ChatGPT model (which was derived from GPT-3.5), and go on to explain why a key-value cache is needed for efficient inference. For illustrative purposes, I use words to represent tokens. ChatGPT uses a decoder-only Transformer architecture, and is trained to predict the next token (e.g. sub-word, punctuation) given a context (i.e. message). I start by illustrating the multi-head self-attention mechanism using a single sentence as an example. Assume that we are training an LLM with a context window of 10 words; a 10-word sentence in the trai...
Read more