LoRA linear layer

Problem

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.

1. lora_forward(x, W, b, A, B, alpha) returns the output of the layer with the adapter.
2. merge_lora(W, A, B, alpha) returns the merged weight.
3. 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

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

• 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

pythonCopyx = 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)

textCopy[[ 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]].

pythonCopymerge_lora(W, A, B, 4.0)

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

pythonCopycount_trainable(4096, 4096, 8)

textCopy65536