Appendix C

Matrix Calculus Cookbook

Vector-Jacobian products for the operations used in this book.

Matrix calculus is the bookkeeping layer behind backpropagation. This appendix is a compact reference for the NumPy operations used in the book: write the forward pass, receive an upstream cotangent, and return vector-Jacobian products (VJPs) with the input shapes. All formulas follow the gradient notation in Section A.4.

C.1 Method

Use index notation when unsure. If $y_j = f_j(\vx)$ and the upstream gradient is $\bar{y}_j = \partial L / \partial y_j$ , then

\bar{x}_i = \sum_j \bar{y}_j\frac{\partial y_j}{\partial x_i} .\tag{C.1}

The shape rule is the guardrail: each returned VJP has the same shape as the primal it belongs to. Broadcasted axes are summed away with unbroadcast. Reductions broadcast the upstream gradient back to the input. VJPs compose backward through the graph; no full Jacobian is materialized.

Listing C.1 Core VJPs in code

def add_forward(x, y):
    return x + y


def add_vjp(grad, x, y):
    return unbroadcast(grad, x.shape), unbroadcast(grad, y.shape)


def multiply_forward(x, y):
    return x * y


def multiply_vjp(grad, x, y):
    return unbroadcast(grad * y, x.shape), unbroadcast(grad * x, y.shape)


def matmul_forward(a, b):
    return a @ b


def matmul_vjp(grad, a, b):
    return grad @ b.T, a.T @ grad

C.2 Cookbook

Let $\bar{Y}$ be the upstream gradient.

Add with broadcasting. $Y=X+B$ . VJP: $\bar{X}=\operatorname{unbroadcast}(\bar{Y}, X)$ , $\bar{B}=\operatorname{unbroadcast}(\bar{Y}, B)$ .
Elementwise multiply. $Y=X \odot B$ . VJP: $\bar{X}=\operatorname{unbroadcast}(\bar{Y}\odot B, X)$ , $\bar{B}=\operatorname{unbroadcast}(\bar{Y}\odot X, B)$ .
Matmul. $Y=AB$ . From $y_{ij}=\sum_k a_{ik}b_{kj}$ , $\bar{A}=\bar{Y}B^\T$ and $\bar{B}=A^\T\bar{Y}$ .
Sum/mean. $Y=\sum_{i\in S} X_i$ copies $\bar{Y}$ to every reduced element. Mean divides that copy by the number of reduced elements.
Exp/log/ReLU/sigmoid/tanh. VJPs multiply $\bar{Y}$ by the scalar derivative: $e^x$ , $1/x$ , $1_{x>0}$ , $\sigma(x)(1-\sigma(x))$ , and $1-\tanh^2(x)$ .
Softmax. For $\vy=\softmax(\vx)$ , $\bar{\vx}=\vy\odot(\bar{\vy}-(\bar{\vy}^\T\vy)\one)$ .
Log-softmax and cross-entropy. If $\ell=\log\softmax(\vx)$ , $\bar{\vx}=\bar{\ell}-\softmax(\vx)\sum_i\bar{\ell}_i$ . For mean cross-entropy, $\bar{Z}=(\softmax(Z)-\operatorname{onehot}(t))/B$ .

Listing C.2 Softmax VJP in code

def softmax_forward(x, axis=-1):
    shifted = x - np.max(x, axis=axis, keepdims=True)
    exp = np.exp(shifted)
    return exp / np.sum(exp, axis=axis, keepdims=True)


def softmax_vjp(grad, y, axis=-1):
    dot = np.sum(grad * y, axis=axis, keepdims=True)
    return y * (grad - dot)

