LoRA: Low-Rank Adaptation

Chapter 1: The Low-Rank Idea

During fine-tuning, each weight matrix W changes by some amount ΔW. Full fine-tuning learns ΔW directly — a dense matrix with the same dimensions as W. For GPT-3's attention layers, W is 12288 × 12288, so ΔW has 150 million entries. Per layer. Times 96 layers.

But Aghajanyan et al. (2021) showed something remarkable: the intrinsic dimensionality of fine-tuning is very low. When you project the weight updates into a random low-dimensional subspace, you can match 90% of full fine-tuning performance with only a few thousand trainable parameters. This means the actual "useful" information in ΔW is far lower-dimensional than its matrix size suggests.

Here's the key equation. Instead of:

LoRA computes:

Where W is frozen (no gradients), B is d × r (initialized to zeros), and A is r × d (initialized from a Gaussian). The product B·A has rank at most r, so the update is inherently low-rank.

python
import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, original_linear, r=8, alpha=16):
        super().__init__()
        d_in = original_linear.in_features    # e.g., 12288
        d_out = original_linear.out_features   # e.g., 12288

        # Freeze the original weight
        self.W = original_linear
        self.W.weight.requires_grad = False

        # LoRA matrices — ONLY these are trained
        self.A = nn.Parameter(torch.randn(r, d_in) * 0.01)  # [r, d_in]
        self.B = nn.Parameter(torch.zeros(d_out, r))            # [d_out, r]
        self.scaling = alpha / r  # scaling factor

    def forward(self, x):
        # Original: W·x, shape [batch, seq, d_out]
        base = self.W(x)
        # LoRA: B·A·x, same shape
        lora = (x @ self.A.T) @ self.B.T * self.scaling
        return base + lora  # h = Wx + BAx

The beauty of this design:

1. At training time: Only A and B receive gradients. W is frozen. Memory for optimizer states drops from O(d²) to O(2·d·r).

2. At inference time: You can merge: W' = W + B·A. The merged weight has exactly the same shape as the original, so there's zero additional latency. LoRA literally disappears at inference.

3. For multiple tasks: Store one base model + a tiny LoRA file per task. Swap LoRA modules to switch tasks. Each LoRA file is ~1-10 MB instead of 350 GB.

Chapter 2: Math of LoRA

Let's derive LoRA's parameter savings precisely and understand the scaling factor.

Parameter count

For a single weight matrix W ∈ R^d×d:

For GPT-3's attention weights (d = 12288, r = 8):

python
d = 12288     # hidden dim
r = 8         # LoRA rank

full_params = d * d                    # 150,994,944
lora_params = 2 * d * r               # 196,608
ratio = full_params / lora_params     # 768x compression

# Applied to all attention Wq and Wv across 96 layers:
total_lora = 2 * lora_params * 96     # 37.7M params
total_full = 175e9                   # 175B params
print(f"LoRA trains {total_lora/total_full*100:.4f}% of params")
# LoRA trains 0.0215% of parameters → ~4650x reduction

The scaling factor α

LoRA introduces a hyperparameter α (alpha) that controls the magnitude of the update:

Why α/r and not just α?

When you change r, the magnitude of B·A changes (more columns in A means a larger product). Dividing by r normalizes this: the update magnitude stays roughly constant regardless of r. This means you can tune α once and then change r freely without retuning the learning rate.

Initialization

LoRA initializes B to zeros and A from a Gaussian. This means at initialization, ΔW = B·A = 0 — the model starts exactly as the pre-trained model. Training only moves the weights away from the pre-trained solution as needed.

This is crucial for stability: the model doesn't "forget" its pre-trained knowledge at the start of fine-tuning. It only learns the minimum necessary adaptation.

Gradient computation

Only A and B receive gradients. The memory savings are substantial:

python
# Memory comparison for Adam optimizer (2 states per param)
d, r = 12288, 8
n_layers = 96
n_matrices = 2  # Wq and Wv

# Full fine-tuning: Adam stores m and v for each param
full_adam = 175e9 * 2 * 4   # 2 states × 4 bytes = 1.4 TB

