Matrix Calculus

Matrix Calculus

Scalar by Matrix Derivatives

Vector by Vector Derivatives

Key Matrix Calculus Rules

Chain Rule for Matrices

Gradient of Loss Functions

Backpropagation Derivation

Hessian Matrices

Python Implementation

import numpy as np

# --- Gradient of f(x) = x^T A x ---
def grad_quadratic(x, A):
    return (A + A.T) @ x

A = np.array([[2, 1], [1, 3]])
x = np.array([1.0, 2.0])
print(f"Gradient of x^T A x: {grad_quadratic(x, A)}")
# Output: [4. 7.]

# --- Hessian of f(x) = x^T A x ---
def hessian_quadratic(A):
    return A + A.T

print(f"Hessian:\n{hessian_quadratic(A)}")
# [[4 2]
#  [2 6]]

# --- Manual Jacobian computation ---
def jacobian(f, x, eps=1e-7):
    """Compute Jacobian via finite differences."""
    n = len(x)
    m = len(f(x))
    J = np.zeros((m, n))
    fx = f(x)
    for j in range(n):
        x_plus = x.copy()
        x_plus[j] += eps
        J[:, j] = (f(x_plus) - fx) / eps
    return J

# Example: f(x) = [x1*x2, sin(x1) + x2^2]
f = lambda x: np.array([x[0]*x[1], np.sin(x[0]) + x[1]**2])
x0 = np.array([1.0, 0.0])
J = jacobian(f, x0)
print(f"Jacobian at (1,0):\n{J}")
# [[0.         1.        ]
#  [0.54030231 0.        ]]

# --- Autograd with JAX ---
import jax
import jax.numpy as jnp
from jax import grad, hessian

def loss_fn(x):
    return jnp.sum(x**2) + jnp.dot(x, jnp.array([1.0, 2.0]))

# Gradient
g = grad(loss_fn)
print(f"JAX gradient: {g(x)}")

# Hessian
H = hessian(loss_fn)
print(f"JAX Hessian:\n{H}")

# --- Backpropagation: 2-layer network ---
def sigmoid(z):
    return 1 / (1 + np.exp(-z))

def sigmoid_prime(z):
    s = sigmoid(z)
    return s * (1 - s)

def forward_pass(x, W1, b1, W2, b2):
    z1 = W1 @ x + b1
    a1 = sigmoid(z1)
    z2 = W2 @ a1 + b2
    return z2, z1, a1

def backward_pass(x, y, z2, z1, a1, W2):
    N = len(y)
    dz2 = (z2 - y) * 2 / N       # dL/dz2
    dW2 = np.outer(dz2, a1)       # dL/dW2
    da1 = W2.T @ dz2              # dL/da1
    dz1 = sigmoid_prime(z1) * da1 # dL/dz1
    dW1 = np.outer(dz1, x)        # dL/dW1
    db1 = dz1                     # dL/db1
    return dW1, db1, dW2

# Simple example
np.random.seed(42)
W1 = np.random.randn(3, 2) * 0.1
b1 = np.zeros(3)
W2 = np.random.randn(1, 3) * 0.1
b2 = np.zeros(1)
x = np.array([1.0, 0.5])
y = np.array([1.0])

z2, z1, a1 = forward_pass(x, W1, b1, W2, b2)
dW1, db1, dW2 = backward_pass(x, y, z2, z1, a1, W2)

print(f"dW1 shape: {dW1.shape}")  # (3, 2)
print(f"dW2 shape: {dW2.shape}")  # (1, 3)
print(f"db1 shape: {db1.shape}")  # (3,)

Applications in AI/ML

Gradient Descent Optimization

The fundamental update rule: $\vec{w}_{t+1} = \vec{w}_t - \eta \nabla_{\vec{w}} \mathcal{L}(\vec{w}_t)$. All variants — SGD, Adam, AdaGrad, RMSProp — modify this basic step using matrix calculus. Adam tracks running averages of both the gradient (first moment) and the gradient squared (second moment), requiring elementwise matrix operations.

Neural Network Training

Each layer performs a linear transformation $W\vec{x} + \vec{b}$ followed by a nonlinear activation $\sigma(\cdot)$. Backpropagation computes $\frac{\partial \mathcal{L}}{\partial W^{(l)}} = \delta^{(l)} (\vec{a}^{(l-1)})^T$ for every layer — this is a rank-1 outer product, computed $L$ times per sample. The total computation for a network with $N$ parameters over $B$ samples is $O(B \cdot N)$ — the backbone of deep learning scalability.

