LoRA (Low-Rank Adaptation) Types

Core data structures for LoRA finetuning of BitNet models.

Overview

LoRA injects trainable low-rank matrices alongside frozen ternary weights:

output = BitLinear(x) + (alpha / rank) * B @ A @ x

Where:

• BitLinear(x): frozen ternary base model output
• A: [rank, inDim] FP32 matrix (Kaiming initialized)
• B: [outDim, rank] FP32 matrix (zero initialized)
• alpha: scaling factor (typically equal to rank)

References

• "LoRA: Low-Rank Adaptation of Large Language Models" (Hu et al., 2021)
• Stanford Alpaca: instruction-following finetuning

structure Hesper.LoRA.Config : Type

LoRA configuration for finetuning

• rank : Nat

  Rank of the low-rank matrices (typical: 4, 8, 16)

• alpha : Float

  Scaling factor: output is multiplied by alpha/rank

• targetModules : List String

  Which attention projections to apply LoRA to

def Hesper.LoRA.instReprConfig.repr : Config → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.LoRA.instReprConfig : Repr Config

Equations

• Hesper.LoRA.instReprConfig = { reprPrec := Hesper.LoRA.instReprConfig.repr }

def Hesper.LoRA.Config.scale (config : Config) : Float

Compute the LoRA scaling factor: alpha / rank

Equations

• config.scale = config.alpha / config.rank.toFloat

structure Hesper.LoRA.Weight : Type

A single LoRA weight pair (A and B matrices) for one projection. Forward: output += scale * B @ (A @ x) A is [rank, inDim], B is [outDim, rank] in row-major FP32.

• a : WebGPU.Buffer

  A matrix: [rank, inDim] FP32, Kaiming initialized

• b : WebGPU.Buffer

  B matrix: [outDim, rank] FP32, zero initialized (so LoRA starts as identity)

• inDim : Nat

  Input dimension

• outDim : Nat

  Output dimension

• rank : Nat

  Rank

structure Hesper.LoRA.WeightGrad : Type

Gradient buffers for a single LoRA weight pair

• dA : WebGPU.Buffer

  Gradient for A: [rank, inDim] FP32

• dB : WebGPU.Buffer

  Gradient for B: [outDim, rank] FP32

structure Hesper.LoRA.AdamState : Type

Adam optimizer state for a single LoRA weight pair

• mA : WebGPU.Buffer

  First moment for A

• vA : WebGPU.Buffer

  Second moment for A

• mB : WebGPU.Buffer

  First moment for B

• vB : WebGPU.Buffer

  Second moment for B

structure Hesper.LoRA.LayerAdapter : Type

LoRA adapter for a single attention layer (Q and V projections)

• loraQ : Weight

  LoRA weights for Q projection

• loraV : Weight

  LoRA weights for V projection

structure Hesper.LoRA.LayerAdapterGrad : Type

Gradient buffers for a single attention layer

• gradQ : WeightGrad
• gradV : WeightGrad

structure Hesper.LoRA.LayerAdapterAdamState : Type

Adam state for a single attention layer

• stateQ : AdamState
• stateV : AdamState

structure Hesper.LoRA.Adapter : Type

Full LoRA adapter for the entire model (all transformer layers)

• config : Config
• layers : Array LayerAdapter

  Per-layer adapter weights, indexed by layer number

structure Hesper.LoRA.AdapterGrad : Type

Full gradient state for the entire model

• layers : Array LayerAdapterGrad

structure Hesper.LoRA.AdapterAdamState : Type

Full Adam optimizer state for the entire model

• layers : Array LayerAdapterAdamState
• step : Nat

  Current optimizer step (for bias correction)

structure Hesper.LoRA.SavedActivations : Type

Saved activations from forward pass, needed for backward. For each LoRA layer, we save the input x and intermediate h = A @ x.

• layers : Array (WebGPU.Buffer × WebGPU.Buffer × WebGPU.Buffer × WebGPU.Buffer)

  Per-layer saved activations: (inputToQ, hQ, inputToV, hV)

structure Hesper.LoRA.TrainConfig : Type

Training configuration

• lr : Float

  Learning rate

• beta1 : Float

  Adam beta1

• beta2 : Float

  Adam beta2

• eps : Float

  Adam epsilon

• epochs : Nat

  Number of training epochs

• logEvery : Nat

  Log every N steps

• maxSeqLen : Nat

  Max sequence length for training

def Hesper.LoRA.instReprTrainConfig.repr : TrainConfig → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.LoRA.instReprTrainConfig : Repr TrainConfig

Equations

• Hesper.LoRA.instReprTrainConfig = { reprPrec := Hesper.LoRA.instReprTrainConfig.repr }