Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

LoRA Implementation › LoRA Module Design

Michael BrenndoerferDecember 2, 202512 min read

Language AI HandbookAll chapters

Part 35 of 67Parameter-Efficient Fine-tuning

Chapter 257LoRA Implementation › LoRA Module Design

Reading section 2 of 11 · LoRA Module Design · 12 minSection 2 of 11 · 12 min

Section 2 of 11 · 12 min read

LoRA Implementation › LoRA Module DesignOne section · Sections in this chapter2 / 11 ↓

Sections in this chapter11

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

The core building block of any LoRA implementation is a module that wraps an existing linear layer and adds the low-rank adaptation path. This wrapper pattern is basic to LoRA's design philosophy: rather than modifying the original model architecture, we augment it with a parallel pathway that learns task-specific adjustments. Think of it as a person who has already internalized a large library of general knowledge, and who adds a small notebook of task-specific reminders on top of that foundation. The notebook doesn't replace the underlying memory, it supplements it, and it can be discarded when working on a different task.

Recall from the previous chapter that LoRA modifies a weight matrix $W$ by adding the product of two smaller matrices. Before writing any code, we should internalize exactly what this formula is and why each component is shaped the way it is.

The core weight update formula is:

W' = W + \frac{\alpha}{r}BA

where:

$W'$ : the modified weight matrix utilized in the forward pass
$W$ : the original pre-trained weight matrix (frozen), with shape $d_{\text{out}} \times d_{\text{in}}$
$B$ : the projection-up matrix ( $\mathbb{R}^{d_{\text{out}} \times r}$ ), initialized to zeros
$A$ : the projection-down matrix ( $\mathbb{R}^{r \times d_{\text{in}}}$ ), initialized with random noise
$r$ : the low rank of the adaptation, typically 4 to 64
$\alpha$ : the scaling factor that controls the adaptation strength

To understand why this decomposition is the right shape, consider what the full weight update $\Delta W$ would look like if we allowed unconstrained fine-tuning. For a layer with input dimension $d_{\text{in}} = 1024$ and output dimension $d_{\text{out}} = 1024$ , the full update matrix has $1024 \times 1024 = 1{,}048{,}576$ parameters. By imposing the constraint that $\Delta W = BA$ where $r = 8$ , we replace that million-parameter update with an 8,192-parameter one: $A \in \mathbb{R}^{8 \times 1024}$ contributes $8 \times 1024 = 8{,}192$ parameters and $B \in \mathbb{R}^{1024 \times 8}$ contributes another $1024 \times 8 = 8{,}192$ , for a total of $16{,}384$ parameters compared to the original million. The rank-8 constraint says: "the update I need can be described by eight basis directions in weight space, not a thousand."

Understanding the role of each component helps clarify why this decomposition works. The original weight $W$ remains frozen throughout training, preserving all the knowledge the model acquired during pre-training. This frozen foundation gives stability and prevents catastrophic forgetting of general capabilities. Meanwhile, the matrices $A$ and $B$ are the only trainable parameters, forming a low-rank "bottleneck" through which all adaptation must pass. The matrix $A$ projects the input down to a much smaller dimension $r$ , capturing the most relevant features for the adaptation task. The matrix $B$ then projects back up to the output dimension, translating these compressed representations into modifications to the layer's output.

This decomposition assumes that weight updates have a low "intrinsic rank." This means a low-rank matrix can approximate the difference between a pre-trained and fine-tuned model. This hypothesis has strong empirical support: we have observed that fine-tuning typically modifies weights in structured ways that don't require the full expressiveness of dense updates. The scaling factor $\alpha/r$ ensures that update magnitudes stay consistent across different ranks. When you increase rank to capture more complex adaptations, the scaling automatically compensates by reducing the per-parameter contribution, maintaining stable learning dynamics regardless of your rank selection.

The key insight is that the product $BA$ always has rank at most $r$ , no matter how many gradient steps we take. This is a basic property of matrix multiplication: the rank of a product is bounded by the minimum rank of its factors. By choosing $r \ll \min(d_{\text{in}}, d_{\text{out}})$ , we force all adaptation to happen within a low-dimensional subspace of the full weight space. The model cannot escape this subspace no matter how long it trains, which is both a constraint and a useful prior.

Let's implement this as a PyTorch module:

In[3]:

Code

import math

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


