Matrix Calculus for Neural Networks: Deriving the MSE Gradient

A linear layer has only a few gradient formulas, yet several implementation errors remain easy to miss: a flattened bias can change the computation without raising an exception; an extra loss divisor changes gradient magnitude; a smaller finite-difference step can make the check worse. This article works through those failures with fixed, inspectable numbers.

Revision, 2026-09-08: the program now uses the same inputs as the hand calculation and checks bias and input gradients as well as weights. Exact rational calculations are separated from floating-point approximations. The old unseeded random example and animation using a different two-input fixture are no longer presented as evidence for this run. Original downloads remain unchanged.

The Run notes above and the program below both use this revision’s linear-layer gradient package. The original deep-learning-math-lab remains a historical reference; do not mix script entry points from the two directories.

1. Fix the Inputs and the Loss Definition

Start with one sample: x has shape 3×1, W is 2×3, and both b and y are 2×1. The objective here is half the sum of squared errors, not an unspecified function called MSE.

W = [[ 0.2, -0.4,  0.1],
     [ 0.7,  0.3, -0.2]]
b = [[ 0.05], [-0.10]]
x = [[ 1.50], [-2.00], [0.50]]
y = [[ 1.00], [-1.00]]

y_hat = W x + b = [[1.20], [0.25]]
e = y_hat - y   = [[0.20], [1.25]]
L = 0.5 * sum(e**2) = 641/800 = 0.80125

For example, the first output is 0.2×1.5 + (-0.4)×(-2) + 0.1×0.5 + 0.05 = 1.2. Every value follows from these inputs. The code does not generate a different random parameter set after the explanation assumes a convenient residual.

2. Read Three Gradients from a Scalar Differential

There is no need to allocate a large Jacobian first. Write the residual in components and collect the coefficient of each small change:

e_i = sum_j W_ij x_j + b_i - y_i
dL  = sum_i e_i de_i
    = sum_i e_i (sum_j x_j dW_ij + sum_j W_ij dx_j + db_i)

dL/dW_ij = e_i x_j
dL/db_i  = e_i
dL/dx_j  = sum_i W_ij e_i

In matrix form, dW = e x^T, db = e, and dx = W^T e. The first two update this layer’s parameters; the last propagates a gradient to an earlier layer. A gradient must have its variable’s shape, but matching shape alone does not establish correct values.

dW = [[ 0.300, -0.400,  0.100],
      [ 1.875, -2.500,  0.625]]
db = [[0.200], [1.250]]
dx = [[0.915], [0.295], [-0.230]]

Increasing W[1,0] changes the second output by 1.5 times that parameter change. Its current residual is 1.25, so the first-order loss coefficient is 1.25×1.5=1.875. This explains the row and column roles in the outer product, rather than relying only on a memorized transpose.

3. What Does the Batch Loss Divide By?

The program stores samples in columns: X=(D,B), Y=(M,B), and b=(M,1). B counts samples and M counts outputs. The objective averages samples, not output coordinates:

E  = W X + b - Y
L  = sum(E**2) / (2 B)
G  = E / B
dW = G X^T
db = sum(G, axis=1, keepdims=True)
dX = W^T G

A shared bias affects every sample, so its contributions must be summed along the sample axis. keepdims=True retains a column shape; it does not choose the correct reduction axis for you.

The package includes this three-output, two-sample control to expose the coincidence in a two-output-only test:

W3 = [[0.2,-0.4,0.1], [0.7,0.3,-0.2], [-0.3,0.5,0.2]]
b3 = [[0.05], [-0.1], [0.1]]
X3 = [[1.5,0.25], [-2.0,1.0], [0.5,-1.5]]
Y3 = [[1.0,-0.5], [-1.0,0.25], [0.2,-0.4]]

half_sample_mean = 1.0696875
element_mean_MSE = 0.713125
element_mean_MSE / half_sample_mean = 2/3

If a framework port changes all gradients by a constant factor, inspect loss reduction first. Immediately compensating with a learning-rate change can hide the fact that the objective itself changed.

4. The Complete NumPy Program

This is the package’s linear_example.py, with no randomness or external dataset. It defaults to float64 and checks all 11 W, b, and X gradient coordinates. Positive and negative perturbations use independent copies, so checking cannot leave the original input perturbed. Y is treated as a fixed target.

"""A fixed linear-layer example; samples occupy columns, loss averages samples."""
import json
import numpy as np


def fixture(dtype=np.float64):
    return tuple(np.array(a, dtype=dtype) for a in (
        [[0.2, -0.4, 0.1], [0.7, 0.3, -0.2]],
        [[0.05], [-0.1]], [[1.5], [-2.0], [0.5]], [[1.0], [-1.0]]))


def evaluate(W, b, X, Y):
    arrays = (W, b, X, Y)
    if any(not isinstance(a, np.ndarray) or a.ndim != 2 for a in arrays):
        raise ValueError("Use two-dimensional NumPy arrays")
    if W.dtype not in (np.dtype('float32'), np.dtype('float64')) or any(a.dtype != W.dtype for a in arrays):
        raise ValueError("Use one common float32 or float64 dtype")
    outputs, inputs = W.shape
    batch = X.shape[1]
    if min(outputs, inputs, batch) == 0 or X.shape[0] != inputs or b.shape != (outputs, 1) or Y.shape != (outputs, batch):
        raise ValueError("Expected W=(M,D), b=(M,1), X=(D,B), Y=(M,B)")
    if any(not np.isfinite(a).all() for a in arrays):
        raise ValueError("Inputs must be finite")
    with np.errstate(over='raise', invalid='raise', divide='raise'):
        error = W @ X + b - Y
        loss = float(np.sum(error * error) * W.dtype.type(0.5 / batch))
        G = error / batch
        gradients = (G @ X.T, G.sum(axis=1, keepdims=True), W.T @ G)
    return loss, gradients


