06. LoRA: Mathematical Foundations

• Primary References:

1. 1
   Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., & Chen, W. (2022). LoRA: Low-Rank Adaptation of Large Language Models. In ICLR 2022.
2. 2
   Aghajanyan, A., Gupta, S., & Zettlemoyer, L. (2021). Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning. In ACL-IJCNLP 2021.
3. 3
   Dettmers, T., Pagnoni, A., Holtzman, A., & Zettlemoyer, L. (2023). QLoRA: Efficient Finetuning of Quantized LLMs. In NeurIPS 2023.
4. 4
   Liu, S.-Y., Wang, C.-Y., Yin, H., Molchanov, P., Wang, Y.-C. F., Cheng, K.-T., & Chen, M.-H. (2024). DoRA: Weight-Decomposed Low-Rank Adaptation. In ICML 2024.

• Why I Wrote This Note: Many engineers treat LoRA as a black-box config (r=16, lora_alpha=32) in HuggingFace peft. I wrote this note to record the exact linear algebra, gradient subspace projection, and scaling dynamics behind LoRA—and how low-rank subspaces can also parameterize cheap Bayesian posteriors (Laplace-LoRA).

• The Intrinsic Rank Hypothesis: Aghajanyan et al. (2021) proved that optimizing within a random low-dimensional subspace \(d_{90}\) recovers \(90\%\) of full fine-tuning accuracy. Hu et al. (2022) showed that the task-specific update matrix \(\Delta W \in \mathbb{R}^{d \times k}\) is strongly rank-deficient, so \(\text{rank}(\Delta W) = r \in [4, 64]\) suffices even when \(d = 4096\).
• Parameter & VRAM Arithmetic: Replacing \(\Delta W \in \mathbb{R}^{d \times k}\) with \(B \in \mathbb{R}^{d \times r}\) and \(A \in \mathbb{R}^{r \times k}\) reduces trainable parameters from \(d \cdot k\) to \(r(d + k)\). For \(d = k = 4096\) and \(r = 16\), trainable parameters drop from \(16.78\text{M}\) to \(131\text{K}\) (\(0.78\%\) of original size), eliminating \(99.2\%\) of Adam optimizer states.
• Why Asymmetric Initialization (\(B = 0, A \sim \mathcal{N}(0, \sigma^2)\)) Is Mandatory: At step \(t = 0\), \(B A = 0 \implies W = W_0\) (preserving the exact pretrained function), while \(\nabla_B \mathcal{L} = (\nabla_h \mathcal{L}) (A x)^\top \neq 0\) allows \(B\) to immediately move off zero on the very first gradient step.
• Zero Inference Latency: Because \(y = W_0 x + \frac{\alpha}{r} B A x = \big(W_0 + \frac{\alpha}{r} B A\big) x\) is strictly linear, we pre-compute \(W_{\text{merged}} \leftarrow W_0 + \frac{\alpha}{r} B A\) before deployment, incurring \(0\) extra FLOPs or kernel launches.

4.1 Low-Rank Update Parameterization and Forward/Backward Mechanics

Consider a frozen pretrained linear projection \(W_0 \in \mathbb{R}^{d \times k}\) mapping input activation \(x \in \mathbb{R}^k\) to output \(h \in \mathbb{R}^d\). LoRA constrains the weight update \(\Delta W \in \mathbb{R}^{d \times k}\) to have rank at most \(r \ll \min(d, k)\) via a two-factor outer product:

where \(A \in \mathbb{R}^{r \times k}\) is the down-projection matrix, \(B \in \mathbb{R}^{d \times r}\) is the up-projection matrix, and \(\alpha > 0\) is a constant scaling hyperparameter.

flowchart BT
    X["Input Activation x ∈ ℝᵏ"] --> W0["Frozen Pretrained Weight W₀ ∈ ℝ^{d×k}<br/>(No Gradients, Optional 4-bit NF4)"]
    X --> A["Trainable Down-Projection A ∈ ℝ^{r×k}<br/>Initialized A ~ N(0, σ²)"]
    A --> BOT["Low-Rank Bottleneck z = Ax ∈ ℝʳ<br/>(r ≪ min(d, k))"]
    BOT --> B["Trainable Up-Projection B ∈ ℝ^{d×r}<br/>Initialized B = 0"]
    W0 --> SUM(("h = W₀x + (α/r)BAx"))
    B -- "Scaled by α / r" --> SUM