class LoRALinear(nn.Module):
    """
    A linear layer with Low-Rank Adaptation (LoRA).

    Wraps an existing linear layer and adds a parallel low-rank path.
    The original weights are frozen, only A and B matrices are trained.
    """

    def __init__(
        self,
        original_layer: nn.Linear,
        rank: int = 4,
        alpha: float = 1.0,
        dropout: float = 0.0,
    ):
        super().__init__()

        self.original_layer = original_layer
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank

        # Get dimensions from the original layer
        self.in_features = original_layer.in_features
        self.out_features = original_layer.out_features

        # Freeze original weights
        for param in self.original_layer.parameters():
            param.requires_grad = False

        # Initialize LoRA matrices
        # A: in_features -> rank (initialized with Kaiming)
        # B: rank -> out_features (initialized with zeros)
        self.lora_A = nn.Parameter(torch.empty(rank, self.in_features))
        self.lora_B = nn.Parameter(torch.zeros(self.out_features, rank))

        # Initialize A with Kaiming uniform
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

        # Optional dropout on LoRA path
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Original path (frozen)
        original_output = self.original_layer(x)

        # LoRA path: x @ A^T @ B^T * scaling
        lora_output = self.dropout(x)
        lora_output = F.linear(lora_output, self.lora_A)  # x @ A^T
        lora_output = F.linear(lora_output, self.lora_B)  # (x @ A^T) @ B^T
        lora_output = lora_output * self.scaling

        return original_output + lora_output

Several design decisions merit careful explanation, as they directly affect the behavior and effectiveness of the adaptation. First, we initialize $B$ with zeros and $A$ with Kaiming uniform initialization. This asymmetric initialization strategy ensures that at the start of training, the LoRA contribution is exactly zero because $BA = 0$ when $B$ contains only zeros. As a result, the model behaves identically to the pre-trained version before any training occurs. This initialization strategy prevents any initial degradation in model quality and gives a clean starting point where all changes come from learned adaptations rather than random noise. The Kaiming initialization for $A$ follows standard practice for layers that will be trained. This keeps appropriate variance for gradient flow.

Second, the scaling factor $\alpha/r$ controls the magnitude of the LoRA updates in a principled way. This scaling emerges from thinking about how the total update magnitude should behave as rank changes. Without scaling, doubling the rank would roughly double the magnitude of updates, since you are summing over more dimensions. The $\alpha/r$ scaling compensates for this effect, which makes it easier to tune other hyperparameters like learning rate without worrying about rank-dependent interactions. In practice, $\alpha$ is often set equal to the rank or to a fixed value like 16 or 32, depending on how aggressively you want the adaptation to influence the model's behavior.

Third, note the separation between original_output and lora_output in the forward pass. These are computed independently, which means the gradient flows through the LoRA path without touching the frozen original layer. PyTorch's autograd engine tracks which tensors have requires_grad=True and only computes gradients for those. Since original_layer.weight has requires_grad=False, no gradient is ever computed for it, saving both memory and computation.

Let's verify our module works correctly:

In[4]:

Code

# Create a base linear layer
base_layer = nn.Linear(768, 768)

# Wrap it with LoRA
lora_layer = LoRALinear(base_layer, rank=8, alpha=8.0)

# Test forward pass
test_input = torch.randn(2, 10, 768)  # batch=2, seq_len=10, hidden=768
output = lora_layer(test_input)

# Calculate parameter counts
base_params = sum(p.numel() for p in base_layer.parameters())
lora_params = lora_layer.lora_A.numel() + lora_layer.lora_B.numel()
percentage = (lora_params / base_params) * 100

Out[5]:

Console

Input shape:  torch.Size([2, 10, 768])
Output shape: torch.Size([2, 10, 768])

Original layer parameters: 590,592
LoRA A parameters: 6,144
LoRA B parameters: 6,144
Total LoRA parameters: 12,288
Parameters added: 2.1%

The parameter counts show LoRA's efficiency and explain why this technique is popular for adapting large models. The LoRA adaptation adds a very small fraction of parameters compared to the original dense layer. To understand why, consider the arithmetic: a dense layer with input and output dimensions of 768 has $768 \times 768 = 589{,}824$ parameters (plus bias). A rank-8 LoRA adds only $768 \times 8 + 8 \times 768 = 12{,}288$ parameters, representing just over 2% of the original. This ratio improves as model dimensions grow, making LoRA especially effective for large language models where hidden dimensions commonly reach 4096 or higher.

Out[6]:

Visualization

Bar chart with four bars showing parameter counts: Original Layer (frozen) at 589,824, LoRA A and LoRA B each at 6,144, and Total LoRA at 12,288 parameters. — Bar chart comparing parameter counts for a standard linear layer versus its LoRA components. The low-rank matrices A and B together add only about 2% of the original layer's parameters, letting extremely efficient fine-tuning.

Verifying Initial BehaviorLink Copied

Since $B$ is initialized to zeros, the LoRA layer should produce identical outputs to the original layer at initialization. This property is needed for so that applying LoRA doesn't accidentally degrade a model's capabilities before training begins. Without this guarantee, applying LoRA would immediately change the model's behavior, which makes it unclear whether any subsequent performance changes were due to fine-tuning or the LoRA insertion itself. The zero-initialization removes this confound entirely.

