Skip to main content
Abhik Sarkar's avatar

LoRA: Low-Rank Adaptation of Large Language Models

How LoRA adapts a frozen large language model by learning a low-rank update ΔW = (α/r)·BA to each weight matrix, training under 1% of the parameters of full fine-tuning and adding zero inference latency once merged.

TL;DR

  • LoRA freezes the pretrained weights W and learns a low-rank update \Delta W = \frac{\alpha}{r} B A, where B \in \mathbb{R}^{d\times r} and A \in \mathbb{R}^{r\times k} with r \ll d,k.
  • This trains well under 1% of the parameters of full fine-tuning, so adaptation fits on a single GPU and each task’s adapter is a few megabytes.
  • At inference, BA is merged back into W, adding zero latency versus the base model.
  • The premise: the weight change needed to adapt a large model has low intrinsic rank.

The low-rank update

Fine-tuning updates every weight in W. LoRA hypothesizes the useful part of that update is low-rank, and parameterizes it as the product of two small matrices:

W' = W + \Delta W = W + \frac{\alpha}{r} B A

Only A and B are trained; W stays frozen. The rank r controls the trade-off between capacity and cost.

Parameters, slashed

Because r is tiny, the trainable count r(d+k) is a rounding error next to the full d \times k.

Merge for free inference

Since \Delta W = BA has the same shape as W, it can be added into the weights once after training. The deployed model runs at exactly the base model’s cost — no extra layers, no added latency.

Implementation

The merge demo closes the loop from a low-rank update to an ordinary weight matrix. Now follow that same loop in a runnable PyTorch example: load pretrained GPT-2, train adapters on Tiny Shakespeare, save them, and merge them for generation.

Each row pairs an original-paper reference with our tensor explanation and the corresponding executable source. The two images are crops from arXiv v2. The Python excerpts come directly from the downloadable example. Tiny Shakespeare is our worked experiment; the paper evaluated other datasets and model sizes.

1. Load GPT-2 and freeze the base

Original paper

§4.1 · Frozen weights · PDF page 4

Our reading: the pretrained matrix stays fixed while the two update matrices learn.

Math and tensor shapes

W_0^{(t+1)}=W_0^{(t)}

The CLI loads openai-community/gpt2 and its tokenizer before installing adapters. Every original parameter is frozen, including embeddings, layer norms, MLPs, biases, and the LM head.

Autograd still follows the computation through frozen layers to reach the trainable adapters.

Python implementation

freeze_model
python
for parameter in model.parameters():
    parameter.requires_grad_(False)

2. Initialize the low-rank branch

Original paper

Figure 1 · PDF page 1

The original diagram separates the frozen path from the trainable A and B path. §4.1 specifies Gaussian A and zero B.

Math and tensor shapes

A\in\mathbb{R}^{r\times k},\quad B\in\mathbb{R}^{d\times r}
\Delta W=\frac{\alpha}{r}BA

Here k is the input width and d the output width. The defaults are r=8 and \alpha=16, giving a scale of 2.

At initialization, B=0, so the adapter contributes zero and preserves the base output. A uses Gaussian noise with standard deviation 0.01.

Python implementation

LoRAConv1D initialization
python
self.base = base
freeze_model(self.base)
self.rank = rank
self.alpha = float(alpha)
self.scale = self.alpha / rank
self.A = nn.Parameter(base.weight.new_empty(rank, in_features))
self.B = nn.Parameter(base.weight.new_zeros(out_features, rank))
nn.init.normal_(self.A, std=0.01)
self.merged = False

3. Place adapters in attention

Original paper

§4.2 · Attention matrices · PDF page 5

Our reading: attention projections are an adaptation target; the paper leaves MLP weights frozen. Our fused-QKV and output targets are an implementation choice.

Math and tensor shapes

In each of GPT-2’s 12 blocks, attn.c_attn projects 768 features to 2,304 fused query, key, and value features. attn.c_proj maps 768 features back to 768.

\begin{aligned}N_{\mathrm{adapter}}&=12r(3072+1536)\\&=442{,}368\quad(r=8)\end{aligned}

At rank 8, the example has 442,368 trainable adapter parameters. The original base has 124,439,808 parameters; with adapters, the total is 124,882,176. Each projection gets its own A and B.

The similarly named MLP c_proj is untouched: the loop visits block.attn only.

Python implementation

install_lora
python
freeze_model(model)
for block in model.transformer.h:
    for name in ("c_attn", "c_proj"):
        projection = getattr(block.attn, name)
        if not isinstance(projection, Conv1D):
            raise ValueError("expected GPT-2 attention Conv1D projections")
        setattr(block.attn, name, LoRAConv1D(projection, rank, alpha))

4. Follow both paths through the forward pass

Original paper

Equation 3 · §4.1 · PDF page 4

Equation 3 uses a column-vector input and prints h = W₀x + BAx. The paragraph below introduces α/r scaling.

Math and tensor shapes

z=xA^\top,\quad u=zB^\top

With tokens stored as rows, the shapes flow from [batch, tokens, k] through [batch, tokens, r] to [batch, tokens, d].

h=xW_0^\top+b_0+\frac{\alpha}{r}u

