Implementing LoRA From Scratch in PyTorch

While open-source implementations of LoRA are readily available (like the HuggingFace PEFT library shown elsewhere in this topic), implementing it from scratch in raw PyTorch builds a much better understanding of the practical mechanics.

A typical LoRA layer comprises two matrices, A and B, implemented in a LoRAWeights class alongside its forward pass. The class decomposes a matrix of dimensionality d×k into two smaller matrices A and B, accepting four parameters: d (rows of W), k (columns of W), r (the rank hyperparameter), and alpha (a scaling parameter controlling the strength of the adaptation). Both self.A and self.B are learnable parameters of the module. Matrix A is initialized from a Gaussian distribution (optionally scaled down so initial values aren't too large); matrix B is initialized as a zero matrix. This specific initialization ensures the product A×B is zero when fine-tuning begins — meaning the original model's weights are exactly retained until any fine-tuning actually happens. In the forward method, the input x is multiplied by A and B, then scaled by alpha/r, and returned as the module's output. A higher alpha means more pronounced adjustments from the LoRA layer; a lower alpha means more subtle changes.

To actually use this on a real network: say you have an existing trained network with layers fc1, fc2, fc3, fc4. The goal is to attach a LoRAWeights instance to each targeted layer — you don't necessarily need to attach one to every single layer (the original LoRA paper, for instance, limited adaptation to attention weights, freezing the feed-forward units entirely for parameter efficiency). In this example, fc4 could be frozen entirely since it's not enormously large compared to the other layers. The network is trained as usual, except the pre-trained base model is frozen and only the A/B matrices in each LoRAWeights layer receive gradient updates. A wrapper module (MyNeuralNetworkWithLoRA) creates one LoRAWeights layer per targeted fully connected layer (loralayer1, loralayer2, loralayer3 for fc1, fc2, fc3), and in its forward method, passes input through each original fully connected layer and adds the corresponding LoRA layer's output to it, before applying the activation function — repeating this for each targeted layer and returning the final output from the last (frozen) fully connected layer. This resulting model can then be trained exactly like any other neural network, with the pre-trained weights guaranteed to stay frozen and only the LoRAWeights parameters actually learning.

import torch
import torch.nn as nn

class LoRAWeights(nn.Module):
    def __init__(self, d: int, k: int, r: int, alpha: float):
        super().__init__()
        self.r = r
        self.alpha = alpha
        self.A = nn.Parameter(torch.randn(d, r) * 0.01)  # Gaussian init, scaled down
        self.B = nn.Parameter(torch.zeros(r, k))         # zero init

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # A @ B is exactly zero at the start of training, since B is zero-initialized
        return (self.alpha / self.r) * (x @ self.A @ self.B)

class MyNeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 10)

class MyNeuralNetworkWithLoRA(nn.Module):
    def __init__(self, model: MyNeuralNetwork, r: int = 8, alpha: float = 16):
        super().__init__()
        self.model = model
        for p in self.model.parameters():
            p.requires_grad = False  # freeze the pre-trained base model

        self.loralayer1 = LoRAWeights(784, 256, r, alpha)
        self.loralayer2 = LoRAWeights(256, 128, r, alpha)
        self.loralayer3 = LoRAWeights(128, 64, r, alpha)
        # fc4 is left frozen with no LoRA layer, as it's comparatively small

    def forward(self, x):
        x = torch.relu(self.model.fc1(x) + self.loralayer1(x))
        x = torch.relu(self.model.fc2(x) + self.loralayer2(x))
        x = torch.relu(self.model.fc3(x) + self.loralayer3(x))
        return self.model.fc4(x)

# Only self.loralayer1/2/3's A and B matrices will receive gradient updates —
# self.model's weights (fc1-fc4) stay exactly frozen throughout training.