# LoRA: Adam only for A and B
lora_adam = n_layers * n_matrices * 2 * d * r * 2 * 4
print(f"Full Adam: {full_adam/1e9:.0f} GB")   # 1400 GB
print(f"LoRA Adam: {lora_adam/1e6:.0f} MB")   # 302 MB

Chapter 3: Where to Apply

A Transformer has four types of weight matrices per attention layer: W_q, W_k, W_v, W_o (query, key, value, output projections). Plus the FFN's up and down projections. Where should you put LoRA?

Hu et al. ran ablations on GPT-3 175B:

Why not FFN layers?

The original LoRA paper focused on attention layers and didn't apply LoRA to FFN. Later work (QLoRA, LoRA+) found that including FFN gives small additional gains, especially for knowledge-intensive tasks. The FFN stores factual knowledge, so adapting it helps when the target task requires domain-specific facts not in the pre-training data.

One rank for all layers?

Hu et al. used the same rank r for all layers, but subsequent work showed that different layers may benefit from different ranks. Lower layers (closer to input) tend to capture general linguistic features and need less adaptation. Upper layers (closer to output) are more task-specific and benefit from higher ranks.

python
# Practical LoRA setup for a Transformer
def add_lora(model, r=8, alpha=16, target_modules=['q_proj', 'v_proj']):
    """Add LoRA to specified attention projections."""
    for name, module in model.named_modules():
        if any(t in name for t in target_modules):
            if isinstance(module, nn.Linear):
                lora_module = LoRALinear(module, r=r, alpha=alpha)
                # Replace the original module
                parent = get_parent(model, name)
                setattr(parent, name.split('.')[-1], lora_module)

    # Freeze everything except LoRA parameters
    for n, p in model.named_parameters():
        p.requires_grad = 'lora' in n  # only LoRA params train

Chapter 4: Rank Selection

The rank r is LoRA's most important hyperparameter. It controls the expressiveness of the adaptation — higher rank means more capacity to learn, but also more parameters and compute.

Hu et al. found that surprisingly low ranks work well:

Why does low rank work?

The weight update ΔW during fine-tuning is approximately low-rank because:

1. Task similarity: Most downstream tasks are related to the pre-training distribution (both involve language). The adaptation is a small perturbation, not a complete rewrite. Small perturbations tend to be low-rank.

2. Intrinsic dimensionality: Aghajanyan et al. showed that the intrinsic dimensionality of fine-tuning for GPT-3 is only ~200-1000 across all layers combined. Distributed across 192 matrices (2 per layer × 96 layers), that's only ~1-5 effective dimensions per matrix — consistent with r ≈ 4-8.

3. Redundancy in attention: Many attention heads learn similar patterns. The adaptation can often be expressed as changing a few "modes" of attention behavior, not independently modifying every head.

Rank selection guidelines

Chapter 5: Results & Comparisons

How does LoRA compare to other parameter-efficient methods and full fine-tuning? Hu et al. ran comprehensive comparisons on GPT-3 175B.

Training efficiency

Beyond parameter count, LoRA saves substantial compute during training:

Note: GPU memory savings are "only" 3.4x (not 37,000x) because you still need to load the full model weights for the forward pass. The savings come from not storing optimizer states and gradients for the frozen weights.

RoBERTa results

LoRA also works well on encoder models. On RoBERTa-Base (125M params) with the GLUE benchmark:

On RoBERTa, LoRA actually exceeds full fine-tuning on some tasks (SST-2: 95.1 vs 94.8). This suggests that the low-rank constraint acts as a beneficial regularizer, preventing overfitting on small datasets.

Chapter 6: LoRA Explorer

Time to see LoRA in action. This explorer lets you configure LoRA parameters and watch how they affect the weight update decomposition, parameter efficiency, and task performance.

Chapter 7: Connections

LoRA sits at the intersection of parameter efficiency, matrix theory, and practical ML engineering. Its connections reveal both its intellectual roots and its enormous impact.

What came before

What came after

The bigger picture

"Simplicity is the ultimate sophistication." — Leonardo da Vinci (and LoRA is as simple as W + BA)