F.linear(x, A) computes xA^\top. GPT-2’s Conv1D stores the base matrix as [k, d], so self.base(inputs) already uses the required orientation and preserves its bias.

Python implementation

LoRAConv1D.forward
python
def forward(self, inputs):
    """Adds the low-rank branch until its weights have been merged."""
    result = self.base(inputs)
    if not self.merged:
        result = result + F.linear(F.linear(inputs, self.A), self.B) * self.scale
    return result

5. Train by predicting the next token

Original paper

§2 · Equation 2 · PDF page 3

Our reading: the language-model objective is optimized over a small set of adaptation parameters. Tiny Shakespeare and this training loop are our example.

Math and tensor shapes

\mathcal{L}=-\frac{1}{T-1}\sum_{t=1}^{T-1}\log p_\theta(x_{t+1}\mid x_{\leq t})

Pass the same unshifted input_ids as inputs and labels. Hugging Face GPT-2 shifts them internally once, so position t predicts token t + 1.

The optimizer receives only A and B. The script splits the corpus contiguously 90/10, samples training windows only from the first split, and evaluates fixed held-out blocks with dropout disabled.

Python implementation

language_model_loss
python
def language_model_loss(model, input_ids):
    """Lets GPT-2 shift labels internally: position t predicts token t + 1."""
    return model(input_ids=input_ids, labels=input_ids, use_cache=False).loss

6. Save adapters and reload the same base

Original paper

§4.2 · Storage and task switching · PDF page 5

Our reading: separate update matrices let tasks share a pretrained model. The adapter.pt and metadata.json format below belongs to our example.

Math and tensor shapes

\theta_{\mathrm{task}}=\{A_i,B_i\}_{i=1}^{24}

adapter.pt holds only A and B tensors. metadata.json records rank, alpha, target paths, architecture, the exact pretrained model revision, and training/data provenance.

Generation reloads that recorded base revision, installs matching adapters, validates the tensor keys and shapes, and loads the adapter weights. This saves an inference checkpoint, without optimizer state for training resume.

Python implementation

save_adapter
python
state = {f"{name}.{key}": getattr(module, key).detach().cpu()
         for name, module in modules.items() for key in ("A", "B")}
directory = Path(directory)
directory.mkdir(parents=True, exist_ok=True)
torch.save(state, directory / "adapter.pt")
(directory / "metadata.json").write_text(
    json.dumps(metadata, indent=2) + "\n", encoding="utf-8")

7. Merge once, then generate

Original paper

§4.1 · Merging for inference · PDF page 4

Our reading: the learned matrix product can be added to the pretrained matrix before deployment, removing the separate adapter computation.

Math and tensor shapes

W_{\mathrm{stored}}\leftarrow W_{\mathrm{stored}}+\frac{\alpha}{r}(BA)^\top

The transpose matches GPT-2’s [k, d] storage. The merged guard prevents adding the update twice, and the forward pass skips the adapter branch afterward.

The generation command compares logits before and after merging with dropout disabled, then uses greedy decoding. It reports the measured numerical difference.

Python implementation

LoRAConv1D.merge
python
@torch.no_grad()
def merge(self):
    """Merges once for inference, transposing BA into Conv1D storage."""
    if not self.merged:
        self.base.weight.add_((self.B @ self.A).T * self.scale)
        self.merged = True
        self.A.requires_grad_(False)
        self.B.requires_grad_(False)

Run the complete example

Download lora_gpt2.py and requirements-lora.txt into the same directory. Create a Python environment, then run these commands there. CUDA is the recommended path; use --device cpu in both commands to check the example on CPU.

Install, train, reload, and generate
bash
python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-lora.txt

python lora_gpt2.py train --device cuda --steps 100 --seq-len 128 --batch-size 2 --seed 42
python lora_gpt2.py generate --device cuda --adapter ~/.cache/abhik-lora-gpt2/adapter --merge

The first training run downloads pretrained GPT-2 and the Tiny Shakespeare corpus. Model, data, and adapter files default to ~/.cache/abhik-lora-gpt2. Run python lora_gpt2.py train --help for rank, alpha, learning rate, cache, and output options, or python lora_gpt2.py generate --help to choose a prompt and generation length.

These seven steps make the practical idea concrete: retain a shared base, learn small task-specific matrices, and fold their update into the base when it is time to serve that task.

Why it mattered

LoRA made adapting frontier models accessible. Instead of copying and retraining an entire model per task, you keep one frozen base and swap in small adapters. It became the backbone of the parameter-efficient fine-tuning ecosystem (PEFT, QLoRA), powering everything from instruction tuning to the thousands of community fine-tunes shared as lightweight adapter files.

  • Attention Is All You Need — LoRA most often adapts the attention projection matrices of this architecture
  • BERT — the pretrain-then-fine-tune paradigm LoRA makes dramatically cheaper
  • Switch Transformer — a different efficiency axis: sparse scaling of one big model vs cheap adaptation of a frozen one
  • FlashAttention — efficiency at the kernel level, complementary to LoRA’s efficiency at the fine-tuning level

Related Papers

If you found this paper review helpful, consider sharing it with others.

Mastodon