LoRA Merging & Magnitude Pruning

The problem

Two ways to change a trained weight matrix $W$ of shape $(d_{\text{in}}, d_{\text{out}})$ cheaply: add a small learned correction with LoRA (Low-Rank Adaptation), or remove the weights that matter least (pruning).

LoRA

Fine-tuning every weight of a model is expensive. LoRA freezes $W$ and learns a low-rank update: two thin matrices $A$ of shape $(d_{\text{in}}, r)$ and $B$ of shape $(r, d_{\text{out}})$, with a small $r$ (the rank, such as 8), scaled by $\alpha / r$:

where $x$ is the input and $\alpha$ (alpha) a scaling constant. The thin matrices are tiny compared with $W$, and the layer can be used with them as they are, or merged: $W' = W + \frac{\alpha}{r} AB$ computed once, so the layer then costs exactly what the original cost.

Example. $W = \begin{pmatrix} 1 & 2 \\ 3 & 4 \end{pmatrix}$, $A = \begin{pmatrix} 1 \\ 0 \end{pmatrix}$, $B = \begin{pmatrix} 0.5 & 0.5 \end{pmatrix}$, $\alpha = 2$, $r = 1$ (the number of columns of $A$).

• $AB = \begin{pmatrix} 0.5 & 0.5 \\ 0 & 0 \end{pmatrix}$, so $W' = W + 2 \cdot AB = \begin{pmatrix} 2 & 3 \\ 3 & 4 \end{pmatrix}$.
• For $x = (1, 1)$: $xW = (4, 6)$; $xA = 1$, then $(xA)B = (0.5, 0.5)$, times $2$ is $(1, 1)$: the adapted layer gives $(5, 7)$.
• The merged layer: $xW' = (5, 7)$. The same.

Why compute (x @ A) @ B and not x @ (A @ B)? $xA$ has only $r$ numbers, so the whole adapted layer costs $O(r(d_{\text{in}} + d_{\text{out}}))$ extra, not a product of two big matrices.

Magnitude pruning

Set the fraction sparsity of weights with the smallest absolute value $|w|$ to zero.

Example. $W = \begin{pmatrix} 0.5 & -0.1 & 0.3 \\ 0.05 & -0.8 & 0.2 \end{pmatrix}$ with sparsity = 0.5: int(0.5 * 6) = 3 weights are zeroed, the ones with the smallest $|w|$: $0.05$, $0.1$ and $0.2$. The result is $\begin{pmatrix} 0.5 & 0 & 0.3 \\ 0 & -0.8 & 0 \end{pmatrix}$.

What to implement

• lora_forward(x, W, A, B, alpha): the adapted layer, without merging. The rank $r$ is A.shape[1].
• merge_lora(W, A, B, alpha): the merged matrix $W'$.
• magnitude_prune(W, sparsity): a copy of W in which the int(sparsity * W.size) entries with the smallest absolute values are 0.

Requirements

• merge_lora must not modify W, and x @ merge_lora(...) must equal lora_forward(...).
• In lora_forward, compute (x @ A) @ B, never the full $AB$: that is what makes LoRA training cheap.
• magnitude_prune zeroes exactly int(sparsity * W.size) entries (break ties by position: np.argsort(np.abs(W), axis=None, kind="stable")).

Hint

np.argsort(..., axis=None) sorts the flattened array and returns the positions in the flat order; out.flat[positions] = 0.0 writes zeros there.