LoRA: train a small update to a large model

Low-rank adaptation (LoRA) fine-tunes a model by freezing its original weights and learning small additional matrices. The base model remains available; the adapter records the change for a task or style.

Consider one linear layer with a weight matrix W. Full fine-tuning can change every entry of W. LoRA represents an update as the product of two matrices, A and B, with a narrow intermediate dimension called the rank. A smaller rank gives fewer trainable parameters and less room for the update to vary.

The forward pass

With row-vector inputs, W has shape input × output, A is input × rank, and B is rank × output. This function accepts those matrices and combines the frozen layer's output with its learned update. Training the adapter is a separate step.

import numpy as np

def lora_linear(x, weight, a, b, alpha):
    rank = a.shape[1]
    return x @ weight + (x @ a) @ b * (alpha / rank)

For a square layer of width 4,096 and rank 16, the adapter has 131,072 parameters instead of the base matrix's 16.8 million. During training, optimizer state and gradients are needed for the adapter parameters rather than the frozen base weights. Activations and the base model still occupy memory.

Keep the adapter, or merge it

The two paths can be evaluated separately, or the scaled product A @ B can be added to W before inference. Merging removes the extra adapter operations for that selected update. Keeping adapters separate makes switching between them convenient.

LoRA reduces the cost of adaptation. Quantization reduces numerical storage precision; QLoRA combines a quantized frozen base with trainable adapters. These address different parts of the memory bill.