LoRA: Low-Rank Adapter Matrices Bolted Onto Part 1's Linear Layer
The problem, from first principles
rl-alignment-qlora-quantized-lora used LoRA's forward pass as one ingredient in a bigger quantization story. This question builds LoRA itself, in full, from scratch: a way to fine-tune a huge frozen Linear layer (linear-regression-hypothesis-function's function, unchanged) by training only a tiny low-rank ADDITION to its output: including deriving the backward pass, which QLoRA's exercise didn't need.
From theory to code
Implement lora_forward (base output plus a scaled low-rank delta) and lora_backward (gradients for ONLY the trainable LoRA matrices: the base layer stays completely frozen).
Constraints
lora_forward(input, base_weight, base_bias, lora_A, lora_B, alpha, rank)returns(output, cache), whereoutput = linear(input, base_weight, base_bias) + (alpha/rank) * (input @ lora_A.T) @ lora_B.T.lora_Ahas shape(rank, in_features);lora_Bhas shape(out_features, rank).lora_backward(grad_output, cache, lora_B, alpha, rank)returns(grad_A, grad_B)ONLY: nevergrad_weight/grad_bias, since the base layer is frozen.- Setting
lora_A/lora_Bto all zeros must makelora_forwardreduce exactly to the plain frozenlinearcall.
Hints
Open one at a time. Each gives away a little more than the last.
Hint 1
Think of the LoRA path as two chained Linear-like operations with no bias: intermediate = input @ lora_A.T (projecting down to rank dimensions), then delta = scale * (intermediate @ lora_B.T) (projecting back up to out_features): backprop through it exactly the way you'd backprop through two stacked linear layers.
Hint 2
grad_B follows the same grad_weight = grad_output.T @ input pattern 03-gradient-checkpointing's linear_backward used, applied to the intermediate -> delta step; grad_A follows the same pattern applied to the input -> intermediate step, using the gradient THROUGH lora_B (grad_output @ lora_B) as its own "grad_output".
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.