LoRA linear layer · AI Engineering

LoRA fine-tunes a large model without a change to its weights. It keeps the weight W of a linear layer frozen and adds a small trainable path through two low rank matrices A and B. After training, you can merge the path into W, so inference costs the same as the base model.

Task

Write three functions.

lora_forward(x, W, b, A, B, alpha) returns the output of the layer with the adapter.
merge_lora(W, A, B, alpha) returns the merged weight.
count_trainable(in_features, out_features, rank) returns the number of trainable values in A and B.

The rank r is the number of rows of A. The output of the layer is

y = x W^T + b + \frac{\alpha}{r} \, x A^T B^T.

The merged weight is W' = W + \frac{\alpha}{r} B A. With the merged weight, x W'^T + b equals the output of lora_forward.

Only A and B train. W and b are frozen and do not count.

Input and output

Name	Type	Shape	Meaning
`x`	float array	`(N, in_features)`	one input row for each of the `N` examples
`W`	float array	`(out_features, in_features)`	the frozen weight
`b`	float array	`(out_features,)`	the frozen bias
`A`	float array	`(r, in_features)`	the down projection of the adapter
`B`	float array	`(out_features, r)`	the up projection of the adapter
`alpha`	float	scalar	the LoRA scale numerator
`in_features`, `out_features`, `rank`	int	scalar	the layer sizes and the rank

lora_forward returns a float array of shape (N, out_features).
merge_lora returns a new float array of shape (out_features, in_features).
count_trainable returns a Python int.

Constraints

1 <= N <= 64, 1 <= in_features, out_features <= 4096 and 1 <= r <= 64.
The rank can be larger than in_features or out_features.
alpha > 0.
Do not change the input arrays.
Use numpy only.

Examples

python

x = np.array([[1.0, 2.0, 0.0], [0.0, -1.0, 1.0]])
W = np.array([[0.5, -0.5, 1.0], [1.0, 0.0, 0.0]])
b = np.array([0.1, -0.2])
A = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
B = np.array([[1.0, 0.5], [-1.0, 0.0]])
lora_forward(x, W, b, A, B, 4.0)

text

[[ 3.6 -1.2]
 [ 2.6 -2.2]]

With B = np.zeros((2, 2)), the same call returns the base output [[-0.4, 0.8], [1.6, -0.2]].

python

merge_lora(W, A, B, 4.0)

text

[[ 2.5  0.5  3. ]
 [-1.   0.  -2. ]]

python

count_trainable(4096, 4096, 8)

text