Low-Rank Adapters

LoRA represents a linear-weight update as two smaller matrices. For a base weight $W : \mathbb{R}^{d_{in}\times d_{out}}$, an adapter of rank $r$ uses $A : \mathbb{R}^{d_{in}\times r}$ and $B : \mathbb{R}^{r\times d_{out}}$:

$$W_{eff}=W+sAB.$$

The matrix orientation agrees with TorchLean's row-batch linear layers. This module defines the typed update and its action on a batch; the training code decides which parameters to optimize.

Reference: Hu et al., “LoRA: Low-Rank Adaptation of Large Language Models” (2021), https://arxiv.org/abs/2106.09685.

structure TorchLean.Adapters.LoRA.Params (α : Type) (inDim rank outDim : ℕ) : Type

LoRA factors for a linear weight of shape inDim × outDim.

• A : Spec.Tensor α (Spec.Shape.dim inDim (Spec.Shape.dim rank Spec.Shape.scalar))

  Projection from the input dimension to the adapter rank.

• B : Spec.Tensor α (Spec.Shape.dim rank (Spec.Shape.dim outDim Spec.Shape.scalar))

  Projection from the adapter rank to the output dimension.

def TorchLean.Adapters.LoRA.delta {α : Type} [Add α] [Mul α] [Zero α] {inDim rank outDim : ℕ} (p : Params α inDim rank outDim) (scale : α) : Spec.Tensor α (Spec.Shape.dim inDim (Spec.Shape.dim outDim Spec.Shape.scalar))

The scaled low-rank update $sAB$.

def TorchLean.Adapters.LoRA.effectiveWeight {α : Type} [Add α] [Mul α] [Sub α] [Zero α] {inDim rank outDim : ℕ} (base : Spec.Tensor α (Spec.Shape.dim inDim (Spec.Shape.dim outDim Spec.Shape.scalar))) (p : Params α inDim rank outDim) (scale : α) : Spec.Tensor α (Spec.Shape.dim inDim (Spec.Shape.dim outDim Spec.Shape.scalar))

Add a LoRA update to a base linear weight.

def TorchLean.Adapters.LoRA.linear {α : Type} [Add α] [Mul α] [Sub α] [Zero α] {batch inDim rank outDim : ℕ} (x : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim inDim Spec.Shape.scalar))) (base : Spec.Tensor α (Spec.Shape.dim inDim (Spec.Shape.dim outDim Spec.Shape.scalar))) (p : Params α inDim rank outDim) (scale : α) : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim outDim Spec.Shape.scalar))

Apply a linear map whose weight is augmented by a LoRA update.