Let us examine the backward pass gradients with respect to the upstream loss gradient \(g_h = \frac{\partial \mathcal{L}}{\partial h} \in \mathbb{R}^d\). By the chain rule, the full fine-tuning gradient would be the rank-\(1\) outer product (summed over a batch):

Under LoRA, we never materialize \(\frac{\partial \mathcal{L}}{\partial W} \in \mathbb{R}^{d \times k}\) in optimizer memory! Instead, gradients project directly through the low-rank factors \(A\) and \(B\):

4.2 Mathematical Justification of Initialization & \(\alpha/r\) Scaling

1. 1
   Why Not Initialize Both \(A = 0\) and \(B = 0\)?

Inspecting the gradient equations above, if \(A_0 = 0\) and \(B_0 = 0\) simultaneously, then at step \(t=0\):

Both matrices are trapped at a saddle point forever! Conversely, if both \(A_0, B_0 \sim \mathcal{N}(0, \sigma^2)\), then \(\Delta W_0 = \frac{\alpha}{r} B_0 A_0 \neq 0\), injecting random Gaussian noise into every Transformer layer at step \(0\) anddestroying pretrained feature alignment. Setting \(B_0 = 0\) and \(A_0 \sim \mathcal{N}(0, \sigma^2)\) (or Kaiming uniform) guarantees \(\Delta W_0 = 0\) while ensuring \(\frac{\partial \mathcal{L}}{\partial B}\big|_{t=0} \neq 0\).

1. 1
   Why Scale by \(\alpha / r\) (or \(\alpha / \sqrt{r}\) in rsLoRA)?

When tuning rank \(r\) (e.g., sweeping \(r \in \{8, 16, 32, 64\}\)), the inner product \((B A)_{ij} = \sum_{m=1}^{r} B_{im} A_{mj}\) sums \(r\) terms. Under Adam optimization after the first step, the magnitude of \((BA)x\) grows proportionally to \(r\). Dividing by \(r\) keeps the effective update magnitude \(\left\|\frac{\alpha}{r} B A x\right\|_2\) roughly invariant to \(r\), allowing engineers to change \(r\) without re-tuning the learning rate \(\eta\) from scratch. (Kalajdzievski, 2023 later showed that for large ranks \(r \ge 64\), Rank-Stabilized LoRA (rsLoRA) scaling by \(\alpha / \sqrt{r}\) prevents rank collapse via the Central Limit Theorem).

4.3 Subspace Similarity & Weight-Decomposed Adaptation (DoRA)

To verify whether two adapters \(A_{r=8}\) and \(A_{r=64}\) learn the same underlying subspace, Hu et al. (2022) compute the normalized Grassmann projection similarity between the top-\(i\) right singular vectors \(U_A^i\) and \(U_B^j\) of \(\Delta W\):

Empirically, \(\phi > 0.5\) for the top singular direction (\(i=j=1\)), confirming that a tiny rank \(r \in [4, 8]\) applied across all attention and MLP matrices (\(W_q, W_k, W_v, W_o, W_{\text{up}}, W_{\text{down}}\)) outperforms a huge rank \(r=64\) applied only to \(W_q, W_v\).

4.4 PyTorch Implementation of Mergeable LoRALinear

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

