LoRA

The hypothesis is that the weight change induced by adaptation has low intrinsic rank even when the weight itself does not. Parameterise that change as a product of two thin matrices and train only those.

$A$ is initialised Gaussian and $B$ at zero, so $\Delta W$ starts at exactly zero and the adapted model begins as the base model. The $\alpha/r$ factor keeps the update scale roughly constant as $r$ is varied, which is what makes the learning rate transferable across ranks.

The $\alpha/r$ scaling is why rank and learning rate interact less than you would expect — but it also means $\alpha$ is not a free parameter once you have tuned the learning rate. Fixing $\alpha = 2r$ is the usual convention; rsLoRA uses $\alpha/\sqrt{r}$ instead.

import math
import torch
from torch import Tensor, nn


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

    def forward(self, x: Tensor) -> Tensor:
        return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale

    @torch.no_grad()
    def merge_(self) -> nn.Linear:
        self.base.weight += (self.B @ self.A) * self.scale
        return self.base

Note the association order in the forward: x @ A.T @ B.T costs $O(r(d+k))$ per token, whereas forming B @ A first costs $O(dkr)$ and throws away the entire point.