Let's verify this guarantee holds in practice:

In[7]:

Code

# Compare outputs before any training
original_output = base_layer(test_input)
lora_output = lora_layer(test_input)
difference = (original_output - lora_output).abs().max()

Out[8]:

Console

Maximum difference at initialization: 0.00e+00

The negligible difference, caused only by floating-point precision limits, confirms that our zero-initialization strategy works as designed. The model will produce exactly the same outputs as before LoRA was applied, giving you confidence that you can add LoRA to any pre-trained model without risking immediate performance degradation. This clean starting point is one of the most underappreciated aspects of LoRA's design. It means you can instrument your model with adapters before deciding which task to fine-tune for, then begin training without any warmup period to recover initial performance.

Worked Example: Manual Forward PassLink Copied

To build intuition for what the LoRA forward pass computes, let's trace through a minimal example by hand. Suppose we have a layer with $d_{\text{in}} = 4$ , $d_{\text{out}} = 4$ , and rank $r = 2$ , and consider a single input vector (ignoring batches for simplicity).

Setup. The original weight matrix $W \in \mathbb{R}^{4 \times 4}$ is frozen. We have learned adaptation matrices:

A = \begin{pmatrix} 0.2 & -0.1 & 0.3 & 0.0 \\ -0.1 & 0.4 & 0.1 & 0.2 \end{pmatrix} \in \mathbb{R}^{2 \times 4}

B = \begin{pmatrix} 0.5 & 0.1 \\ -0.3 & 0.2 \\ 0.1 & -0.4 \\ 0.2 & 0.3 \end{pmatrix} \in \mathbb{R}^{4 \times 2}

With $\alpha = 1$ and $r = 2$ , the scaling factor is $\alpha/r = 0.5$ .

Input. Consider the input vector $\mathbf{x} = [1.0,\ 0.5,\ -1.0,\ 0.0]^T$ .

Step 1: Project down with $A$ . We compute $A\mathbf{x} \in \mathbb{R}^2$ :

\begin{aligned} (A\mathbf{x})_1 &= 0.2(1.0) + (-0.1)(0.5) + 0.3(-1.0) + 0.0(0.0) = 0.2 - 0.05 - 0.3 = -0.15 \\ (A\mathbf{x})_2 &= (-0.1)(1.0) + 0.4(0.5) + 0.1(-1.0) + 0.2(0.0) = -0.1 + 0.2 - 0.1 = 0.0 \end{aligned}

So $A\mathbf{x} = [-0.15,\ 0.0]^T$ . Notice that this intermediate vector lives in $\mathbb{R}^2$ , not $\mathbb{R}^4$ : the rank-2 bottleneck is already doing its job, compressing the 4-dimensional input into 2 dimensions.

Step 2: Project back up with $B$ . We compute $B(A\mathbf{x}) \in \mathbb{R}^4$ :

\begin{aligned} (BA\mathbf{x})_1 &= 0.5(-0.15) + 0.1(0.0) = -0.075 \\ (BA\mathbf{x})_2 &= (-0.3)(-0.15) + 0.2(0.0) = 0.045 \\ (BA\mathbf{x})_3 &= 0.1(-0.15) + (-0.4)(0.0) = -0.015 \\ (BA\mathbf{x})_4 &= 0.2(-0.15) + 0.3(0.0) = -0.030 \end{aligned}

So $BA\mathbf{x} = [-0.075,\ 0.045,\ -0.015,\ -0.030]^T$ .

Step 3: Apply scaling. We multiply by $\alpha/r = 0.5$ :

\frac{\alpha}{r} BA\mathbf{x} = [-0.0375,\ 0.0225,\ -0.0075,\ -0.015]^T

Step 4: Add to original output. We add this to the output of the frozen layer $W\mathbf{x}$ . Suppose the frozen layer produces $W\mathbf{x} = [0.5,\ -0.2,\ 0.8,\ 0.1]^T$ . The final output is:

\mathbf{y} = W\mathbf{x} + \frac{\alpha}{r}BA\mathbf{x} = [0.5 - 0.0375,\ -0.2 + 0.0225,\ 0.8 - 0.0075,\ 0.1 - 0.015]^T

\mathbf{y} = [0.4625,\ -0.1775,\ 0.7925,\ 0.085]^T

The adaptation is small but structured: it shifts the output in a direction determined by the learned matrices $A$ and $B$ . As training proceeds, these matrices adjust to push outputs in whatever direction reduces the task loss, always constrained to the rank-2 subspace. You can see intuitively why higher rank lets richer adaptations: a rank-8 bottleneck gives the model eight independent directions in which to modify outputs, while rank-1 restricts it to a single direction.