Fixing A to Hit b: Backpropagation, One Matrix at a Time

You trust the input \(\mathbf{x}\). The matrix \(A\) is only approximately right, so \(A\mathbf{x}\) misses the target \(\mathbf{b}\). Backpropagation is the rule for changing \(A\) so the miss shrinks. Forward quantities are in blue; the error flowing backward is in orange.

1The setup

\[ \color{#2456c8}{\mathbf{y} = A\mathbf{x}}, \qquad \color{#b6461c}{\mathbf{r} = A\mathbf{x} - \mathbf{b}}, \qquad L = \tfrac12 \lVert A\mathbf{x} - \mathbf{b} \rVert^2 = \tfrac12\, \mathbf{r}^{\mathsf T}\mathbf{r} \]

\(A\) is \(k\times m\), \(\mathbf{x}\in\mathbb{R}^m\), \(\mathbf{b}\in\mathbb{R}^k\). The residual \(\mathbf{r}\) says how far, and in which direction, each output is off. The loss \(L\) is half its squared length. We change \(A\), never \(\mathbf{x}\).

2The gradient is an outer product

One output is \(y_i = \sum_j A_{ij}\, x_j\), so \(\partial y_i / \partial A_{ij} = x_j\). The chain rule gives every entry of the gradient at once:

\[ \frac{\partial L}{\partial A_{ij}} = (y_i - b_i)\, x_j = r_i\, x_j \qquad\Longrightarrow\qquad \boxed{\;\nabla_{\!A} L = \mathbf{r}\,\mathbf{x}^{\mathsf T}\;} \]

The update to a weight matrix is the error coming back times the input going forward. Row \(i\) is scaled by how wrong output \(i\) was; column \(j\) by how much input \(j\) contributed. The result has rank one.

3One step, and what it does to the error

\[ A \;\leftarrow\; A - \eta\, \mathbf{r}\,\mathbf{x}^{\mathsf T} \] \[ \mathbf{r}' = (A - \eta\,\mathbf{r}\mathbf{x}^{\mathsf T})\,\mathbf{x} - \mathbf{b} = \mathbf{r} - \eta\,\mathbf{r}\,(\mathbf{x}^{\mathsf T}\mathbf{x}) = \bigl(1 - \eta \lVert\mathbf{x}\rVert^2\bigr)\,\mathbf{r} \]

Every step multiplies the residual by the same factor \(\rho = 1 - \eta\lVert\mathbf{x}\rVert^2\). The learning rate decides everything:

learning rate	factor \(\rho\)	what happens
\(0 < \eta < 1/\lVert\mathbf{x}\rVert^2\)	\(0 < \rho < 1\)	error shrinks smoothly
\(\eta = 1/\lVert\mathbf{x}\rVert^2\)	\(\rho = 0\)	error gone in one step
\(1/\lVert\mathbf{x}\rVert^2 < \eta < 2/\lVert\mathbf{x}\rVert^2\)	\(-1 < \rho < 0\)	overshoots, flips sign, still shrinks
\(\eta > 2/\lVert\mathbf{x}\rVert^2\)	\(\lvert\rho\rvert > 1\)	grows every step: diverges

The safe learning rate depends on \(\lVert\mathbf{x}\rVert^2\). That is why networks normalise their inputs and activations: it keeps one learning rate safe for every layer.

4The one-step fix is the smallest possible change

With \(\eta = 1/\lVert\mathbf{x}\rVert^2\) the update is

\[ \Delta A = \frac{(\mathbf{b} - A\mathbf{x})\,\mathbf{x}^{\mathsf T}}{\mathbf{x}^{\mathsf T}\mathbf{x}} \]

It solves \(\min \lVert \Delta A \rVert_F\) subject to \((A + \Delta A)\,\mathbf{x} = \mathbf{b}\). With a Lagrange multiplier \(\boldsymbol\lambda\), stationarity gives \(\Delta A = \boldsymbol\lambda\mathbf{x}^{\mathsf T}\); substituting into the constraint gives \(\boldsymbol\lambda = (\mathbf{b} - A\mathbf{x})/\lVert\mathbf{x}\rVert^2\). The change lives only along \(\mathbf{x}\): for any input \(\mathbf{u}\perp\mathbf{x}\), \(\Delta A\,\mathbf{u} = 0\), so nothing the matrix already did on other directions is disturbed. This is the normalised LMS (Widrow–Hoff) rule, also known as a Kaczmarz projection.

5A worked example

\[ \mathbf{x} = \begin{bmatrix}1\\2\end{bmatrix},\quad A = \begin{bmatrix}1&0\\0&1\end{bmatrix},\quad \mathbf{b} = \begin{bmatrix}2\\3\end{bmatrix} \;\Rightarrow\; A\mathbf{x} = \begin{bmatrix}1\\2\end{bmatrix},\quad \color{#b6461c}{\mathbf{r} = \begin{bmatrix}-1\\-1\end{bmatrix}},\quad \lVert\mathbf{x}\rVert^2 = 5 \] \[ \mathbf{r}\mathbf{x}^{\mathsf T} = \begin{bmatrix}-1&-2\\-1&-2\end{bmatrix},\qquad A' = A - \tfrac15\,\mathbf{r}\mathbf{x}^{\mathsf T} = \begin{bmatrix}1.2&0.4\\0.2&1.4\end{bmatrix},\qquad A'\mathbf{x} = \begin{bmatrix}2.0\\3.0\end{bmatrix} = \mathbf{b}\;\checkmark \]

Column 2 of \(A\) moved twice as much as column 1, because \(x_2 = 2\) is twice \(x_1\). The input that contributed more takes more of the blame.

Try it a 2 × 2 example, preloaded with section 5

A (starting guess)

b (target)

learning rate η, as a multiple of 1/‖x‖²

current A

A x

residual r = A x − b

‖r‖, step

6Many examples at once: least squares

Stack \(n\) inputs as the columns of \(X\) (\(m\times n\)) and their targets as the columns of \(B\) (\(k\times n\)):

\[ L = \tfrac12 \lVert AX - B \rVert_F^2, \qquad \nabla_{\!A} L = (AX - B)\,X^{\mathsf T} = \sum_{t=1}^{n} \mathbf{r}_t\,\mathbf{x}_t^{\mathsf T} \] \[ \nabla_{\!A} L = 0 \;\Longrightarrow\; A^{*} X X^{\mathsf T} = B X^{\mathsf T} \;\Longrightarrow\; A^{*} = B X^{\mathsf T}(XX^{\mathsf T})^{-1} = B X^{+} \]

The gradient is the sum of the per-example outer products from section 2. When examples disagree no matrix satisfies them all, and \(A^{*}\) is the best compromise in the squared-error sense: the normal equations, solved with the pseudo-inverse \(X^{+}\). Gradient descent, or stochastic gradient descent one example at a time, walks toward \(A^{*}\) without ever forming that inverse, which is how networks with millions of rows train.

7Two layers: what "backward" means

Put a nonlinearity \(\sigma\) between two matrices:

\[ \color{#2456c8}{\mathbf{z}_1 = A_1\mathbf{x},\quad \mathbf{h}_1 = \sigma(\mathbf{z}_1),\quad \mathbf{y} = A_2\mathbf{h}_1}, \qquad L = \tfrac12\lVert\mathbf{y}-\mathbf{b}\rVert^2 \] \[ \begin{aligned} \color{#b6461c}{\boldsymbol\delta_2} &= \mathbf{y} - \mathbf{b} & \nabla_{\!A_2} L &= \boldsymbol\delta_2\,\mathbf{h}_1^{\mathsf T} \\ \color{#b6461c}{\boldsymbol\delta_1} &= \bigl(A_2^{\mathsf T}\boldsymbol\delta_2\bigr) \odot \sigma'(\mathbf{z}_1) & \nabla_{\!A_1} L &= \boldsymbol\delta_1\,\mathbf{x}^{\mathsf T} \end{aligned} \]

Transposes carry the error backward. Forward, \(A_2\) maps hidden units to outputs; backward, \(A_2^{\mathsf T}\) maps output errors to hidden-unit errors. Each hidden unit is blamed in proportion to the weights through which it pushed the outputs wrong.
\(\odot\,\sigma'(\mathbf{z}_1)\) gates the blame. A unit that was switched off (zero slope) passes no error back and its row of \(A_1\) does not change.
Every layer repeats section 2: its weight update is its own error times its own input.

8Why it matters in practice

The cost is matrix products. The backward pass does one product for the error (\(A^{\mathsf T}\boldsymbol\delta\)) and one for the gradient (\(\boldsymbol\delta\,\mathbf{x}^{\mathsf T}\)) per layer, the same kind of work as the forward pass. A training step therefore costs about three forward passes and is compute-bound, like prefill.
Calibration is the same exercise. Fitting a bandwidth-efficiency factor so a performance model's predicted tokens per second match measured ones is a least-squares fit: the residual drives the correction, as in section 6.
Fine-tuning is section 4 with a rank budget. A single example's correction has rank one. LoRA assumes the correction a new task needs, summed over many examples, still has low rank, and learns \(\Delta A = UV^{\mathsf T}\) with rank 8 to 64 instead of a full matrix.