Medium End-to-End

Train LoRA Adapter End-to-End

Train a low-rank correction to a fixed linear regression model. The base matrix stays unchanged while two smaller matrices, a and b, learn a correction a @ b. Use explicit gradients and full-batch SGD, without an optimizer or automatic differentiation.

Inputs and forward pass

Implement train_lora(x, y, w_base, a0, b0, lr, n_steps):

x has shape (N, d_in) and y has shape (N, d_out).
w_base has shape (d_in, d_out).
a0 has shape (d_in, r) and b0 has shape (r, d_out).
lr is a nonnegative learning rate; n_steps is a nonnegative integer.

All matrix dimensions are positive. Start from the supplied a0 and b0. This exercise uses row-vector orientation and a correction scale of one:

effective_w = w_base + a @ b
prediction = x @ effective_w
loss = mean((prediction - y)**2)

There is no bias, dropout or extra alpha / r factor. The correction has rank at most r; choosing a small r reduces the number of trainable parameters. Neither w_base nor any input array should be modified.

Gradients and update order

The mean includes all N * d_out squared errors. At each step, compute:

G = (2 / (N * d_out)) * x.T @ (prediction - y)
grad_a = G @ b.T
grad_b = a.T @ G

G is the gradient with respect to the effective matrix, with shape (d_in, d_out). The chain rule gives grad_a of shape (d_in, r) and grad_b of shape (r, d_out).

Compute both gradients from the same current a and b, then update:

a = a - lr * grad_a
b = b - lr * grad_b

Do not use the updated a to calculate grad_b. Repeat exactly n_steps times, recomputing predictions and gradients each time. Do not stop early.

Output and checks

Return one vector: all entries of final a, flattened in row-major order, followed by all entries of final b in the same order. Its length is d_in*r + r*d_out. Do not return the effective matrix or the loss.

For zero steps or zero learning rate, return the initial factors in that format. If both initial factors are zero, both gradients stay zero. Use the supplied initialization as-is; this edge case is not a recommended way to initialize a trainable correction.

For a one-step check, take x = [[1, 0], [0, 1]], y = [[1, 0], [0, 1]], zero w_base, a0 = [[1], [0.5]], b0 = [[0.2, -0.4]], and lr = 0.1. The gradients are grad_a = [[0], [0.25]] and grad_b = [[-0.375, -0.5]]. Return [1, 0.475, 0.2375, -0.35].

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…