Low-Rank Adaptation¶

An important paradigm for large language models is the idea of pre-training on large-scale general domain data followed by fine tuning on task-specific data. However, fine tuning retrains all model parameters which can become costly. Low-Rank Adaptation freezes the pre-trained model weights and injects trainable rank decomposition matrices into each layer, which drastically reduces the memory requirements.

Consider a language modeling problem where the goal is to maximize the conditional probabilities given a task-specific prompt. Given a pre-trained autoregressive model $P_{\Phi }\left(y|x\right)$, each downstream task is represented by a training dataset of context-target pairs of token sequences $Z=\left\{\right(x_1,y_1),...,(x_n,y_n\left)\right\}$.

Traditionally the model is initialized to pre-trained weights ${\Phi }_0$ and updated to ${\Phi }_0+\Delta \Phi$ with gradient descent on a conditional language modeling objective.

$$\underset{\Phi }{max}\sum _{(x,y)\in Z}\sum _{t=1}^{|y|}log\left(P_{\Phi }\right(y_t|x,y_{<t}\left)\right)$$

Training on each downstream tasks results in a different set of parameters $\Delta \Phi$ with the same dimensions as ${\Phi }_0$. Low-Rank Adaptation encodes these task-specific models with a much smaller sized set of parameters $\Theta$.

$$\underset{\Theta }{max}\sum _{(x,y)\in Z}\sum _{t=1}^{|y|}log\left(p_{{\Phi }_0+\Delta \Phi \left(\Theta \right)}\right(y_t|x,y_{<t}\left)\right)$$

For a pre-trained weight matrix $W_0\in R^{d\times k}$, the weight update is constrained through a low-rank matrix decomposition $W_0+\Delta W=W_0+BA$ where $B\in R^{d\times r}$, $A\in R^{r\times k}$ and the rank $r\ll min(d,k)$. $A$ is randomly initialized with a Gaussian and B is initialized with zeros such that $\Delta W=BA=0$ at the beginning of training. The modified forward pass yields:

$$h=W_{0}x+\Delta Wx=W_{0}x+BAx$$