Medium Research

LoRA Update

LoRA, or Low-Rank Adaptation, represents a weight correction as the product of two smaller matrices. During adaptation, the base weights stay fixed while the smaller matrices are trained. Here, compute a forward pass using supplied parameters; do not perform a training update.

Your task

Implement lora_forward(x, W, A, B, alpha, r):

x is a floating-point vector of shape (d_in,).
W has shape (d_out, d_in) and contains the base weights.
A has shape (r, d_in) and maps the input to r features.
B has shape (d_out, r) and maps those features to the output width.
r is a positive integer matching the inner dimensions of A and B.
alpha is a scalar controlling the correction's strength.

Use matrix-times-vector orientation, with W @ x as the base output:

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

Return a vector of shape (d_out,). The scale applies only to the correction, not to W @ x. Use ordinary division, with no rounding of alpha / r. Add no bias, activation or dropout, and preserve all inputs.

Equivalently, form delta_W = scale * (B @ A) and return (W + delta_W) @ x. Both forms describe the same calculation; the first avoids constructing a full (d_out, d_in) correction matrix.

Example

For x = [2, 3], an identity W, A = [[1, 0]], B = [[1], [0]], alpha = 1 and r = 1, the base output is [2, 3]. The correction is [2, 0], so return [4, 3].

When alpha = 0, return W @ x. The input and output widths can differ, so the base path is not generally an identity or a direct addition of x.

What does low rank mean?

The matrix B @ A has rank at most r; it need not have rank exactly r. The two factors contain r * (d_in + d_out) parameters instead of d_in * d_out for an unrestricted correction. Choosing a sufficiently small r reduces the number of trainable parameters. This describes the correction, not the rank of the full effective matrix W + delta_W.

Background: Hu et al., LoRA: Low-Rank Adaptation of Large Language Models.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Train a smaller set of parameters

Low-rank adaptation, or LoRA, learns a change to a weight matrix using two smaller matrices. For row-vector inputs and weight shape (d_in, d_out), one convention is

$$W_{effective}=W+\frac{\alpha}{r}AB,$$

where A has shape (d_in, r) and B has shape (r, d_out). The base weight W can remain fixed while A and B learn.

For d_in = d_out = 100 and r = 4, the factors contain 400 + 400 = 800 parameters, compared with 10,000 in the full matrix. The saving comes with a constraint: the update has rank at most r. Follow the task's matrix orientation and scale, since implementations use different conventions.

Start with an initialization that allows learning

If both factors are zero, both receive zero gradients through their product. A common initialization makes one factor random and the other zero. The initial weight change is still zero, but the zero factor can receive a nonzero gradient and start learning.

Distinguish other adaptation methods

An adapter applies a small network to intermediate activations, often reducing then restoring the feature width in a residual branch. Prefix tuning learns context such as key/value vectors that tokens can attend to. These approaches modify different parts of the model and are not substitutes for one another.

To freeze base weights, ensure the optimizer cannot move them through momentum or weight decay as well as through current gradients. Check that base values stay unchanged after an update, while the loss still has a gradient path to the trainable factors or prefix.

Detecting runtime…