LoRA fine tuning fundamentals

A pretrained language model already contains useful transformations. Adapting it to a task means changing those transformations so that the desired responses become more likely. LoRA lets us learn those changes through small pairs of matrices while keeping the pretrained weights frozen.

The central idea is to limit the rank of the weight update. The original weight matrix can retain its full rank, and the complete model still uses its original layers. The restriction applies to the correction we learn for each selected matrix.

We can understand that restriction with a two-dimensional example, then use exactly the same construction inside a transformer. Along the way, we will implement a small LoRA layer in PyTorch and connect its dimensions to the rank-8 adapters in our JSON fine-tuning experiment.

Fine tuning learns a change in weights

Start with a linear projection:

The input $h$ has $d_{\mathrm{in}}$ values, and the output has $d_{\mathrm{out}}$ values. Each row of $W_0$ reads a pattern from the input through a dot product. In a transformer, query, key, value, and feed-forward projections all perform operations of this kind.

Full fine-tuning can update the original matrix. We can describe the result as $W_0+\Delta W$, where $\Delta W$ is the total change accumulated over training. A full update has a separate adjustable entry for every weight.

LoRA keeps $W_0$ fixed and represents the correction as

Here, $s$ is a scale factor and $r$ controls the width of the adapter. Training adjusts $A$ and $B$. To see what this buys us, we first need a picture of rank.

Rank counts independent output directions

A matrix transforms an input by taking a linear combination of its columns. The column space is the set of all outputs that transformation can produce. Its dimension is the matrix’s rank.

Consider this correction:

Its second column is twice its first. Both columns point along the direction $(1,2)$, so every possible output lies on the same line. It has four nonzero entries but only one independent output direction: its rank is 1.

For an input $h=(h_1,h_2)$,

The scalar changes with the input, while the available output direction stays the same. Rank 0 would produce only the zero vector. A rank-2 matrix in this two-dimensional example can produce outputs across the plane. The matrix rank article develops this geometry and its connection to singular values in more detail.

Build that correction from two matrices

The expression above already gives us a factorization:

$A$ first reads a scalar from the input: $Ah=h_1+2h_2$. Then $B$ turns that scalar into an output correction along $(1,2)$. With $h=(3,1)$, this becomes

This is a rank-1 outer product. For a wider adapter, each row $a_i^\top$ of $A$ reads one scalar coefficient, and each column $b_i$ of $B$ supplies an output direction:

The correction always lies in the span of $B$‘s columns, so

The bound can be strict: columns can be dependent, a factor can contain zeros, or training can use fewer directions than the adapter allows. Rank 8 means the correction has rank at most 8. The rows of $A$ and columns of $B$ are learned directions; their coordinates can involve many input and output features.

Read with a key, write with a value

Each term $b_i(a_i^\top h)$ has two roles: the row $a_i^\top$ reads a scalar from the input, and the column $b_i$ supplies the output direction scaled by that scalar. We can think of them as a key and value. A negative score reverses the contribution; there is no softmax or expert selection in ordinary LoRA.

The standalone article Feature detectors: how neurons read and write develops the dot products, scaled output vectors, and outer products behind this mechanism with an interactive example. These mathematical roles alone do not guarantee that a key corresponds to a named topic.

A rank 1 update can preserve a rank 2 layer

Set the frozen base to the identity matrix, $W_0=I$, and use scale $s=1$. The adapted matrix is

Its determinant is $2\cdot5-2\cdot2=6$, so it has rank 2. The correction maps inputs onto a line, while the complete layer still maps across a plane. For our selected input, the original output $(3,1)$ and correction $(5,10)$ add to $(8,11)$.

Try No update, Rank 1, and Rank 2 below. Move the input sliders to watch $Ah$ and the selected output change. These are hand-chosen matrices, so we can inspect their ranks exactly.

In Rank 2, we add a second independent outer product:

The correction can now span the plane too. In all three presets, the complete layer retains rank 2. This is the distinction we need for LoRA: restricting each weight correction does not force the full pretrained matrix to have that rank. Across many adapted layers and nonlinear operations, the model’s behavior can change substantially.

The parameter savings appear at larger dimensions

A dense update to a projection needs $d_{\mathrm{out}}d_{\mathrm{in}}$ trainable values. The two LoRA factors need

For a Qwen query projection with 896 inputs and 896 outputs:

That is 56 times fewer trainable values for this projection. The original 802,816 matrix weights remain present as frozen weights.

The two-dimensional example is deliberately small: a rank-1 factorization uses $1\cdot(2+2)=4$ parameters, the same count as a dense $2\times2$ matrix. It teaches the rank constraint; the large projection shows the storage advantage when $r$ is small compared with the matrix dimensions.

The original LoRA paper motivates low-rank adaptation with the hypothesis that useful task-specific weight changes can occupy a much smaller space than the pretrained weights. How restrictive that choice is depends on the task and selected layers. A higher rank allows a richer correction, but better validation performance still has to be measured.