class LoRALinear(nn.Module):
    # Linear layer augmented with a mergeable low-rank adaptation branch ΔW = (alpha / r) * B @ A
    def __init__(self, base_layer: nn.Linear, r: int = 8, lora_alpha: float = 16.0, lora_dropout: float = 0.05):
        super().__init__()
        self.in_features = base_layer.in_features
        self.out_features = base_layer.out_features
        self.r = r
        self.scaling = lora_alpha / float(r)
        self.merged = False

        # Freeze pretrained weights W_0
        self.weight = base_layer.weight
        self.weight.requires_grad = False
        self.bias = base_layer.bias
        if self.bias is not None:
            self.bias.requires_grad = False

        # Low-rank factors A in R^{r x in} and B in R^{out x r}
        self.lora_A = nn.Parameter(torch.empty((r, self.in_features), device=self.weight.device, dtype=self.weight.dtype))
        self.lora_B = nn.Parameter(torch.zeros((self.out_features, r), device=self.weight.device, dtype=self.weight.dtype))
        self.dropout = nn.Dropout(p=lora_dropout) if lora_dropout > 0.0 else nn.Identity()

        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

    def merge_weights(self) -> None:
        # Fold (alpha / r) * B @ A into W_0 for zero-latency deployment
        if not self.merged:
            with torch.no_grad():
                self.weight.data += self.scaling * (self.lora_B @ self.lora_A)
            self.merged = True

    def unmerge_weights(self) -> None:
        # Restore W_0 to swap another task adapter dynamically
        if self.merged:
            with torch.no_grad():
                self.weight.data -= self.scaling * (self.lora_B @ self.lora_A)
            self.merged = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.merged:
            return F.linear(x, self.weight, self.bias)
        base_out = F.linear(x, self.weight, self.bias)
        lora_out = (self.dropout(x) @ self.lora_A.t() @ self.lora_B.t()) * self.scaling
        return base_out + lora_out

1. 1
   Bayesian LoRA for LLM & Vision Uncertainty: One of the most exciting intersections with my uncertainty research (Understanding Predictive Uncertainty and Deep Ensembles vs MC Dropout) is Laplace-LoRA / LoRA Ensembles (Yang et al., ICLR 2024; Wang et al., 2023). Computing a Kronecker-factored Laplace approximation or training an \(M=5\) ensemble over \(7\text{B}\) parameters is impossible, but training \(M=5\) independent LoRA adapters \(\{B_m, A_m\}_{m=1}^{5}\) costs \(< 4\%\) extra VRAM and captures high-quality epistemic uncertainty!
2. 2
   Domain Adaptation in Medical Vision: When adapting pretrained backbones to chest X-ray tuberculosis screening (ICEEICT 2024) or lung segmentation (ICICT 2025) on small hospital cohorts (\(N < 2,000\)), full fine-tuning rapidly overfits and destroys calibration. Low-rank adaptation acts as a structural regularizer that preserves pretrained edge/texture priors.

• Adapt All Linear Layers, Not Just \(W_q, W_v\): The original 2022 LoRA paper focused on attention \(W_q, W_v\) for Transformer encoders/decoders, but Dettmers et al. (QLoRA, 2023) conclusively showed that applying LoRA to all linear layers (including MLP \(W_{\text{gate}}, W_{\text{up}}, W_{\text{down}}\)) is far more critical for accuracy than increasing rank \(r\) from \(8\) to \(64\).
• Magnitude-Direction Coupled Updates (Why DoRA Helps): Full fine-tuning often alters weight column directions slightly while scaling magnitudes significantly (or vice versa). Standard LoRA couples magnitude and direction changes linearly, which is why DoRA (\(W = m \frac{W_0 + BA}{\|W_0 + BA\|_c}\)) converges faster at tiny ranks (\(r=4\)).

• Implement self-contained LoRALinear with exact zero-latency merge_weights() and unmerge_weights().Verified
• Implement a 5-Member LoRA Ensemble (MultiLoRALinear) that runs \(1\) shared base matmul \(W_0 x\) and \(5\) batched low-rank matmuls \(\frac{\alpha}{r} B_m (A_m x)\) in a single forward pass to compute epistemic mutual information \(\mathcal{I}[y; \theta \mid x]\) at \(\sim 1.05\times\) latency.Pending
• Compare standard \(\alpha/r\) scaling against rsLoRA (\(\alpha/\sqrt{r}\)) and DoRA on medical VQA and clinical text classification.Pending

"Pretraining requires high rank to build the coordinate system of representation space; fine-tuning only requires low rank to rotate and align existing axes toward the target task."

Multi-Adapter Batched Epistemic Inference Identity:

• Connected Research Notes:
• Deep Ensembles vs MC Dropout — Replacing \(M\) full models with \(M\) LoRA adapters solves the \(M\times\) VRAM bottleneck of Deep Ensembles.
• Mixture-of-Experts Architecture — Routing tokens across a bank of specialized LoRA experts (MoLoRA / LoRAHub).
• Test-Time Adaptation — Using ultra-low-rank (\(r=1, 2\)) adapters for fast test-time domain adaptation under covariate shift.