LayerNorm. Over width $m$ , $\mu=\operatorname{mean}(X)$ , $\hat{X}=(X-\mu)/\sqrt{\operatorname{var}(X)\epsilon}], and stem:[Y=\gamma\hat{X}\beta$ . With $G=\bar{Y}\odot\gamma$ , $\bar{X}=s(G-\operatorname{mean}G- \hat{X}\operatorname{mean}(G\odot\hat{X}))$ , where $s=1/\sqrt{\operatorname{var}(X)+\epsilon}$ . Also $\bar{\gamma}=\sum\bar{Y}\odot\hat{X}$ , $\bar{\beta}=\sum\bar{Y}$ .
RMSNorm. $Y=\gamma X / \sqrt{\operatorname{mean}(X^2)+\epsilon}$ . The VJP is like LayerNorm without centering; the code uses $\bar{X}=G s - Xs^3\operatorname{mean}(G\odot X)$ .
Embedding gather. $Y=W[I]$ . Scatter-add each upstream row into $\bar{W}$ at its selected index; repeated IDs add.
Reshape/transpose. Reshape sends $\bar{Y}$ to the original shape. Transpose applies the inverse permutation.
Scaled dot-product attention. $S=QK^\T/\sqrt{d}$ , $P=\softmax(S)$ , $O=PV$ . VJPs: $\bar{V}=P^\T\bar{O}$ , $\bar{P}=\bar{O}V^\T$ , $\bar{S}=\operatorname{softmaxVJP}(\bar{P},P)$ , $\bar{Q}=\bar{S}K/\sqrt d$ , $\bar{K}=\bar{S}^\T Q/\sqrt d$ .

Listing C.3 Attention VJP in code

def scaled_dot_product_attention(q, k, v):
    scale = 1 / np.sqrt(q.shape[-1])
    scores = (q @ np.swapaxes(k, -1, -2)) * scale
    probabilities = softmax_forward(scores, axis=-1)
    return probabilities @ v


def attention_vjp(grad, q, k, v):
    scale = 1 / np.sqrt(q.shape[-1])
    scores = (q @ np.swapaxes(k, -1, -2)) * scale
    probabilities = softmax_forward(scores, axis=-1)
    grad_v = np.swapaxes(probabilities, -1, -2) @ grad
    grad_prob = grad @ np.swapaxes(v, -1, -2)
    grad_scores = softmax_vjp(grad_prob, probabilities, axis=-1)
    grad_q = (grad_scores @ k) * scale
    grad_k = (np.swapaxes(grad_scores, -1, -2) @ q) * scale
    return grad_q, grad_k, grad_v

In practice

Autodiff systems implement these VJPs as kernels or kernel graphs, then gradient-check tricky new ops with finite differences [baydin2015automatic]. Transformers depend especially on LayerNorm [ba2016layer], RMSNorm variants [zhang2019root], and scaled dot-product attention [vaswani2017attention]. Efficient systems save or recompute only the tensors needed by these VJPs; checkpointing trades extra forward compute for lower activation memory [griewank2008].

Key equations

\bar{x}_i = \sum_j \bar{y}_j\frac{\partial y_j}{\partial x_i}

Y=AB,\qquad \bar{A}=\bar{Y}B^\T,\quad \bar{B}=A^\T\bar{Y}

\bar{\vx}=\vy\odot(\bar{\vy}-(\bar{\vy}^\T\vy)\one),\quad \vy=\softmax(\vx)

\bar{Z}=\frac{\softmax(Z)-\operatorname{onehot}(t)}{B}

O=\softmax(QK^\T/\sqrt d)V

C.3 Teach it

The one-sentence version. Backprop is shape-preserving bookkeeping: each operation receives an upstream gradient and returns one gradient per input.

An analogy. A VJP is an expense report. The loss sends one bill downstream; each operation splits that bill among the inputs that caused it.

At the board.

Write $y_{ij}=\sum_k a_{ik}b_{kj}$ and derive the two matmul VJPs by summing over the repeated index.
Broadcast a bias across a batch, then sum the batch axis to get the bias gradient.
Show softmax as "subtract expected upstream under $\vy$ ".
Finish with attention as matmul, softmax, matmul in reverse.

Misconceptions to address. A gradient is not allowed to keep the output shape if the input shape was different. Broadcasting in the forward pass means summing in the backward pass. Softmax does not need a dense Jacobian.

Check for understanding. If a bias $b\in\R^d$ is added to every row of $X\in\R^{B\times d}$ , why is $\bar{b}$ a sum over $B$ rows?

C.4 Exercises

Exercise C.1 ★ Bias broadcasting

For $Y=X+b$ with $X\in\R^{B\times d}$ and $b\in\R^d$ , derive the VJP for $b$ .

Solution in Appendix D

Exercise C.2 ★★ Matmul by indices

Starting from $y_{ij}=\sum_k a_{ik}b_{kj}$ , derive $\bar{A}=\bar{Y}B^\T$ and $\bar{B}=A^\T\bar{Y}$ .

Solution in Appendix D

Exercise C.3 ★★ Softmax VJP

Show that the softmax VJP can be computed as $\vy\odot(\bar{\vy}-(\bar{\vy}^\T\vy)\one)$ without forming the Jacobian.

Solution in Appendix D

Exercise C.4 ★★★ Attention gradient check

Use the chapter code to compute the query VJP of scaled dot-product attention and finite-difference check it on a tiny seeded tensor.

Solution in Appendix D

References

[parr2018] T. Parr and J. Howard. The matrix calculus you need for deep learning. 2018. arXiv:1802.01528
[griewank2008] A. Griewank and A. Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd edition. SIAM, 2008.
[ba2016layer] J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer Normalization. 2016. arXiv:1607.06450
[baydin2015automatic] A. G. Baydin et al. Automatic Differentiation in Machine Learning: a Survey. 2015. arXiv:1502.05767
[vaswani2017attention] A. Vaswani et al. Attention Is All You Need. 2017. arXiv:1706.03762
[zhang2019root] B. Zhang and R. Sennrich. Root Mean Square Layer Normalization. 2019. arXiv:1910.07467