Connect LoRA to singular value decomposition

Any rank-$r$ matrix has an SVD representation $U_r\Sigma_rV_r^\top$. One valid factorization is $B=U_r\Sigma_r$ and $A=V_r^\top$. This gives another way to see why two narrow matrices can represent a rank-$r$ correction.

Default LoRA training learns the factors directly from the task loss. It does not first train a dense update and then truncate its SVD, and it does not replace the original pretrained matrix with a low-rank approximation. Those would be different procedures.

The factors are also not unique. For an invertible $r\times r$ matrix $C$, $(BC)(C^{-1}A)=BA$. Changing the intermediate coordinate system can preserve the same correction. This is one reason to interpret the combined update and its effect on inputs, rather than assigning a fixed semantic meaning to one adapter coordinate.

Add the adapter alongside a frozen projection

The complete forward pass is

Both branches receive the same input. The base branch preserves the pretrained transformation; the adapter branch supplies the learned correction. There is no activation between $A$ and $B$ in this construction, so their product is a linear update to the selected projection.

In standard LoRA, $s=\alpha/r$. Our transformer example uses rank 8 and alpha 16, giving $s=2$. Alpha controls this explicit scale; it is separate from the optimizer’s learning rate. Other variants can use a different scaling rule, so the formula here describes standard LoRA in PEFT.

Start with the pretrained behavior

PEFT’s default linear-layer initialization gives $A$ random values and sets $B$ to zero. The adapter initially contributes zero, so the layer starts with its pretrained output. During the first backward pass, $B$ can receive a gradient through the random $A$. The task-loss gradient for $A$ is initially zero because it passes through $B$. Once $B$ changes, gradients can train both factors. Initializing both factors to zero would block those task-loss gradients. The PEFT initialization guide describes the default and alternative choices.

Frozen weights still take part in the computation. Backpropagation must carry the learning signal through the operations needed to reach earlier adapters. Freezing a weight removes its parameter gradient and optimizer updates; it does not remove the layer from the model.

Implement and train a small adapter in PyTorch

The example below learns our rank-1 correction on CPU, with no pretrained model download. We generate input vectors and ask the layer to match $(I+\Delta W)h$. Mean squared error fits this small regression problem. The purpose is to observe the factors learning while the original identity matrix stays unchanged.

nn.Linear stores weights with shape (out_features, in_features). For a batch of row vectors, PyTorch computes x @ weight.T; the nested F.linear calls below implement the column-vector expression $B(Ah)$ with that convention.

Use Expand code for the full-screen view. The code shown here is the same file as examples/lora_linear.py beside this article.

"""Learn a rank-1 correction to a frozen linear layer on CPU. No LLM download."""

import copy
import json
import math

import torch
from torch import nn
from torch.nn import functional as F


class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, rank: int, alpha: float):
        super().__init__()
        self.base = base.requires_grad_(False)
        self.scale = alpha / rank
        self.A = nn.Parameter(torch.empty(rank, base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, rank))
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))

    def forward(self, x):
        # F.linear(x, W) computes x @ W.T for a batch of row vectors.
        correction = F.linear(F.linear(x, self.A), self.B)
        return self.base(x) + self.scale * correction

    @torch.no_grad()
    def merged_copy(self):
        merged = copy.deepcopy(self.base)
        merged.weight.add_(self.scale * (self.B @ self.A))
        return merged


