Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

LoRA Mathematics › Code Implementation

Michael BrenndoerferDecember 1, 20259 min read

Language AI HandbookAll chapters

Part 35 of 67Parameter-Efficient Fine-tuning

Chapter 256LoRA Mathematics › Code Implementation

Reading section 8 of 12 · Code Implementation · 9 minSection 8 of 12 · 9 min

Section 8 of 12 · 9 min read

LoRA Mathematics › Code ImplementationOne section · Sections in this chapter8 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We will implement the LoRA mathematics in PyTorch. This shows the forward pass, gradient computation, and initialization. The implementation verifies our theoretical predictions empirically, connecting the mathematical derivations to runnable code that you can experiment with and extend.

In[7]:

Code

import torch

torch.manual_seed(42)

In[8]:

Code

import math

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


class LoRALinear(nn.Module):
    """Linear layer with LoRA adaptation"""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 4,
        alpha: float = 1.0,
        pretrained_weight: torch.Tensor = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.scaling = alpha / rank

        # Frozen pre-trained weight
        if pretrained_weight is None:
            pretrained_weight = torch.randn(out_features, in_features) * 0.02
        self.register_buffer("W0", pretrained_weight)

        # LoRA matrices (trainable)
        self.A = nn.Parameter(torch.randn(rank, in_features) / math.sqrt(rank))
        self.B = nn.Parameter(torch.zeros(out_features, rank))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass with LoRA adaptation."""
        # Original forward pass
        h0 = F.linear(x, self.W0)

        # LoRA path: x -> A -> B (with scaling)
        z = F.linear(x, self.A)  # Shape: (batch, rank)
        delta_h = F.linear(z, self.B)  # Shape: (batch, out_features)

        return h0 + self.scaling * delta_h

    def get_merged_weight(self) -> "torch.Tensor":
        """Return W0 + (alpha/r) * BA for inference."""
        return self.W0 + self.scaling * (self.B @ self.A)

Let's verify the initialization produces zero updates:

In[9]:

Code

# Create a LoRA layer
layer = LoRALinear(in_features=64, out_features=128, rank=8, alpha=8.0)

# Check that BA = 0 at initialization
BA = layer.B @ layer.A
shape_a = layer.A.shape
shape_b = layer.B.shape
shape_ba = BA.shape
max_ba = BA.abs().max().item()

Out[10]:

Console

Shape of A: torch.Size([8, 64])
Shape of B: torch.Size([128, 8])
Shape of BA: torch.Size([128, 64])
Max absolute value in BA: 0.000e+00

The product $BA$ is exactly zero because $B$ is initialized to zero, confirming our model starts with pretrained behavior. Matrix $A$ has dimensions (rank, in_features) for the down-projection, while $B$ has dimensions (out_features, rank) for the up-projection. Their product $BA$ maintains the full weight matrix dimensions but contains only zeros at initialization, meaning the LoRA path contributes nothing to the output until training begins.

Verifying Gradient FlowLink Copied

Let's trace gradients through a forward-backward pass to confirm the asymmetric gradient behavior we derived mathematically:

In[11]:

Code

# Create layer and input
layer = LoRALinear(in_features=8, out_features=4, rank=2, alpha=2.0)
x = torch.randn(1, 8, requires_grad=True)

# Forward pass
h = layer(x)
input_shape = x.shape
output_shape = h.shape

# Backward pass with a simple loss
loss = h.sum()
loss.backward()

grad_a = layer.A.grad
grad_b = layer.B.grad

Out[12]:

Console

Input shape: torch.Size([1, 8])
Output shape: torch.Size([1, 4])
Gradient norm of A: 0.000e+00
Gradient norm of B: 7.864e+00

As predicted by our mathematical analysis, matrix $A$ 's gradient is zero at initialization because $B = 0$ in the gradient formula $\frac{\partial \mathcal{L}}{\partial A} = \frac{\alpha}{r} B^T g x^T$ , while $B$ receives non-zero gradients because $\frac{\partial \mathcal{L}}{\partial B} = \frac{\alpha}{r} g (Ax)^T$ depends on $A$ (which is non-zero) rather than $B$ . This asymmetry demonstrates the coupled dynamics derived mathematically: matrix $A$ begins learning only after $B$ has moved away from zero, creating a brief initial phase where $B$ adapts alone before both matrices co-evolve.

Parameter Efficiency CalculationLink Copied

Let's compute the parameter savings for a realistic scenario corresponding to a real large language model:

In[13]:

Code

def calculate_lora_efficiency(d: int, k: int, r: int, num_layers: int) -> dict:
    """Calculate parameter counts and savings."""
    full_params = d * k * num_layers
    lora_params = r * (d + k) * num_layers
    ratio = lora_params / full_params

    return {
        "full_params": full_params,
        "lora_params": lora_params,
        "ratio_percent": ratio * 100,
        "compression": 1 / ratio,
    }


# LLaMA-7B attention projections: d=k=4096, 32 layers, 4 matrices per layer
results = calculate_lora_efficiency(d=4096, k=4096, r=16, num_layers=32 * 4)

Out[14]:

Console

LLaMA-7B Attention LoRA Analysis (r=16):
  Full fine-tuning params: 2,147,483,648
  LoRA params: 16,777,216
  Percentage of full: 0.781%
  Compression ratio: 128x

The 128x compression achieves dramatic parameter reduction while maintaining strong adaptation capability. Full finetuning of these attention layers requires 2.1 billion parameters, while LoRA needs only 16.8 million trainable parameters, making finetuning practical on consumer GPUs with 16-24GB of memory. This dramatic reduction enables efficient multi-task serving where different LoRA adapters can be swapped without duplicating the base model weights.

Visualizing Rank EffectsLink Copied

Let's visualize how different ranks affect the approximation capacity, connecting the theoretical analysis of SVD truncation to the practical question of rank selection:

In[15]:

Code

# Create a target matrix that simulates a fine-tuning update
# (low intrinsic rank with some noise)
d, k = 64, 64
true_rank = 8
U_true = torch.randn(d, true_rank)
V_true = torch.randn(true_rank, k)
target_delta_W = U_true @ V_true + 0.1 * torch.randn(d, k)

# Try different LoRA ranks
ranks = [1, 2, 4, 8, 16, 32, 64]
errors = []

for r in ranks:
    # Best rank-r approximation via SVD
    U, S, Vh = torch.linalg.svd(target_delta_W, full_matrices=False)
    approx = (U[:, :r] * S[:r]) @ Vh[:r, :]
    error = torch.norm(target_delta_W - approx).item()
    errors.append(error)

# Store true_rank for use in visualization
true_rank_value = true_rank

Out[16]:

Visualization

Line plot showing reconstruction error decreasing as LoRA rank increases from 1 to 64. — Reconstruction error for a 64x64 fine-tuning update matrix with intrinsic rank 8, computed via SVD truncation. Error decreases sharply until rank 8, then levels off at higher ranks. This pattern demonstrates the matrix's true intrinsic rank and explains why LoRA with r=8 to r=16 captures nearly all useful information while achieving strong parameter compression for typical fine-tuning tasks.

The error drops sharply until we reach the true intrinsic rank of the target matrix (rank 8 in this example), then decreases much more slowly. This illustrates the "knee" in the error curve that practitioners look for when choosing rank: the point at which increasing $r$ provides diminishing returns in approximation quality. For this synthetic matrix, $r = 8$ captures nearly all the useful structure, while $r = 32$ or $r = 64$ adds only marginal improvement at much higher parameter cost.

Gradient Magnitude AnalysisLink Copied

Let's examine how gradients evolve during the first few training steps, confirming the bootstrap dynamics we analyzed theoretically:

In[17]:

Code

# Track gradient magnitudes over steps
layer = LoRALinear(in_features=32, out_features=32, rank=4, alpha=4.0)
optimizer = torch.optim.Adam([layer.A, layer.B], lr=0.01)

grad_A_norms = []
grad_B_norms = []
B_norms = []

for step in range(20):
    optimizer.zero_grad()

    x = torch.randn(16, 32)
    target = torch.randn(16, 32)

    output = layer(x)
    loss = F.mse_loss(output, target)
    loss.backward()

    grad_A_norms.append(layer.A.grad.norm().item())
    grad_B_norms.append(layer.B.grad.norm().item())
    B_norms.append(layer.B.data.norm().item())

    optimizer.step()

Out[18]:

Visualization

Line plot showing gradient norms for matrices A and B over 20 training steps, with A starting at zero and rising as B moves away from zero, showing the asymmetric learning dynamics of LoRA initialization. — Gradient magnitudes for matrices A and B during 20 training steps reveal asymmetric learning dynamics. Matrix B receives non-zero gradients from step 1, while matrix A's gradient is zero initially. As B evolves, gradient flow to A enables both matrices to co-develop. This coupled dynamic explains why B must adapt first before joint optimization begins.

Line plot showing the Frobenius norm of matrix B growing from zero to approximately 0.6 over 20 training steps, illustrating stable weight growth after LoRA's zero initialization of B. — Matrix B Frobenius norm grows smoothly from zero to approximately 0.6 over 20 training steps, letting gradient flow to A through coupled dynamics. This stable, predictable growth shows how zero initialization preserves pretrained behavior while transitioning to full two-matrix adaptation.

The plot confirms the mathematical analysis. Matrix $A$ 's gradient starts at zero and grows as $B$ moves away from zero. The coupled dynamics quickly bring both matrices into active training, with both gradient norms reaching comparable magnitudes after just a handful of steps. The smooth, monotonic growth of $\|B\|$ shows that the zero initialization does not cause instability or slow convergence.

Weight Merging for InferenceLink Copied

LoRA's trained adapters can be merged into the base weights for inference, eliminating any computational overhead from the low-rank decomposition:

In[19]:

Code

# After training, we can merge LoRA into the base weights
layer = LoRALinear(in_features=64, out_features=64, rank=8, alpha=8.0)

# Simulate some training
for _ in range(10):
    x = torch.randn(16, 64)
    loss = layer(x).sum()
    loss.backward()
    with torch.no_grad():
        layer.A -= 0.01 * layer.A.grad
        layer.B -= 0.01 * layer.B.grad
        layer.A.grad.zero_()
        layer.B.grad.zero_()

# Verify equivalence
x_test = torch.randn(4, 64)
output_separate = layer(x_test)
merged_weight = layer.get_merged_weight()
output_merged = F.linear(x_test, merged_weight)

max_diff = (output_separate - output_merged).abs().max().item()
are_equivalent = torch.allclose(output_separate, output_merged)

Out[20]:

Console

Max difference: 2.432e-05
Outputs are equivalent: True

After merging, the LoRA matrices can be discarded and replaced with a single weight matrix that has zero inference overhead compared to the original model. The negligible difference confirms that merging is mathematically exact within floating-point accuracy. This property matters enormously for deployment: during development, separate $A$ and $B$ matrices provide flexibility for experimentation, while production uses merged weights to eliminate computational overhead from the low-rank decomposition. A model fine-tuned with LoRA at rank 16 runs at exactly the same speed as the original model after merging, with no additional latency.