def finite_difference(arrays, slot, h=1e-5):
    if slot not in (0, 1, 2) or not np.isscalar(h) or not np.isfinite(h) or h <= 0:
        raise ValueError("Choose W/b/X (0/1/2) and a finite positive step")
    evaluate(*arrays)
    numerical = np.empty(arrays[slot].shape, dtype=np.float64)
    for index in np.ndindex(numerical.shape):
        plus = [a.copy() for a in arrays]
        minus = [a.copy() for a in arrays]
        plus[slot][index] += h
        minus[slot][index] -= h
        numerical[index] = (evaluate(*plus)[0] - evaluate(*minus)[0]) / (2 * h)
    return numerical


def main():
    W, b, X, Y = arrays = fixture()
    loss, gradients = evaluate(*arrays)
    checks = {}
    for slot, name in enumerate(('W', 'b', 'X')):
        numerical = finite_difference(arrays, slot)
        np.testing.assert_allclose(gradients[slot], numerical, atol=1e-8, rtol=1e-8)
        checks[name] = float(np.max(np.abs(gradients[slot] - numerical)))
    next_loss, _ = evaluate(W - 0.1 * gradients[0], b - 0.1 * gradients[1], X, Y)
    print(json.dumps({'loss': loss, 'error': (W @ X + b - Y).tolist(),
        'dW': gradients[0].tolist(), 'db': gradients[1].tolist(), 'dX': gradients[2].tolist(),
        'max_abs_difference': checks, 'loss_after_one_update': next_loss}, indent=2))
    print('LINEAR_EXAMPLE_OK')


if __name__ == '__main__':
    main()

The printed value 0.8012499999999998 differs from decimal 0.80125 through binary floating-point rounding. Comparisons use atol=1e-8, rtol=1e-8 for the float64 fixtures, not equality of every printed digit.

The final update changes only W and b, keeping X and Y fixed. With rate 0.1, the new loss is 0.050078125. For this single sample, the updated residual is one quarter of the original, making the loss one sixteenth as large. That is not a promise of descent for arbitrary networks or rates.

5. What Was Actually Checked?

Three fixed fixtures were run on 2026-09-08 using arm64 macOS, Python 3.13.9, and NumPy 2.3.5. In addition to central differences, a separate Python Fraction implementation computes losses and derivatives with scalar loops rather than reusing NumPy matrix operations as its reference.

There are 46 gradient coordinates in total. Each coordinate is also checked with exact rational central differences at steps 1/10 and 1/100000, giving 92 exact-equality comparisons. Maximum absolute difference between NumPy analytical gradients and the rational reference is approximately 4.44e-16. These are tests of specified fixtures, not a proof for all possible inputs.

The full gradient CSV retains floating-point discrepancies, and the exact derivative CSV retains fractions. The old CSV contains only six numerical W checks rounded to 0.000000 at six decimal places; numerical bias fields are empty. The new package does not mislabel those empty fields as verified gradients.

6. Why Can a Smaller Step Make the Check Worse?

Central differences use [L(θ+h)-L(θ-h)]/(2h). A general smooth nonlinear function has a trade-off between truncation and floating-point roundoff. Here, however, fixing the other variables and perturbing one coordinate makes the squared loss a quadratic in that coordinate. Its central-difference truncation error vanishes in exact arithmetic. A familiar U-shaped error curve must not be assumed for this particular function.

View the full-size error plot; download all 24 sweep rows.

Finite differences are a numerical diagnostic, not absolute ground truth. Nonsmooth points, random operations, rounding, and shared storage can affect interpretation. The PyTorch gradcheck documentation likewise notes that default tolerances target double precision and that nondifferentiable points can cause failures. PyTorch was not installed or run for this audit; NumPy checks must not be described as cross-framework validation.

7. Test Broadcasting with a Failing Input

Flattening b from (2,1) to (2,) lets W @ X + b.ravel() - Y run while producing a (2,2) error array. Loss changes from the correct 0.80125 to 1.7825. That is a reproducible semantic error without claiming that this tiny example must exhaust GPU memory.

NumPy broadcasting rules compare dimensions from the right; they must match or one must be 1. Thus (64,1)+(64,) produces a 64×64 result. Its float32 payload is 16,384 bytes, or 16 KiB. This small result does not itself demonstrate GPU OOM; memory consequences at larger scales require separate measurement.

These controls establish that specific mistakes are detectable, not that a test count proves correctness. The differential derivation, shape contract, independent reference, and roundoff analysis serve different purposes.

8. Download and Continue Checking

Download the linear-layer gradient lab and run these commands from the extracted directory:

python3 -m venv .venv
.venv/bin/python -m pip install -r requirements.txt
.venv/bin/python linear_example.py
.venv/bin/python audit_matrix.py --output reproduced

The success markers are LINEAR_EXAMPLE_OK and MATRIX_GRADIENT_AUDIT_OK. The result JSON records environment, source hashes, every fixture, negative controls, and limitations. Plotting is optional; its commands and dependencies are documented in the package README.

No GPU throughput, memory-bandwidth, distributed-training, or mixed-precision performance benchmark was run. Finite-difference h and mixed-precision loss scaling are different mechanisms; changing loss scale does not replace checking whether a perturbation is representable.

Continue with two-layer MLP backpropagation and gradient checking to see how dX becomes an earlier layer’s upstream gradient. Then compare fixed experiments with gradient descent, Momentum, and Adam to distinguish correct gradient computation from the behavior of a chosen update rule.