def main():
    torch.manual_seed(42)
    torch.set_num_threads(1)
    base = nn.Linear(2, 2, bias=False)
    with torch.no_grad():
        base.weight.copy_(torch.eye(2))
    layer = LoRALinear(base, rank=1, alpha=1.0)
    original_weight = base.weight.detach().clone()

    x = torch.randn(256, 2)
    desired_delta = torch.tensor([[1.0, 2.0], [2.0, 4.0]])
    targets = F.linear(x, original_weight + desired_delta)
    torch.testing.assert_close(layer(x), base(x))  # B starts at zero.
    parameters = [p for p in layer.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(parameters, lr=0.05, weight_decay=0.0)
    initial_loss = F.mse_loss(layer(x), targets).item()

    for step in range(500):
        optimizer.zero_grad(set_to_none=True)
        loss = F.mse_loss(layer(x), targets)
        loss.backward()
        if step == 0:
            first_A_gradient_norm = layer.A.grad.norm().item()
            first_B_gradient_norm = layer.B.grad.norm().item()
            assert first_A_gradient_norm == 0.0
            assert first_B_gradient_norm > 0.0
        optimizer.step()

    with torch.no_grad():
        delta = layer.scale * (layer.B @ layer.A)
        adapted = original_weight + delta
        final_loss = F.mse_loss(layer(x), targets).item()
        merged = layer.merged_copy()
        torch.testing.assert_close(delta, desired_delta, atol=1e-4, rtol=1e-4)
        torch.testing.assert_close(merged(x), layer(x), atol=1e-5, rtol=1e-5)
        assert torch.equal(base.weight, original_weight)
        assert base.weight.grad is None
        assert torch.linalg.matrix_rank(delta, atol=1e-5).item() == 1
        assert torch.linalg.matrix_rank(adapted, atol=1e-5).item() == 2
        print(json.dumps({
            "device": str(x.device), "steps": 500,
            "initial_mse": initial_loss, "final_mse": final_loss,
            "first_A_gradient_norm": first_A_gradient_norm,
            "first_B_gradient_norm": first_B_gradient_norm,
            "base_unchanged": torch.equal(base.weight, original_weight),
            "update_rank": 1, "adapted_weight_rank": 2,
            "learned_delta": delta.tolist(),
        }, indent=2))


if __name__ == "__main__":
    main()

From this article’s directory, use uv to prepare a CPU-only PyTorch environment and run the example:

uv run --no-project --python 3.12 \
  --with torch==2.7.1 \
  --index https://download.pytorch.org/whl/cpu \
  examples/lora_linear.py

The script checks that the initial adapter preserves the original output, the first gradient for $A$ is zero, the first gradient for $B$ is nonzero, and the frozen base never changes. After training, it checks that the learned correction matches our chosen matrix, has rank 1, and leaves the complete weight matrix with rank 2.

In our seed-42 run, mean squared error fell from about 10.94 to $5.9\times10^{-14}$. The learned correction matched the matrix above to FP32 precision. The useful checks are the unchanged base, the rank of the correction, and agreement between the separate and merged computations.

This toy uses FP32 CPU tensors and has no dropout or quantization. In the LLM experiment, we use PEFT to attach adapters to the actual projections and manage the model integration.

The loss trains the task through the adapters

LoRA specifies how we parameterize weight changes. The task still supplies the training objective. In supervised language-model fine-tuning, that is usually next-token cross entropy:

For our JSON task, the target is the correct assistant response. Prompt and padding positions are ignored, and loss is averaged over response tokens. Backpropagation computes gradients for the adapter factors, and AdamW updates those trainable parameters. There is no separate loss asking a matrix to become low rank: its factorized shape already enforces the bound.

The CPU example uses mean squared error because its targets are numeric vectors. The JSON experiment uses cross entropy because its targets are token IDs. The adapter construction is compatible with both objectives.

Put the same construction inside transformer blocks

Our pinned Qwen2.5-0.5B-Instruct checkpoint has 24 transformer blocks and hidden vectors with 896 values. Each block contains attention projections and a feed-forward network. We adapt the query and value projections inside every block:

The value projection is narrower because this model uses grouped-query attention: 14 query heads share two key/value heads, with 64 values per head. The query output is $14\cdot64=896$ values; the value output is $2\cdot64=128$.

Across the model, this gives $24\cdot(14{,}336+8{,}192)=540{,}672$ trainable parameters. There are 48 adapted projections, each with its own pair of matrices. Their corrections affect computation inside the transformer blocks. The original attention weights, feed-forward weights, embeddings, and normalization weights remain frozen in this recipe.

The corresponding PEFT configuration is short:

LoraConfig(
    task_type="CAUSAL_LM",
    r=8,
    lora_alpha=16,
    lora_dropout=0.0,
    target_modules=["q_proj", "v_proj"],
    bias="none",
)

Choosing query and value projections is our experiment’s setting. LoRA can target other linear projections, including feed-forward layers. Rank controls the size of each selected correction; the target-module list controls how many corrections we add and where they act.

Save the correction or merge it for inference

With dropout disabled, the two branches can be combined algebraically:

The CPU script creates a merged copy and checks that it gives the same outputs up to floating-point tolerance. It leaves the original frozen base intact. In a supported unquantized PEFT model, merge_and_unload() combines the adapter with the base for inference, removing the separate adapter branch.

Saving an unmerged adapter keeps the correction small and lets us load it alongside the compatible pretrained checkpoint. The adapter does not contain the full base model; its configuration and the correct base revision are part of what we need to reproduce the adapted model. See the PEFT checkpoint format.

What the smaller parameter count buys us

Fewer trainable weights mean fewer parameter gradients and optimizer-state entries, and a smaller task-specific checkpoint. The base weights and the activations needed for training still consume memory. Our 540,672 FP32 adapter weights alone occupy about 2.06 MiB, while the measured Qwen microbatch-4 sweep used about 3.82 GiB of peak allocated GPU memory for the complete workload.

Rank and batch size control different resources. Rank changes the adapter’s shape and representational capacity. Microbatch size changes how many examples we process together and how much activation memory we need. Increasing rank does not automatically improve held-out accuracy, and reducing trainable parameters does not imply the same reduction in training time.

For this method, the useful picture is a pretrained transformation plus a correction built from a few input measurements and output directions. We preserve the original matrix, learn the factors, and evaluate whether that constrained change solves the task. The JSON fine-tuning experiment puts those pieces together with real requests, token loss, saved adapters, and GPU measurements.