LoRA Linear Layer

Write the forward pass of a linear layer with a LoRA adapter, and the function that merges the adapter back into the weights.

Fine-tuning every weight of a large model is expensive. LoRA (low-rank adaptation) freezes the pretrained weight W and learns a small update in the form of two thin matrices: A of shape (d_in, r) and B of shape (r, d_out), with a small rank r. The update is scaled by alpha / r:

y = x @ W + (alpha / r) * (x @ A) @ B

After training, the update can be folded into the weight, so inference costs nothing extra:

W_merged = W + (alpha / r) * A @ B

B usually starts at zero, so at the start of training the layer is exactly the pretrained one.

Function signature

def lora_forward(x, W, A, B, alpha):
def lora_merge(W, A, B, alpha):

Example

Input

x = np.array([[1.0, 2.0]])
W = np.eye(2)
A = np.array([[1.0], [0.0]])
B = np.array([[0.0, 3.0]])
lora_forward(x, W, A, B, alpha=4.0)

Output

array([[ 1., 14.]])

Reasoning

The frozen path gives x @ W = [1, 2]. The adapter has rank 1, so the scale is 4 / 1 = 4. x @ A = [1], then @ B gives [0, 3], scaled to [0, 12]. The sum is [1, 14].

Constraints and notes

• x has shape (n, d_in) and W has shape (d_in, d_out).
• Read the rank r from the shape of A.
• Compute (x @ A) @ B, not x @ (A @ B). Both are correct, but the first never builds a full (d_in, d_out) matrix.
• lora_forward(x, W, A, B, alpha) must equal x @ lora_merge(W, A, B, alpha).