Attention Mechanisms (Transformers)

The attention function $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$ requires computing the gradient through softmax (a matrix function) and matrix multiplication. The derivative of softmax produces a full Jacobian matrix, not a diagonal one, making this one of the more complex backprop computations in modern architectures.

Regularization

L2 regularization adds $\frac{\lambda}{2}\|W\|_F^2$ to the loss, contributing $\lambda W$ to the weight gradient — a direct application of $\frac{\partial}{\partial W}\|W\|_F^2 = 2W$. L1 regularization uses the subgradient of $\|W\|_1$, leading to sparse solutions.

Generative Models

In GANs, the discriminator's loss involves gradients through a minimax objective. The generator's gradient $\nabla_{\vec{z}} D(G(\vec{z}))$ flows through both the generator and discriminator networks, requiring careful application of the chain rule through both networks simultaneously.

Common Mistakes

Interview Questions

Q1: What is the gradient of $f(\vec{x}) = \vec{x}^T A \vec{x}$?

$(A + A^T)\vec{x}$. If $A$ is symmetric, this simplifies to $2A\vec{x}$. Derivation: expand $f = \sum_{i,j} a_{ij} x_i x_j$, differentiate with respect to $x_k$, and collect terms to get $(A\vec{x})_k + (A^T\vec{x})_k$.

Q2: Why is backpropagation $O(N)$ where $N$ is the number of parameters?

Each layer's gradient is an outer product $\delta^{(l)} (\vec{a}^{(l-1)})^T$, which costs $O(n_l \cdot n_{l-1})$ — exactly the number of parameters in $W^{(l)}$. Summing over all $L$ layers gives $O(\sum_l n_l \cdot n_{l-1}) = O(N)$. The backward pass has the same time complexity as the forward pass.

Q3: What is the Jacobian of the softmax function?

For $\text{softmax}(\vec{z})_i = \frac{e^{z_i}}{\sum_k e^{z_k}}$, the Jacobian is:

where $\delta_{ij}$ is the Kronecker delta. In matrix form: $J = \text{diag}(\vec{p}) - \vec{p}\vec{p}^T$. This is NOT a diagonal matrix — it's a rank-1 update to a diagonal matrix.

Q4: How does the Hessian relate to Newton's method in optimization?

Newton's method uses the Hessian to compute the optimal step: $\vec{x}_{\text{new}} = \vec{x}_{\text{old}} - H^{-1}\nabla f$. If $H \succ 0$, the step always points toward the minimum. However, computing $H$ costs $O(n^2)$ space and $O(n^3)$ for the inverse, so quasi-Newton methods (BFGS, L-BFGS) approximate $H^{-1}$ instead.

Q5: Derive the gradient of the cross-entropy loss with respect to the logits.

For softmax output $\hat{p}_k = \frac{e^{z_k}}{\sum_j e^{z_j}}$ and cross-entropy $\mathcal{L} = -\sum_k y_k \log(\hat{p}_k)$:

(using $\sum_j y_j = 1$ for one-hot targets). This is why cross-entropy + softmax gives such clean gradients.

Q6: What is the difference between the Jacobian and the Hessian?

The Jacobian $J \in \mathbb{R}^{m \times n}$ is the matrix of first derivatives of a vector function $\vec{f}: \mathbb{R}^n \to \mathbb{R}^m$. The Hessian $H \in \mathbb{R}^{n \times n}$ is the matrix of second derivatives of a scalar function $f: \mathbb{R}^n \to \mathbb{R}$. The Hessian is the Jacobian of the gradient: $H = J_{\nabla f}$.

Practice Problems

Quick Reference

Cross-References

• Previous topic: 020 - Linear Algebra Basis and Dimension
• Next topic: 022 - Linear Algebra Advanced
• Related concepts:
  • Matrix Multiplication — matrix calculus operations rely on matrix multiplication
  • Eigenvalues and Eigenvectors — Hessian eigenvectors define principal curvature directions
  • Calculus Chain Rule — scalar chain rule; matrix chain rule generalizes it
  • Partial Derivatives — foundation of all matrix calculus
  • Optimization — gradient descent and Newton's method